Run TLS destructors at process exit on all platforms

library/std/src/sys/thread_local/exit/unix.rs

@@ -0,0 +1,10 @@
+use crate::mem;
+
+pub unsafe fn at_process_exit(cb: unsafe extern "C" fn()) {
+    // Miri does not support atexit.
+    #[cfg(not(miri))]
+    assert_eq!(unsafe { libc::atexit(mem::transmute(cb)) }, 0);
+
+    #[cfg(miri)]
+    let _ = cb;
+}

library/std/src/sys/thread_local/guard/key.rs

@@ -3,21 +3,40 @@
 //! that will run all native TLS destructors in the destructor list.
 
 use crate::ptr;
+use crate::sync::atomic::{AtomicBool, Ordering};
+use crate::sys::thread_local::exit::at_process_exit;
 use crate::sys::thread_local::key::{LazyKey, set};
 
 #[cfg(target_thread_local)]
 pub fn enable() {
-    use crate::sys::thread_local::destructors;
+    fn enable_thread() {
+        static DTORS: LazyKey = LazyKey::new(Some(run_thread));
 
-    static DTORS: LazyKey = LazyKey::new(Some(run));
+        // Setting the key value to something other than NULL will result in the
+        // destructor being run at thread exit.
+        unsafe {
+            set(DTORS.force(), ptr::without_provenance_mut(1));
+        }
+
+        unsafe extern "C" fn run_thread(_: *mut u8) {
+            run()
+        }
+    }
 
-    // Setting the key value to something other than NULL will result in the
-    // destructor being run at thread exit.
-    unsafe {
-        set(DTORS.force(), ptr::without_provenance_mut(1));
+    fn enable_process() {
+        static REGISTERED: AtomicBool = AtomicBool::new(false);
+        if !REGISTERED.swap(true, Ordering::Relaxed) {
+            unsafe { at_process_exit(run_process) };
+        }
+
+        unsafe extern "C" fn run_process() {
+            run()
+        }
     }
 
-    unsafe extern "C" fn run(_: *mut u8) {
+    fn run() {
+        use crate::sys::thread_local::destructors;
+
         unsafe {
             destructors::run();
             // On platforms with `__cxa_thread_atexit_impl`, `destructors::run`
@@ -28,33 +47,55 @@ pub fn enable() {
             crate::rt::thread_cleanup();
         }
     }
+
+    enable_thread();
+    enable_process();
 }
 
 /// On platforms with key-based TLS, the system runs the destructors for us.
 /// We still have to make sure that [`crate::rt::thread_cleanup`] is called,
 /// however. This is done by defering the execution of a TLS destructor to
 /// the next round of destruction inside the TLS destructors.
+///
+/// POSIX systems do not run TLS destructors at process exit.
+/// Thus we register our own callback to invoke them in that case.
 #[cfg(not(target_thread_local))]
 pub fn enable() {
-    const DEFER: *mut u8 = ptr::without_provenance_mut(1);
-    const RUN: *mut u8 = ptr::without_provenance_mut(2);
-
-    static CLEANUP: LazyKey = LazyKey::new(Some(run));
-
-    unsafe { set(CLEANUP.force(), DEFER) }
-
-    unsafe extern "C" fn run(state: *mut u8) {
-        if state == DEFER {
-            // Make sure that this function is run again in the next round of
-            // TLS destruction. If there is no futher round, there will be leaks,
-            // but that's okay, `thread_cleanup` is not guaranteed to be called.
-            unsafe { set(CLEANUP.force(), RUN) }
-        } else {
-            debug_assert_eq!(state, RUN);
-            // If the state is still RUN in the next round of TLS destruction,
-            // it means that no other TLS destructors defined by this runtime
-            // have been run, as they would have set the state to DEFER.
-            crate::rt::thread_cleanup();
+    fn enable_thread() {
+        const DEFER: *mut u8 = ptr::without_provenance_mut(1);
+        const RUN: *mut u8 = ptr::without_provenance_mut(2);
+
+        static CLEANUP: LazyKey = LazyKey::new(Some(run_thread));
+
+        unsafe { set(CLEANUP.force(), DEFER) }
+
+        unsafe extern "C" fn run_thread(state: *mut u8) {
+            if state == DEFER {
+                // Make sure that this function is run again in the next round of
+                // TLS destruction. If there is no futher round, there will be leaks,
+                // but that's okay, `thread_cleanup` is not guaranteed to be called.
+                unsafe { set(CLEANUP.force(), RUN) }
+            } else {
+                debug_assert_eq!(state, RUN);
+                // If the state is still RUN in the next round of TLS destruction,
+                // it means that no other TLS destructors defined by this runtime
+                // have been run, as they would have set the state to DEFER.
+                crate::rt::thread_cleanup();
+            }
+        }
+    }
+
+    fn enable_process() {
+        static REGISTERED: AtomicBool = AtomicBool::new(false);
+        if !REGISTERED.swap(true, Ordering::Relaxed) {
+            unsafe { at_process_exit(run_process) };
+        }
+
+        unsafe extern "C" fn run_process() {
+            unsafe { crate::sys::thread_local::key::run_dtors() };
         }
     }
+
+    enable_thread();
+    enable_process();
 }

library/std/src/sys/thread_local/guard/statik.rs

@@ -0,0 +1,24 @@
+//! The platform has no threads, so we just need to register
+//! a process exit callback.
+
+use crate::cell::Cell;
+use crate::sys::thread_local::exit::at_process_exit;
+use crate::sys::thread_local::statik::run_dtors;
+
+pub fn enable() {
+    struct Registered(Cell<bool>);
+    // SAFETY: the target doesn't have threads.
+    unsafe impl Sync for Registered {}
+
+    static REGISTERED: Registered = Registered(Cell::new(false));
+
+    if !REGISTERED.0.get() {
+        REGISTERED.0.set(true);
+        unsafe { at_process_exit(run_process) };
+    }
+
+    unsafe extern "C" fn run_process() {
+        unsafe { run_dtors() };
+        crate::rt::thread_cleanup();
+    }
+}

library/std/src/sys/thread_local/key/racy.rs

@@ -78,4 +78,66 @@ impl LazyKey {
             },
         }
     }
+
+    /// Registers destructor to run at process exit.
+    #[cfg(not(target_thread_local))]
+    pub fn register_process_dtor(&'static self) {
+        if self.dtor.is_none() {
+            return;
+        }
+
+        crate::sys::thread_local::guard::enable();
+        lazy_keys().borrow_mut().push(self);
+    }
+}
+
+/// POSIX does not run TLS destructors on process exit.
+/// Thus we keep our own thread-local list for that purpose.
+#[cfg(not(target_thread_local))]
+fn lazy_keys() -> &'static crate::cell::RefCell<Vec<&'static LazyKey>> {
+    static KEY: LazyKey = LazyKey::new(Some(drop_lazy_keys));
+
+    unsafe extern "C" fn drop_lazy_keys(ptr: *mut u8) {
+        let ptr = ptr as *mut crate::cell::RefCell<Vec<&'static LazyKey>>;
+        drop(unsafe { Box::from_raw(ptr) });
+    }
+
+    let key = KEY.force();
+    let mut ptr = unsafe { super::get(key) as *const crate::cell::RefCell<Vec<&'static LazyKey>> };
+    if ptr.is_null() {
+        let list = Box::new(crate::cell::RefCell::new(Vec::new()));
+        ptr = Box::into_raw(list);
+        unsafe { super::set(key, ptr as _) };
+    }
+
+    unsafe { &*ptr }
+}
+
+/// Run destructors at process exit.
+///
+/// SAFETY: This will and must only be run by the destructor callback in [`guard`].
+#[cfg(not(target_thread_local))]
+pub unsafe fn run_dtors() {
+    let lazy_keys_cell = lazy_keys();
+
+    for _ in 0..5 {
+        let mut any_run = false;
+
+        for lazy_key in lazy_keys_cell.take() {
+            let key = lazy_key.force();
+            let ptr = unsafe { super::get(key) };
+            if !ptr.is_null() {
+                // SAFETY: only keys with destructors are registered.
+                unsafe {
+                    let Some(dtor) = &lazy_key.dtor else { crate::hint::unreachable_unchecked() };
+                    dtor(ptr);
+                }
+                any_run = true;
+            }
+        }
+
+        if !any_run {
+            break;
+        }
+    }
 }

library/std/src/sys/thread_local/key/windows.rs

@@ -126,6 +126,10 @@ impl LazyKey {
             }
         }
     }
+
+    pub fn register_process_dtor(&'static self) {
+        // On Windows destructor registration is performed in LazyKey::init.
+    }
 }
 
 unsafe impl Send for LazyKey {}

