Treat thread-local statics on main thread as static roots for leakage analysis

src/concurrency/thread.rs

@@ -33,6 +33,15 @@ enum SchedulingAction {
     Sleep(Duration),
 }
 
+/// What to do with TLS allocations from terminated threads
+pub enum TlsAllocAction {
+    /// Deallocate backing memory of thread-local statics as usual
+    Deallocate,
+    /// Skip deallocating backing memory of thread-local statics and consider all memory reachable
+    /// from them as "allowed to leak" (like global `static`s).
+    Leak,
+}
+
 /// Trait for callbacks that can be executed when some event happens, such as after a timeout.
 pub trait MachineCallback<'mir, 'tcx>: VisitTags {
     fn call(&self, ecx: &mut InterpCx<'mir, 'tcx, MiriMachine<'mir, 'tcx>>) -> InterpResult<'tcx>;
@@ -1051,7 +1060,8 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                         // See if this thread can do something else.
                         match this.run_on_stack_empty()? {
                             Poll::Pending => {} // keep going
-                            Poll::Ready(()) => this.terminate_active_thread()?,
+                            Poll::Ready(()) =>
+                                this.terminate_active_thread(TlsAllocAction::Deallocate)?,
                         }
                     }
                 }
@@ -1066,21 +1076,29 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
     }
 
     /// Handles thread termination of the active thread: wakes up threads joining on this one,
-    /// and deallocated thread-local statics.
+    /// and deals with the thread's thread-local statics according to `tls_alloc_action`.
     ///
     /// This is called by the eval loop when a thread's on_stack_empty returns `Ready`.
     #[inline]
-    fn terminate_active_thread(&mut self) -> InterpResult<'tcx> {
+    fn terminate_active_thread(&mut self, tls_alloc_action: TlsAllocAction) -> InterpResult<'tcx> {
         let this = self.eval_context_mut();
         let thread = this.active_thread_mut();
         assert!(thread.stack.is_empty(), "only threads with an empty stack can be terminated");
         thread.state = ThreadState::Terminated;
 
         let current_span = this.machine.current_span();
-        for ptr in
-            this.machine.threads.thread_terminated(this.machine.data_race.as_mut(), current_span)
-        {
-            this.deallocate_ptr(ptr.into(), None, MiriMemoryKind::Tls.into())?;
+        let thread_local_allocations =
+            this.machine.threads.thread_terminated(this.machine.data_race.as_mut(), current_span);
+        for ptr in thread_local_allocations {
+            match tls_alloc_action {
+                TlsAllocAction::Deallocate =>
+                    this.deallocate_ptr(ptr.into(), None, MiriMemoryKind::Tls.into())?,
+                TlsAllocAction::Leak =>
+                    if let Some(alloc) = ptr.provenance.get_alloc_id() {
+                        trace!("Thread-local static leaked and stored as static root: {:?}", alloc);
+                        this.machine.static_roots.push(alloc);
+                    },
+            }
         }
         Ok(())
     }

src/eval.rs

@@ -11,6 +11,7 @@ use log::info;
 use rustc_middle::ty::Ty;
 
 use crate::borrow_tracker::RetagFields;
+use crate::concurrency::thread::TlsAllocAction;
 use crate::diagnostics::report_leaks;
 use rustc_data_structures::fx::FxHashSet;
 use rustc_hir::def::Namespace;
@@ -244,9 +245,9 @@ impl MainThreadState {
                 // Figure out exit code.
                 let ret_place = this.machine.main_fn_ret_place.clone().unwrap();
                 let exit_code = this.read_target_isize(&ret_place)?;
-                // Need to call this ourselves since we are not going to return to the scheduler
-                // loop, and we want the main thread TLS to not show up as memory leaks.
-                this.terminate_active_thread()?;
+                // Deal with our thread-local memory. We do *not* want to actually free it, instead we consider TLS
+                // to be like a global `static`, so that all memory reached by it is considered to "not leak".
+                this.terminate_active_thread(TlsAllocAction::Leak)?;
                 // Stop interpreter loop.
                 throw_machine_stop!(TerminationInfo::Exit { code: exit_code, leak_check: true });
             }

tests/fail/leak_in_lib_tls.rs

@@ -0,0 +1,21 @@
+//@error-in-other-file: memory leaked
+//@normalize-stderr-test: ".*│.*" -> "$$stripped$$"
+
+use std::cell::Cell;
+
+pub fn main() {
+    thread_local! {
+        static TLS: Cell<Option<&'static i32>> = Cell::new(None);
+    }
+
+    std::thread::spawn(|| {
+        TLS.with(|cell| {
+            cell.set(Some(Box::leak(Box::new(123))));
+        });
+    })
+    .join()
+    .unwrap();
+
+    // Imagine the program running for a long time while the thread is gone
+    // and this memory still sits around, unused -- leaked.
+}

tests/fail/leak_in_lib_tls.stderr

