3 files changed, 88 insertions, 141 deletions

diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs
index 3f93eae6f..80bd49bad 100644
--- a/embassy-executor/src/raw/mod.rs
+++ b/embassy-executor/src/raw/mod.rs
@@ -17,7 +17,7 @@ mod run_queue;
 mod state;
 
 #[cfg(feature = "integrated-timers")]
-mod timer_queue;
+pub mod timer_queue;
 #[cfg(feature = "trace")]
 mod trace;
 pub(crate) mod util;
@@ -31,9 +31,6 @@ use core::pin::Pin;
 use core::ptr::NonNull;
 use core::task::{Context, Poll};
 
-#[cfg(feature = "integrated-timers")]
-use embassy_time_driver::AlarmHandle;
-
 use self::run_queue::{RunQueue, RunQueueItem};
 use self::state::State;
 use self::util::{SyncUnsafeCell, UninitCell};
@@ -47,8 +44,7 @@ pub(crate) struct TaskHeader {
     pub(crate) executor: SyncUnsafeCell<Option<&'static SyncExecutor>>,
     poll_fn: SyncUnsafeCell<Option<unsafe fn(TaskRef)>>,
 
-    #[cfg(feature = "integrated-timers")]
-    pub(crate) expires_at: SyncUnsafeCell<u64>,
+    /// Integrated timer queue storage. This field should not be accessed outside of the timer queue.
     #[cfg(feature = "integrated-timers")]
     pub(crate) timer_queue_item: timer_queue::TimerQueueItem,
 }
@@ -80,6 +76,12 @@ impl TaskRef {
         unsafe { self.ptr.as_ref() }
     }
 
+    /// Returns a reference to the executor that the task is currently running on.
+    #[cfg(feature = "integrated-timers")]
+    pub unsafe fn executor(self) -> Option<&'static Executor> {
+        self.header().executor.get().map(|e| Executor::wrap(e))
+    }
+
     /// The returned pointer is valid for the entire TaskStorage.
     pub(crate) fn as_ptr(self) -> *const TaskHeader {
         self.ptr.as_ptr()
@@ -121,8 +123,6 @@ impl<F: Future + 'static> TaskStorage<F> {
                 poll_fn: SyncUnsafeCell::new(None),
 
                 #[cfg(feature = "integrated-timers")]
-                expires_at: SyncUnsafeCell::new(0),
-                #[cfg(feature = "integrated-timers")]
                 timer_queue_item: timer_queue::TimerQueueItem::new(),
             },
             future: UninitCell::uninit(),
@@ -160,9 +160,6 @@ impl<F: Future + 'static> TaskStorage<F> {
             Poll::Ready(_) => {
                 this.future.drop_in_place();
                 this.raw.state.despawn();
-
-                #[cfg(feature = "integrated-timers")]
-                this.raw.expires_at.set(u64::MAX);
             }
             Poll::Pending => {}
         }
