1 files changed, 57 insertions, 32 deletions

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();
+                }
             }
         }
     }