rt(alt): fix memory leak and increase max preemption when running Loom CI tests (#5911)

The memory leak was caused by a bug during shutdown where some state was leaked.

.github/workflows/loom.yml

@@ -120,5 +120,3 @@ jobs:
         working-directory: tokio
         env:
           SCOPE: ${{ matrix.scope }}
-          # TODO: remove this before stabilizing
-          LOOM_MAX_PREEMPTIONS: 1

tokio/src/runtime/scheduler/multi_thread_alt/handle.rs

@@ -39,7 +39,7 @@ impl Handle {
     }
 
     pub(crate) fn shutdown(&self) {
-        self.shared.close();
+        self.shared.close(self);
         self.driver.unpark();
     }
 

tokio/src/runtime/scheduler/multi_thread_alt/idle.rs

@@ -4,7 +4,7 @@
 
 use crate::loom::sync::atomic::{AtomicBool, AtomicUsize};
 use crate::loom::sync::MutexGuard;
-use crate::runtime::scheduler::multi_thread_alt::{worker, Core, Shared};
+use crate::runtime::scheduler::multi_thread_alt::{worker, Core, Handle, Shared};
 
 use std::sync::atomic::Ordering::{AcqRel, Acquire, Release};
 
@@ -146,23 +146,13 @@ impl Idle {
         // Find a sleeping worker
         if let Some(worker) = synced.idle.sleepers.pop() {
             // Find an available core
-            if let Some(mut core) = synced.idle.available_cores.pop() {
+            if let Some(mut core) = self.try_acquire_available_core(&mut synced.idle) {
                 debug_assert!(!core.is_searching);
                 core.is_searching = true;
 
-                self.idle_map.unset(core.index);
-                debug_assert!(self.idle_map.matches(&synced.idle.available_cores));
-
                 // Assign the core to the worker
                 synced.assigned_cores[worker] = Some(core);
 
-                let num_idle = synced.idle.available_cores.len();
-                #[cfg(not(loom))]
-                debug_assert_eq!(num_idle, self.num_idle.load(Acquire) - 1);
-
-                // Update the number of sleeping workers
-                self.num_idle.store(num_idle, Release);
-
                 // Drop the lock before notifying the condvar.
                 drop(synced);
 
@@ -198,6 +188,7 @@ impl Idle {
 
         for _ in 0..num {
             if let Some(worker) = synced.idle.sleepers.pop() {
+                // TODO: can this be switched to use next_available_core?
                 if let Some(core) = synced.idle.available_cores.pop() {
                     debug_assert!(!core.is_searching);
 
@@ -236,15 +227,10 @@ impl Idle {
         // eventually find the cores and shut them down.
         while !synced.idle.sleepers.is_empty() && !synced.idle.available_cores.is_empty() {
             let worker = synced.idle.sleepers.pop().unwrap();
-            let core = synced.idle.available_cores.pop().unwrap();
+            let core = self.try_acquire_available_core(&mut synced.idle).unwrap();
-
-            self.idle_map.unset(core.index);
 
             synced.assigned_cores[worker] = Some(core);
             shared.condvars[worker].notify_one();
-
-            self.num_idle
-                .store(synced.idle.available_cores.len(), Release);
         }
 
         debug_assert!(self.idle_map.matches(&synced.idle.available_cores));
@@ -255,6 +241,18 @@ impl Idle {
         }
     }
 
+    pub(super) fn shutdown_unassigned_cores(&self, handle: &Handle, shared: &Shared) {
+        // If there are any remaining cores, shut them down here.
+        //
+        // This code is a bit convoluted to avoid lock-reentry.
+        while let Some(core) = {
+            let mut synced = shared.synced.lock();
+            self.try_acquire_available_core(&mut synced.idle)
+        } {
+            shared.shutdown_core(handle, core);
+        }
+    }
+
     /// The worker releases the given core, making it available to other workers
     /// that are waiting.
     pub(super) fn release_core(&self, synced: &mut worker::Synced, core: Box<Core>) {

tokio/src/runtime/scheduler/multi_thread_alt/worker.rs

@@ -560,7 +560,7 @@ impl Worker {
                 // First try to acquire an available core
                 if let Some(core) = self.try_acquire_available_core(cx, &mut synced) {
                     // Try to poll a task from the global queue
-                    let maybe_task = self.next_remote_task_synced(cx, &mut synced);
+                    let maybe_task = cx.shared().next_remote_task_synced(&mut synced);
                     (maybe_task, core)
                 } else {
                     // block the thread to wait for a core to be assinged to us
@@ -589,10 +589,7 @@ impl Worker {
             }
         }
 
-        self.pre_shutdown(cx, &mut core);
+        cx.shared().shutdown_core(&cx.handle, core);
-
-        // Signal shutdown
-        self.shutdown_core(cx, core);
 
         // It is possible that tasks wake others during drop, so we need to
         // clear the defer list.
@@ -746,7 +743,7 @@ impl Worker {
             }
         }
 
-        if let Some(task) = self.next_local_task(&mut core) {
+        if let Some(task) = core.next_local_task() {
             return Ok((Some(task), core));
         }
 
@@ -759,12 +756,7 @@ impl Worker {
         }
 
         let mut synced = cx.shared().synced.lock();
-        self.next_remote_task_synced(cx, &mut synced)
+        cx.shared().next_remote_task_synced(&mut synced)
-    }
-
-    fn next_remote_task_synced(&self, cx: &Context, synced: &mut Synced) -> Option<Notified> {
-        // safety: we only have access to a valid `Synced` in this file.
-        unsafe { cx.shared().inject.pop(&mut synced.inject) }
     }
 
     fn next_remote_task_batch(&self, cx: &Context, mut core: Box<Core>) -> NextTaskResult {
@@ -818,14 +810,6 @@ impl Worker {
         ret
     }
 
-    fn next_local_task(&self, core: &mut Core) -> Option<Notified> {
-        self.next_lifo_task(core).or_else(|| core.run_queue.pop())
-    }
-
-    fn next_lifo_task(&self, core: &mut Core) -> Option<Notified> {
-        core.lifo_slot.take()
-    }
-
     /// Function responsible for stealing tasks from another worker
     ///
     /// Note: Only if less than half the workers are searching for tasks to steal
@@ -948,7 +932,7 @@ impl Worker {
                 };
 
                 // Check for a task in the LIFO slot
-                let task = match self.next_lifo_task(&mut core) {
+                let task = match core.next_lifo_task() {
                     Some(task) => task,
                     None => {
                         self.reset_lifo_enabled(cx);
@@ -1229,7 +1213,7 @@ impl Worker {
             if cx.shared().inject.is_closed(&mut synced.inject) {
                 synced.shutdown_driver = Some(driver);
                 self.shutdown_clear_defer(cx);
-                self.shutdown_finalize(cx, synced);
+                cx.shared().shutdown_finalize(&cx.handle, &mut synced);
                 return Err(());
             }
 
@@ -1281,61 +1265,6 @@ impl Worker {
         core.lifo_slot.is_some() || !core.run_queue.is_empty()
     }
 
-    /// Signals all tasks to shut down, and waits for them to complete. Must run
-    /// before we enter the single-threaded phase of shutdown processing.
-    fn pre_shutdown(&self, cx: &Context, core: &mut Core) {
-        // Signal to all tasks to shut down.
-        cx.shared().owned.close_and_shutdown_all();
-
-        core.stats.submit(&cx.shared().worker_metrics[core.index]);
-    }
-
-    /// Signals that a worker has observed the shutdown signal and has replaced
-    /// its core back into its handle.
-    ///
-    /// If all workers have reached this point, the final cleanup is performed.
-    fn shutdown_core(&self, cx: &Context, core: Box<Core>) {
-        let mut synced = cx.shared().synced.lock();
-        synced.shutdown_cores.push(core);
-
-        self.shutdown_finalize(cx, synced);
-    }
-
-    fn shutdown_finalize(&self, cx: &Context, mut synced: MutexGuard<'_, Synced>) {
-        // Wait for all cores
-        if synced.shutdown_cores.len() != cx.shared().remotes.len() {
-            return;
-        }
-
-        let driver = synced.shutdown_driver.take();
-
-        if cx.shared().driver_enabled() && driver.is_none() {
-            return;
-        }
-
-        debug_assert!(cx.shared().owned.is_empty());
-
-        for mut core in synced.shutdown_cores.drain(..) {
-            // Drain tasks from the local queue
-            while self.next_local_task(&mut core).is_some() {}
-        }
-
-        // Shutdown the driver
-        if let Some(mut driver) = driver {
-            driver.shutdown(&cx.handle.driver);
-        }
-
-        // Drain the injection queue
-        //
-        // We already shut down every task, so we can simply drop the tasks. We
-        // cannot call `next_remote_task()` because we already hold the lock.
-        //
-        // safety: passing in correct `idle::Synced`
-        while let Some(task) = self.next_remote_task_synced(cx, &mut synced) {
-            drop(task);
-        }
-    }
-
     fn reset_lifo_enabled(&self, cx: &Context) {
         cx.lifo_enabled
             .set(!cx.handle.shared.config.disable_lifo_slot);
@@ -1379,7 +1308,22 @@ impl Context {
     }
 }
 
+impl Core {
+    fn next_local_task(&mut self) -> Option<Notified> {
+        self.next_lifo_task().or_else(|| self.run_queue.pop())
+    }
+
+    fn next_lifo_task(&mut self) -> Option<Notified> {
+        self.lifo_slot.take()
+    }
+}
+
 impl Shared {
+    fn next_remote_task_synced(&self, synced: &mut Synced) -> Option<Notified> {
+        // safety: we only have access to a valid `Synced` in this file.
+        unsafe { self.inject.pop(&mut synced.inject) }
+    }
+
     pub(super) fn schedule_task(&self, task: Notified, is_yield: bool) {
         use std::ptr;
 
@@ -1449,17 +1393,25 @@ impl Shared {
         self.idle.notify_remote(synced, self);
     }
 
-    pub(super) fn close(&self) {
+    pub(super) fn close(&self, handle: &Handle) {
-        let mut synced = self.synced.lock();
+        {
+            let mut synced = self.synced.lock();
 
-        if let Some(driver) = self.driver.take() {
+            if let Some(driver) = self.driver.take() {
-            synced.shutdown_driver = Some(driver);
+                synced.shutdown_driver = Some(driver);
-        }
+            }
+
+            if !self.inject.close(&mut synced.inject) {
+                return;
+            }
 
-        if self.inject.close(&mut synced.inject) {
             // Set the shutdown flag on all available cores
             self.idle.shutdown(&mut synced, self);
         }
+
+        // Any unassigned cores need to be shutdown, but we have to first drop
+        // the lock
+        self.idle.shutdown_unassigned_cores(handle, self);
     }
 
     fn push_remote_task(&self, synced: &mut Synced, task: Notified) {
@@ -1498,6 +1450,52 @@ impl Shared {
     fn driver_enabled(&self) -> bool {
         self.condvars.len() > self.remotes.len()
     }
+
+    pub(super) fn shutdown_core(&self, handle: &Handle, mut core: Box<Core>) {
+        self.owned.close_and_shutdown_all();
+
+        core.stats.submit(&self.worker_metrics[core.index]);
+
+        let mut synced = self.synced.lock();
+        synced.shutdown_cores.push(core);
+
+        self.shutdown_finalize(handle, &mut synced);
+    }
+
+    pub(super) fn shutdown_finalize(&self, handle: &Handle, synced: &mut Synced) {
+        // Wait for all cores
+        if synced.shutdown_cores.len() != self.remotes.len() {
+            return;
+        }
+
+        let driver = synced.shutdown_driver.take();
+
+        if self.driver_enabled() && driver.is_none() {
+            return;
+        }
+
+        debug_assert!(self.owned.is_empty());
+
+        for mut core in synced.shutdown_cores.drain(..) {
+            // Drain tasks from the local queue
+            while core.next_local_task().is_some() {}
+        }
+
+        // Shutdown the driver
+        if let Some(mut driver) = driver {
+            driver.shutdown(&handle.driver);
+        }
+
+        // Drain the injection queue
+        //
+        // We already shut down every task, so we can simply drop the tasks. We
+        // cannot call `next_remote_task()` because we already hold the lock.
+        //
+        // safety: passing in correct `idle::Synced`
+        while let Some(task) = self.next_remote_task_synced(synced) {
+            drop(task);
+        }
+    }
 }
 
 impl Overflow<Arc<Handle>> for Shared {
@@ -1514,10 +1512,20 @@ impl Overflow<Arc<Handle>> for Shared {
 }
 
 impl<'a> Lock<inject::Synced> for &'a Shared {
-    type Handle = InjectGuard<'a>;
+    type Handle = SyncedGuard<'a>;
 
     fn lock(self) -> Self::Handle {
-        InjectGuard {
+        SyncedGuard {
+            lock: self.synced.lock(),
+        }
+    }
+}
+
+impl<'a> Lock<Synced> for &'a Shared {
+    type Handle = SyncedGuard<'a>;
+
+    fn lock(self) -> Self::Handle {
+        SyncedGuard {
             lock: self.synced.lock(),
         }
     }
@@ -1537,16 +1545,28 @@ impl task::Schedule for Arc<Handle> {
     }
 }
 
-pub(crate) struct InjectGuard<'a> {
+impl AsMut<Synced> for Synced {
+    fn as_mut(&mut self) -> &mut Synced {
+        self
+    }
+}
+
+pub(crate) struct SyncedGuard<'a> {
     lock: crate::loom::sync::MutexGuard<'a, Synced>,
 }
 
-impl<'a> AsMut<inject::Synced> for InjectGuard<'a> {
+impl<'a> AsMut<inject::Synced> for SyncedGuard<'a> {
     fn as_mut(&mut self) -> &mut inject::Synced {
         &mut self.lock.inject
     }
 }
 
+impl<'a> AsMut<Synced> for SyncedGuard<'a> {
+    fn as_mut(&mut self) -> &mut Synced {
+        &mut self.lock
+    }
+}
+
 #[track_caller]
 fn with_current<R>(f: impl FnOnce(Option<&Context>) -> R) -> R {
     use scheduler::Context::MultiThreadAlt;