2 files changed, 161 insertions, 0 deletions

diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs
index e38a2af66..73f4f00ee 100644
--- a/embassy-executor/src/raw/mod.rs
+++ b/embassy-executor/src/raw/mod.rs
@@ -138,6 +138,12 @@ impl TaskRef {
     pub(crate) fn as_ptr(self) -> *const TaskHeader {
         self.ptr.as_ptr()
     }
+
+    /// Get the ID for a task
+    #[cfg(feature = "trace")]
+    pub fn as_id(self) -> u32 {
+        self.ptr.as_ptr() as u32
+    }
 }
 
 /// Raw storage in which a task can be spawned.
@@ -224,6 +230,9 @@ impl<F: Future + 'static> TaskStorage<F> {
                 // Make sure we despawn last, so that other threads can only spawn the task
                 // after we're done with it.
                 this.raw.state.despawn();
+
+                #[cfg(feature = "trace")]
+                trace::task_end(self, &task);
             }
             Poll::Pending => {}
         }
@@ -420,6 +429,9 @@ 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")]
+        trace::poll_start(self);
+
         self.run_queue.dequeue_all(|p| {
             let task = p.header();
 
diff --git a/embassy-executor/src/raw/trace.rs b/embassy-executor/src/raw/trace.rs
index b34387b58..57222d60b 100644
--- a/embassy-executor/src/raw/trace.rs
+++ b/embassy-executor/src/raw/trace.rs
@@ -1,16 +1,157 @@
 #![allow(unused)]
 use crate::raw::{SyncExecutor, TaskRef};
 
+//! # Tracing
+//!
+//! The `trace` feature enables a number of callbacks that can be used to track the
+//! lifecycle of tasks and/or executors.
+//!
+//! Callbacks will have one or both of the following IDs passed to them:
+//!
+//! 1. A `task_id`, a `u32` value unique to a task for the duration of the time it is valid
+//! 2. An `executor_id`, a `u32` value unique to an executor for the duration of the time it is
+//!    valid
+//!
+//! Today, both `task_id` and `executor_id` are u32s containing the least significant 32 bits of
+//! the address of the task or executor, however this is NOT a stable guarantee, and MAY change
+//! at any time.
+//!
+//! IDs are only guaranteed to be unique for the duration of time the item is valid. If a task
+//! ends, and is respond, it MAY or MAY NOT have the same ID. For tasks, this time is defined
+//! as the time between `_embassy_trace_task_new` and `_embassy_trace_task_end` for a given task.
+//! For executors, this time is not defined, but is often "forever" for practical embedded
+//! programs.
+//!
+//! Callbacks can be used by enabling the `trace` feature, and providing implementations of the
+//! `extern "Rust"` functions below. All callbacks must be implemented.
+//!
+//! ## Stability
+//!
+//! The `trace` interface is considered unstable. Callbacks may change, be added, or be removed
+//! in any minor or trivial release.
+//!
+//! ## Task Tracing lifecycle
+//!
+//! ```text
+//! ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─
+//!        │(1)                                            │
+//! │      │
+//!   ╔════▼════╗ (2) ┌─────────┐ (3) ┌─────────┐          │
+//! │ ║ SPAWNED ║────▶│ WAITING │────▶│ RUNNING │
+//!   ╚═════════╝     └─────────┘     └─────────┘          │
+//! │                 ▲         ▲     │    │    │
+//!                   │           (4)      │    │(6)       │
+//! │                 │         └ ─ ─ ┘    │    │
+//!                   │                    │    │          │
+//! │             ┌──────┐             (5) │    │  ┌─────┐
+//!               │ IDLE │◀────────────────┘    └─▶│ END │ │
+//! │             └──────┘                         └─────┘
+//!   ┌──────────────────────┐                             │
+//! └ ┤ Task Trace Lifecycle │─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─
+//!   └──────────────────────┘
+//! ```
+//!
+//! 1. A task is spawned. `_embassy_trace_task_new` is called
+//! 2. A task is enqueued for the first time, `_embassy_trace_task_ready_begin` is called
+//! 3. A task is polled, `_embassy_trace_task_exec_begin` is called
+//! 4. WHILE a task is polled, the task is re-awoken, and `_embassy_trace_task_ready_begin` is
+//!      called. The task does not IMMEDIATELY move state, until polling is complete and the
+//!      RUNNING state is existed. `_embassy_trace_task_exec_end` is called when polling is
+//!      complete, marking the transition to WAITING
+//! 5. Polling is complete, `_embassy_trace_task_exec_end` is called
+//! 6. The task has completed, and `_embassy_trace_task_end` is called.
+//!
+//! ## Executor Tracing lifecycle
+//!
+//! ```text
+//! ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─
+//!       │(1)                                             │
+//! │     │
+//!   ╔═══▼══╗   (2)     ┌────────────┐  (3)  ┌─────────┐  │
+//! │ ║ IDLE ║──────────▶│ SCHEDULING │──────▶│ POLLING │
+//!   ╚══════╝           └────────────┘       └─────────┘  │
+//! │     ▲              │            ▲            │
+//!       │      (5)     │            │  (4)       │       │
+//! │     └──────────────┘            └────────────┘
+//!   ┌──────────────────────────┐                         │
+//! └ ┤ Executor Trace Lifecycle │─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─
+//!   └──────────────────────────┘
+//! ```
+//!
+//! 1. The executor is started (no associated trace)
+//! 2. A task on this executor is awoken. `_embassy_trace_task_ready_begin` is called
+//!      when this occurs, and `_embassy_trace_poll_start` is called when the executor
+//!      actually begins running.
+//! 3. The executor has decided a task to poll. `_embassy_trace_task_exec_begin` is called.
+//! 4. The executor finishes polling the task. `_embassy_trace_task_exec_end` is called.
+//! 5. The executor has finished polling tasks. `_embassy_trace_executor_idle` is called.
+
 #[cfg(not(feature = "rtos-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`.
+    ///
+    /// This marks the EXECUTOR state transition from IDLE -> SCHEDULING.
+    fn _embassy_trace_poll_start(executor_id: u32);
+
+    /// This callback is called AFTER a task is initialized/allocated, and BEFORE
+    /// it is enqueued to run for the first time. If the task ends (and does not
+    /// loop "forever"), there will be a matching call to `_embassy_trace_task_end`.
+    ///
+    /// Tasks start life in the SPAWNED state.
     fn _embassy_trace_task_new(executor_id: u32, task_id: u32);
+
+    /// This callback is called AFTER a task is destructed/freed. This will always
+    /// have a prior matching call to `_embassy_trace_task_new`.
+    fn _embassy_trace_task_end(executor_id: u32, task_id: u32);
+
+    /// This callback is called AFTER a task has been dequeued from the runqueue,
+    /// and BEFORE the task is polled. There will always be a matching call to
+    /// `_embassy_trace_task_exec_end`.
+    ///
+    /// This marks the TASK state transition from WAITING -> RUNNING
+    /// This marks the EXECUTOR state transition from SCHEDULING -> POLLING
     fn _embassy_trace_task_exec_begin(executor_id: u32, task_id: u32);
+
+    /// This callback is called AFTER a task has completed polling. There will
+    /// always be a matching call to `_embassy_trace_task_exec_begin`.
+    ///
+    /// This marks the TASK state transition from either:
+    /// * RUNNING -> IDLE - if there were no `_embassy_trace_task_ready_begin` events
+    ///     for this task since the last `_embassy_trace_task_exec_begin` for THIS task
+    /// * RUNNING -> WAITING - if there WAS a `_embassy_trace_task_ready_begin` event
+    ///     for this task since the last `_embassy_trace_task_exec_begin` for THIS task
+    ///
+    /// This marks the EXECUTOR state transition from POLLING -> SCHEDULING
     fn _embassy_trace_task_exec_end(excutor_id: u32, task_id: u32);
+
+    /// This callback is called AFTER the waker for a task is awoken, and BEFORE it
+    /// is added to the run queue.
+    ///
+    /// If the given task is currently RUNNING, this marks no state change, BUT the
+    /// RUNNING task will then move to the WAITING stage when polling is complete.
+    ///
+    /// If the given task is currently IDLE, this marks the TASK state transition
+    /// from IDLE -> WAITING.
     fn _embassy_trace_task_ready_begin(executor_id: u32, task_id: u32);
+
+    /// This callback is called AFTER all dequeued tasks in a single call to poll
+    /// have been processed. This will always be paired with a call to
+    /// `_embassy_trace_executor_idle`.
+    ///
+    /// This marks the EXECUTOR state transition from
     fn _embassy_trace_executor_idle(executor_id: u32);
 }
 
 #[inline]
+pub(crate) fn poll_start(executor: &SyncExecutor) {
+    #[cfg(not(feature = "rtos-trace"))]
+    unsafe {
+        _embassy_trace_poll_start(executor as *const _ as u32)
+    }
+}
+
+#[inline]
 pub(crate) fn task_new(executor: &SyncExecutor, task: &TaskRef) {
     #[cfg(not(feature = "rtos-trace"))]
     unsafe {
@@ -22,6 +163,14 @@ pub(crate) fn task_new(executor: &SyncExecutor, task: &TaskRef) {
 }
 
 #[inline]
+pub(crate) fn task_end(executor: &SyncExecutor, task: &TaskRef) {
+    #[cfg(not(feature = "rtos-trace"))]
+    unsafe {
+        _embassy_trace_task_end(executor as *const _ as u32, task.as_ptr() as u32)
+    }
+}
+
+#[inline]
 pub(crate) fn task_ready_begin(executor: &SyncExecutor, task: &TaskRef) {
     #[cfg(not(feature = "rtos-trace"))]
     unsafe {