@@ -0,0 +1,32 @@
+error: memory leaked: ALLOC (Rust heap, size: 4, align: 4), allocated here:
+  --> RUSTLIB/alloc/src/alloc.rs:LL:CC
+   |
+LL |         __rust_alloc(layout.size(), layout.align())
+   |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   |
+   = note: inside `std::alloc::alloc` at RUSTLIB/alloc/src/alloc.rs:LL:CC
+   = note: inside `std::alloc::Global::alloc_impl` at RUSTLIB/alloc/src/alloc.rs:LL:CC
+   = note: inside `<std::alloc::Global as std::alloc::Allocator>::allocate` at RUSTLIB/alloc/src/alloc.rs:LL:CC
+   = note: inside `alloc::alloc::exchange_malloc` at RUSTLIB/alloc/src/alloc.rs:LL:CC
+   = note: inside `std::boxed::Box::<i32>::new` at RUSTLIB/alloc/src/boxed.rs:LL:CC
+note: inside closure
+  --> $DIR/leak_in_lib_tls.rs:LL:CC
+   |
+LL |             cell.set(Some(Box::leak(Box::new(123))));
+   |                                     ^^^^^^^^^^^^^
+   = note: inside `std::thread::LocalKey::<std::cell::Cell<std::option::Option<&i32>>>::try_with::<{closure@$DIR/leak_in_lib_tls.rs:LL:CC}, ()>` at RUSTLIB/std/src/thread/local.rs:LL:CC
+   = note: inside `std::thread::LocalKey::<std::cell::Cell<std::option::Option<&i32>>>::with::<{closure@$DIR/leak_in_lib_tls.rs:LL:CC}, ()>` at RUSTLIB/std/src/thread/local.rs:LL:CC
+note: inside closure
+  --> $DIR/leak_in_lib_tls.rs:LL:CC
+   |
+LL | /         TLS.with(|cell| {
+LL | |             cell.set(Some(Box::leak(Box::new(123))));
+LL | |         });
+   | |__________^
+
+note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` for a verbose backtrace
+
+note: the evaluated program leaked memory, pass `-Zmiri-ignore-leaks` to disable this check
+
+error: aborting due to previous error
+

tests/fail/leak_in_static_tls.rs

@@ -0,0 +1,22 @@
+//@error-in-other-file: memory leaked
+//@normalize-stderr-test: ".*│.*" -> "$$stripped$$"
+
+#![feature(thread_local)]
+
+use std::cell::Cell;
+
+/// Ensure that leaks through `thread_local` statics *not in the main thread*
+/// are detected.
+pub fn main() {
+    #[thread_local]
+    static TLS: Cell<Option<&'static i32>> = Cell::new(None);
+
+    std::thread::spawn(|| {
+        TLS.set(Some(Box::leak(Box::new(123))));
+    })
+    .join()
+    .unwrap();
+
+    // Imagine the program running for a long time while the thread is gone
+    // and this memory still sits around, unused -- leaked.
+}

tests/fail/leak_in_static_tls.stderr

@@ -0,0 +1,23 @@
+error: memory leaked: ALLOC (Rust heap, size: 4, align: 4), allocated here:
+  --> RUSTLIB/alloc/src/alloc.rs:LL:CC
+   |
+LL |         __rust_alloc(layout.size(), layout.align())
+   |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   |
+   = note: inside `std::alloc::alloc` at RUSTLIB/alloc/src/alloc.rs:LL:CC
+   = note: inside `std::alloc::Global::alloc_impl` at RUSTLIB/alloc/src/alloc.rs:LL:CC
+   = note: inside `<std::alloc::Global as std::alloc::Allocator>::allocate` at RUSTLIB/alloc/src/alloc.rs:LL:CC
+   = note: inside `alloc::alloc::exchange_malloc` at RUSTLIB/alloc/src/alloc.rs:LL:CC
+   = note: inside `std::boxed::Box::<i32>::new` at RUSTLIB/alloc/src/boxed.rs:LL:CC
+note: inside closure
+  --> $DIR/leak_in_static_tls.rs:LL:CC
+   |
+LL |         TLS.set(Some(Box::leak(Box::new(123))));
+   |                                ^^^^^^^^^^^^^
+
+note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` for a verbose backtrace
+
+note: the evaluated program leaked memory, pass `-Zmiri-ignore-leaks` to disable this check
+
+error: aborting due to previous error
+

tests/pass/concurrency/tls_leak_main_thread_allowed.rs

@@ -0,0 +1,21 @@
+//@ignore-target-windows: Windows uses a different mechanism for `thread_local!`
+#![feature(thread_local)]
+
+use std::cell::Cell;
+
+// Thread-local variables in the main thread are basically like `static` (they live
+// as long as the program does), so make sure we treat them the same for leak purposes.
+pub fn main() {
+    thread_local! {
+        static TLS_KEY: Cell<Option<&'static i32>> = Cell::new(None);
+    }
+
+    TLS_KEY.with(|cell| {
+        cell.set(Some(Box::leak(Box::new(123))));
+    });
+
+    #[thread_local]
+    static TLS: Cell<Option<&'static i32>> = Cell::new(None);
+
+    TLS.set(Some(Box::leak(Box::new(123))));
+}