3 files changed, 111 insertions, 54 deletions

diff --git a/embassy-macros/src/macros/task.rs b/embassy-macros/src/macros/task.rs
index 396ce18f2..e48de3d63 100644
--- a/embassy-macros/src/macros/task.rs
+++ b/embassy-macros/src/macros/task.rs
@@ -73,26 +73,10 @@ pub fn run(args: syn::AttributeArgs, f: syn::ItemFn) -> Result<TokenStream, Toke
         // in the user's code.
         #task_inner
 
-        #visibility fn #task_ident(#fargs) -> #embassy_path::executor::SpawnToken<impl ::core::future::Future + 'static> {
-            use ::core::future::Future;
-            use #embassy_path::executor::SpawnToken;
-            use #embassy_path::executor::raw::TaskStorage;
-
-            type Fut = impl Future + 'static;
-
-            #[allow(clippy::declare_interior_mutable_const)]
-            const NEW_TS: TaskStorage<Fut> = TaskStorage::new();
-
-            static POOL: [TaskStorage<Fut>; #pool_size] = [NEW_TS; #pool_size];
-
-            // Opaque type laundering, to obscure its origin!
-            // Workaround for "opaque type's hidden type cannot be another opaque type from the same scope"
-            // https://github.com/rust-lang/rust/issues/96406
-            fn launder_tait(token: SpawnToken<impl Future+'static>) -> SpawnToken<impl Future+'static> {
-                token
-            }
-
-            launder_tait(unsafe { TaskStorage::spawn_pool(&POOL, move || #task_inner_ident(#(#arg_names,)*)) })
+        #visibility fn #task_ident(#fargs) -> #embassy_path::executor::SpawnToken<impl Sized> {
+            type Fut = impl ::core::future::Future + 'static;
+            static POOL: #embassy_path::executor::raw::TaskPool<Fut, #pool_size> = #embassy_path::executor::raw::TaskPool::new();
+            unsafe { POOL._spawn_async_fn(move || #task_inner_ident(#(#arg_names,)*)) }
         }
     };
 
diff --git a/embassy/src/executor/raw/mod.rs b/embassy/src/executor/raw/mod.rs
index 3ae52ae31..5cf399cdf 100644
--- a/embassy/src/executor/raw/mod.rs
+++ b/embassy/src/executor/raw/mod.rs
@@ -110,8 +110,8 @@ impl TaskHeader {
 /// Raw storage in which a task can be spawned.
 ///
 /// This struct holds the necessary memory to spawn one task whose future is `F`.
-/// At a given time, the `Task` may be in spawned or not-spawned state. You may spawn it
-/// with [`Task::spawn()`], which will fail if it is already spawned.
+/// At a given time, the `TaskStorage` may be in spawned or not-spawned state. You
+/// may spawn it with [`TaskStorage::spawn()`], which will fail if it is already spawned.
 ///
 /// A `TaskStorage` must live forever, it may not be deallocated even after the task has finished
 /// running. Hence the relevant methods require `&'static self`. It may be reused, however.
@@ -121,7 +121,7 @@ impl TaskHeader {
 /// the memory for the task is allocated: on the stack, or on the heap with e.g. `Box::leak`, etc.
 
 // repr(C) is needed to guarantee that the Task is located at offset 0
-// This makes it safe to cast between Task and Task pointers.
+// This makes it safe to cast between TaskHeader and TaskStorage pointers.
 #[repr(C)]
 pub struct TaskStorage<F: Future + 'static> {
     raw: TaskHeader,
@@ -129,6 +129,9 @@ pub struct TaskStorage<F: Future + 'static> {
 }
 
 impl<F: Future + 'static> TaskStorage<F> {
+    #[cfg(feature = "nightly")]
+    const NEW: Self = Self::new();
+
     /// Create a new TaskStorage, in not-spawned state.
     #[cfg(feature = "nightly")]
     pub const fn new() -> Self {
@@ -147,22 +150,6 @@ impl<F: Future + 'static> TaskStorage<F> {
         }
     }
 
-    /// Try to spawn a task in a pool.
-    ///
-    /// See [`Self::spawn()`] for details.
-    ///
-    /// This will loop over the pool and spawn the task in the first storage that
-    /// is currently free. If none is free,
-    pub fn spawn_pool(pool: &'static [Self], future: impl FnOnce() -> F) -> SpawnToken<F> {
-        for task in pool {
-            if task.spawn_allocate() {
-                return unsafe { task.spawn_initialize(future) };
-            }
-        }
-
-        SpawnToken::new_failed()
-    }
-
     /// Try to spawn the task.
     ///
     /// The `future` closure constructs the future. It's only called if spawning is
@@ -172,15 +159,15 @@ impl<F: Future + 'static> TaskStorage<F> {
     ///
     /// This function will fail if the task is already spawned and has not finished running.
     /// In this case, the error is delayed: a "poisoned" SpawnToken is returned, which will
-    /// cause [`Executor::spawn()`] to return the error.
+    /// cause [`Spawner::spawn()`] to return the error.
     ///
     /// Once the task has finished running, you may spawn it again. It is allowed to spawn it
     /// on a different executor.
-    pub fn spawn(&'static self, future: impl FnOnce() -> F) -> SpawnToken<F> {
+    pub fn spawn(&'static self, future: impl FnOnce() -> F) -> SpawnToken<impl Sized> {
         if self.spawn_allocate() {
-            unsafe { self.spawn_initialize(future) }
+            unsafe { SpawnToken::<F>::new(self.spawn_initialize(future)) }
         } else {
-            SpawnToken::new_failed()
+            SpawnToken::<F>::new_failed()
         }
     }
 
@@ -192,12 +179,11 @@ impl<F: Future + 'static> TaskStorage<F> {
             .is_ok()
     }
 
-    unsafe fn spawn_initialize(&'static self, future: impl FnOnce() -> F) -> SpawnToken<F> {
+    unsafe fn spawn_initialize(&'static self, future: impl FnOnce() -> F) -> NonNull<TaskHeader> {
         // Initialize the task
         self.raw.poll_fn.write(Self::poll);
         self.future.write(future());
-
-        SpawnToken::new(NonNull::new_unchecked(&self.raw as *const TaskHeader as _))
+        NonNull::new_unchecked(&self.raw as *const TaskHeader as *mut TaskHeader)
     }
 
     unsafe fn poll(p: NonNull<TaskHeader>) {
@@ -222,6 +208,89 @@ impl<F: Future + 'static> TaskStorage<F> {
 
 unsafe impl<F: Future + 'static> Sync for TaskStorage<F> {}
 
+/// Raw storage that can hold up to N tasks of the same type.
+///
+/// This is essentially a `[TaskStorage<F>; N]`.
+#[cfg(feature = "nightly")]
+pub struct TaskPool<F: Future + 'static, const N: usize> {
+    pool: [TaskStorage<F>; N],
+}
+
+#[cfg(feature = "nightly")]
+impl<F: Future + 'static, const N: usize> TaskPool<F, N> {
+    /// Create a new TaskPool, with all tasks in non-spawned state.
+    pub const fn new() -> Self {
+        Self {
+            pool: [TaskStorage::NEW; N],
+        }
+    }
+
+    /// Try to spawn a task in the pool.
+    ///
+    /// See [`TaskStorage::spawn()`] for details.
+    ///
+    /// This will loop over the pool and spawn the task in the first storage that
+    /// is currently free. If none is free, a "poisoned" SpawnToken is returned,
+    /// which will cause [`Spawner::spawn()`] to return the error.
+    pub fn spawn(&'static self, future: impl FnOnce() -> F) -> SpawnToken<impl Sized> {
+        for task in &self.pool {
+            if task.spawn_allocate() {
+                return unsafe { SpawnToken::<F>::new(task.spawn_initialize(future)) };
+            }
+        }
+
+        SpawnToken::<F>::new_failed()
+    }
+
+    /// Like spawn(), but allows the task to be send-spawned if the args are Send even if
+    /// the future is !Send.
+    ///
+    /// Not covered by semver guarantees. DO NOT call this directly. Intended to be used
+    /// by the Embassy macros ONLY.
+    ///
+    /// 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 _spawn_async_fn<FutFn>(&'static self, future: FutFn) -> SpawnToken<impl Sized>
+    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>`.
+
+        for task in &self.pool {
+            if task.spawn_allocate() {
+                return SpawnToken::<FutFn>::new(task.spawn_initialize(future));
+            }
+        }
+
+        SpawnToken::<FutFn>::new_failed()
+    }
+}
+
 /// Raw executor.
 ///
 /// This is the core of the Embassy executor. It is low-level, requiring manual
diff --git a/embassy/src/executor/spawner.rs b/embassy/src/executor/spawner.rs
index e6770e299..73c1f786f 100644
--- a/embassy/src/executor/spawner.rs
+++ b/embassy/src/executor/spawner.rs
@@ -12,17 +12,21 @@ use super::raw;
 /// value is a `SpawnToken` that represents an instance of the task, ready to spawn. You must
 /// then spawn it into an executor, typically with [`Spawner::spawn()`].
 ///
+/// The generic parameter `S` determines whether the task can be spawned in executors
+/// in other threads or not. If `S: Send`, it can, which allows spawning it into a [`SendSpawner`].
+/// If not, it can't, so it can only be spawned into the current thread's executor, with [`Spawner`].
+///
 /// # Panics
 ///
 /// Dropping a SpawnToken instance panics. You may not "abort" spawning a task in this way.
 /// Once you've invoked a task function and obtained a SpawnToken, you *must* spawn it.
 #[must_use = "Calling a task function does nothing on its own. You must spawn the returned SpawnToken, typically with Spawner::spawn()"]
-pub struct SpawnToken<F> {
+pub struct SpawnToken<S> {
     raw_task: Option<NonNull<raw::TaskHeader>>,
-    phantom: PhantomData<*mut F>,
+    phantom: PhantomData<*mut S>,
 }
 
-impl<F> SpawnToken<F> {
+impl<S> SpawnToken<S> {
     pub(crate) unsafe fn new(raw_task: NonNull<raw::TaskHeader>) -> Self {
         Self {
             raw_task: Some(raw_task),
@@ -38,7 +42,7 @@ impl<F> SpawnToken<F> {
     }
 }
 
-impl<F> Drop for SpawnToken<F> {
+impl<S> Drop for SpawnToken<S> {
     fn drop(&mut self) {
         // TODO deallocate the task instead.
         panic!("SpawnToken instances may not be dropped. You must pass them to Spawner::spawn()")
@@ -97,7 +101,7 @@ impl Spawner {
     /// Spawn a task into an executor.
     ///
     /// You obtain the `token` by calling a task function (i.e. one marked with `#[embassy::task]`).
-    pub fn spawn<F>(&self, token: SpawnToken<F>) -> Result<(), SpawnError> {
+    pub fn spawn<S>(&self, token: SpawnToken<S>) -> Result<(), SpawnError> {
         let task = token.raw_task;
         mem::forget(token);
 
@@ -119,7 +123,7 @@ impl Spawner {
     /// # Panics
     ///
     /// Panics if the spawning fails.
-    pub fn must_spawn<F>(&self, token: SpawnToken<F>) {
+    pub fn must_spawn<S>(&self, token: SpawnToken<S>) {
         unwrap!(self.spawn(token));
     }
 
@@ -173,7 +177,7 @@ impl SendSpawner {
     /// Spawn a task into an executor.
     ///
     /// You obtain the `token` by calling a task function (i.e. one marked with `#[embassy::task]`).
-    pub fn spawn<F: Send>(&self, token: SpawnToken<F>) -> Result<(), SpawnError> {
+    pub fn spawn<S: Send>(&self, token: SpawnToken<S>) -> Result<(), SpawnError> {
         let header = token.raw_task;
         mem::forget(token);
 
@@ -191,7 +195,7 @@ impl SendSpawner {
     /// # Panics
     ///
     /// Panics if the spawning fails.
-    pub fn must_spawn<F: Send>(&self, token: SpawnToken<F>) {
+    pub fn must_spawn<S: Send>(&self, token: SpawnToken<S>) {
         unwrap!(self.spawn(token));
     }
 }