Run TLS destructors at process exit on all platforms

This calls TLS destructors for the thread that initiates
a process exit. This is done by registering a process-wide
exit callback.

Previously UNIX platforms other than Linux and Apple
did not destruct TLS variables on the thread that
initiated the process exit (either by returning from
main or calling std::process::exit).

Author: surban

library/std/src/lib.rs

@@ -280,6 +280,7 @@
 #![feature(asm_experimental_arch)]
 #![feature(autodiff)]
 #![feature(cfg_sanitizer_cfi)]
+#![feature(cfg_target_has_atomic)]
 #![feature(cfg_target_thread_local)]
 #![feature(cfi_encoding)]
 #![feature(concat_idents)]

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

@@ -17,6 +17,9 @@ pub struct LazyKey {
     key: AtomicUsize,
     /// Destructor for the TLS value.
     dtor: Option<unsafe extern "C" fn(*mut u8)>,
+    /// Next element of process-wide destructor list.
+    #[cfg(not(target_thread_local))]
+    next: atomic::AtomicPtr<LazyKey>,
 }
 
 // Define a sentinel value that is likely not to be returned
@@ -31,18 +34,23 @@ const KEY_SENTVAL: usize = libc::PTHREAD_KEYS_MAX + 1;
 
 impl LazyKey {
     pub const fn new(dtor: Option<unsafe extern "C" fn(*mut u8)>) -> LazyKey {
-        LazyKey { key: atomic::AtomicUsize::new(KEY_SENTVAL), dtor }
+        LazyKey {
+            key: atomic::AtomicUsize::new(KEY_SENTVAL),
+            dtor,
+            #[cfg(not(target_thread_local))]
+            next: atomic::AtomicPtr::new(crate::ptr::null_mut()),
+        }
     }
 
     #[inline]
-    pub fn force(&self) -> super::Key {
+    pub fn force(&'static self) -> super::Key {
         match self.key.load(Ordering::Acquire) {
             KEY_SENTVAL => self.lazy_init() as super::Key,
             n => n as super::Key,
         }
     }
 
-    fn lazy_init(&self) -> usize {
+    fn lazy_init(&'static self) -> usize {
         // POSIX allows the key created here to be KEY_SENTVAL, but the compare_exchange
         // below relies on using KEY_SENTVAL as a sentinel value to check who won the
         // race to set the shared TLS key. As far as I know, there is no
@@ -70,7 +78,13 @@ impl LazyKey {
             Ordering::Acquire,
         ) {
             // The CAS succeeded, so we've created the actual key
-            Ok(_) => key as usize,
+            Ok(_) => {
+                #[cfg(not(target_thread_local))]
+                if self.dtor.is_some() {
+                    unsafe { register_dtor(self) };
+                }
+                key as usize
+            }
             // If someone beat us to the punch, use their key instead
             Err(n) => unsafe {
                 super::destroy(key);
@@ -79,3 +93,57 @@ impl LazyKey {
         }
     }
 }
+
+/// POSIX does not run TLS destructors on process exit.
+/// Thus we keep our own global list for that purpose.
+#[cfg(not(target_thread_local))]
+static DTORS: atomic::AtomicPtr<LazyKey> = atomic::AtomicPtr::new(crate::ptr::null_mut());
+
+/// Registers destructor to run at process exit.
+#[cfg(not(target_thread_local))]
+unsafe fn register_dtor(key: &'static LazyKey) {
+    crate::sys::thread_local::guard::enable();
+
+    let this = <*const LazyKey>::cast_mut(key);
+    // Use acquire ordering to pass along the changes done by the previously
+    // registered keys when we store the new head with release ordering.
+    let mut head = DTORS.load(Ordering::Acquire);
+    loop {
+        key.next.store(head, Ordering::Relaxed);
+        match DTORS.compare_exchange_weak(head, this, Ordering::Release, Ordering::Acquire) {
+            Ok(_) => break,
+            Err(new) => head = new,
+        }
+    }
+}
+
+/// 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() {
+    for _ in 0..5 {
+        let mut any_run = false;
+
+        // Use acquire ordering to observe key initialization.
+        let mut cur = DTORS.load(Ordering::Acquire);
+        while !cur.is_null() {
+            let key = unsafe { (*cur).key.load(Ordering::Acquire) };
+            let dtor = unsafe { (*cur).dtor.unwrap() };
+            cur = unsafe { (*cur).next.load(Ordering::Relaxed) };
+
+            let ptr = unsafe { super::get(key as _) };
+            if !ptr.is_null() {
+                unsafe {
+                    super::set(key as _, crate::ptr::null_mut());
+                    dtor(ptr as *mut _);
+                    any_run = true;
+                }
+            }
+        }
+
+        if !any_run {
+            break;
+        }
+    }
+}

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.