3 files changed, 107 insertions, 41 deletions

diff --git a/embassy-executor/Cargo.toml b/embassy-executor/Cargo.toml
index 7c930b658..0a5360e5d 100644
--- a/embassy-executor/Cargo.toml
+++ b/embassy-executor/Cargo.toml
@@ -31,7 +31,7 @@ features = ["defmt", "arch-cortex-m", "executor-thread", "executor-interrupt"]
 [dependencies]
 defmt = { version = "0.3", optional = true }
 log = { version = "0.4.14", optional = true }
-rtos-trace = { version = "0.1.2", optional = true }
+rtos-trace = { version = "0.1.3", optional = true }
 
 embassy-executor-macros = { version = "0.6.2", path = "../embassy-executor-macros" }
 embassy-time-driver = { version = "0.1.0", path = "../embassy-time-driver", optional = true }
@@ -91,6 +91,10 @@ arch-spin = ["_arch"]
 executor-thread = []
 ## Enable the interrupt-mode executor (available in Cortex-M only)
 executor-interrupt = []
+## Enable tracing support (adds some overhead)
+trace = []
+## Enable support for rtos-trace framework
+rtos-trace = ["dep:rtos-trace", "trace"]
 
 #! ### Task Arena Size
 #! Sets the [task arena](#task-arena) size. Necessary if you’re not using `nightly`.
diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs
index ebabee1ba..3f93eae6f 100644
--- a/embassy-executor/src/raw/mod.rs
+++ b/embassy-executor/src/raw/mod.rs
@@ -18,6 +18,8 @@ mod state;
 
 #[cfg(feature = "integrated-timers")]
 mod timer_queue;
+#[cfg(feature = "trace")]
+mod trace;
 pub(crate) mod util;
 #[cfg_attr(feature = "turbowakers", path = "waker_turbo.rs")]
 mod waker;
@@ -31,8 +33,6 @@ use core::task::{Context, Poll};
 
 #[cfg(feature = "integrated-timers")]
 use embassy_time_driver::AlarmHandle;
-#[cfg(feature = "rtos-trace")]
-use rtos_trace::trace;
 
 use self::run_queue::{RunQueue, RunQueueItem};
 use self::state::State;
