Fix task_end callback

2 files changed, 7 insertions, 3 deletions

diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs
index 73f4f00ee..5b1f33a0e 100644
--- a/embassy-executor/src/raw/mod.rs
+++ b/embassy-executor/src/raw/mod.rs
@@ -219,6 +219,9 @@ impl<F: Future + 'static> TaskStorage<F> {
         let mut cx = Context::from_waker(&waker);
         match future.poll(&mut cx) {
             Poll::Ready(_) => {
+                #[cfg(feature = "trace")]
+                let exec_ptr: *const SyncExecutor = this.raw.executor.load(Ordering::Relaxed);
+
                 // As the future has finished and this function will not be called
                 // again, we can safely drop the future here.
                 this.future.drop_in_place();
@@ -232,7 +235,7 @@ impl<F: Future + 'static> TaskStorage<F> {
                 this.raw.state.despawn();
 
                 #[cfg(feature = "trace")]
-                trace::task_end(self, &task);
+                trace::task_end(exec_ptr, &p);
             }
             Poll::Pending => {}
         }
diff --git a/embassy-executor/src/raw/trace.rs b/embassy-executor/src/raw/trace.rs
index 0952a9bc0..56724b0bb 100644
--- a/embassy-executor/src/raw/trace.rs
+++ b/embassy-executor/src/raw/trace.rs
@@ -84,6 +84,7 @@
 //! 5. The executor has finished polling tasks. `_embassy_trace_executor_idle` is called.
 
 #![allow(unused)]
+
 use crate::raw::{SyncExecutor, TaskRef};
 
 #[cfg(not(feature = "rtos-trace"))]
@@ -163,10 +164,10 @@ pub(crate) fn task_new(executor: &SyncExecutor, task: &TaskRef) {
 }
 
 #[inline]
-pub(crate) fn task_end(executor: &SyncExecutor, task: &TaskRef) {
+pub(crate) fn task_end(executor: *const SyncExecutor, task: &TaskRef) {
     #[cfg(not(feature = "rtos-trace"))]
     unsafe {
-        _embassy_trace_task_end(executor as *const _ as u32, task.as_ptr() as u32)
+        _embassy_trace_task_end(executor as u32, task.as_ptr() as u32)
     }
 }