2 files changed, 54 insertions, 85 deletions

diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs
index 4b17d4982..bdaa32951 100644
--- a/embassy-executor/src/raw/mod.rs
+++ b/embassy-executor/src/raw/mod.rs
@@ -16,7 +16,7 @@ mod run_queue;
 #[cfg_attr(not(target_has_atomic = "8"), path = "state_critical_section.rs")]
 mod state;
 
-#[cfg(feature = "trace")]
+#[cfg(feature = "_any_trace")]
 pub mod trace;
 pub(crate) mod util;
 #[cfg_attr(feature = "turbowakers", path = "waker_turbo.rs")]
@@ -41,6 +41,7 @@ use self::state::State;
 use self::util::{SyncUnsafeCell, UninitCell};
 pub use self::waker::task_from_waker;
 use super::SpawnToken;
+use crate::{Metadata, SpawnError};
 
 #[no_mangle]
 extern "Rust" fn __embassy_time_queue_item_from_waker(waker: &Waker) -> &'static mut TimerQueueItem {
@@ -94,11 +95,10 @@ pub(crate) struct TaskHeader {
 
     /// Integrated timer queue storage. This field should not be accessed outside of the timer queue.
     pub(crate) timer_queue_item: TimerQueueItem,
-    #[cfg(feature = "trace")]
-    pub(crate) name: Option<&'static str>,
-    #[cfg(feature = "trace")]
-    pub(crate) id: u32,
-    #[cfg(feature = "trace")]
+
+    pub(crate) metadata: Metadata,
+
+    #[cfg(feature = "rtos-trace")]
     all_tasks_next: AtomicPtr<TaskHeader>,
 }
 
@@ -129,6 +129,10 @@ impl TaskRef {
         unsafe { self.ptr.as_ref() }
     }
 
+    pub(crate) fn metadata(self) -> &'static Metadata {
+        unsafe { &self.ptr.as_ref().metadata }
+    }
+
     /// Returns a reference to the executor that the task is currently running on.
     pub unsafe fn executor(self) -> Option<&'static Executor> {
         let executor = self.header().executor.load(Ordering::Relaxed);
@@ -148,6 +152,12 @@ impl TaskRef {
     pub(crate) fn as_ptr(self) -> *const TaskHeader {
         self.ptr.as_ptr()
     }
+
+    /// Returns the task ID.
+    /// This can be used in combination with rtos-trace to match task names with IDs
+    pub fn id(&self) -> u32 {
+        self.as_ptr() as u32
+    }
 }
 
 /// Raw storage in which a task can be spawned.
@@ -189,11 +199,8 @@ impl<F: Future + 'static> TaskStorage<F> {
                 poll_fn: SyncUnsafeCell::new(None),
 
                 timer_queue_item: TimerQueueItem::new(),
-                #[cfg(feature = "trace")]
-                name: None,
-                #[cfg(feature = "trace")]
-                id: 0,
-                #[cfg(feature = "trace")]
+                metadata: Metadata::new(),
+                #[cfg(feature = "rtos-trace")]
                 all_tasks_next: AtomicPtr::new(core::ptr::null_mut()),
             },
             future: UninitCell::uninit(),
@@ -213,11 +220,11 @@ impl<F: Future + 'static> TaskStorage<F> {
     ///
     /// 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<impl Sized> {
+    pub fn spawn(&'static self, future: impl FnOnce() -> F) -> Result<SpawnToken<impl Sized>, SpawnError> {
         let task = AvailableTask::claim(self);
         match task {
-            Some(task) => task.initialize(future),
-            None => SpawnToken::new_failed(),
+            Some(task) => Ok(task.initialize(future)),
+            None => Err(SpawnError::Busy),
         }
     }
 
