implement support for multiple TLS destructors on macOS

4 files changed, 87 insertions, 30 deletions

diff --git a/src/tools/miri/src/shims/tls.rs b/src/tools/miri/src/shims/tls.rs
index c91386fa877..87074468789 100644
--- a/src/tools/miri/src/shims/tls.rs
+++ b/src/tools/miri/src/shims/tls.rs
@@ -36,9 +36,9 @@ pub struct TlsData<'tcx> {
     /// pthreads-style thread-local storage.
     keys: BTreeMap<TlsKey, TlsEntry<'tcx>>,
 
-    /// A single per thread destructor of the thread local storage (that's how
-    /// things work on macOS) with a data argument.
-    macos_thread_dtors: BTreeMap<ThreadId, (ty::Instance<'tcx>, Scalar)>,
+    /// On macOS, each thread holds a list of destructor functions with their
+    /// respective data arguments.
+    macos_thread_dtors: BTreeMap<ThreadId, Vec<(ty::Instance<'tcx>, Scalar)>>,
 }
 
 impl<'tcx> Default for TlsData<'tcx> {
@@ -119,26 +119,15 @@ impl<'tcx> TlsData<'tcx> {
         }
     }
 
-    /// Set the thread wide destructor of the thread local storage for the given
-    /// thread. This function is used to implement `_tlv_atexit` shim on MacOS.
-    ///
-    /// Thread wide dtors are available only on MacOS. There is one destructor
-    /// per thread as can be guessed from the following comment in the
-    /// [`_tlv_atexit`
-    /// implementation](https://github.com/opensource-apple/dyld/blob/195030646877261f0c8c7ad8b001f52d6a26f514/src/threadLocalVariables.c#L389):
-    ///
-    /// NOTE: this does not need locks because it only operates on current thread data
-    pub fn set_macos_thread_dtor(
+    /// Add a thread local storage destructor for the given thread. This function
+    /// is used to implement the `_tlv_atexit` shim on MacOS.
+    pub fn add_macos_thread_dtor(
         &mut self,
         thread: ThreadId,
         dtor: ty::Instance<'tcx>,
         data: Scalar,
     ) -> InterpResult<'tcx> {
-        if self.macos_thread_dtors.insert(thread, (dtor, data)).is_some() {
-            throw_unsup_format!(
-                "setting more than one thread local storage destructor for the same thread is not supported"
-            );
-        }
+        self.macos_thread_dtors.entry(thread).or_default().push((dtor, data));
         Ok(())
     }
 
@@ -202,6 +191,10 @@ impl<'tcx> TlsData<'tcx> {
         for TlsEntry { data, .. } in self.keys.values_mut() {
             data.remove(&thread_id);
         }
+
+        if let Some(dtors) = self.macos_thread_dtors.remove(&thread_id) {
+            assert!(dtors.is_empty(), "the destructors should have already been run");
+        }
     }
 }
 
@@ -212,7 +205,7 @@ impl VisitProvenance for TlsData<'_> {
         for scalar in keys.values().flat_map(|v| v.data.values()) {
             scalar.visit_provenance(visit);
         }
-        for (_, scalar) in macos_thread_dtors.values() {
+        for (_, scalar) in macos_thread_dtors.values().flatten() {
             scalar.visit_provenance(visit);
         }
     }
@@ -225,6 +218,7 @@ pub struct TlsDtorsState<'tcx>(TlsDtorsStatePriv<'tcx>);
 enum TlsDtorsStatePriv<'tcx> {
     #[default]
     Init,
+    MacOsDtors,
     PthreadDtors(RunningDtorState),
     /// For Windows Dtors, we store the list of functions that we still have to call.
     /// These are functions from the magic `.CRT$XLB` linker section.
