fix(runtime): store active task's Waker in Scheduler (#494)

Author: Paraworker

Cargo.toml

@@ -52,6 +52,7 @@ nix = "0.30.1"
 once_cell = "1.18.0"
 os_pipe = "1.1.4"
 paste = "1.0.14"
+pin-project-lite = "0.2.16"
 rand = "0.9.0"
 rustls = { version = "0.23.1", default-features = false }
 rustls-native-certs = "0.8.0"

compio-io/Cargo.toml

@@ -15,7 +15,7 @@ compio-buf = { workspace = true, features = ["arrayvec", "bytes"] }
 futures-util = { workspace = true, features = ["sink"] }
 paste = { workspace = true }
 thiserror = { workspace = true, optional = true }
-pin-project-lite = { version = "0.2.14", optional = true }
+pin-project-lite = { workspace = true, optional = true }
 serde = { version = "1.0.219", optional = true }
 serde_json = { version = "1.0.140", optional = true }
 

compio-runtime/Cargo.toml

@@ -41,8 +41,9 @@ crossbeam-queue = { workspace = true }
 futures-util = { workspace = true }
 once_cell = { workspace = true }
 scoped-tls = "1.0.1"
-slab = { workspace = true, optional = true }
+slab = { workspace = true }
 socket2 = { workspace = true }
+pin-project-lite = { workspace = true }
 
 # Windows specific dependencies
 [target.'cfg(windows)'.dependencies]
@@ -61,11 +62,19 @@ block2 = "0.6.0"
 
 [features]
 event = ["dep:cfg-if", "compio-buf/arrayvec"]
-time = ["dep:slab"]
+time = []
 
 # Enable it to always notify the driver when a task schedules.
 notify-always = []
 
+[[test]]
+name = "custom_loop"
+required-features = ["event", "time"]
+
 [[test]]
 name = "event"
 required-features = ["event"]
+
+[[test]]
+name = "drop"
+required-features = ["time"]

compio-runtime/src/runtime/mod.rs

