feat(maitake): dequeue one task at a time in scheduler

maitake/src/scheduler.rs

@@ -181,7 +181,7 @@ use crate::{
 };
 use core::{future::Future, marker::PhantomData, ptr};
 
-use cordyceps::mpsc_queue::MpscQueue;
+use cordyceps::mpsc_queue::{MpscQueue, TryDequeueError};
 
 #[cfg(any(feature = "tracing-01", feature = "tracing-02", test))]
 use mycelium_util::fmt;
@@ -623,7 +623,7 @@ impl StaticScheduler {
     ///
     /// Only a single CPU core/thread may tick a given scheduler at a time. If
     /// another call to `tick` is in progress on a different core, this method
-    /// will wait until that call to `tick` completes before ticking the scheduler.
+    /// will immediately return.
     ///
     /// See [the module-level documentation][run-loops] for more information on
     /// using this function to implement a system's run loop.
@@ -716,7 +716,7 @@ impl LocalStaticScheduler {
     ///
     /// To spawn `!`[`Send`] tasks using a [`Builder`](task::Builder), use the
     /// [`Builder::spawn_local`](task::Builder::spawn_local) method.
-    /// 
+    ///
     /// [task `Builder`]: task::Builder
     #[must_use]
     pub fn build_task<'a>(&'static self) -> task::Builder<'a, &'static Self> {
@@ -746,7 +746,7 @@ impl LocalStaticScheduler {
     ///
     /// Only a single CPU core/thread may tick a given scheduler at a time. If
     /// another call to `tick` is in progress on a different core, this method
-    /// will wait until that call to `tick` completes before ticking the scheduler.
+    /// will immediately return.
     ///
     /// See [the module-level documentation][run-loops] for more information on
     /// using this function to implement a system's run loop.
@@ -888,7 +888,20 @@ impl Core {
             has_remaining: false,
         };
 
-        for task in self.run_queue.consume().take(n) {
+        while tick.polled < n {
+            let task = match self.run_queue.try_dequeue() {
+                Ok(task) => task,
+                // If inconsistent, just try again.
+                Err(TryDequeueError::Inconsistent) => {
+                    core::hint::spin_loop();
+                    continue;
+                }
+                // Queue is empty or busy (in use by something else), bail out.
+                Err(TryDequeueError::Busy | TryDequeueError::Empty) => {
+                    break;
+                }
+            };
+
             self.queued.fetch_sub(1, Relaxed);
             let _span = trace_span!(
                 "poll",
@@ -1211,8 +1224,7 @@ feature! {
         ///
         /// Only a single CPU core/thread may tick a given scheduler at a time. If
         /// another call to `tick` is in progress on a different core, this method
-        /// will wait until that call to `tick` completes before ticking the
-        /// scheduler.
+        /// will immediately return.
         ///
         /// See [the module-level documentation][run-loops] for more information on
         /// using this function to implement a system's run loop.
@@ -1555,8 +1567,7 @@ feature! {
         ///
         /// Only a single CPU core/thread may tick a given scheduler at a time. If
         /// another call to `tick` is in progress on a different core, this method
-        /// will wait until that call to `tick` completes before ticking the
-        /// scheduler.
+        /// will immediately return.
         ///
         /// See [the module-level documentation][run-loops] for more information on
         /// using this function to implement a system's run loop.

maitake/tests/scheduler/alloc.rs

@@ -1,5 +1,5 @@
 use super::*;
-use mycelium_util::sync::Lazy;
+use mycelium_util::sync::{Lazy, spin::Mutex};
 
 #[test]
 fn basically_works() {
@@ -66,3 +66,50 @@ fn many_yields() {
     assert_eq!(COMPLETED.load(Ordering::SeqCst), TASKS);
     assert!(!tick.has_remaining);
 }
+
+#[test]
+fn steal_blocked() {
+    static SCHEDULER_1: Lazy<StaticScheduler> = Lazy::new(StaticScheduler::new);
+    static SCHEDULER_2: Lazy<StaticScheduler> = Lazy::new(StaticScheduler::new);
+    static MUTEX: Mutex<()> = Mutex::new(());
+    static READY: AtomicBool = AtomicBool::new(false);
+    static IT_WORKED: AtomicBool = AtomicBool::new(false);
+
+    util::trace_init();
+
+    let guard = MUTEX.lock();
+
+    let thread = std::thread::spawn(|| {
+        SCHEDULER_1.spawn(async {
+            READY.store(true, Ordering::Release);
+
+            // block this thread
+            let _guard = MUTEX.lock();
+        });
+
+        SCHEDULER_1.spawn(async {
+            IT_WORKED.store(true, Ordering::Release);
+        });
+
+        SCHEDULER_1.tick()
+    });
+
+    while !READY.load(Ordering::Acquire) {
+        core::hint::spin_loop();
+    }
+
+    assert!(SCHEDULER_1.current_task().is_some());
+
+    let stolen = SCHEDULER_1.try_steal().unwrap().spawn_n(&SCHEDULER_2.get(), 1);
+    assert_eq!(stolen, 1);
+
+    let tick = SCHEDULER_2.tick();
+    assert!(IT_WORKED.load(Ordering::Acquire));
+    assert_eq!(tick.polled, 1);
+    assert_eq!(tick.completed, 1);
+
+    drop(guard);
+    let tick = thread.join().unwrap();
+    assert_eq!(tick.polled, 1);
+    assert_eq!(tick.completed, 1);
+}