@@ -229,7 +236,7 @@ impl<F: Future + 'static> TaskStorage<F> {
         let mut cx = Context::from_waker(&waker);
         match future.poll(&mut cx) {
             Poll::Ready(_) => {
-                #[cfg(feature = "trace")]
+                #[cfg(feature = "_any_trace")]
                 let exec_ptr: *const SyncExecutor = this.raw.executor.load(Ordering::Relaxed);
 
                 // As the future has finished and this function will not be called
@@ -244,7 +251,7 @@ impl<F: Future + 'static> TaskStorage<F> {
                 // after we're done with it.
                 this.raw.state.despawn();
 
-                #[cfg(feature = "trace")]
+                #[cfg(feature = "_any_trace")]
                 trace::task_end(exec_ptr, &p);
             }
             Poll::Pending => {}
@@ -279,6 +286,7 @@ impl<F: Future + 'static> AvailableTask<F> {
 
     fn initialize_impl<S>(self, future: impl FnOnce() -> F) -> SpawnToken<S> {
         unsafe {
+            self.task.raw.metadata.reset();
             self.task.raw.poll_fn.set(Some(TaskStorage::<F>::poll));
             self.task.future.write_in_place(future);
 
@@ -345,10 +353,10 @@ impl<F: Future + 'static, const N: usize> TaskPool<F, N> {
         }
     }
 
-    fn spawn_impl<T>(&'static self, future: impl FnOnce() -> F) -> SpawnToken<T> {
+    fn spawn_impl<T>(&'static self, future: impl FnOnce() -> F) -> Result<SpawnToken<T>, SpawnError> {
         match self.pool.iter().find_map(AvailableTask::claim) {
-            Some(task) => task.initialize_impl::<T>(future),
-            None => SpawnToken::new_failed(),
+            Some(task) => Ok(task.initialize_impl::<T>(future)),
+            None => Err(SpawnError::Busy),
         }
     }
 
@@ -359,7 +367,7 @@ impl<F: Future + 'static, const N: usize> TaskPool<F, N> {
     /// 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()`](super::Spawner::spawn) to return the error.
-    pub fn spawn(&'static self, future: impl FnOnce() -> F) -> SpawnToken<impl Sized> {
+    pub fn spawn(&'static self, future: impl FnOnce() -> F) -> Result<SpawnToken<impl Sized>, SpawnError> {
         self.spawn_impl::<F>(future)
     }
 
@@ -372,7 +380,7 @@ impl<F: Future + 'static, const N: usize> TaskPool<F, N> {
     /// 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>
+    pub unsafe fn _spawn_async_fn<FutFn>(&'static self, future: FutFn) -> Result<SpawnToken<impl Sized>, SpawnError>
     where
         FutFn: FnOnce() -> F,
     {
@@ -417,7 +425,7 @@ impl SyncExecutor {
     /// - `task` must NOT be already enqueued (in this executor or another one).
     #[inline(always)]
     unsafe fn enqueue(&self, task: TaskRef, l: state::Token) {
-        #[cfg(feature = "trace")]
+        #[cfg(feature = "_any_trace")]
         trace::task_ready_begin(self, &task);
 
         if self.run_queue.enqueue(task, l) {
@@ -430,7 +438,7 @@ impl SyncExecutor {
             .executor
             .store((self as *const Self).cast_mut(), Ordering::Relaxed);
 
-        #[cfg(feature = "trace")]
+        #[cfg(feature = "_any_trace")]
         trace::task_new(self, &task);
 
         state::locked(|l| {
@@ -442,23 +450,23 @@ impl SyncExecutor {
     ///
     /// Same as [`Executor::poll`], plus you must only call this on the thread this executor was created.
     pub(crate) unsafe fn poll(&'static self) {
-        #[cfg(feature = "trace")]
+        #[cfg(feature = "_any_trace")]
         trace::poll_start(self);
 
         self.run_queue.dequeue_all(|p| {
             let task = p.header();
 
-            #[cfg(feature = "trace")]
+            #[cfg(feature = "_any_trace")]
             trace::task_exec_begin(self, &p);
 
             // Run the task
             task.poll_fn.get().unwrap_unchecked()(p);
 
-            #[cfg(feature = "trace")]
+            #[cfg(feature = "_any_trace")]
             trace::task_exec_end(self, &p);
         });
 
-        #[cfg(feature = "trace")]
+        #[cfg(feature = "_any_trace")]
         trace::executor_idle(self)
     }
 }
diff --git a/embassy-executor/src/raw/trace.rs b/embassy-executor/src/raw/trace.rs
index f484abf58..b3086948c 100644
--- a/embassy-executor/src/raw/trace.rs
+++ b/embassy-executor/src/raw/trace.rs
@@ -95,17 +95,20 @@ use crate::spawner::{SpawnError, SpawnToken, Spawner};
 /// This static provides access to the global task tracker which maintains
 /// a list of all tasks in the system. It's automatically updated by the
 /// task lifecycle hooks in the trace module.
-pub static TASK_TRACKER: TaskTracker = TaskTracker::new();
+#[cfg(feature = "rtos-trace")]
+pub(crate) static TASK_TRACKER: TaskTracker = TaskTracker::new();
 
 /// A thread-safe tracker for all tasks in the system
 ///
 /// This struct uses an intrusive linked list approach to track all tasks
 /// without additional memory allocations. It maintains a global list of
 /// tasks that can be traversed to find all currently existing tasks.
-pub struct TaskTracker {
+#[cfg(feature = "rtos-trace")]
+pub(crate) struct TaskTracker {
     head: AtomicPtr<TaskHeader>,
 }
 
+#[cfg(feature = "rtos-trace")]
 impl TaskTracker {
     /// Creates a new empty task tracker
     ///
@@ -125,7 +128,7 @@ impl TaskTracker {
     /// # Arguments
     /// * `task` - The task reference to add to the tracker
     pub fn add(&self, task: TaskRef) {
-        let task_ptr = task.as_ptr() as *mut TaskHeader;
+        let task_ptr = task.as_ptr();
 
         loop {
             let current_head = self.head.load(Ordering::Acquire);
@@ -135,7 +138,7 @@ impl TaskTracker {
 
             if self
                 .head
-                .compare_exchange(current_head, task_ptr, Ordering::Release, Ordering::Relaxed)
+                .compare_exchange(current_head, task_ptr.cast_mut(), Ordering::Release, Ordering::Relaxed)
                 .is_ok()
             {
                 break;
@@ -165,50 +168,7 @@ impl TaskTracker {
     }
 }
 
-/// Extension trait for `TaskRef` that provides tracing functionality.
-///
-/// This trait is only available when the `trace` feature is enabled.
-/// It extends `TaskRef` with methods for accessing and modifying task identifiers
-/// and names, which are useful for debugging, logging, and performance analysis.
-pub trait TaskRefTrace {
-    /// Get the name for a task
-    fn name(&self) -> Option<&'static str>;
-
-    /// Set the name for a task
-    fn set_name(&self, name: Option<&'static str>);
-
-    /// Get the ID for a task
-    fn id(&self) -> u32;
-
-    /// Set the ID for a task
-    fn set_id(&self, id: u32);
-}
-
-impl TaskRefTrace for TaskRef {
-    fn name(&self) -> Option<&'static str> {
-        self.header().name
-    }
-
-    fn set_name(&self, name: Option<&'static str>) {
-        unsafe {
-            let header_ptr = self.ptr.as_ptr() as *mut TaskHeader;
-            (*header_ptr).name = name;
-        }
-    }
-
-    fn id(&self) -> u32 {
-        self.header().id
-    }
-
-    fn set_id(&self, id: u32) {
-        unsafe {
-            let header_ptr = self.ptr.as_ptr() as *mut TaskHeader;
-            (*header_ptr).id = id;
-        }
-    }
-}
-
-#[cfg(not(feature = "rtos-trace"))]
+#[cfg(feature = "trace")]
 extern "Rust" {
     /// This callback is called when the executor begins polling. This will always
     /// be paired with a later call to `_embassy_trace_executor_idle`.
@@ -270,7 +230,7 @@ extern "Rust" {
 
 #[inline]
 pub(crate) fn poll_start(executor: &SyncExecutor) {
-    #[cfg(not(feature = "rtos-trace"))]
+    #[cfg(feature = "trace")]
     unsafe {
         _embassy_trace_poll_start(executor as *const _ as u32)
     }
@@ -278,7 +238,7 @@ pub(crate) fn poll_start(executor: &SyncExecutor) {
 
 #[inline]
 pub(crate) fn task_new(executor: &SyncExecutor, task: &TaskRef) {
-    #[cfg(not(feature = "rtos-trace"))]
+    #[cfg(feature = "trace")]
     unsafe {
         _embassy_trace_task_new(executor as *const _ as u32, task.as_ptr() as u32)
     }
@@ -286,7 +246,7 @@ pub(crate) fn task_new(executor: &SyncExecutor, task: &TaskRef) {
     #[cfg(feature = "rtos-trace")]
     {
         rtos_trace::trace::task_new(task.as_ptr() as u32);
-        let name = task.name().unwrap_or("unnamed task\0");
+        let name = task.metadata().name().unwrap_or("unnamed task\0");
         let info = rtos_trace::TaskInfo {
             name,
             priority: 0,
@@ -302,7 +262,7 @@ pub(crate) fn task_new(executor: &SyncExecutor, task: &TaskRef) {
 
 #[inline]
 pub(crate) fn task_end(executor: *const SyncExecutor, task: &TaskRef) {
-    #[cfg(not(feature = "rtos-trace"))]
+    #[cfg(feature = "trace")]
     unsafe {
         _embassy_trace_task_end(executor as u32, task.as_ptr() as u32)
     }
@@ -310,7 +270,7 @@ pub(crate) fn task_end(executor: *const SyncExecutor, task: &TaskRef) {
 
 #[inline]
 pub(crate) fn task_ready_begin(executor: &SyncExecutor, task: &TaskRef) {
-    #[cfg(not(feature = "rtos-trace"))]
+    #[cfg(feature = "trace")]
     unsafe {
         _embassy_trace_task_ready_begin(executor as *const _ as u32, task.as_ptr() as u32)
     }
@@ -320,7 +280,7 @@ pub(crate) fn task_ready_begin(executor: &SyncExecutor, task: &TaskRef) {
 
 #[inline]
 pub(crate) fn task_exec_begin(executor: &SyncExecutor, task: &TaskRef) {
-    #[cfg(not(feature = "rtos-trace"))]
+    #[cfg(feature = "trace")]
     unsafe {
         _embassy_trace_task_exec_begin(executor as *const _ as u32, task.as_ptr() as u32)
     }
@@ -330,7 +290,7 @@ pub(crate) fn task_exec_begin(executor: &SyncExecutor, task: &TaskRef) {
 
 #[inline]
 pub(crate) fn task_exec_end(executor: &SyncExecutor, task: &TaskRef) {
-    #[cfg(not(feature = "rtos-trace"))]
+    #[cfg(feature = "trace")]
     unsafe {
         _embassy_trace_task_exec_end(executor as *const _ as u32, task.as_ptr() as u32)
     }
@@ -340,7 +300,7 @@ pub(crate) fn task_exec_end(executor: &SyncExecutor, task: &TaskRef) {
 
 #[inline]
 pub(crate) fn executor_idle(executor: &SyncExecutor) {
-    #[cfg(not(feature = "rtos-trace"))]
+    #[cfg(feature = "trace")]
     unsafe {
         _embassy_trace_executor_idle(executor as *const _ as u32)
     }
@@ -356,6 +316,7 @@ pub(crate) fn executor_idle(executor: &SyncExecutor) {
 ///
 /// # Returns
 /// An iterator that yields `TaskRef` items for each task
+#[cfg(feature = "rtos-trace")]
 fn get_all_active_tasks() -> impl Iterator<Item = TaskRef> + 'static {
     struct TaskIterator<'a> {
         tracker: &'a TaskTracker,
@@ -384,6 +345,7 @@ fn get_all_active_tasks() -> impl Iterator<Item = TaskRef> + 'static {
 }
 
 /// Perform an action on each active task
+#[cfg(feature = "rtos-trace")]
 fn with_all_active_tasks<F>(f: F)
 where
     F: FnMut(TaskRef),
@@ -395,9 +357,8 @@ where
 impl rtos_trace::RtosTraceOSCallbacks for crate::raw::SyncExecutor {
     fn task_list() {
         with_all_active_tasks(|task| {
-            let name = task.name().unwrap_or("unnamed task\0");
             let info = rtos_trace::TaskInfo {
-                name,
+                name: task.metadata().name().unwrap_or("unnamed task\0"),
                 priority: 0,
                 stack_base: 0,
                 stack_size: 0,