library/std/src/sys/thread_local/mod.rs

@@ -85,18 +85,16 @@ pub(crate) mod guard {
         } else if #[cfg(target_os = "windows")] {
             mod windows;
             pub(crate) use windows::enable;
-        } else if #[cfg(any(
-            all(target_family = "wasm", not(
-                all(target_os = "wasi", target_env = "p1", target_feature = "atomics")
-            )),
-            target_os = "uefi",
-            target_os = "zkvm",
+        } else if #[cfg(all(
+            target_family = "wasm",
+            target_feature = "atomics",
+            not(all(target_os = "wasi", target_env = "p1"))
         ))] {
             pub(crate) fn enable() {
                 // FIXME: Right now there is no concept of "thread exit" on
-                // wasm, but this is likely going to show up at some point in
-                // the form of an exported symbol that the wasm runtime is going
-                // to be expected to call. For now we just leak everything, but
+                // wasm-unknown-unknown, but this is likely going to show up at some
+                // point in the form of an exported symbol that the wasm runtime is
+                // going to be expected to call. For now we just leak everything, but
                 // if such a function starts to exist it will probably need to
                 // iterate the destructor list with these functions:
                 #[cfg(all(target_family = "wasm", target_feature = "atomics"))]
@@ -115,6 +113,13 @@ pub(crate) mod guard {
         } else if #[cfg(target_os = "solid_asp3")] {
             mod solid;
             pub(crate) use solid::enable;
+        } else if #[cfg(any(
+            all(target_family = "wasm", not(target_feature = "atomics")),
+            target_os = "uefi",
+            target_os = "zkvm",
+        ))] {
+            mod statik;
+            pub(crate) use statik::enable;
         } else {
             mod key;
             pub(crate) use key::enable;
@@ -144,6 +149,8 @@ pub(crate) mod key {
             #[cfg(test)]
             mod tests;
             pub(super) use racy::LazyKey;
+            #[cfg(not(target_thread_local))]
+            pub(super) use racy::run_dtors;
             pub(super) use unix::{Key, set};
             #[cfg(any(not(target_thread_local), test))]
             pub(super) use unix::get;
@@ -158,7 +165,7 @@ pub(crate) mod key {
             mod sgx;
             #[cfg(test)]
             mod tests;
-            pub(super) use racy::LazyKey;
+            pub(super) use racy::{LazyKey, run_dtors};
             pub(super) use sgx::{Key, get, set};
             use sgx::{create, destroy};
         } else if #[cfg(target_os = "xous")] {
@@ -174,6 +181,31 @@ pub(crate) mod key {
     }
 }
 
+/// Process exit callback.
+///
+/// Some platforms (POSIX) do not run TLS destructors at process exit.
+/// Thus we need to register an exit callback to run them in that case.
+pub(crate) mod exit {
+    cfg_if::cfg_if! {
+        if #[cfg(any(
+            all(
+                not(target_vendor = "apple"),
+                not(target_family = "wasm"),
+                target_family = "unix",
+            ),
+            target_os = "teeos",
+            all(target_os = "wasi", target_env = "p1"),
+        ))] {
+            mod unix;
+            pub(super) use unix::at_process_exit;
+        } else if #[cfg(target_family = "wasm")] {
+            pub unsafe fn at_process_exit(cb: unsafe extern "C" fn()) {
+                let _ = cb;
+            }
+        }
+    }
+}
+
 /// Run a callback in a scenario which must not unwind (such as a `extern "C"
 /// fn` declared in a user crate). If the callback unwinds anyway, then
 /// `rtabort` with a message about thread local panicking on drop.

library/std/src/sys/thread_local/os.rs

@@ -79,7 +79,11 @@ impl<T: 'static> Storage<T> {
             unsafe { &(*ptr).value }
         } else {
             // SAFETY: trivially correct.
-            unsafe { Self::try_initialize(key, ptr, i, f) }
+            let (ptr, was_null) = unsafe { Self::try_initialize(key, ptr, i, f) };
+            if was_null {
+                self.key.register_process_dtor();
+            }
+            ptr
         }
     }
 
@@ -91,10 +95,10 @@ impl<T: 'static> Storage<T> {
         ptr: *mut Value<T>,
         i: Option<&mut Option<T>>,
         f: impl FnOnce() -> T,
-    ) -> *const T {
+    ) -> (*const T, bool) {
         if ptr.addr() == 1 {
             // destructor is running
-            return ptr::null();
+            return (ptr::null(), false);
         }
 
         let value = Box::new(Value { value: i.and_then(Option::take).unwrap_or_else(f), key });
@@ -120,7 +124,7 @@ impl<T: 'static> Storage<T> {
         }
 
         // SAFETY: We just created this value above.
-        unsafe { &(*ptr).value }
+        (unsafe { &(*ptr).value }, old.is_null())
     }
 }