From 6ec9bcb1c4dfbe5fc5365d93e75c516bb03bf9fc Mon Sep 17 00:00:00 2001
From: Dario Nieuwenhuis <dirbaio@dirbaio.net>
Date: Thu, 11 Sep 2025 16:15:27 +0200
Subject: executor: add priority scheduler.

---
 embassy-executor/src/metadata.rs      | 26 ++++++++++++++++++++++++++
 embassy-executor/src/raw/mod.rs       |  5 -----
 embassy-executor/src/raw/run_queue.rs | 29 ++++++++++++++++++++++++-----
 3 files changed, 50 insertions(+), 10 deletions(-)

(limited to 'embassy-executor/src')

diff --git a/embassy-executor/src/metadata.rs b/embassy-executor/src/metadata.rs
index 4220048a6..bc0df0f83 100644
--- a/embassy-executor/src/metadata.rs
+++ b/embassy-executor/src/metadata.rs
@@ -1,6 +1,8 @@
 #[cfg(feature = "metadata-name")]
 use core::cell::Cell;
 use core::future::{poll_fn, Future};
+#[cfg(feature = "scheduler-priority")]
+use core::sync::atomic::{AtomicU8, Ordering};
 use core::task::Poll;
 
 #[cfg(feature = "metadata-name")]
@@ -14,6 +16,8 @@ use crate::raw::Deadline;
 pub struct Metadata {
     #[cfg(feature = "metadata-name")]
     name: Mutex<Cell<Option<&'static str>>>,
+    #[cfg(feature = "scheduler-priority")]
+    priority: AtomicU8,
     #[cfg(feature = "scheduler-deadline")]
     deadline: raw::Deadline,
 }
@@ -23,6 +27,8 @@ impl Metadata {
         Self {
             #[cfg(feature = "metadata-name")]
             name: Mutex::new(Cell::new(None)),
+            #[cfg(feature = "scheduler-priority")]
+            priority: AtomicU8::new(0),
             // NOTE: The deadline is set to zero to allow the initializer to reside in `.bss`. This
             // will be lazily initalized in `initialize_impl`
             #[cfg(feature = "scheduler-deadline")]
@@ -33,6 +39,14 @@ impl Metadata {
     pub(crate) fn reset(&self) {
         #[cfg(feature = "metadata-name")]
         critical_section::with(|cs| self.name.borrow(cs).set(None));
+
+        #[cfg(feature = "scheduler-priority")]
+        self.set_priority(0);
+
+        // By default, deadlines are set to the maximum value, so that any task WITH
+        // a set deadline will ALWAYS be scheduled BEFORE a task WITHOUT a set deadline
+        #[cfg(feature = "scheduler-deadline")]
+        self.unset_deadline();
     }
 
     /// Get the metadata for the current task.
@@ -61,6 +75,18 @@ impl Metadata {
         critical_section::with(|cs| self.name.borrow(cs).set(Some(name)))
     }
 
+    /// Get this task's priority.
+    #[cfg(feature = "scheduler-priority")]
+    pub fn priority(&self) -> u8 {
+        self.priority.load(Ordering::Relaxed)
+    }
+
+    /// Set this task's priority.
+    #[cfg(feature = "scheduler-priority")]
+    pub fn set_priority(&self, priority: u8) {
+        self.priority.store(priority, Ordering::Relaxed)
+    }
+
     /// Get this task's deadline.
     #[cfg(feature = "scheduler-deadline")]
     pub fn deadline(&self) -> u64 {
diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs
index 51a363385..9f36c60bc 100644
--- a/embassy-executor/src/raw/mod.rs
+++ b/embassy-executor/src/raw/mod.rs
@@ -300,11 +300,6 @@ impl<F: Future + 'static> AvailableTask<F> {
             self.task.raw.poll_fn.set(Some(TaskStorage::<F>::poll));
             self.task.future.write_in_place(future);
 
-            // By default, deadlines are set to the maximum value, so that any task WITH
-            // a set deadline will ALWAYS be scheduled BEFORE a task WITHOUT a set deadline
-            #[cfg(feature = "scheduler-deadline")]
-            self.task.raw.metadata.unset_deadline();
-
             let task = TaskRef::new(self.task);
 
             SpawnToken::new(task)
diff --git a/embassy-executor/src/raw/run_queue.rs b/embassy-executor/src/raw/run_queue.rs
index d98c26f73..b8b052310 100644
--- a/embassy-executor/src/raw/run_queue.rs
+++ b/embassy-executor/src/raw/run_queue.rs
@@ -2,7 +2,7 @@ use core::ptr::{addr_of_mut, NonNull};
 
 use cordyceps::sorted_list::Links;
 use cordyceps::Linked;
-#[cfg(feature = "scheduler-deadline")]
+#[cfg(any(feature = "scheduler-priority", feature = "scheduler-deadline"))]
 use cordyceps::SortedList;
 
 #[cfg(target_has_atomic = "ptr")]
@@ -83,7 +83,7 @@ impl RunQueue {
     /// Empty the queue, then call `on_task` for each task that was in the queue.
     /// NOTE: It is OK for `on_task` to enqueue more tasks. In this case they're left in the queue
     /// and will be processed by the *next* call to `dequeue_all`, *not* the current one.
-    #[cfg(not(feature = "scheduler-deadline"))]
+    #[cfg(not(any(feature = "scheduler-priority", feature = "scheduler-deadline")))]
     pub(crate) fn dequeue_all(&self, on_task: impl Fn(TaskRef)) {
         let taken = self.stack.take_all();
         for taskref in taken {
@@ -106,10 +106,29 @@ impl RunQueue {
     ///
     /// This process will repeat until the local `sorted` queue AND the global
     /// runqueue are both empty, at which point this function will return.
-    #[cfg(feature = "scheduler-deadline")]
+    #[cfg(any(feature = "scheduler-priority", feature = "scheduler-deadline"))]
     pub(crate) fn dequeue_all(&self, on_task: impl Fn(TaskRef)) {
-        let mut sorted =
-            SortedList::<TaskHeader>::new_with_cmp(|lhs, rhs| lhs.metadata.deadline().cmp(&rhs.metadata.deadline()));
+        let mut sorted = SortedList::<TaskHeader>::new_with_cmp(|lhs, rhs| {
+            // compare by priority first
+            #[cfg(feature = "scheduler-priority")]
+            {
+                let lp = lhs.metadata.priority();
+                let rp = rhs.metadata.priority();
+                if lp != rp {
+                    return lp.cmp(&rp).reverse();
+                }
+            }
+            // compare deadlines in case of tie.
+            #[cfg(feature = "scheduler-deadline")]
+            {
+                let ld = lhs.metadata.deadline();
+                let rd = rhs.metadata.deadline();
+                if ld != rd {
+                    return ld.cmp(&rd);
+                }
+            }
+            core::cmp::Ordering::Equal
+        });
 
         loop {
             // For each loop, grab any newly pended items
-- 
cgit