Use Droppable Future instead of TaskMetadata (#4)

- TaskMetadata does not Drop when using tokio::select
- Added tokio integration tests for join! and select!

Author: coder137

Cargo.toml

@@ -5,6 +5,7 @@ edition = "2021"
 
 [dependencies]
 async-task = "4.7"
+pin-project = "1"
 
 [dev-dependencies]
 tokio = { version = "1", features = ["full"] }

src/droppable_future.rs

@@ -0,0 +1,51 @@
+use std::{future::Future, pin::Pin};
+
+use pin_project::{pin_project, pinned_drop};
+
+#[pin_project(PinnedDrop)]
+pub struct DroppableFuture<F, D>
+where
+    F: Future,
+    D: Fn(),
+{
+    #[pin]
+    future: F,
+    on_drop: D,
+}
+
+impl<F, D> DroppableFuture<F, D>
+where
+    F: Future,
+    D: Fn(),
+{
+    pub fn new(future: F, on_drop: D) -> Self {
+        Self { future, on_drop }
+    }
+}
+
+impl<F, D> Future for DroppableFuture<F, D>
+where
+    F: Future,
+    D: Fn(),
+{
+    type Output = F::Output;
+
+    fn poll(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Self::Output> {
+        let this = self.project();
+        this.future.poll(cx)
+    }
+}
+
+#[pinned_drop]
+impl<F, D> PinnedDrop for DroppableFuture<F, D>
+where
+    F: Future,
+    D: Fn(),
+{
+    fn drop(self: Pin<&mut Self>) {
+        (self.on_drop)();
+    }
+}

src/lib.rs

@@ -1,3 +1,6 @@
+mod droppable_future;
+use droppable_future::*;
+
 mod task_identifier;
 pub use task_identifier::*;
 

src/ticked_async_executor.rs

@@ -6,7 +6,7 @@ use std::{
     },
 };
 
-use crate::TaskIdentifier;
+use crate::{DroppableFuture, TaskIdentifier};
 
 #[derive(Debug)]
 pub enum TaskState {
@@ -16,37 +16,11 @@ pub enum TaskState {
     Drop(TaskIdentifier),
 }
 
-pub type Task<T, O> = async_task::Task<T, TaskMetadata<O>>;
-type TaskRunnable<O> = async_task::Runnable<TaskMetadata<O>>;
-type Payload<O> = (TaskIdentifier, TaskRunnable<O>);
+pub type Task<T> = async_task::Task<T>;
+type Payload = (TaskIdentifier, async_task::Runnable);
 
-/// Task Metadata associated with TickedAsyncExecutor
-///
-/// Primarily used to track when the Task is completed/cancelled
-pub struct TaskMetadata<O>
-where
-    O: Fn(TaskState) + Send + Sync + 'static,
-{
-    num_spawned_tasks: Arc<AtomicUsize>,
-    identifier: TaskIdentifier,
-    observer: O,
-}
-
-impl<O> Drop for TaskMetadata<O>
-where
-    O: Fn(TaskState) + Send + Sync + 'static,
-{
-    fn drop(&mut self) {
-        self.num_spawned_tasks.fetch_sub(1, Ordering::Relaxed);
-        (self.observer)(TaskState::Drop(self.identifier.clone()));
-    }
-}
-
-pub struct TickedAsyncExecutor<O>
-where
-    O: Fn(TaskState) + Send + Sync + 'static,
-{
-    channel: (mpsc::Sender<Payload<O>>, mpsc::Receiver<Payload<O>>),
+pub struct TickedAsyncExecutor<O> {
+    channel: (mpsc::Sender<Payload>, mpsc::Receiver<Payload>),
     num_woken_tasks: Arc<AtomicUsize>,
     num_spawned_tasks: Arc<AtomicUsize>,
 
@@ -79,22 +53,14 @@ where
         &self,
         identifier: impl Into<TaskIdentifier>,
         future: impl Future<Output = T> + Send + 'static,
-    ) -> Task<T, O>
+    ) -> Task<T>
     where
         T: Send + 'static,
     {
         let identifier = identifier.into();
-        self.num_spawned_tasks.fetch_add(1, Ordering::Relaxed);
-        (self.observer)(TaskState::Spawn(identifier.clone()));
-
-        let schedule = self.runnable_schedule_cb(identifier.clone());
-        let (runnable, task) = async_task::Builder::new()
-            .metadata(TaskMetadata {
-                num_spawned_tasks: self.num_spawned_tasks.clone(),
-                identifier,
-                observer: self.observer.clone(),
-            })
-            .spawn(|_m| future, schedule);
+        let future = self.droppable_future(identifier.clone(), future);
+        let schedule = self.runnable_schedule_cb(identifier);
+        let (runnable, task) = async_task::spawn(future, schedule);
         runnable.schedule();
         task
     }
@@ -103,22 +69,14 @@ where
         &self,
         identifier: impl Into<TaskIdentifier>,
         future: impl Future<Output = T> + 'static,
-    ) -> Task<T, O>
+    ) -> Task<T>
     where
         T: 'static,
     {
         let identifier = identifier.into();
-        self.num_spawned_tasks.fetch_add(1, Ordering::Relaxed);
-        (self.observer)(TaskState::Spawn(identifier.clone()));
-
-        let schedule = self.runnable_schedule_cb(identifier.clone());
-        let (runnable, task) = async_task::Builder::new()
-            .metadata(TaskMetadata {
-                num_spawned_tasks: self.num_spawned_tasks.clone(),
-                identifier,
-                observer: self.observer.clone(),
-            })
-            .spawn_local(move |_m| future, schedule);
+        let future = self.droppable_future(identifier.clone(), future);
+        let schedule = self.runnable_schedule_cb(identifier);
+        let (runnable, task) = async_task::spawn_local(future, schedule);
         runnable.schedule();
         task
     }
