Added droppable future for Ticked Async Executor

- Gives visibility into number of spawned/dropped tasks

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,99 +1,5 @@
-use std::{
-    future::Future,
-    sync::{
-        atomic::{AtomicUsize, Ordering},
-        mpsc, Arc,
-    },
-};
+mod droppable_future;
+use droppable_future::*;
 
-use async_task::{Runnable, Task};
-
-pub struct TickedAsyncLocalExecutor {
-    channel: (mpsc::Sender<Runnable>, mpsc::Receiver<Runnable>),
-    num_woken_tasks: Arc<AtomicUsize>,
-}
-
-impl Default for TickedAsyncLocalExecutor {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-impl TickedAsyncLocalExecutor {
-    pub fn new() -> Self {
-        Self {
-            channel: mpsc::channel(),
-            num_woken_tasks: Arc::new(AtomicUsize::new(0)),
-        }
-    }
-
-    pub fn spawn_local<T>(&self, future: impl Future<Output = T> + 'static) -> Task<T>
-    where
-        T: 'static,
-    {
-        let sender = self.channel.0.clone();
-        let num_woken_tasks = self.num_woken_tasks.clone();
-        let schedule = move |runnable| {
-            sender.send(runnable).unwrap_or(());
-            num_woken_tasks.fetch_add(1, Ordering::Relaxed);
-        };
-        let (runnable, task) = async_task::spawn_local(future, schedule);
-        runnable.schedule();
-        task
-    }
-
-    /// Run the woken tasks once
-    ///
-    /// NOTE: Will not run tasks that are woken/scheduled immediately after `Runnable::run`
-    pub fn tick(&self) {
-        let num_woken_tasks = self.num_woken_tasks.load(Ordering::Relaxed);
-        self.channel
-            .1
-            .try_iter()
-            .take(num_woken_tasks)
-            .for_each(|runnable| {
-                runnable.run();
-            });
-        self.num_woken_tasks
-            .fetch_sub(num_woken_tasks, Ordering::Relaxed);
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn test_multiple_tasks() {
-        let executor = TickedAsyncLocalExecutor::new();
-        executor
-            .spawn_local(async move {
-                println!("A: Start");
-                tokio::task::yield_now().await;
-                println!("A: End");
-            })
-            .detach();
-
-        executor
-            .spawn_local(async move {
-                println!("B: Start");
-                tokio::task::yield_now().await;
-                println!("B: End");
-            })
-            .detach();
-
-        executor
-            .spawn_local(async move {
-                println!("C: Start");
-                tokio::task::yield_now().await;
-                println!("C: End");
-            })
-            .detach();
-
-        // A, B, C: Start
-        executor.tick();
-
-        // A, B, C: End
-        executor.tick();
-    }
-}
+mod ticked_async_executor;
+pub use ticked_async_executor::*;

src/ticked_async_executor.rs

@@ -0,0 +1,134 @@
+use std::{
+    future::Future,
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        mpsc, Arc,
+    },
+};
+
+use async_task::{Runnable, Task};
+
+use crate::DroppableFuture;
+
+pub struct TickedAsyncExecutor {
+    channel: (mpsc::Sender<Runnable>, mpsc::Receiver<Runnable>),
+    num_woken_tasks: Arc<AtomicUsize>,
+    num_spawned_tasks: Arc<AtomicUsize>,
+}
+
+impl Default for TickedAsyncExecutor {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl TickedAsyncExecutor {
+    pub fn new() -> Self {
+        Self {
+            channel: mpsc::channel(),
+            num_woken_tasks: Arc::new(AtomicUsize::new(0)),
+            num_spawned_tasks: Arc::new(AtomicUsize::new(0)),
+        }
+    }
+
+    pub fn spawn<T>(&self, future: impl Future<Output = T> + Send + 'static) -> Task<T>
+    where
+        T: Send + 'static,
+    {
+        let future = self.droppable_future(future);
+        let schedule = self.runnable_schedule_cb();
+        let (runnable, task) = async_task::spawn(future, schedule);
+        runnable.schedule();
+        task
+    }
+
+    pub fn spawn_local<T>(&self, future: impl Future<Output = T> + 'static) -> Task<T>
+    where
+        T: 'static,
+    {
+        let future = self.droppable_future(future);
+        let schedule = self.runnable_schedule_cb();
+        let (runnable, task) = async_task::spawn_local(future, schedule);
+        runnable.schedule();
+        task
+    }
+
+    /// Run the woken tasks once
+    ///
+    /// NOTE: Will not run tasks that are woken/scheduled immediately after `Runnable::run`
+    pub fn tick(&self) {
+        let num_woken_tasks = self.num_woken_tasks.load(Ordering::Relaxed);
+        self.channel
+            .1
+            .try_iter()
+            .take(num_woken_tasks)
+            .for_each(|runnable| {
+                runnable.run();
+            });
+        self.num_woken_tasks
+            .fetch_sub(num_woken_tasks, Ordering::Relaxed);
+    }
+
+    fn droppable_future<F>(&self, future: F) -> DroppableFuture<F, impl Fn()>
+    where
+        F: Future,
+    {
+        self.num_spawned_tasks.fetch_add(1, Ordering::Relaxed);
+        let num_spawned_tasks = self.num_spawned_tasks.clone();
+        DroppableFuture::new(future, move || {
+            num_spawned_tasks.fetch_sub(1, Ordering::Relaxed);
+        })
+    }
+
+    fn runnable_schedule_cb(&self) -> impl Fn(Runnable) {
+        let sender = self.channel.0.clone();
+        let num_woken_tasks = self.num_woken_tasks.clone();
+        move |runnable| {
+            sender.send(runnable).unwrap_or(());
+            num_woken_tasks.fetch_add(1, Ordering::Relaxed);
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_multiple_tasks() {
+        let executor = TickedAsyncExecutor::new();
+        executor
+            .spawn_local(async move {
+                println!("A: Start");
+                tokio::task::yield_now().await;
+                println!("A: End");
+            })
+            .detach();
+
+        executor
+            .spawn_local(async move {
+                println!("B: Start");
+                tokio::task::yield_now().await;
+                println!("B: End");
+            })
+            .detach();
+
+        executor
+            .spawn_local(async move {
+                println!("C: Start");
+                tokio::task::yield_now().await;
+                println!("C: End");
+            })
+            .detach();
+
+        // A, B, C: Start
+        executor.tick();
+
+        // A, B, C: End
+        executor.tick();
+    }
+
+    // TODO, Test Task cancellation
+    // TODO, Test FallibleTasks
+    // TODO, Test Edge cases
+}