From 462d04c6d5a0fc6072cf9bdb0faa60da74ff46d2 Mon Sep 17 00:00:00 2001
From: Kat Perez <katperez@microsoft.com>
Date: Thu, 8 May 2025 13:34:32 -0400
Subject: move TaskTracker to trace

---
 embassy-executor/src/raw/trace.rs | 74 +++++++++++++++++++++++++++++++++++++--
 1 file changed, 72 insertions(+), 2 deletions(-)

(limited to 'embassy-executor/src/raw/trace.rs')

diff --git a/embassy-executor/src/raw/trace.rs b/embassy-executor/src/raw/trace.rs
index eb960f721..b59da0526 100644
--- a/embassy-executor/src/raw/trace.rs
+++ b/embassy-executor/src/raw/trace.rs
@@ -82,11 +82,11 @@
 #![allow(unused)]
 
 use core::cell::UnsafeCell;
-use core::sync::atomic::{AtomicUsize, Ordering};
+use core::sync::atomic::{AtomicPtr, AtomicUsize, Ordering};
 
 use rtos_trace::TaskInfo;
 
-use crate::raw::{SyncExecutor, TaskHeader, TaskRef, TaskTracker};
+use crate::raw::{SyncExecutor, TaskHeader, TaskRef};
 use crate::spawner::{SpawnError, SpawnToken, Spawner};
 
 /// Extension trait adding tracing capabilities to the Spawner
@@ -139,6 +139,76 @@ impl TraceExt for Spawner {
 #[cfg(feature = "trace")]
 pub 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.
+#[cfg(feature = "trace")]
+pub struct TaskTracker {
+    head: AtomicPtr<TaskHeader>,
+}
+
+#[cfg(feature = "trace")]
+impl TaskTracker {
+    /// Creates a new empty task tracker
+    ///
+    /// Initializes a tracker with no tasks in its list.
+    pub const fn new() -> Self {
+        Self {
+            head: AtomicPtr::new(core::ptr::null_mut()),
+        }
+    }
+
+    /// Adds a task to the tracker
+    ///
+    /// This method inserts a task at the head of the intrusive linked list.
+    /// The operation is thread-safe and lock-free, using atomic operations
+    /// to ensure consistency even when called from different contexts.
+    ///
+    /// # 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;
+
+        loop {
+            let current_head = self.head.load(Ordering::Acquire);
+            unsafe {
+                (*task_ptr).all_tasks_next.store(current_head, Ordering::Relaxed);
+            }
+
+            if self
+                .head
+                .compare_exchange(current_head, task_ptr, Ordering::Release, Ordering::Relaxed)
+                .is_ok()
+            {
+                break;
+            }
+        }
+    }
+
+    /// Performs an operation on each task in the tracker
+    ///
+    /// This method traverses the entire list of tasks and calls the provided
+    /// function for each task. This allows inspecting or processing all tasks
+    /// in the system without modifying the tracker's structure.
+    ///
+    /// # Arguments
+    /// * `f` - A function to call for each task in the tracker
+    pub fn for_each<F>(&self, mut f: F)
+    where
+        F: FnMut(TaskRef),
+    {
+        let mut current = self.head.load(Ordering::Acquire);
+        while !current.is_null() {
+            let task = unsafe { TaskRef::from_ptr(current) };
+            f(task);
+
+            current = unsafe { (*current).all_tasks_next.load(Ordering::Acquire) };
+        }
+    }
+}
+
 #[cfg(not(feature = "rtos-trace"))]
 extern "Rust" {
     /// This callback is called when the executor begins polling. This will always
-- 
cgit