refactor scheduling of TLS dtors

Author: RalfJung

src/tools/miri/src/concurrency/thread.rs

@@ -3,6 +3,7 @@
 use std::cell::RefCell;
 use std::collections::hash_map::Entry;
 use std::num::TryFromIntError;
+use std::task::Poll;
 use std::time::{Duration, SystemTime};
 
 use log::trace;
@@ -16,6 +17,7 @@ use rustc_target::spec::abi::Abi;
 
 use crate::concurrency::data_race;
 use crate::concurrency::sync::SynchronizationState;
+use crate::shims::tls;
 use crate::*;
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -24,10 +26,8 @@ pub enum SchedulingAction {
     ExecuteStep,
     /// Execute a timeout callback.
     ExecuteTimeoutCallback,
-    /// Execute destructors of the active thread.
-    ExecuteDtors,
-    /// Stop the program.
-    Stop,
+    /// Wait for a bit, until there is a timeout to be called.
+    Sleep(Duration),
 }
 
 /// Trait for callbacks that can be executed when some event happens, such as after a timeout.
@@ -41,9 +41,6 @@ type TimeoutCallback<'mir, 'tcx> = Box<dyn MachineCallback<'mir, 'tcx> + 'tcx>;
 #[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq, Hash)]
 pub struct ThreadId(u32);
 