@@ -154,7 +154,7 @@ impl Runtime {
 
     /// Block on the future till it completes.
     pub fn block_on<F: Future>(&self, future: F) -> F::Output {
-        CURRENT_RUNTIME.set(self, || {
+        self.enter(|| {
             let mut result = None;
             unsafe { self.spawn_unchecked(async { result = Some(future.await) }) }.detach();
             loop {
@@ -348,6 +348,14 @@ impl Runtime {
     }
 }
 
+impl Drop for Runtime {
+    fn drop(&mut self) {
+        self.enter(|| {
+            self.scheduler.clear();
+        })
+    }
+}
+
 impl AsRawFd for Runtime {
     fn as_raw_fd(&self) -> RawFd {
         self.driver.borrow().as_raw_fd()

compio-runtime/src/runtime/scheduler/drop_hook.rs

@@ -0,0 +1,43 @@
+use std::{
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+use pin_project_lite::pin_project;
+
+/// Calls a function when dropped.
+struct Defer<F: FnMut()>(F);
+
+impl<F: FnMut()> Drop for Defer<F> {
+    fn drop(&mut self) {
+        (self.0)();
+    }
+}
+
+pin_project! {
+    /// A future wrapper that runs a hook when dropped.
+    pub(crate) struct DropHook<Fut, Hook: FnMut()> {
+        #[pin]
+        future: Fut,
+        _hook: Defer<Hook>,
+    }
+}
+
+impl<Fut, Hook: FnMut()> DropHook<Fut, Hook> {
+    /// Creates a new [`DropHook`].
+    pub(crate) fn new(future: Fut, hook: Hook) -> Self {
+        Self {
+            future,
+            _hook: Defer(hook),
+        }
+    }
+}
+
+impl<Fut: Future, Hook: FnMut()> Future for DropHook<Fut, Hook> {
+    type Output = Fut::Output;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        self.project().future.poll(cx)
+    }
+}

compio-runtime/src/runtime/scheduler/mod.rs

@@ -1,9 +1,13 @@
-use crate::runtime::scheduler::{local_queue::LocalQueue, send_wrapper::SendWrapper};
+use crate::runtime::scheduler::{
+    drop_hook::DropHook, local_queue::LocalQueue, send_wrapper::SendWrapper,
+};
 use async_task::{Runnable, Task};
 use compio_driver::NotifyHandle;
 use crossbeam_queue::SegQueue;
-use std::{future::Future, marker::PhantomData, sync::Arc};
+use slab::Slab;
+use std::{cell::RefCell, future::Future, marker::PhantomData, rc::Rc, sync::Arc, task::Waker};
 
+mod drop_hook;
 mod local_queue;
 mod send_wrapper;
 
@@ -68,13 +72,38 @@ impl TaskQueue {
         let local_queue = unsafe { self.local_queue.get_unchecked() };
         local_queue.is_empty() && self.sync_queue.is_empty()
     }
+
+    /// Clears both queues.
+    ///
+    /// # Safety
+    ///
+    /// Call this method in the same thread as the creator.
+    unsafe fn clear(&self) {
+        // SAFETY: See the safety comment of this method.
+        let local_queue = unsafe { self.local_queue.get_unchecked() };
+
+        while let Some(item) = local_queue.pop() {
+            drop(item);
+        }
+
+        while let Some(item) = self.sync_queue.pop() {
+            drop(item);
+        }
+    }
 }
 
 /// A scheduler for managing and executing tasks.
 pub(crate) struct Scheduler {
+    /// Queue for scheduled tasks.
     task_queue: Arc<TaskQueue>,
+
+    /// `Waker` of active tasks.
+    active_tasks: Rc<RefCell<Slab<Waker>>>,
+
+    /// Number of scheduler ticks for each `run` invocation.
     event_interval: usize,
-    // `Scheduler` is `!Send` and `!Sync`.
+
+    /// Makes this type `!Send` and `!Sync`.
     _local_marker: PhantomData<*const ()>,
 }
 
@@ -83,6 +112,7 @@ impl Scheduler {
     pub(crate) fn new(event_interval: usize) -> Self {
         Self {
             task_queue: Arc::new(TaskQueue::new()),
+            active_tasks: Rc::new(RefCell::new(Slab::new())),
             event_interval,
             _local_marker: PhantomData,
         }
@@ -101,20 +131,39 @@ impl Scheduler {
     where
         F: Future,
     {
+        let mut active_tasks = self.active_tasks.borrow_mut();
+        let task_entry = active_tasks.vacant_entry();
+
+        let future = {
+            let active_tasks = self.active_tasks.clone();
+            let index = task_entry.key();
+
+            // Wrap the future with a drop hook to remove the waker upon completion.
+            DropHook::new(future, move || {
+                active_tasks.borrow_mut().try_remove(index);
+            })
+        };
+
         let schedule = {
-            // Use `Weak` to break reference cycle.
-            // `TaskQueue` -> `Runnable` -> `TaskQueue`
+            // The schedule closure is managed by the `Waker` and may be dropped on another thread,
+            // so use `Weak` to ensure the `TaskQueue` is always dropped on the creator thread.
             let task_queue = Arc::downgrade(&self.task_queue);
 
             move |runnable| {
-                if let Some(task_queue) = task_queue.upgrade() {
-                    task_queue.push(runnable, &notify);
-                }
+                // The `upgrade()` never fails because all tasks are dropped when the `Scheduler` is dropped,
+                // if a `Waker` is used after that, the schedule closure will never be called.
+                task_queue.upgrade().unwrap().push(runnable, &notify);
             }
         };
 
         let (runnable, task) = async_task::spawn_unchecked(future, schedule);
+
+        // Store the waker.
+        task_entry.insert(runnable.waker());
+
+        // Schedule the task for execution.
         runnable.schedule();
+
         task
     }
 
@@ -124,11 +173,13 @@ impl Scheduler {
     pub(crate) fn run(&self) -> bool {
         for _ in 0..self.event_interval {
             // SAFETY:
-            // `Scheduler` is `!Send` and `!Sync`, so this method is only called
-            // on `TaskQueue`'s creator thread.
+            // This method is only called on `TaskQueue`'s creator thread
+            // because `Scheduler` is `!Send` and `!Sync`.
             let tasks = unsafe { self.task_queue.pop() };
 
             // Run the tasks, which will poll the futures.
+            //
+            // SAFETY:
             // Since spawned tasks are not required to be `Send`, they must always be polled
             // on the same thread. Because `Scheduler` is `!Send` and `!Sync`, this is safe.
             match tasks {
@@ -147,8 +198,29 @@ impl Scheduler {
         }
 
         // SAFETY:
-        // `Scheduler` is `!Send` and `!Sync`, so this method is only called
-        // on `TaskQueue`'s creator thread.
+        // This method is only called on `TaskQueue`'s creator thread
+        // because `Scheduler` is `!Send` and `!Sync`.
         !unsafe { self.task_queue.is_empty() }
     }
+
+    /// Clears all active tasks.
+    ///
+    /// This method **must** be called before the scheduler is dropped.
+    pub(crate) fn clear(&self) {
+        // Drain and wake all wakers, which will schedule all active tasks.
+        self.active_tasks
+            .borrow_mut()
+            .drain()
+            .for_each(|waker| waker.wake());
+
+        // Then drop all scheduled tasks, which will drop all futures.
+        //
+        // SAFETY:
+        // Since spawned tasks are not required to be `Send`, they must always be dropped
+        // on the same thread. Because `Scheduler` is `!Send` and `!Sync`, this is safe.
+        //
+        // This method is only called on `TaskQueue`'s creator thread
+        // because `Scheduler` is `!Send` and `!Sync`.
+        unsafe { self.task_queue.clear() };
+    }
 }

compio-runtime/tests/drop.rs

@@ -0,0 +1,89 @@
+use futures_util::task::AtomicWaker;
+use std::{
+    future::Future,
+    pin::Pin,
+    sync::Arc,
+    task::{Context, Poll},
+    thread::{self, ThreadId},
+};
+
+struct DropWatcher {
+    waker: Arc<AtomicWaker>,
+    thread_id: ThreadId,
+}
+
+impl DropWatcher {
+    fn new(waker: Arc<AtomicWaker>) -> Self {
+        Self {
+            waker,
+            thread_id: thread::current().id(),
+        }
+    }
+}
+
+impl Future for DropWatcher {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
+        self.waker.register(cx.waker());
+        Poll::Pending
+    }
+}
+
+impl Drop for DropWatcher {
+    fn drop(&mut self) {
+        if self.thread_id != thread::current().id() {
+            panic!("DropWatcher dropped on a different thread!!!");
+        }
+    }
+}
+
+#[test]
+fn test_drop_with_timer() {
+    compio_runtime::Runtime::new().unwrap().block_on(async {
+        compio_runtime::spawn(async {
+            loop {
+                compio_runtime::time::sleep(std::time::Duration::from_secs(1)).await;
+            }
+        })
+        .detach();
+    })
+}
+
+#[test]
+fn test_wake_after_runtime_drop() {
+    let waker = Arc::new(AtomicWaker::new());
+    let waker_clone = waker.clone();
+
+    let rt = compio_runtime::Runtime::new().unwrap();
+
+    rt.block_on(async move {
+        compio_runtime::spawn(DropWatcher::new(waker_clone)).detach();
+    });
+
+    drop(rt);
+
+    // Use `unwrap()` to ensure there is a waker stored.
+    waker.take().unwrap().wake();
+}
+
+#[test]
+fn test_wake_from_another_thread_after_runtime_drop() {
+    let waker = Arc::new(AtomicWaker::new());
+    let waker_clone = waker.clone();
+
+    let rt = compio_runtime::Runtime::new().unwrap();
+
+    rt.block_on(async move {
+        compio_runtime::spawn(DropWatcher::new(waker_clone)).detach();
+    });
+
+    drop(rt);
+
+    thread::spawn(move || {
+        // Use `unwrap()` to ensure there is a waker stored.
+        waker.take().unwrap().wake();
+    })
+    .join()
+    .unwrap();
+}