Rtos: avoid scanning timer/run qeues (#4301)

* Avoid scanning queues

* Fix warning

Author: bugadani

esp-rtos/src/run_queue.rs

@@ -172,7 +172,6 @@ impl RunQueue {
                 containing_queue.as_mut().remove(ready_task);
             }
         }
-        self.ready_tasks[priority_n].remove(ready_task);
         self.ready_tasks[priority_n].push(ready_task);
 
         cfg_if::cfg_if! {

esp-rtos/src/scheduler.rs

@@ -80,6 +80,8 @@ impl CpuSchedulerState {
                 pinned_to: None,
 
                 wakeup_at: 0,
+                run_queued: false,
+                timer_queued: false,
 
                 alloc_list_item: TaskListItem::None,
                 ready_queue_item: TaskListItem::None,

esp-rtos/src/task/mod.rs

@@ -42,12 +42,24 @@ pub(crate) type TaskListItem = Option<TaskPtr>;
 
 /// An abstraction that allows the task to contain multiple different queue pointers.
 pub(crate) trait TaskListElement: Default {
+    /// Returns the pointer to the next element in the list.
     fn next(task: TaskPtr) -> Option<TaskPtr>;
+
+    /// Sets the pointer to the next element in the list.
     fn set_next(task: TaskPtr, next: Option<TaskPtr>);
+
+    /// Returns whether the task is in the list. If this function returns `None`, we don't know.
+    fn is_in_queue(_task: TaskPtr) -> Option<bool> {
+        // By default we don't store this information, so we return "Don't know".
+        None
+    }
+
+    /// Marks whether the task is in the list.
+    fn mark_in_queue(_task: TaskPtr, _in_queue: bool) {}
 }
 
 macro_rules! task_list_item {
-    ($struct:ident, $field:ident) => {
+    ($struct:ident, $field:ident $(, $in_queue_field:ident)?) => {
         #[derive(Default)]
         pub(crate) struct $struct;
         impl TaskListElement for $struct {
@@ -60,14 +72,27 @@ macro_rules! task_list_item {
                     task.as_mut().$field = next;
                 }
             }
+
+            $(
+                fn is_in_queue(task: TaskPtr) -> Option<bool> {
+                    Some(unsafe { task.as_ref().$in_queue_field })
+                }
+
+                fn mark_in_queue(mut task: TaskPtr, in_queue: bool) {
+                    unsafe {
+                        task.as_mut().$in_queue_field = in_queue;
+                    }
+                }
+            )?
         }
     };
 }
 
+task_list_item!(TaskReadyQueueElement, ready_queue_item, run_queued);
+task_list_item!(TaskTimerQueueElement, timer_queue_item, timer_queued);
+// These aren't perf critical, no need to waste memory on caching list status:
 task_list_item!(TaskAllocListElement, alloc_list_item);
-task_list_item!(TaskReadyQueueElement, ready_queue_item);
 task_list_item!(TaskDeleteListElement, delete_list_item);
-task_list_item!(TaskTimerQueueElement, timer_queue_item);
 
 /// Extension trait for common task operations. These should be inherent methods but we can't
 /// implement stuff for NonNull.
@@ -153,6 +178,11 @@ impl<E: TaskListElement> TaskList<E> {
     }
 
     pub fn push(&mut self, task: TaskPtr) {
+        if E::is_in_queue(task) == Some(true) {
+            return;
+        }
+        E::mark_in_queue(task, true);
+
         debug_assert!(E::next(task).is_none());
         E::set_next(task, self.head);
         self.head = Some(task);
@@ -164,12 +194,18 @@ impl<E: TaskListElement> TaskList<E> {
         if let Some(task) = popped {
             self.head = E::next(task);
             E::set_next(task, None);
+            E::mark_in_queue(task, false);
         }
 
         popped
     }
 
     pub fn remove(&mut self, task: TaskPtr) {
+        if E::is_in_queue(task) == Some(false) {
+            return;
+        }
+        E::mark_in_queue(task, false);
+
         // TODO: maybe this (and TaskQueue::remove) may prove too expensive.
         let mut list = core::mem::take(self);
         while let Some(popped) = list.pop() {
@@ -213,6 +249,11 @@ impl<E: TaskListElement> TaskQueue<E> {
     }
 
     pub fn push(&mut self, task: TaskPtr) {
+        if E::is_in_queue(task) == Some(true) {
+            return;
+        }
+        E::mark_in_queue(task, true);
+
         debug_assert!(E::next(task).is_none());
         if let Some(tail) = self.tail {
             E::set_next(tail, Some(task));
@@ -231,41 +272,39 @@ impl<E: TaskListElement> TaskQueue<E> {
             if self.head.is_none() {
                 self.tail = None;
             }
+            E::mark_in_queue(task, false);
         }
 
         popped
     }
 
     #[cfg(multi_core)]
     pub fn pop_if(&mut self, cond: impl Fn(&Task) -> bool) -> Option<TaskPtr> {
-        let mut head = self.head.take();
-        self.tail = None;
-
         let mut popped = None;
-        while let Some(task) = head {
-            head = E::next(task);
-            E::set_next(task, None);
-            if cond(unsafe { task.as_ref() }) {
+
+        let mut list = core::mem::take(self);
+        while let Some(task) = list.pop() {
+            if popped.is_none() && cond(unsafe { task.as_ref() }) {
+                E::mark_in_queue(task, false);
                 popped = Some(task);
-                break;
             } else {
                 self.push(task);
             }
         }
 
-        while let Some(task) = head {
-            head = E::next(task);
-            E::set_next(task, None);
-            self.push(task);
-        }
-
         popped
     }
 
     pub fn remove(&mut self, task: TaskPtr) {
+        if E::is_in_queue(task) == Some(false) {
+            return;
+        }
+
         let mut list = core::mem::take(self);
         while let Some(popped) = list.pop() {
-            if popped != task {
+            if popped == task {
+                E::mark_in_queue(task, false);
+            } else {
                 self.push(popped);
             }
         }
@@ -372,6 +411,11 @@ pub(crate) struct Task {
 
     pub wakeup_at: u64,
 
+    /// Whether the task is currently queued in the run queue.
+    pub run_queued: bool,
+    /// Whether the task is currently queued in the timer queue.
+    pub timer_queued: bool,
+
     /// The current wait queue this task is in.
     pub(crate) current_queue: Option<NonNull<WaitQueue>>,
 
@@ -463,6 +507,8 @@ impl Task {
             pinned_to,
 
             wakeup_at: 0,
+            timer_queued: false,
+            run_queued: false,
 
             alloc_list_item: TaskListItem::None,
             ready_queue_item: TaskListItem::None,
@@ -555,6 +601,8 @@ pub(super) fn allocate_main_task(
     scheduler.per_cpu[current_cpu].main_task.priority = Priority::ZERO;
     scheduler.per_cpu[current_cpu].main_task.state = TaskState::Ready;
     scheduler.per_cpu[current_cpu].main_task.stack = stack;
+    scheduler.per_cpu[current_cpu].main_task.run_queued = false;
+    scheduler.per_cpu[current_cpu].main_task.timer_queued = false;
     #[cfg(multi_core)]
     {
         scheduler.per_cpu[current_cpu].main_task.pinned_to = Some(cpu);

esp-rtos/src/timer/embassy.rs

@@ -1,18 +1,19 @@
 use core::task::Waker;
 
+use embassy_time_queue_utils::Queue;
 use esp_sync::NonReentrantMutex;
 
 use crate::SCHEDULER;
 
 pub(super) struct TimerQueueInner {
-    queue: embassy_time_queue_utils::Queue,
+    queue: Queue,
     pub next_wakeup: u64,
 }
 
 impl TimerQueueInner {
     const fn new() -> Self {
         Self {
-            queue: embassy_time_queue_utils::Queue::new(),
+            queue: Queue::new(),
             next_wakeup: u64::MAX,
         }
     }
@@ -70,13 +71,13 @@ impl embassy_time_driver::Driver for TimerQueue {
 
     #[inline]
     fn schedule_wake(&self, at: u64, waker: &Waker) {
+        // Note that we don't put the thread to sleep here, as other embassy tasks may be
+        // ready. The thread will go to sleep when it can.
         if self.inner.with(|inner| inner.schedule_wake(at, waker)) {
             // Next wakeup time became shorter, re-arm the timer.
             // FIXME: avoid two separate critical sections.
-            // FIXME: this likely interferes with time slicing - we just keep pushing the time slice
-            // out on the other core, if active.
-            SCHEDULER.with(|s| {
-                unwrap!(s.time_driver.as_mut()).arm_next_wakeup(crate::now());
+            SCHEDULER.with(|scheduler| {
+                unwrap!(scheduler.time_driver.as_mut()).arm_next_wakeup(crate::now());
             });
         }
     }

esp-sync/src/lib.rs

@@ -142,11 +142,6 @@ mod multi_core {
             }
         }
 
-        #[inline]
-        fn is_owned_by(&self, thread: usize) -> bool {
-            self.owner.load(Ordering::Relaxed) == thread
-        }
-
         #[inline]
         pub fn lock(&self, lock: &impl crate::RawLock) -> crate::RestoreState {
             // We acquire the lock inside an interrupt-free context to prevent a subtle
@@ -199,7 +194,7 @@ mod multi_core {
         #[inline]
         pub unsafe fn unlock(&self) {
             #[cfg(debug_assertions)]
-            if !self.is_owned_by(thread_id()) {
+            if self.owner.load(Ordering::Relaxed) != thread_id() {
                 panic_attempt_unlock_not_owned();
             }
             self.owner.store(UNUSED_THREAD_ID_VALUE, Ordering::Release);