Make AvailableTask public, deduplicate

2 files changed, 65 insertions, 51 deletions

diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs
index 7caa3302f..c1d82e18a 100644
--- a/embassy-executor/src/raw/mod.rs
+++ b/embassy-executor/src/raw/mod.rs
@@ -147,10 +147,7 @@ impl<F: Future + 'static> TaskStorage<F> {
     pub fn spawn(&'static self, future: impl FnOnce() -> F) -> SpawnToken<impl Sized> {
         let task = AvailableTask::claim(self);
         match task {
-            Some(task) => {
-                let task = task.initialize(future);
-                unsafe { SpawnToken::<F>::new(task) }
-            }
+            Some(task) => task.initialize(future),
             None => SpawnToken::new_failed(),
         }
     }
@@ -186,12 +183,16 @@ impl<F: Future + 'static> TaskStorage<F> {
     }
 }
 
-struct AvailableTask<F: Future + 'static> {
+/// An uninitialized [`TaskStorage`].
+pub struct AvailableTask<F: Future + 'static> {
     task: &'static TaskStorage<F>,
 }
 
 impl<F: Future + 'static> AvailableTask<F> {
-    fn claim(task: &'static TaskStorage<F>) -> Option<Self> {
+    /// Try to claim a [`TaskStorage`].
+    ///
+    /// This function returns `None` if a task has already been spawned and has not finished running.
+    pub fn claim(task: &'static TaskStorage<F>) -> Option<Self> {
         task.raw
             .state
             .compare_exchange(0, STATE_SPAWNED | STATE_RUN_QUEUED, Ordering::AcqRel, Ordering::Acquire)
@@ -199,12 +200,56 @@ impl<F: Future + 'static> AvailableTask<F> {
             .map(|_| Self { task })
     }
 
-    fn initialize(self, future: impl FnOnce() -> F) -> TaskRef {
+    fn initialize_impl<S>(self, future: impl FnOnce() -> F) -> SpawnToken<S> {
         unsafe {
             self.task.raw.poll_fn.set(Some(TaskStorage::<F>::poll));
             self.task.future.write(future());
+
+            let task = TaskRef::new(self.task);
+
+            SpawnToken::new(task)
         }
-        TaskRef::new(self.task)
+    }
+
+    /// Initialize the [`TaskStorage`] to run the given future.
+    pub fn initialize(self, future: impl FnOnce() -> F) -> SpawnToken<F> {
+        self.initialize_impl::<F>(future)
+    }
+
+    /// Initialize the [`TaskStorage`] to run the given future.
+    ///
+    /// # Safety
+    ///
+    /// `future` must be a closure of the form `move || my_async_fn(args)`, where `my_async_fn`
+    /// is an `async fn`, NOT a hand-written `Future`.
+    #[doc(hidden)]
+    pub unsafe fn __initialize_async_fn<FutFn>(self, future: impl FnOnce() -> F) -> SpawnToken<FutFn> {
+        // When send-spawning a task, we construct the future in this thread, and effectively
+        // "send" it to the executor thread by enqueuing it in its queue. Therefore, in theory,
+        // send-spawning should require the future `F` to be `Send`.
+        //
+        // The problem is this is more restrictive than needed. Once the future is executing,
+        // it is never sent to another thread. It is only sent when spawning. It should be
+        // enough for the task's arguments to be Send. (and in practice it's super easy to
+        // accidentally make your futures !Send, for example by holding an `Rc` or a `&RefCell` across an `.await`.)
+        //
+        // We can do it by sending the task args and constructing the future in the executor thread
+        // on first poll. However, this cannot be done in-place, so it'll waste stack space for a copy
+        // of the args.
+        //
+        // Luckily, an `async fn` future contains just the args when freshly constructed. So, if the
+        // args are Send, it's OK to send a !Send future, as long as we do it before first polling it.
+        //
+        // (Note: this is how the generators are implemented today, it's not officially guaranteed yet,
+        // but it's possible it'll be guaranteed in the future. See zulip thread:
+        // https://rust-lang.zulipchat.com/#narrow/stream/187312-wg-async/topic/.22only.20before.20poll.22.20Send.20futures )
+        //
+        // The `FutFn` captures all the args, so if it's Send, the task can be send-spawned.
+        // This is why we return `SpawnToken<FutFn>` below.
+        //
+        // This ONLY holds for `async fn` futures. The other `spawn` methods can be called directly
+        // by the user, with arbitrary hand-implemented futures. This is why these return `SpawnToken<F>`.
+        self.initialize_impl::<FutFn>(future)
     }
 }
 
@@ -223,6 +268,13 @@ impl<F: Future + 'static, const N: usize> TaskPool<F, N> {
         }
     }
 
