feat(maitake): implement JoinHandles

maitake/Cargo.toml

@@ -54,9 +54,10 @@ optional = true
 [dev-dependencies]
 futures-util = "0.3"
 futures = "0.3"
+tokio-test = "0.4"
 
 [target.'cfg(not(loom))'.dev-dependencies]
-tracing-subscriber = { git = "https://github.com/tokio-rs/tracing", features = ["ansi", "fmt"] }
+tracing-subscriber = { git = "https://github.com/tokio-rs/tracing", default-features = false, features = ["ansi", "fmt"] }
 tracing-02 = { package = "tracing", git = "https://github.com/tokio-rs/tracing", default-features = false }
 console-subscriber = "0.1.6"
 

maitake/src/loom.rs

@@ -6,7 +6,55 @@ mod inner {
     #![allow(dead_code)]
 
     #[cfg(feature = "alloc")]
-    pub(crate) use loom::alloc;
+    pub(crate) mod alloc {
+        use super::sync::Arc;
+        use core::{
+            future::Future,
+            pin::Pin,
+            task::{Context, Poll},
+        };
+        pub(crate) use loom::alloc::*;
+        #[pin_project::pin_project]
+        pub(crate) struct TrackFuture<F> {
+            #[pin]
+            inner: F,
+            track: Arc<()>,
+        }
+
+        impl<F: Future> Future for TrackFuture<F> {
+            type Output = TrackFuture<F::Output>;
+            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+                let this = self.project();
+                this.inner.poll(cx).map(|inner| TrackFuture {
+                    inner,
+                    track: this.track.clone(),
+                })
+            }
+        }
+
+        impl<F> TrackFuture<F> {
+            /// Wrap a `Future` in a `TrackFuture` that participates in Loom's
+            /// leak checking.
+            #[track_caller]
+            pub(crate) fn new(inner: F) -> Self {
+                Self {
+                    inner,
+                    track: Arc::new(()),
+                }
+            }
+
+            /// Stop tracking this future, and return the inner value.
+            pub(crate) fn into_inner(self) -> F {
+                self.inner
+            }
+        }
+
+        #[track_caller]
+        pub(crate) fn track_future<F: Future>(inner: F) -> TrackFuture<F> {
+            TrackFuture::new(inner)
+        }
+    }
+
     pub(crate) use loom::{cell, future, model, thread};
 
     pub(crate) mod sync {

maitake/src/scheduler.rs

@@ -1,4 +1,4 @@
-use crate::task::{self, Header, Storage, TaskRef};
+use crate::task::{self, Header, JoinHandle, Storage, TaskRef};
 use core::{future::Future, pin::Pin};
 
 use cordyceps::mpsc_queue::MpscQueue;
@@ -87,13 +87,15 @@ impl StaticScheduler {
     /// [`Storage`]: crate::task::Storage
     #[inline]
     #[track_caller]
-    pub fn spawn_allocated<F, STO>(&'static self, task: STO::StoredTask)
+    pub fn spawn_allocated<F, STO>(&'static self, task: STO::StoredTask) -> JoinHandle<F::Output>
     where
         F: Future + 'static,
+        F::Output: 'static,
         STO: Storage<&'static Self, F>,
     {
-        let tr = TaskRef::new_allocated::<&'static Self, F, STO>(task);
-        self.schedule(tr);
+        let (task, join) = TaskRef::new_allocated::<&'static Self, F, STO>(task);
+        self.schedule(task);
+        join
     }
 
     /// Returns a new [task `Builder`] for configuring tasks prior to spawning
@@ -245,18 +247,26 @@ feature! {
 
         #[inline]
         #[track_caller]
-        pub fn spawn(&self, future: impl Future + 'static) {
-            self.schedule(TaskRef::new(self.clone(), future));
+        pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
+        where
+            F: Future + 'static,
+            F::Output: 'static,
+        {
+            let (task, join) = TaskRef::new(self.clone(), future);
+            self.schedule(task);
+            join
         }
 
         #[inline]
         #[track_caller]
-        pub fn spawn_allocated<F>(&'static self, task: Box<Task<Self, F, BoxStorage>>)
+        pub fn spawn_allocated<F>(&'static self, task: Box<Task<Self, F, BoxStorage>>) -> JoinHandle<F::Output>
         where
             F: Future + 'static,
+            F::Output: 'static,
         {
-            let tr = TaskRef::new_allocated::<Self, F, BoxStorage>(task);
-            self.schedule(tr);
+            let (task, join) = TaskRef::new_allocated::<Self, F, BoxStorage>(task);
+            self.schedule(task);
+            join
         }
 
         pub fn tick(&self) -> Tick {
@@ -278,14 +288,20 @@ feature! {
 
         #[inline]
         #[track_caller]
-        pub fn spawn(&'static self, future: impl Future + 'static) {
-            self.schedule(TaskRef::new(self, future));
+        pub fn spawn<F>(&'static self, future: F) -> JoinHandle<F::Output>
+        where
+            F: Future + 'static,
+            F::Output: 'static,
+        {
+            let (task, join) = TaskRef::new(self, future);
+            self.schedule(task);
+            join
         }
     }
 
     impl Core {
         fn new() -> Self {
-            let stub_task = TaskRef::new(Stub, Stub);
+            let (stub_task, _) = TaskRef::new(Stub, Stub);
             Self {
                 run_queue: MpscQueue::new_with_stub(test_dbg!(stub_task)),
             }

maitake/src/scheduler/tests.rs

@@ -87,8 +87,12 @@ mod custom_storage {
         }
     }
 
-    impl<F: Future + 'static> MyBoxTask<&'static StaticScheduler, F> {
-        fn spawn(scheduler: &'static StaticScheduler, future: F) {
+    impl<F> MyBoxTask<&'static StaticScheduler, F>
+    where
+        F: Future + 'static,
+        F::Output: 'static,
+    {
+        fn spawn(scheduler: &'static StaticScheduler, future: F) -> task::JoinHandle<F::Output> {
             let task = MyBoxTask(Box::new(Task::new(scheduler, future)));
             scheduler.spawn_allocated::<F, MyBoxStorage>(task)
         }
@@ -155,39 +159,10 @@ mod loom {
     use super::*;
     use crate::loom::{
         self,
+        alloc::track_future,
         sync::{atomic::AtomicBool, Arc},
         thread,
     };
-    use core::{
-        future::Future,
-        pin::Pin,
-        task::{Context, Poll},
-    };
-
-    #[pin_project::pin_project]
-    struct TrackFuture<F> {
-        #[pin]
-        inner: F,
-        track: Arc<()>,
-    }
-
-    impl<F: Future> Future for TrackFuture<F> {
-        type Output = TrackFuture<F::Output>;
-        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-            let this = self.project();
-            this.inner.poll(cx).map(|inner| TrackFuture {
-                inner,
-                track: this.track.clone(),
-            })
-        }
-    }
-
-    fn track_future<F: Future>(inner: F) -> TrackFuture<F> {
-        TrackFuture {
-            inner,
-            track: Arc::new(()),
-        }
-    }
 
     #[test]
     fn basically_works() {