-/// The main thread. When it terminates, the whole application terminates.
-const MAIN_THREAD: ThreadId = ThreadId(0);
-
 impl ThreadId {
     pub fn to_u32(self) -> u32 {
         self.0
@@ -118,6 +115,12 @@ pub struct Thread<'mir, 'tcx> {
     /// The virtual call stack.
     stack: Vec<Frame<'mir, 'tcx, Provenance, FrameData<'tcx>>>,
 
+    /// The function to call when the stack ran empty, to figure out what to do next.
+    /// Conceptually, this is the interpreter implementation of the things that happen 'after' the
+    /// Rust language entry point for this thread returns (usually implemented by the C or OS runtime).
+    /// (`None` is an error, it means the callback has not been set up yet or is actively running.)
+    pub(crate) on_stack_empty: Option<StackEmptyCallback<'mir, 'tcx>>,
+
     /// The index of the topmost user-relevant frame in `stack`. This field must contain
     /// the value produced by `get_top_user_relevant_frame`.
     /// The `None` state here represents
@@ -137,19 +140,10 @@ pub struct Thread<'mir, 'tcx> {
     pub(crate) last_error: Option<MPlaceTy<'tcx, Provenance>>,
 }
 
-impl<'mir, 'tcx> Thread<'mir, 'tcx> {
-    /// Check if the thread is done executing (no more stack frames). If yes,
-    /// change the state to terminated and return `true`.
-    fn check_terminated(&mut self) -> bool {
-        if self.state == ThreadState::Enabled {
-            if self.stack.is_empty() {
-                self.state = ThreadState::Terminated;
-                return true;
-            }
-        }
-        false
-    }
+pub type StackEmptyCallback<'mir, 'tcx> =
+    Box<dyn FnMut(&mut MiriInterpCx<'mir, 'tcx>) -> InterpResult<'tcx, Poll<()>>>;
 
+impl<'mir, 'tcx> Thread<'mir, 'tcx> {
     /// Get the name of the current thread, or `<unnamed>` if it was not set.
     fn thread_name(&self) -> &[u8] {
         if let Some(ref thread_name) = self.thread_name { thread_name } else { b"<unnamed>" }
@@ -202,28 +196,21 @@ impl<'mir, 'tcx> std::fmt::Debug for Thread<'mir, 'tcx> {
     }
 }
 
-impl<'mir, 'tcx> Default for Thread<'mir, 'tcx> {
-    fn default() -> Self {
+impl<'mir, 'tcx> Thread<'mir, 'tcx> {
+    fn new(name: Option<&str>, on_stack_empty: Option<StackEmptyCallback<'mir, 'tcx>>) -> Self {
         Self {
             state: ThreadState::Enabled,
-            thread_name: None,
+            thread_name: name.map(|name| Vec::from(name.as_bytes())),
             stack: Vec::new(),
             top_user_relevant_frame: None,
             join_status: ThreadJoinStatus::Joinable,
             panic_payload: None,
             last_error: None,
+            on_stack_empty,
         }
     }
 }
 
-impl<'mir, 'tcx> Thread<'mir, 'tcx> {
-    fn new(name: &str) -> Self {
-        let mut thread = Thread::default();
-        thread.thread_name = Some(Vec::from(name.as_bytes()));
-        thread
-    }
-}
-
 impl VisitTags for Thread<'_, '_> {
     fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
         let Thread {
@@ -234,6 +221,7 @@ impl VisitTags for Thread<'_, '_> {
             state: _,
             thread_name: _,
             join_status: _,
+            on_stack_empty: _, // we assume the closure captures no GC-relevant state
         } = self;
 
         panic_payload.visit_tags(visit);
@@ -327,22 +315,6 @@ pub struct ThreadManager<'mir, 'tcx> {
     timeout_callbacks: FxHashMap<ThreadId, TimeoutCallbackInfo<'mir, 'tcx>>,
 }
 
-impl<'mir, 'tcx> Default for ThreadManager<'mir, 'tcx> {
-    fn default() -> Self {
-        let mut threads = IndexVec::new();
-        // Create the main thread and add it to the list of threads.
-        threads.push(Thread::new("main"));
-        Self {
-            active_thread: ThreadId::new(0),
-            threads,
-            sync: SynchronizationState::default(),
-            thread_local_alloc_ids: Default::default(),
-            yield_active_thread: false,
-            timeout_callbacks: FxHashMap::default(),
-        }
-    }
-}
-
 impl VisitTags for ThreadManager<'_, '_> {
     fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
         let ThreadManager {
@@ -367,8 +339,28 @@ impl VisitTags for ThreadManager<'_, '_> {
     }
 }
 
+impl<'mir, 'tcx> Default for ThreadManager<'mir, 'tcx> {
+    fn default() -> Self {
+        let mut threads = IndexVec::new();
+        // Create the main thread and add it to the list of threads.
+        threads.push(Thread::new(Some("main"), None));
+        Self {
+            active_thread: ThreadId::new(0),
+            threads,
+            sync: SynchronizationState::default(),
+            thread_local_alloc_ids: Default::default(),
+            yield_active_thread: false,
+            timeout_callbacks: FxHashMap::default(),
+        }
+    }
+}
+
 impl<'mir, 'tcx: 'mir> ThreadManager<'mir, 'tcx> {
-    pub(crate) fn init(ecx: &mut MiriInterpCx<'mir, 'tcx>) {
+    pub(crate) fn init(
+        ecx: &mut MiriInterpCx<'mir, 'tcx>,
+        on_main_stack_empty: StackEmptyCallback<'mir, 'tcx>,
+    ) {
+        ecx.machine.threads.threads[ThreadId::new(0)].on_stack_empty = Some(on_main_stack_empty);
         if ecx.tcx.sess.target.os.as_ref() != "windows" {
             // The main thread can *not* be joined on except on windows.
             ecx.machine.threads.threads[ThreadId::new(0)].join_status = ThreadJoinStatus::Detached;
@@ -411,9 +403,9 @@ impl<'mir, 'tcx: 'mir> ThreadManager<'mir, 'tcx> {
     }
 
     /// Create a new thread and returns its id.
-    fn create_thread(&mut self) -> ThreadId {
+    fn create_thread(&mut self, on_stack_empty: StackEmptyCallback<'mir, 'tcx>) -> ThreadId {
         let new_thread_id = ThreadId::new(self.threads.len());
-        self.threads.push(Default::default());
+        self.threads.push(Thread::new(None, Some(on_stack_empty)));
         new_thread_id
     }
 
@@ -458,6 +450,7 @@ impl<'mir, 'tcx: 'mir> ThreadManager<'mir, 'tcx> {
     }
 
     /// Get a mutable borrow of the currently active thread.
+    /// (Private for a bit of protection.)
     fn active_thread_mut(&mut self) -> &mut Thread<'mir, 'tcx> {
         &mut self.threads[self.active_thread]
     }
@@ -669,37 +662,25 @@ impl<'mir, 'tcx: 'mir> ThreadManager<'mir, 'tcx> {
     /// long as we can and switch only when we have to (the active thread was
     /// blocked, terminated, or has explicitly asked to be preempted).
     fn schedule(&mut self, clock: &Clock) -> InterpResult<'tcx, SchedulingAction> {
-        // Check whether the thread has **just** terminated (`check_terminated`
-        // checks whether the thread has popped all its stack and if yes, sets
-        // the thread state to terminated).
-        if self.threads[self.active_thread].check_terminated() {
-            return Ok(SchedulingAction::ExecuteDtors);
-        }
-        // If we get here again and the thread is *still* terminated, there are no more dtors to run.
-        if self.threads[MAIN_THREAD].state == ThreadState::Terminated {
-            // The main thread terminated; stop the program.
-            // We do *not* run TLS dtors of remaining threads, which seems to match rustc behavior.
-            return Ok(SchedulingAction::Stop);
-        }
         // This thread and the program can keep going.
         if self.threads[self.active_thread].state == ThreadState::Enabled
             && !self.yield_active_thread
         {
             // The currently active thread is still enabled, just continue with it.
             return Ok(SchedulingAction::ExecuteStep);
         }
-        // The active thread yielded. Let's see if there are any timeouts to take care of. We do
-        // this *before* running any other thread, to ensure that timeouts "in the past" fire before
-        // any other thread can take an action. This ensures that for `pthread_cond_timedwait`, "an
-        // error is returned if [...] the absolute time specified by abstime has already been passed
-        // at the time of the call".
+        // The active thread yielded or got terminated. Let's see if there are any timeouts to take
+        // care of. We do this *before* running any other thread, to ensure that timeouts "in the
+        // past" fire before any other thread can take an action. This ensures that for
+        // `pthread_cond_timedwait`, "an error is returned if [...] the absolute time specified by
+        // abstime has already been passed at the time of the call".
         // <https://pubs.opengroup.org/onlinepubs/9699919799/functions/pthread_cond_timedwait.html>
         let potential_sleep_time =
             self.timeout_callbacks.values().map(|info| info.call_time.get_wait_time(clock)).min();
         if potential_sleep_time == Some(Duration::new(0, 0)) {
             return Ok(SchedulingAction::ExecuteTimeoutCallback);
         }
-        // No callbacks scheduled, pick a regular thread to execute.
+        // No callbacks immediately scheduled, pick a regular thread to execute.
         // The active thread blocked or yielded. So we go search for another enabled thread.
         // Crucially, we start searching at the current active thread ID, rather than at 0, since we
         // want to avoid always scheduling threads 0 and 1 without ever making progress in thread 2.
@@ -730,9 +711,7 @@ impl<'mir, 'tcx: 'mir> ThreadManager<'mir, 'tcx> {
             // All threads are currently blocked, but we have unexecuted
             // timeout_callbacks, which may unblock some of the threads. Hence,
             // sleep until the first callback.
-
-            clock.sleep(sleep_time);
-            Ok(SchedulingAction::ExecuteTimeoutCallback)
+            Ok(SchedulingAction::Sleep(sleep_time))
         } else {
             throw_machine_stop!(TerminationInfo::Deadlock);
         }
@@ -773,18 +752,9 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         }
     }
 
+    /// Start a regular (non-main) thread.
     #[inline]
-    fn create_thread(&mut self) -> ThreadId {
-        let this = self.eval_context_mut();
-        let id = this.machine.threads.create_thread();
-        if let Some(data_race) = &mut this.machine.data_race {
-            data_race.thread_created(&this.machine.threads, id);
-        }
-        id
-    }
-
-    #[inline]
-    fn start_thread(
+    fn start_regular_thread(
         &mut self,
         thread: Option<MPlaceTy<'tcx, Provenance>>,
         start_routine: Pointer<Option<Provenance>>,
@@ -795,7 +765,13 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         let this = self.eval_context_mut();
 
         // Create the new thread
-        let new_thread_id = this.create_thread();
+        let new_thread_id = this.machine.threads.create_thread({
+            let mut state = tls::TlsDtorsState::default();
+            Box::new(move |m| state.on_stack_empty(m))
+        });
+        if let Some(data_race) = &mut this.machine.data_race {
+            data_race.thread_created(&this.machine.threads, new_thread_id);
+        }
 
         // Write the current thread-id, switch to the next thread later
         // to treat this write operation as occuring on the current thread.
@@ -888,12 +864,6 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         this.machine.threads.get_total_thread_count()
     }
 
-    #[inline]
-    fn has_terminated(&self, thread_id: ThreadId) -> bool {
-        let this = self.eval_context_ref();
-        this.machine.threads.has_terminated(thread_id)
-    }
-
     #[inline]
     fn have_all_terminated(&self) -> bool {
         let this = self.eval_context_ref();
@@ -943,26 +913,22 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
     where
         'mir: 'c,
     {
-        let this = self.eval_context_ref();
-        this.machine.threads.get_thread_name(thread)
+        self.eval_context_ref().machine.threads.get_thread_name(thread)
     }
 
     #[inline]
     fn block_thread(&mut self, thread: ThreadId) {
-        let this = self.eval_context_mut();
-        this.machine.threads.block_thread(thread);
+        self.eval_context_mut().machine.threads.block_thread(thread);
     }
 
     #[inline]
     fn unblock_thread(&mut self, thread: ThreadId) {
-        let this = self.eval_context_mut();
-        this.machine.threads.unblock_thread(thread);
+        self.eval_context_mut().machine.threads.unblock_thread(thread);
     }
 
     #[inline]
     fn yield_active_thread(&mut self) {
-        let this = self.eval_context_mut();
-        this.machine.threads.yield_active_thread();
+        self.eval_context_mut().machine.threads.yield_active_thread();
     }
 
     #[inline]
@@ -1024,6 +990,19 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         Ok(())
     }
 
+    #[inline]
+    fn run_on_stack_empty(&mut self) -> InterpResult<'tcx, Poll<()>> {
+        let this = self.eval_context_mut();
+        let mut callback = this
+            .active_thread_mut()
+            .on_stack_empty
+            .take()
+            .expect("`on_stack_empty` not set up, or already running");
+        let res = callback(this)?;
+        this.active_thread_mut().on_stack_empty = Some(callback);
+        Ok(res)
+    }
+
     /// Decide which action to take next and on which thread.
     #[inline]
     fn schedule(&mut self) -> InterpResult<'tcx, SchedulingAction> {
@@ -1034,10 +1013,14 @@ 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.
     ///
-    /// This is called from `tls.rs` after handling the TLS dtors.
+    /// This is called by the eval loop when a thread's on_stack_empty returns `Ready`.
     #[inline]
     fn thread_terminated(&mut self) -> 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;
+
         for ptr in this.machine.threads.thread_terminated(this.machine.data_race.as_mut()) {
             this.deallocate_ptr(ptr.into(), None, MiriMemoryKind::Tls.into())?;
         }