@@ -352,8 +352,8 @@ impl SyncExecutor {
     /// - `task` must NOT be already enqueued (in this executor or another one).
     #[inline(always)]
     unsafe fn enqueue(&self, task: TaskRef) {
-        #[cfg(feature = "rtos-trace")]
-        trace::task_ready_begin(task.as_ptr() as u32);
+        #[cfg(feature = "trace")]
+        trace::task_ready_begin(self, &task);
 
         if self.run_queue.enqueue(task) {
             self.pender.pend();
@@ -369,8 +369,8 @@ impl SyncExecutor {
     pub(super) unsafe fn spawn(&'static self, task: TaskRef) {
         task.header().executor.set(Some(self));
 
-        #[cfg(feature = "rtos-trace")]
-        trace::task_new(task.as_ptr() as u32);
+        #[cfg(feature = "trace")]
+        trace::task_new(self, &task);
 
         self.enqueue(task);
     }
@@ -400,14 +400,14 @@ impl SyncExecutor {
                     return;
                 }
 
-                #[cfg(feature = "rtos-trace")]
-                trace::task_exec_begin(p.as_ptr() as u32);
+                #[cfg(feature = "trace")]
+                trace::task_exec_begin(self, &p);
 
                 // Run the task
                 task.poll_fn.get().unwrap_unchecked()(p);
 
-                #[cfg(feature = "rtos-trace")]
-                trace::task_exec_end();
+                #[cfg(feature = "trace")]
+                trace::task_exec_end(self, &p);
 
                 // Enqueue or update into timer_queue
                 #[cfg(feature = "integrated-timers")]
@@ -430,8 +430,8 @@ impl SyncExecutor {
             }
         }
 
-        #[cfg(feature = "rtos-trace")]
-        trace::system_idle();
+        #[cfg(feature = "trace")]
+        trace::executor_idle(self)
     }
 }
 
@@ -593,31 +593,3 @@ impl embassy_time_queue_driver::TimerQueue for TimerQueue {
 
 #[cfg(feature = "integrated-timers")]
 embassy_time_queue_driver::timer_queue_impl!(static TIMER_QUEUE: TimerQueue = TimerQueue);
-
-#[cfg(all(feature = "rtos-trace", feature = "integrated-timers"))]
-const fn gcd(a: u64, b: u64) -> u64 {
-    if b == 0 {
-        a
-    } else {
-        gcd(b, a % b)
-    }
-}
-
-#[cfg(feature = "rtos-trace")]
-impl rtos_trace::RtosTraceOSCallbacks for Executor {
-    fn task_list() {
-        // We don't know what tasks exist, so we can't send them.
-    }
-    #[cfg(feature = "integrated-timers")]
-    fn time() -> u64 {
-        const GCD_1M: u64 = gcd(embassy_time_driver::TICK_HZ, 1_000_000);
-        embassy_time_driver::now() * (1_000_000 / GCD_1M) / (embassy_time_driver::TICK_HZ / GCD_1M)
-    }
-    #[cfg(not(feature = "integrated-timers"))]
-    fn time() -> u64 {
-        0
-    }
-}
-
-#[cfg(feature = "rtos-trace")]
-rtos_trace::global_os_callbacks! {Executor}
diff --git a/embassy-executor/src/raw/trace.rs b/embassy-executor/src/raw/trace.rs
new file mode 100644
index 000000000..c7bcf9c11
--- /dev/null
+++ b/embassy-executor/src/raw/trace.rs
@@ -0,0 +1,90 @@
+#![allow(unused)]
+use crate::raw::{SyncExecutor, TaskRef};
+
+#[cfg(not(feature = "rtos-trace"))]
+extern "Rust" {
+    fn _embassy_trace_task_new(executor_id: u32, task_id: u32);
+    fn _embassy_trace_task_exec_begin(executor_id: u32, task_id: u32);
+    fn _embassy_trace_task_exec_end(excutor_id: u32, task_id: u32);
+    fn _embassy_trace_task_ready_begin(executor_id: u32, task_id: u32);
+    fn _embassy_trace_executor_idle(executor_id: u32);
+}
+
+#[inline]
+pub(crate) fn task_new(executor: &SyncExecutor, task: &TaskRef) {
+    #[cfg(not(feature = "rtos-trace"))]
+    unsafe {
+        _embassy_trace_task_new(executor as *const _ as u32, task.as_ptr() as u32)
+    }
+
+    #[cfg(feature = "rtos-trace")]
+    rtos_trace::trace::task_new(task.as_ptr() as u32);
+}
+
+#[inline]
+pub(crate) fn task_ready_begin(executor: &SyncExecutor, task: &TaskRef) {
+    #[cfg(not(feature = "rtos-trace"))]
+    unsafe {
+        _embassy_trace_task_ready_begin(executor as *const _ as u32, task.as_ptr() as u32)
+    }
+    #[cfg(feature = "rtos-trace")]
+    rtos_trace::trace::task_ready_begin(task.as_ptr() as u32);
+}
+
+#[inline]
+pub(crate) fn task_exec_begin(executor: &SyncExecutor, task: &TaskRef) {
+    #[cfg(not(feature = "rtos-trace"))]
+    unsafe {
+        _embassy_trace_task_exec_begin(executor as *const _ as u32, task.as_ptr() as u32)
+    }
+    #[cfg(feature = "rtos-trace")]
+    rtos_trace::trace::task_exec_begin(task.as_ptr() as u32);
+}
+
+#[inline]
+pub(crate) fn task_exec_end(executor: &SyncExecutor, task: &TaskRef) {
+    #[cfg(not(feature = "rtos-trace"))]
+    unsafe {
+        _embassy_trace_task_exec_end(executor as *const _ as u32, task.as_ptr() as u32)
+    }
+    #[cfg(feature = "rtos-trace")]
+    rtos_trace::trace::task_exec_end();
+}
+
+#[inline]
+pub(crate) fn executor_idle(executor: &SyncExecutor) {
+    #[cfg(not(feature = "rtos-trace"))]
+    unsafe {
+        _embassy_trace_executor_idle(executor as *const _ as u32)
+    }
+    #[cfg(feature = "rtos-trace")]
+    rtos_trace::trace::system_idle();
+}
+
+#[cfg(all(feature = "rtos-trace", feature = "integrated-timers"))]
+const fn gcd(a: u64, b: u64) -> u64 {
+    if b == 0 {
+        a
+    } else {
+        gcd(b, a % b)
+    }
+}
+
+#[cfg(feature = "rtos-trace")]
+impl rtos_trace::RtosTraceOSCallbacks for crate::raw::SyncExecutor {
+    fn task_list() {
+        // We don't know what tasks exist, so we can't send them.
+    }
+    #[cfg(feature = "integrated-timers")]
+    fn time() -> u64 {
+        const GCD_1M: u64 = gcd(embassy_time_driver::TICK_HZ, 1_000_000);
+        embassy_time_driver::now() * (1_000_000 / GCD_1M) / (embassy_time_driver::TICK_HZ / GCD_1M)
+    }
+    #[cfg(not(feature = "integrated-timers"))]
+    fn time() -> u64 {
+        0
+    }
+}
+
+#[cfg(feature = "rtos-trace")]
+rtos_trace::global_os_callbacks! {SyncExecutor}