Run TLS destructors for main thread on UNIX platforms

library/std/src/lib.rs

@@ -282,6 +282,7 @@
 #![feature(allow_internal_unstable)]
 #![feature(asm_experimental_arch)]
 #![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/guard/key.rs

@@ -7,17 +7,38 @@ 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()
+        }
+    }
+
+    #[cfg(target_has_atomic_load_store = "8")]
+    fn enable_process() {
+        use crate::sync::atomic::{AtomicBool, Ordering};
+        use crate::sys::thread_local::key::at_process_exit;
 
-    // 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));
+        static REGISTERED: AtomicBool = AtomicBool::new(false);
+        if !REGISTERED.swap(true, Ordering::AcqRel) {
+            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,6 +49,11 @@ pub fn enable() {
             crate::rt::thread_cleanup();
         }
     }
+
+    enable_thread();
+
+    #[cfg(target_has_atomic_load_store = "8")]
+    enable_process();
 }
 
 /// On platforms with key-based TLS, the system runs the destructors for us.

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

@@ -21,3 +21,7 @@ pub unsafe fn get(key: Key) -> *mut u8 {
 pub unsafe fn destroy(key: Key) {
     Tls::destroy(AbiKey::from_usize(key))
 }
+
+pub unsafe fn at_process_exit(cb: unsafe extern "C" fn()) {
+    let _ = cb;
+}

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

@@ -26,3 +26,13 @@ pub unsafe fn destroy(key: Key) {
     let r = unsafe { libc::pthread_key_delete(key) };
     debug_assert_eq!(r, 0);
 }
+
+#[inline]
+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/mod.rs

@@ -141,7 +141,7 @@ pub(crate) mod key {
             #[cfg(test)]
             mod tests;
             pub(super) use racy::LazyKey;
-            pub(super) use unix::{Key, set};
+            pub(super) use unix::{Key, set, at_process_exit};
             #[cfg(any(not(target_thread_local), test))]
             pub(super) use unix::get;
             use unix::{create, destroy};
@@ -156,7 +156,7 @@ pub(crate) mod key {
             #[cfg(test)]
             mod tests;
             pub(super) use racy::LazyKey;
-            pub(super) use sgx::{Key, get, set};
+            pub(super) use sgx::{Key, get, set, at_process_exit};
             use sgx::{create, destroy};
         } else if #[cfg(target_os = "xous")] {
             mod racy;