Do not scan for new expiration if nothing changes

embassy-rs · Nov 28, 2024 · 04c1419 · 04c1419
1 parent 91e1b3f
commit 04c1419
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 1 deletion.
diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs
@@ -170,7 +170,12 @@ impl<F: Future + 'static> TaskStorage<F> {
                 this.raw.state.despawn();
 
                 #[cfg(feature = "integrated-timers")]
-                this.raw.next_expiration.set(u64::MAX);
+                this.raw
+                    .executor
+                    .get()
+                    .unwrap_unchecked()
+                    .timer_queue
+                    .notify_task_exited(p);
             }
             Poll::Pending => {}
         }
@@ -407,6 +412,8 @@ impl SyncExecutor {
                     //   - 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.
+                    #[cfg(feature = "integrated-timers")]
+                    self.timer_queue.notify_task_exited(p);
                     return;
                 }
 

diff --git a/embassy-executor/src/raw/timer_queue.rs b/embassy-executor/src/raw/timer_queue.rs
@@ -19,20 +19,31 @@ impl TimerQueueItem {
 
 pub(crate) struct TimerQueue {
     head: Mutex<Cell<Option<TaskRef>>>,
+    rescan: Mutex<Cell<bool>>,
 }
 
 impl TimerQueue {
     pub const fn new() -> Self {
         Self {
             head: Mutex::new(Cell::new(None)),
+            rescan: Mutex::new(Cell::new(true)),
         }
     }
 
+    pub(crate) unsafe fn notify_task_exited(&self, p: TaskRef) {
+        let task = p.header();
+        task.next_expiration.set(u64::MAX);
+        critical_section::with(|cs| {
+            self.rescan.borrow(cs).set(true);
+        });
+    }
+
     pub(crate) unsafe fn update(&self, p: TaskRef, at: u64) {
         let task = p.header();
         if at < task.next_expiration.get() {
             task.next_expiration.set(at);
             critical_section::with(|cs| {
+                self.rescan.borrow(cs).set(true);
                 if task.state.timer_enqueue() {
                     let prev = self.head.borrow(cs).replace(Some(p));
                     task.timer_queue_item.next.borrow(cs).set(prev);
@@ -44,6 +55,10 @@ impl TimerQueue {
     pub(crate) unsafe fn next_expiration(&self) -> u64 {
         let mut res = u64::MAX;
         critical_section::with(|cs| {
+            let rescan = self.rescan.borrow(cs).replace(false);
+            if !rescan {
+                return;
+            }
             self.retain(cs, |p| {
                 let task = p.header();
                 let expires = task.next_expiration.get();
@@ -57,15 +72,20 @@ impl TimerQueue {
 
     pub(crate) unsafe fn dequeue_expired(&self, now: u64, on_task: impl Fn(TaskRef)) {
         critical_section::with(|cs| {
+            let mut changed = false;
             self.retain(cs, |p| {
                 let task = p.header();
                 if task.expires_at.borrow(cs).get() <= now {
                     on_task(p);
+                    changed = true;
                     false
                 } else {
                     true
                 }
             });
+            if changed {
+                self.rescan.borrow(cs).set(true);
+            }
         });
     }