@@ -243,11 +237,10 @@ impl<'tcx> TlsDtorsState<'tcx> {
                 Init => {
                     match this.tcx.sess.target.os.as_ref() {
                         "macos" => {
-                            // The macOS thread wide destructor runs "before any TLS slots get
-                            // freed", so do that first.
-                            this.schedule_macos_tls_dtor()?;
-                            // When that destructor is done, go on with the pthread dtors.
-                            break 'new_state PthreadDtors(Default::default());
+                            // macOS has a _tlv_atexit function that allows
+                            // registering destructors without associated keys.
+                            // These are run first.
+                            break 'new_state MacOsDtors;
                         }
                         _ if this.target_os_is_unix() => {
                             // All other Unixes directly jump to running the pthread dtors.
@@ -266,6 +259,14 @@ impl<'tcx> TlsDtorsState<'tcx> {
                         }
                     }
                 }
+                MacOsDtors => {
+                    match this.schedule_macos_tls_dtor()? {
+                        Poll::Pending => return Ok(Poll::Pending),
+                        // After all macOS destructors are run, the system switches
+                        // to destroying the pthread destructors.
+                        Poll::Ready(()) => break 'new_state PthreadDtors(Default::default()),
+                    }
+                }
                 PthreadDtors(state) => {
                     match this.schedule_next_pthread_tls_dtor(state)? {
                         Poll::Pending => return Ok(Poll::Pending), // just keep going
@@ -328,12 +329,15 @@ trait EvalContextPrivExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         Ok(())
     }
 
-    /// Schedule the MacOS thread destructor of the thread local storage to be
-    /// executed.
-    fn schedule_macos_tls_dtor(&mut self) -> InterpResult<'tcx> {
+    /// Schedule the macOS thread local storage destructors to be executed.
+    fn schedule_macos_tls_dtor(&mut self) -> InterpResult<'tcx, Poll<()>> {
         let this = self.eval_context_mut();
         let thread_id = this.active_thread();
-        if let Some((instance, data)) = this.machine.tls.macos_thread_dtors.remove(&thread_id) {
+        // macOS keeps track of TLS destructors in a stack. If a destructor
+        // registers another destructor, it will be run next.
+        // See https://github.com/apple-oss-distributions/dyld/blob/d552c40cd1de105f0ec95008e0e0c0972de43456/dyld/DyldRuntimeState.cpp#L2277
+        let dtor = this.machine.tls.macos_thread_dtors.get_mut(&thread_id).and_then(Vec::pop);
+        if let Some((instance, data)) = dtor {
             trace!("Running macos dtor {:?} on {:?} at {:?}", instance, data, thread_id);
 
             this.call_function(
@@ -343,8 +347,11 @@ trait EvalContextPrivExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 None,
                 StackPopCleanup::Root { cleanup: true },
             )?;
+
+            return Ok(Poll::Pending);
         }
-        Ok(())
+
+        Ok(Poll::Ready(()))
     }
 
     /// Schedule a pthread TLS destructor. Returns `true` if found
diff --git a/src/tools/miri/src/shims/unix/macos/foreign_items.rs b/src/tools/miri/src/shims/unix/macos/foreign_items.rs
index 25002f0a611..aefb5b2de56 100644
--- a/src/tools/miri/src/shims/unix/macos/foreign_items.rs
+++ b/src/tools/miri/src/shims/unix/macos/foreign_items.rs
@@ -132,7 +132,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 let dtor = this.get_ptr_fn(dtor)?.as_instance()?;
                 let data = this.read_scalar(data)?;
                 let active_thread = this.active_thread();
-                this.machine.tls.set_macos_thread_dtor(active_thread, dtor, data)?;
+                this.machine.tls.add_macos_thread_dtor(active_thread, dtor, data)?;
             }
 
             // Querying system information
diff --git a/src/tools/miri/tests/pass/tls/macos_tlv_atexit.rs b/src/tools/miri/tests/pass/tls/macos_tlv_atexit.rs
new file mode 100644
index 00000000000..845f50c1eba
--- /dev/null
+++ b/src/tools/miri/tests/pass/tls/macos_tlv_atexit.rs
@@ -0,0 +1,43 @@
+//@only-target-darwin
+
+use std::thread;
+
+extern "C" {
+    fn _tlv_atexit(dtor: unsafe extern "C" fn(*mut u8), arg: *mut u8);
+}
+
+fn register<F>(f: F)
+where
+    F: FnOnce() + 'static,
+{
+    // This will receive the pointer passed into `_tlv_atexit`, which is the
+    // original `f` but boxed up.
+    unsafe extern "C" fn run<F>(ptr: *mut u8)
+    where
+        F: FnOnce() + 'static,
+    {
+        let f = unsafe { Box::from_raw(ptr as *mut F) };
+        f()
+    }
+
+    unsafe {
+        _tlv_atexit(run::<F>, Box::into_raw(Box::new(f)) as *mut u8);
+    }
+}
+
+fn main() {
+    thread::spawn(|| {
+        register(|| println!("dtor 2"));
+        register(|| println!("dtor 1"));
+        println!("exiting thread");
+    })
+    .join()
+    .unwrap();
+
+    println!("exiting main");
+    register(|| println!("dtor 5"));
+    register(|| {
+        println!("registering dtor in dtor 3");
+        register(|| println!("dtor 4"));
+    });
+}
diff --git a/src/tools/miri/tests/pass/tls/macos_tlv_atexit.stdout b/src/tools/miri/tests/pass/tls/macos_tlv_atexit.stdout
new file mode 100644
index 00000000000..89d6ca25935
--- /dev/null
+++ b/src/tools/miri/tests/pass/tls/macos_tlv_atexit.stdout
@@ -0,0 +1,7 @@
+exiting thread
+dtor 1
+dtor 2
+exiting main
+registering dtor in dtor 3
+dtor 4
+dtor 5