Yield from the main thread a few times at shutdown

src/concurrency/thread.rs

@@ -44,6 +44,11 @@ pub struct ThreadId(u32);
 /// The main thread. When it terminates, the whole application terminates.
 const MAIN_THREAD: ThreadId = ThreadId(0);
 
+/// When the main thread would exit, we will yield to any other thread that is ready to execute.
+/// But we must only do that a finite number of times, or a background thread running `loop {}`
+/// will hang the program.
+const MAIN_THREAD_YIELDS_AT_SHUTDOWN: u64 = 1_000;
+
 impl ThreadId {
     pub fn to_u32(self) -> u32 {
         self.0
@@ -131,16 +136,9 @@ pub struct Thread<'mir, 'tcx> {
 }
 
 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
+    /// Check if the thread is just done executing (still enabled, but no more stack frames).
+    fn should_terminate(&self) -> bool {
+        self.state == ThreadState::Enabled && self.stack.is_empty()
     }
 
     /// Get the name of the current thread, or `<unnamed>` if it was not set.
@@ -276,6 +274,9 @@ pub struct ThreadManager<'mir, 'tcx> {
     yield_active_thread: bool,
     /// Callbacks that are called once the specified time passes.
     timeout_callbacks: FxHashMap<ThreadId, TimeoutCallbackInfo<'mir, 'tcx>>,
+    /// When the main thread is about to exit, we give other threads a few chances to finish up
+    /// whatever they are doing before we consider them leaked.
+    main_thread_yields_at_shutdown_remaining: u64,
 }
 
 impl<'mir, 'tcx> Default for ThreadManager<'mir, 'tcx> {
@@ -290,6 +291,7 @@ impl<'mir, 'tcx> Default for ThreadManager<'mir, 'tcx> {
             thread_local_alloc_ids: Default::default(),
             yield_active_thread: false,
             timeout_callbacks: FxHashMap::default(),
+            main_thread_yields_at_shutdown_remaining: MAIN_THREAD_YIELDS_AT_SHUTDOWN,
         }
     }
 }
@@ -302,6 +304,7 @@ impl VisitTags for ThreadManager<'_, '_> {
             timeout_callbacks,
             active_thread: _,
             yield_active_thread: _,
+            main_thread_yields_at_shutdown_remaining: _,
             sync,
         } = self;
 
@@ -620,17 +623,34 @@ 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() {
+        // Check whether the thread ought to terminate, and if so start executing TLS
+        // destructors.
+        if self.threads[self.active_thread].should_terminate() {
+            self.threads[self.active_thread].state = ThreadState::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);
+            // If we are the main thread and we are Terminated, we ought to stop.
+            // But, if there are any other threads which can execute, yield to them instead of falling
+            // through to the termination state. This is to give them a chance to clean up and quit.
+            let active_thread = &self.threads[self.active_thread];
+            if self.active_thread == MAIN_THREAD
+                && active_thread.state == ThreadState::Terminated
+                && self
+                    .threads
+                    .iter()
+                    .any(|t| t.state == ThreadState::Enabled && !t.stack.is_empty())
+                && self.main_thread_yields_at_shutdown_remaining > 0
+            {
+                self.yield_active_thread = true;
+                self.main_thread_yields_at_shutdown_remaining -= 1;
+            } else {
+                // 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
@@ -671,7 +691,17 @@ impl<'mir, 'tcx: 'mir> ThreadManager<'mir, 'tcx> {
             }
         }
         self.yield_active_thread = false;
+
         if self.threads[self.active_thread].state == ThreadState::Enabled {
+            // The main thread is special, it is allowed to yield when it is shutting down and has
+            // nothing to execute. So if another thread yields back to the main thread when it is in
+            // this partly-shut-down state, we make another pass through the scheduler to either yield
+            // back off them main thread, or enter its shutdown sequence.
+            if self.active_thread == MAIN_THREAD
+                && self.threads[self.active_thread].stack.is_empty()
+            {
+                return self.schedule(clock);
+            }
             return Ok(SchedulingAction::ExecuteStep);
         }
         // We have not found a thread to execute.