@@ -316,34 +313,16 @@ impl Pender {
 pub(crate) struct SyncExecutor {
     run_queue: RunQueue,
     pender: Pender,
-
-    #[cfg(feature = "integrated-timers")]
-    pub(crate) timer_queue: timer_queue::TimerQueue,
-    #[cfg(feature = "integrated-timers")]
-    alarm: AlarmHandle,
 }
 
 impl SyncExecutor {
     pub(crate) fn new(pender: Pender) -> Self {
-        #[cfg(feature = "integrated-timers")]
-        let alarm = unsafe { unwrap!(embassy_time_driver::allocate_alarm()) };
-
         Self {
             run_queue: RunQueue::new(),
             pender,
-
-            #[cfg(feature = "integrated-timers")]
-            timer_queue: timer_queue::TimerQueue::new(),
-            #[cfg(feature = "integrated-timers")]
-            alarm,
         }
     }
 
-    pub(crate) unsafe fn initialize(&'static self) {
-        #[cfg(feature = "integrated-timers")]
-        embassy_time_driver::set_alarm_callback(self.alarm, Self::alarm_callback, self as *const _ as *mut ());
-    }
-
     /// Enqueue a task in the task queue
     ///
     /// # Safety
@@ -360,12 +339,6 @@ impl SyncExecutor {
         }
     }
 
-    #[cfg(feature = "integrated-timers")]
-    fn alarm_callback(ctx: *mut ()) {
-        let this: &Self = unsafe { &*(ctx as *const Self) };
-        this.pender.pend();
-    }
-
     pub(super) unsafe fn spawn(&'static self, task: TaskRef) {
         task.header().executor.set(Some(self));
 
@@ -379,56 +352,27 @@ 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) {
-        #[allow(clippy::never_loop)]
-        loop {
-            #[cfg(feature = "integrated-timers")]
-            self.timer_queue
-                .dequeue_expired(embassy_time_driver::now(), wake_task_no_pend);
-
-            self.run_queue.dequeue_all(|p| {
-                let task = p.header();
-
-                #[cfg(feature = "integrated-timers")]
-                task.expires_at.set(u64::MAX);
-
-                if !task.state.run_dequeue() {
-                    // If task is not running, ignore it. This can happen in the following scenario:
-                    //   - Task gets dequeued, poll starts
-                    //   - While task is being polled, it gets woken. It gets placed in the queue.
-                    //   - Task poll finishes, returning done=true
-                    //   - RUNNING bit is cleared, but the task is already in the queue.
-                    return;
-                }
-
-                #[cfg(feature = "trace")]
-                trace::task_exec_begin(self, &p);
+        self.run_queue.dequeue_all(|p| {
+            let task = p.header();
+
+            if !task.state.run_dequeue() {
+                // If task is not running, ignore it. This can happen in the following scenario:
+                //   - Task gets dequeued, poll starts
+                //   - While task is being polled, it gets woken. It gets placed in the queue.
+                //   - Task poll finishes, returning done=true
+                //   - RUNNING bit is cleared, but the task is already in the queue.
+                return;
+            }
 
-                // Run the task
-                task.poll_fn.get().unwrap_unchecked()(p);
+            #[cfg(feature = "trace")]
+            trace::task_exec_begin(self, &p);
 
-                #[cfg(feature = "trace")]
-                trace::task_exec_end(self, &p);
+            // Run the task
+            task.poll_fn.get().unwrap_unchecked()(p);
 
-                // Enqueue or update into timer_queue
-                #[cfg(feature = "integrated-timers")]
-                self.timer_queue.update(p);
-            });
-
-            #[cfg(feature = "integrated-timers")]
-            {
-                // If this is already in the past, set_alarm might return false
-                // In that case do another poll loop iteration.
-                let next_expiration = self.timer_queue.next_expiration();
-                if embassy_time_driver::set_alarm(self.alarm, next_expiration) {
-                    break;
-                }
-            }
-
-            #[cfg(not(feature = "integrated-timers"))]
-            {
-                break;
-            }
-        }
+            #[cfg(feature = "trace")]
+            trace::task_exec_end(self, &p);
+        });
 
         #[cfg(feature = "trace")]
         trace::executor_idle(self)
@@ -494,15 +438,6 @@ impl Executor {
         }
     }
 
-    /// Initializes the executor.
-    ///
-    /// # Safety
-    ///
-    /// This function must be called once before any other method is called.
-    pub unsafe fn initialize(&'static self) {
-        self.inner.initialize();
-    }
-
     /// Spawn a task in this executor.
     ///
     /// # Safety
@@ -575,21 +510,3 @@ pub fn wake_task_no_pend(task: TaskRef) {
         }
     }
 }
-
-#[cfg(feature = "integrated-timers")]
-struct TimerQueue;
-
-#[cfg(feature = "integrated-timers")]
-impl embassy_time_queue_driver::TimerQueue for TimerQueue {
-    fn schedule_wake(&'static self, at: u64, waker: &core::task::Waker) {
-        let task = waker::task_from_waker(waker);
-        let task = task.header();
-        unsafe {
-            let expires_at = task.expires_at.get();
-            task.expires_at.set(expires_at.min(at));
-        }
-    }
-}
-
-#[cfg(feature = "integrated-timers")]
-embassy_time_queue_driver::timer_queue_impl!(static TIMER_QUEUE: TimerQueue = TimerQueue);
diff --git a/embassy-executor/src/raw/timer_queue.rs b/embassy-executor/src/raw/timer_queue.rs
index 94a5f340b..953bf014f 100644
--- a/embassy-executor/src/raw/timer_queue.rs
+++ b/embassy-executor/src/raw/timer_queue.rs
@@ -1,75 +1,100 @@
+//! Timer queue operations.
 use core::cmp::min;
 
+use super::util::SyncUnsafeCell;
 use super::TaskRef;