+    fn spawn_impl<T>(&'static self, future: impl FnOnce() -> F) -> SpawnToken<T> {
+        match self.pool.iter().find_map(AvailableTask::claim) {
+            Some(task) => task.initialize_impl::<T>(future),
+            None => SpawnToken::new_failed(),
+        }
+    }
+
     /// Try to spawn a task in the pool.
     ///
     /// See [`TaskStorage::spawn()`] for details.
@@ -231,14 +283,7 @@ impl<F: Future + 'static, const N: usize> TaskPool<F, N> {
     /// is currently free. If none is free, a "poisoned" SpawnToken is returned,
     /// which will cause [`Spawner::spawn()`](super::Spawner::spawn) to return the error.
     pub fn spawn(&'static self, future: impl FnOnce() -> F) -> SpawnToken<impl Sized> {
-        let task = self.pool.iter().find_map(AvailableTask::claim);
-        match task {
-            Some(task) => {
-                let task = task.initialize(future);
-                unsafe { SpawnToken::<F>::new(task) }
-            }
-            None => SpawnToken::new_failed(),
-        }
+        self.spawn_impl::<F>(future)
     }
 
     /// Like spawn(), but allows the task to be send-spawned if the args are Send even if
@@ -254,40 +299,8 @@ impl<F: Future + 'static, const N: usize> TaskPool<F, N> {
     where
         FutFn: FnOnce() -> F,
     {
-        // When send-spawning a task, we construct the future in this thread, and effectively
-        // "send" it to the executor thread by enqueuing it in its queue. Therefore, in theory,
-        // send-spawning should require the future `F` to be `Send`.
-        //
-        // The problem is this is more restrictive than needed. Once the future is executing,
-        // it is never sent to another thread. It is only sent when spawning. It should be
-        // enough for the task's arguments to be Send. (and in practice it's super easy to
-        // accidentally make your futures !Send, for example by holding an `Rc` or a `&RefCell` across an `.await`.)
-        //
-        // We can do it by sending the task args and constructing the future in the executor thread
-        // on first poll. However, this cannot be done in-place, so it'll waste stack space for a copy
-        // of the args.
-        //
-        // Luckily, an `async fn` future contains just the args when freshly constructed. So, if the
-        // args are Send, it's OK to send a !Send future, as long as we do it before first polling it.
-        //
-        // (Note: this is how the generators are implemented today, it's not officially guaranteed yet,
-        // but it's possible it'll be guaranteed in the future. See zulip thread:
-        // https://rust-lang.zulipchat.com/#narrow/stream/187312-wg-async/topic/.22only.20before.20poll.22.20Send.20futures )
-        //
-        // The `FutFn` captures all the args, so if it's Send, the task can be send-spawned.
-        // This is why we return `SpawnToken<FutFn>` below.
-        //
-        // This ONLY holds for `async fn` futures. The other `spawn` methods can be called directly
-        // by the user, with arbitrary hand-implemented futures. This is why these return `SpawnToken<F>`.
-
-        let task = self.pool.iter().find_map(AvailableTask::claim);
-        match task {
-            Some(task) => {
-                let task = task.initialize(future);
-                unsafe { SpawnToken::<FutFn>::new(task) }
-            }
-            None => SpawnToken::new_failed(),
-        }
+        // See the comment in AvailableTask::__initialize_async_fn for explanation.
+        self.spawn_impl::<FutFn>(future)
     }
 }
 
diff --git a/embassy-executor/src/spawner.rs b/embassy-executor/src/spawner.rs
index 2b6224045..5a3a0dee1 100644
--- a/embassy-executor/src/spawner.rs
+++ b/embassy-executor/src/spawner.rs
@@ -33,7 +33,8 @@ impl<S> SpawnToken<S> {
         }
     }
 
-    pub(crate) fn new_failed() -> Self {
+    /// Return a SpawnToken that represents a failed spawn.
+    pub fn new_failed() -> Self {
         Self {
             raw_task: None,
             phantom: PhantomData,