1 files changed, 109 insertions, 11 deletions

diff --git a/embassy-executor/src/spawner.rs b/embassy-executor/src/spawner.rs
index 271606244..522d97db3 100644
--- a/embassy-executor/src/spawner.rs
+++ b/embassy-executor/src/spawner.rs
@@ -1,9 +1,12 @@
-use core::future::poll_fn;
+use core::future::{poll_fn, Future};
 use core::marker::PhantomData;
 use core::mem;
+use core::sync::atomic::Ordering;
 use core::task::Poll;
 
 use super::raw;
+#[cfg(feature = "trace")]
+use crate::raw::trace::TaskRefTrace;
 
 /// Token to spawn a newly-created task in an executor.
 ///
@@ -21,7 +24,7 @@ use super::raw;
 /// 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<S> {
-    raw_task: Option<raw::TaskRef>,
+    pub(crate) raw_task: Option<raw::TaskRef>,
     phantom: PhantomData<*mut S>,
 }
 
@@ -33,6 +36,15 @@ impl<S> SpawnToken<S> {
         }
     }
 
+    /// Returns the task id if available, otherwise 0
+    /// This can be used in combination with rtos-trace to match task names with id's
+    pub fn id(&self) -> u32 {
+        match self.raw_task {
+            None => 0,
+            Some(t) => t.as_ptr() as u32,
+        }
+    }
+
     /// Return a SpawnToken that represents a failed spawn.
     pub fn new_failed() -> Self {
         Self {
@@ -50,8 +62,7 @@ impl<S> Drop for SpawnToken<S> {
 }
 
 /// Error returned when spawning a task.
-#[derive(Copy, Clone, Debug)]
-#[cfg_attr(feature = "defmt", derive(defmt::Format))]
+#[derive(Copy, Clone)]
 pub enum SpawnError {
     /// Too many instances of this task are already running.
     ///
@@ -61,6 +72,31 @@ pub enum SpawnError {
     Busy,
 }
 
+impl core::fmt::Debug for SpawnError {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        core::fmt::Display::fmt(self, f)
+    }
+}
+
+impl core::fmt::Display for SpawnError {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        match self {
+            SpawnError::Busy => write!(f, "Busy - Too many instances of this task are already running. Check the `pool_size` attribute of the task."),
+        }
+    }
+}
+
+#[cfg(feature = "defmt")]
+impl defmt::Format for SpawnError {
+    fn format(&self, f: defmt::Formatter) {
+        match self {
+            SpawnError::Busy => defmt::write!(f, "Busy - Too many instances of this task are already running. Check the `pool_size` attribute of the task."),
+        }
+    }
+}
+
+impl core::error::Error for SpawnError {}
+
 /// Handle to spawn tasks into an executor.
 ///
 /// This Spawner can spawn any task (Send and non-Send ones), but it can
@@ -69,7 +105,7 @@ pub enum SpawnError {
 /// If you want to spawn tasks from another thread, use [SendSpawner].
 #[derive(Copy, Clone)]
 pub struct Spawner {
-    executor: &'static raw::Executor,
+    pub(crate) executor: &'static raw::Executor,
     not_send: PhantomData<*mut ()>,
 }
 
@@ -89,14 +125,19 @@ impl Spawner {
     /// # Panics
     ///
     /// Panics if the current executor is not an Embassy executor.
-    pub async fn for_current_executor() -> Self {
+    pub fn for_current_executor() -> impl Future<Output = Self> {
         poll_fn(|cx| {
             let task = raw::task_from_waker(cx.waker());
-            let executor = unsafe { task.header().executor.get().unwrap_unchecked() };
+            let executor = unsafe {
+                task.header()
+                    .executor
+                    .load(Ordering::Relaxed)
+                    .as_ref()
+                    .unwrap_unchecked()
+            };
             let executor = unsafe { raw::Executor::wrap(executor) };
             Poll::Ready(Self::new(executor))
         })
-        .await
     }
 
     /// Spawn a task into an executor.
@@ -134,6 +175,58 @@ impl Spawner {
     pub fn make_send(&self) -> SendSpawner {
         SendSpawner::new(&self.executor.inner)
     }
+
+    /// Return the unique ID of this Spawner's Executor.
+    pub fn executor_id(&self) -> usize {
+        self.executor.id()
+    }
+}
+
+/// Extension trait adding tracing capabilities to the Spawner
+///
+/// This trait provides an additional method to spawn tasks with an associated name,
+/// which can be useful for debugging and tracing purposes.
+pub trait SpawnerTraceExt {
+    /// Spawns a new task with a specified name.
+    ///
+    /// # Arguments
+    /// * `name` - Static string name to associate with the task
+    /// * `token` - Token representing the task to spawn
+    ///
+    /// # Returns
+    /// Result indicating whether the spawn was successful
+    fn spawn_named<S>(&self, name: &'static str, token: SpawnToken<S>) -> Result<(), SpawnError>;
+}
+
+/// Implementation of the SpawnerTraceExt trait for Spawner when trace is enabled
+#[cfg(feature = "trace")]
+impl SpawnerTraceExt for Spawner {
+    fn spawn_named<S>(&self, name: &'static str, token: SpawnToken<S>) -> Result<(), SpawnError> {
+        let task = token.raw_task;
+        core::mem::forget(token);
+
+        match task {
+            Some(task) => {
+                // Set the name and ID when trace is enabled
+                task.set_name(Some(name));
+                let task_id = task.as_ptr() as u32;
+                task.set_id(task_id);
+
+                unsafe { self.executor.spawn(task) };
+                Ok(())
+            }
+            None => Err(SpawnError::Busy),
+        }
+    }
+}
+
+/// Implementation of the SpawnerTraceExt trait for Spawner when trace is disabled
+#[cfg(not(feature = "trace"))]
+impl SpawnerTraceExt for Spawner {
+    fn spawn_named<S>(&self, _name: &'static str, token: SpawnToken<S>) -> Result<(), SpawnError> {
+        // When trace is disabled, just forward to regular spawn and ignore the name
+        self.spawn(token)
+    }
 }
 
 /// Handle to spawn tasks into an executor from any thread.
@@ -161,13 +254,18 @@ impl SendSpawner {
     /// # Panics
     ///
     /// Panics if the current executor is not an Embassy executor.
-    pub async fn for_current_executor() -> Self {
+    pub fn for_current_executor() -> impl Future<Output = Self> {
         poll_fn(|cx| {
             let task = raw::task_from_waker(cx.waker());
-            let executor = unsafe { task.header().executor.get().unwrap_unchecked() };
+            let executor = unsafe {
+                task.header()
+                    .executor
+                    .load(Ordering::Relaxed)
+                    .as_ref()
+                    .unwrap_unchecked()
+            };
             Poll::Ready(Self::new(executor))
         })
-        .await
     }
 
     /// Spawn a task into an executor.