@@ -146,7 +104,29 @@ where
             .fetch_sub(num_woken_tasks, Ordering::Relaxed);
     }
 
-    fn runnable_schedule_cb(&self, identifier: TaskIdentifier) -> impl Fn(TaskRunnable<O>) {
+    fn droppable_future<F>(
+        &self,
+        identifier: TaskIdentifier,
+        future: F,
+    ) -> DroppableFuture<F, impl Fn()>
+    where
+        F: Future,
+    {
+        let observer = self.observer.clone();
+
+        // Spawn Task
+        self.num_spawned_tasks.fetch_add(1, Ordering::Relaxed);
+        observer(TaskState::Spawn(identifier.clone()));
+
+        // Droppable Future registering on_drop callback
+        let num_spawned_tasks = self.num_spawned_tasks.clone();
+        DroppableFuture::new(future, move || {
+            num_spawned_tasks.fetch_sub(1, Ordering::Relaxed);
+            observer(TaskState::Drop(identifier.clone()));
+        })
+    }
+
+    fn runnable_schedule_cb(&self, identifier: TaskIdentifier) -> impl Fn(async_task::Runnable) {
         let sender = self.channel.0.clone();
         let num_woken_tasks = self.num_woken_tasks.clone();
         let observer = self.observer.clone();
@@ -160,15 +140,13 @@ where
 
 #[cfg(test)]
 mod tests {
-    use tokio::join;
-
     use super::*;
 
     #[test]
-    fn test_multiple_local_tasks() {
+    fn test_multiple_tasks() {
         let executor = TickedAsyncExecutor::default();
         executor
-            .spawn_local("A", async move {
+            .spawn("A", async move {
                 tokio::task::yield_now().await;
             })
             .detach();
@@ -187,7 +165,7 @@ mod tests {
     }
 
     #[test]
-    fn test_local_tasks_cancellation() {
+    fn test_task_cancellation() {
         let executor = TickedAsyncExecutor::new(|_state| println!("{_state:?}"));
         let task1 = executor.spawn_local("A", async move {
             loop {
@@ -205,39 +183,7 @@ mod tests {
 
         executor
             .spawn_local("CancelTasks", async move {
-                let (t1, t2) = join!(task1.cancel(), task2.cancel());
-                assert_eq!(t1, None);
-                assert_eq!(t2, None);
-            })
-            .detach();
-        assert_eq!(executor.num_tasks(), 3);
-
-        // Since we have cancelled the tasks above, the loops should eventually end
-        while executor.num_tasks() != 0 {
-            executor.tick();
-        }
-    }
-
-    #[test]
-    fn test_tasks_cancellation() {
-        let executor = TickedAsyncExecutor::new(|_state| println!("{_state:?}"));
-        let task1 = executor.spawn("A", async move {
-            loop {
-                tokio::task::yield_now().await;
-            }
-        });
-
-        let task2 = executor.spawn(format!("B"), async move {
-            loop {
-                tokio::task::yield_now().await;
-            }
-        });
-        assert_eq!(executor.num_tasks(), 2);
-        executor.tick();
-
-        executor
-            .spawn_local("CancelTasks", async move {
-                let (t1, t2) = join!(task1.cancel(), task2.cancel());
+                let (t1, t2) = tokio::join!(task1.cancel(), task2.cancel());
                 assert_eq!(t1, None);
                 assert_eq!(t2, None);
             })

tests/tokio_tests.rs

@@ -0,0 +1,79 @@
+use ticked_async_executor::TickedAsyncExecutor;
+
+#[test]
+fn test_tokio_join() {
+    let executor = TickedAsyncExecutor::default();
+
+    let (tx1, mut rx1) = tokio::sync::mpsc::channel::<usize>(1);
+    let (tx2, mut rx2) = tokio::sync::mpsc::channel::<usize>(1);
+    executor
+        .spawn("ThreadedFuture", async move {
+            let (a, b) = tokio::join!(rx1.recv(), rx2.recv());
+            assert_eq!(a.unwrap(), 10);
+            assert_eq!(b.unwrap(), 20);
+        })
+        .detach();
+
+    let (tx3, mut rx3) = tokio::sync::mpsc::channel::<usize>(1);
+    let (tx4, mut rx4) = tokio::sync::mpsc::channel::<usize>(1);
+    executor
+        .spawn("LocalFuture", async move {
+            let (a, b) = tokio::join!(rx3.recv(), rx4.recv());
+            assert_eq!(a.unwrap(), 10);
+            assert_eq!(b.unwrap(), 20);
+        })
+        .detach();
+
+    tx1.try_send(10).unwrap();
+    tx3.try_send(10).unwrap();
+    for _ in 0..10 {
+        executor.tick();
+    }
+    tx2.try_send(20).unwrap();
+    tx4.try_send(20).unwrap();
+
+    while executor.num_tasks() != 0 {
+        executor.tick();
+    }
+}
+
+#[test]
+fn test_tokio_select() {
+    let executor = TickedAsyncExecutor::default();
+
+    let (tx1, mut rx1) = tokio::sync::mpsc::channel::<usize>(1);
+    let (_tx2, mut rx2) = tokio::sync::mpsc::channel::<usize>(1);
+    executor
+        .spawn("ThreadedFuture", async move {
+            tokio::select! {
+                data = rx1.recv() => {
+                    assert_eq!(data.unwrap(), 10);
+                }
+                _ = rx2.recv() => {}
+            }
+        })
+        .detach();
+
+    let (tx3, mut rx3) = tokio::sync::mpsc::channel::<usize>(1);
+    let (_tx4, mut rx4) = tokio::sync::mpsc::channel::<usize>(1);
+    executor
+        .spawn("LocalFuture", async move {
+            tokio::select! {
+                data = rx3.recv() => {
+                    assert_eq!(data.unwrap(), 10);
+                }
+                _ = rx4.recv() => {}
+            }
+        })
+        .detach();
+
+    for _ in 0..10 {
+        executor.tick();
+    }
+
+    tx1.try_send(10).unwrap();
+    tx3.try_send(10).unwrap();
+    while executor.num_tasks() != 0 {
+        executor.tick();
+    }
+}