-use crate::raw::util::SyncUnsafeCell;
 
 pub(crate) struct TimerQueueItem {
     next: SyncUnsafeCell<Option<TaskRef>>,
+    expires_at: SyncUnsafeCell<u64>,
 }
 
 impl TimerQueueItem {
     pub const fn new() -> Self {
         Self {
             next: SyncUnsafeCell::new(None),
+            expires_at: SyncUnsafeCell::new(0),
         }
     }
 }
 
-pub(crate) struct TimerQueue {
+/// A timer queue, with items integrated into tasks.
+pub struct TimerQueue {
     head: SyncUnsafeCell<Option<TaskRef>>,
 }
 
 impl TimerQueue {
+    /// Creates a new timer queue.
     pub const fn new() -> Self {
         Self {
             head: SyncUnsafeCell::new(None),
         }
     }
 
-    pub(crate) unsafe fn update(&self, p: TaskRef) {
-        let task = p.header();
-        if task.expires_at.get() != u64::MAX {
+    /// Schedules a task to run at a specific time.
+    ///
+    /// If this function returns `true`, the called should find the next expiration time and set
+    /// a new alarm for that time.
+    pub fn schedule_wake(&mut self, at: u64, p: TaskRef) -> bool {
+        unsafe {
+            let task = p.header();
+            let item = &task.timer_queue_item;
             if task.state.timer_enqueue() {
-                task.timer_queue_item.next.set(self.head.get());
-                self.head.set(Some(p));
+                // If not in the queue, add it and update.
+                let prev = self.head.replace(Some(p));
+                item.next.set(prev);
+            } else if at <= item.expires_at.get() {
+                // If expiration is sooner than previously set, update.
+            } else {
+                // Task does not need to be updated.
+                return false;
             }
+
+            item.expires_at.set(at);
+            true
         }
     }
 
-    pub(crate) unsafe fn next_expiration(&self) -> u64 {
-        let mut res = u64::MAX;
-        self.retain(|p| {
-            let task = p.header();
-            let expires = task.expires_at.get();
-            res = min(res, expires);
-            expires != u64::MAX
-        });
-        res
-    }
+    /// Dequeues expired timers and returns the next alarm time.
+    ///
+    /// The provided callback will be called for each expired task. Tasks that never expire
+    /// will be removed, but the callback will not be called.
+    pub fn next_expiration(&mut self, now: u64) -> u64 {
+        let mut next_expiration = u64::MAX;
 
-    pub(crate) unsafe fn dequeue_expired(&self, now: u64, on_task: impl Fn(TaskRef)) {
         self.retain(|p| {
             let task = p.header();
-            if task.expires_at.get() <= now {
-                on_task(p);
+            let item = &task.timer_queue_item;
+            let expires = unsafe { item.expires_at.get() };
+
+            if expires <= now {
+                // Timer expired, process task.
+                super::wake_task(p);
                 false
             } else {
-                true
+                // Timer didn't yet expire, or never expires.
+                next_expiration = min(next_expiration, expires);
+                expires != u64::MAX
             }
         });
+
+        next_expiration
     }
 
-    pub(crate) unsafe fn retain(&self, mut f: impl FnMut(TaskRef) -> bool) {
-        let mut prev = &self.head;
-        while let Some(p) = prev.get() {
-            let task = p.header();
-            if f(p) {
-                // Skip to next
-                prev = &task.timer_queue_item.next;
-            } else {
-                // Remove it
-                prev.set(task.timer_queue_item.next.get());
-                task.state.timer_dequeue();
+    fn retain(&self, mut f: impl FnMut(TaskRef) -> bool) {
+        unsafe {
+            let mut prev = &self.head;
+            while let Some(p) = prev.get() {
+                let task = p.header();
+                let item = &task.timer_queue_item;
+                if f(p) {
+                    // Skip to next
+                    prev = &item.next;
+                } else {
+                    // Remove it
+                    prev.set(item.next.get());
+                    task.state.timer_dequeue();
+                }
             }
         }
     }
diff --git a/embassy-executor/src/raw/util.rs b/embassy-executor/src/raw/util.rs
index c46085e45..e2633658a 100644
--- a/embassy-executor/src/raw/util.rs
+++ b/embassy-executor/src/raw/util.rs
@@ -54,4 +54,9 @@ impl<T> SyncUnsafeCell<T> {
     {
         *self.value.get()
     }
+
+    #[cfg(feature = "integrated-timers")]
+    pub unsafe fn replace(&self, value: T) -> T {
+        core::mem::replace(&mut *self.value.get(), value)
+    }
 }