Auto merge of #2931 - max-heller:issue-2881, r=RalfJung

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

Miri currently treats allocations as leaked if they're only referenced in thread-local statics. For threads other than the main thread, this is correct, since the thread can terminate before the program does, but references in the main thread's locals should be treated as living for the duration of the program since the thread lives for the duration of the program.

This PR adds thread-local statics and TLS keys as "static roots" for leakage analysis, but does not yet bless the example program from #2881. See https://github.com/rust-lang/miri/issues/2881#issuecomment-1585666652

Closes #2881

src/tools/miri/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/tools/miri/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 });
             }

src/tools/miri/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.
+}

src/tools/miri/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
+

src/tools/miri/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.
+}

src/tools/miri/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
+

src/tools/miri/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))));
+}