Runtime(fix): store Waker in Scheduler

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

@@ -43,7 +43,7 @@ once_cell = { workspace = true }
 scoped-tls = "1.0.1"
 slab = { workspace = true }
 socket2 = { workspace = true }
-pin-project-lite = "0.2.16"
+pin-project-lite = { workspace = true }
 
 # Windows specific dependencies
 [target.'cfg(windows)'.dependencies]
@@ -67,6 +67,14 @@ 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/tests/drop.rs

@@ -1,6 +1,45 @@
-#[cfg(feature = "time")]
+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() {
+fn test_drop_with_timer() {
     compio_runtime::Runtime::new().unwrap().block_on(async {
         compio_runtime::spawn(async {
             loop {
@@ -10,3 +49,41 @@ fn test_drop() {
         .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();
+}