Merge pull request #4608 from diondokter/upstream-drs-2

[embassy-executor]: Upstream "Earliest Deadline First" Scheduler (version 2)

6 files changed, 316 insertions, 164 deletions

diff --git a/embassy-executor/src/metadata.rs b/embassy-executor/src/metadata.rs
index f92c9b37c..4220048a6 100644
--- a/embassy-executor/src/metadata.rs
+++ b/embassy-executor/src/metadata.rs
@@ -7,11 +7,15 @@ use core::task::Poll;
 use critical_section::Mutex;
 
 use crate::raw;
+#[cfg(feature = "scheduler-deadline")]
+use crate::raw::Deadline;
 
 /// Metadata associated with a task.
 pub struct Metadata {
     #[cfg(feature = "metadata-name")]
     name: Mutex<Cell<Option<&'static str>>>,
+    #[cfg(feature = "scheduler-deadline")]
+    deadline: raw::Deadline,
 }
 
 impl Metadata {
@@ -19,6 +23,10 @@ impl Metadata {
         Self {
             #[cfg(feature = "metadata-name")]
             name: Mutex::new(Cell::new(None)),
+            // 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")]
+            deadline: raw::Deadline::new_unset(),
         }
     }
 
@@ -52,4 +60,63 @@ impl Metadata {
     pub fn set_name(&self, name: &'static str) {
         critical_section::with(|cs| self.name.borrow(cs).set(Some(name)))
     }
+
+    /// Get this task's deadline.
+    #[cfg(feature = "scheduler-deadline")]
+    pub fn deadline(&self) -> u64 {
+        self.deadline.instant_ticks()
+    }
+
+    /// Set this task's deadline.
+    ///
+    /// This method does NOT check whether the deadline has already passed.
+    #[cfg(feature = "scheduler-deadline")]
+    pub fn set_deadline(&self, instant_ticks: u64) {
+        self.deadline.set(instant_ticks);
+    }
+
+    /// Remove this task's deadline.
+    /// This brings it back to the defaul where it's not scheduled ahead of other tasks.
+    #[cfg(feature = "scheduler-deadline")]
+    pub fn unset_deadline(&self) {
+        self.deadline.set(Deadline::UNSET_TICKS);
+    }
+
+    /// Set this task's deadline `duration_ticks` in the future from when
+    /// this future is polled. This deadline is saturated to the max tick value.
+    ///
+    /// Analogous to `Timer::after`.
+    #[cfg(all(feature = "scheduler-deadline", feature = "embassy-time-driver"))]
+    pub fn set_deadline_after(&self, duration_ticks: u64) {
+        let now = embassy_time_driver::now();
+
+        // Since ticks is a u64, saturating add is PROBABLY overly cautious, leave
+        // it for now, we can probably make this wrapping_add for performance
+        // reasons later.
+        let deadline = now.saturating_add(duration_ticks);
+
+        self.set_deadline(deadline);
+    }
+
+    /// Set the this task's deadline `increment_ticks` from the previous deadline.
+    ///
+    /// This deadline is saturated to the max tick value.
+    ///
+    /// Note that by default (unless otherwise set), tasks start life with the deadline
+    /// not set, which means this method will have no effect.
+    ///
+    /// Analogous to one increment of `Ticker::every().next()`.
+    ///
+    /// Returns the deadline that was set.
+    #[cfg(feature = "scheduler-deadline")]
+    pub fn increment_deadline(&self, duration_ticks: u64) {
+        let last = self.deadline();
+
+        // Since ticks is a u64, saturating add is PROBABLY overly cautious, leave
+        // it for now, we can probably make this wrapping_add for performance
+        // reasons later.
+        let deadline = last.saturating_add(duration_ticks);
+
+        self.set_deadline(deadline);
+    }
 }
diff --git a/embassy-executor/src/raw/deadline.rs b/embassy-executor/src/raw/deadline.rs
new file mode 100644
index 000000000..cc89fadb0
--- /dev/null
+++ b/embassy-executor/src/raw/deadline.rs
@@ -0,0 +1,44 @@
+use core::sync::atomic::{AtomicU32, Ordering};
+
+/// A type for interacting with the deadline of the current task
+///
+/// Requires the `scheduler-deadline` feature.
+///
+/// Note: Interacting with the deadline should be done locally in a task.
+/// In theory you could try to set or read the deadline from another task,
+/// but that will result in weird (though not unsound) behavior.
+pub(crate) struct Deadline {
+    instant_ticks_hi: AtomicU32,
+    instant_ticks_lo: AtomicU32,
+}
+
+impl Deadline {
+    pub(crate) const fn new(instant_ticks: u64) -> Self {
+        Self {
+            instant_ticks_hi: AtomicU32::new((instant_ticks >> 32) as u32),
+            instant_ticks_lo: AtomicU32::new(instant_ticks as u32),
+        }
+    }
+
+    pub(crate) const fn new_unset() -> Self {
+        Self::new(Self::UNSET_TICKS)
+    }
+
+    pub(crate) fn set(&self, instant_ticks: u64) {
+        self.instant_ticks_hi
+            .store((instant_ticks >> 32) as u32, Ordering::Relaxed);
+        self.instant_ticks_lo.store(instant_ticks as u32, Ordering::Relaxed);
+    }
+
+    /// Deadline value in ticks, same time base and ticks as `embassy-time`
+    pub(crate) fn instant_ticks(&self) -> u64 {
+        let hi = self.instant_ticks_hi.load(Ordering::Relaxed) as u64;
+        let lo = self.instant_ticks_lo.load(Ordering::Relaxed) as u64;
+
+        (hi << 32) | lo
+    }
+
+    /// Sentinel value representing an "unset" deadline, which has lower priority
+    /// than any other set deadline value
+    pub(crate) const UNSET_TICKS: u64 = u64::MAX;
+}
diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs
index 4280c5750..51a363385 100644
--- a/embassy-executor/src/raw/mod.rs
+++ b/embassy-executor/src/raw/mod.rs
@@ -7,8 +7,6 @@
 //! Using this module requires respecting subtle safety contracts. If you can, prefer using the safe
 //! [executor wrappers](crate::Executor) and the [`embassy_executor::task`](embassy_executor_macros::task) macro, which are fully safe.
 
-#[cfg_attr(target_has_atomic = "ptr", path = "run_queue_atomics.rs")]
-#[cfg_attr(not(target_has_atomic = "ptr"), path = "run_queue_critical_section.rs")]
 mod run_queue;
 
 #[cfg_attr(all(cortex_m, target_has_atomic = "32"), path = "state_atomics_arm.rs")]
@@ -28,6 +26,9 @@ pub(crate) mod util;
 #[cfg_attr(feature = "turbowakers", path = "waker_turbo.rs")]
 mod waker;
 
+#[cfg(feature = "scheduler-deadline")]
+mod deadline;
+
 use core::future::Future;
 use core::marker::PhantomData;
 use core::mem;
@@ -38,6 +39,8 @@ use core::sync::atomic::AtomicPtr;
 use core::sync::atomic::Ordering;
 use core::task::{Context, Poll, Waker};
 
+#[cfg(feature = "scheduler-deadline")]
+pub(crate) use deadline::Deadline;
 use embassy_executor_timer_queue::TimerQueueItem;
 #[cfg(feature = "arch-avr")]
 use portable_atomic::AtomicPtr;
@@ -96,6 +99,7 @@ extern "Rust" fn __embassy_time_queue_item_from_waker(waker: &Waker) -> &'static
 pub(crate) struct TaskHeader {
     pub(crate) state: State,
     pub(crate) run_queue_item: RunQueueItem,
+
     pub(crate) executor: AtomicPtr<SyncExecutor>,
     poll_fn: SyncUnsafeCell<Option<unsafe fn(TaskRef)>>,
 
@@ -296,6 +300,11 @@ 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
new file mode 100644
index 000000000..d98c26f73
--- /dev/null
+++ b/embassy-executor/src/raw/run_queue.rs
@@ -0,0 +1,194 @@
+use core::ptr::{addr_of_mut, NonNull};
+
+use cordyceps::sorted_list::Links;
+use cordyceps::Linked;
+#[cfg(feature = "scheduler-deadline")]
+use cordyceps::SortedList;
+
+#[cfg(target_has_atomic = "ptr")]
+type TransferStack<T> = cordyceps::TransferStack<T>;
+
+#[cfg(not(target_has_atomic = "ptr"))]
+type TransferStack<T> = MutexTransferStack<T>;
+
+use super::{TaskHeader, TaskRef};
+
+/// Use `cordyceps::sorted_list::Links` as the singly linked list
+/// for RunQueueItems.
+pub(crate) type RunQueueItem = Links<TaskHeader>;
+
+/// Implements the `Linked` trait, allowing for singly linked list usage
+/// of any of cordyceps' `TransferStack` (used for the atomic runqueue),
+/// `SortedList` (used with the DRS scheduler), or `Stack`, which is
+/// popped atomically from the `TransferStack`.
+unsafe impl Linked<Links<TaskHeader>> for TaskHeader {
+    type Handle = TaskRef;
+
+    // Convert a TaskRef into a TaskHeader ptr
+    fn into_ptr(r: TaskRef) -> NonNull<TaskHeader> {
+        r.ptr
+    }
+
+    // Convert a TaskHeader into a TaskRef
+    unsafe fn from_ptr(ptr: NonNull<TaskHeader>) -> TaskRef {
+        TaskRef { ptr }
+    }
+
+    // Given a pointer to a TaskHeader, obtain a pointer to the Links structure,
+    // which can be used to traverse to other TaskHeader nodes in the linked list
+    unsafe fn links(ptr: NonNull<TaskHeader>) -> NonNull<Links<TaskHeader>> {
+        let ptr: *mut TaskHeader = ptr.as_ptr();
+        NonNull::new_unchecked(addr_of_mut!((*ptr).run_queue_item))
+    }
+}
+
+/// Atomic task queue using a very, very simple lock-free linked-list queue:
+///
+/// To enqueue a task, task.next is set to the old head, and head is atomically set to task.
+///
+/// Dequeuing is done in batches: the queue is emptied by atomically replacing head with
+/// null. Then the batch is iterated following the next pointers until null is reached.
+///
+/// Note that batches will be iterated in the reverse order as they were enqueued. This is OK
+/// for our purposes: it can't create fairness problems since the next batch won't run until the
+/// current batch is completely processed, so even if a task enqueues itself instantly (for example
+/// by waking its own waker) can't prevent other tasks from running.
+pub(crate) struct RunQueue {
+    stack: TransferStack<TaskHeader>,
+}
+
+impl RunQueue {
+    pub const fn new() -> Self {
+        Self {
+            stack: TransferStack::new(),
+        }
+    }
+
+    /// Enqueues an item. Returns true if the queue was empty.
+    ///
+    /// # Safety
+    ///
+    /// `item` must NOT be already enqueued in any queue.
+    #[inline(always)]
+    pub(crate) unsafe fn enqueue(&self, task: TaskRef, _tok: super::state::Token) -> bool {
+        self.stack.push_was_empty(
+            task,
+            #[cfg(not(target_has_atomic = "ptr"))]
+            _tok,
+        )
+    }
+
+    /// # Standard atomic 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"))]
+    pub(crate) fn dequeue_all(&self, on_task: impl Fn(TaskRef)) {
+        let taken = self.stack.take_all();
+        for taskref in taken {
+            run_dequeue(&taskref);
+            on_task(taskref);
+        }
+    }
+
+    /// # Earliest Deadline First Scheduler
+    ///
+    /// This algorithm will loop until all enqueued tasks are processed.
+    ///
+    /// Before polling a task, all currently enqueued tasks will be popped from the
+    /// runqueue, and will be added to the working `sorted` list, a linked-list that
+    /// sorts tasks by their deadline, with nearest deadline items in the front, and
+    /// furthest deadline items in the back.
+    ///
+    /// After popping and sorting all pending tasks, the SOONEST task will be popped
+    /// from the front of the queue, and polled by calling `on_task` on it.
+    ///
+    /// 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")]
+    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()));
+
+        loop {
+            // For each loop, grab any newly pended items
+            let taken = self.stack.take_all();
+
+            // Sort these into the list - this is potentially expensive! We do an
+            // insertion sort of new items, which iterates the linked list.
+            //
+            // Something on the order of `O(n * m)`, where `n` is the number
+            // of new tasks, and `m` is the number of already pending tasks.
+            sorted.extend(taken);
+
+            // Pop the task with the SOONEST deadline. If there are no tasks
+            // pending, then we are done.
+            let Some(taskref) = sorted.pop_front() else {
+                return;
+            };
+
+            // We got one task, mark it as dequeued, and process the task.
+            run_dequeue(&taskref);
+            on_task(taskref);
+        }
+    }
+}
+
+/// atomic state does not require a cs...
+#[cfg(target_has_atomic = "ptr")]
+#[inline(always)]
+fn run_dequeue(taskref: &TaskRef) {
+    taskref.header().state.run_dequeue();
+}
+
+/// ...while non-atomic state does
+#[cfg(not(target_has_atomic = "ptr"))]
+#[inline(always)]
+fn run_dequeue(taskref: &TaskRef) {
+    critical_section::with(|cs| {
+        taskref.header().state.run_dequeue(cs);
+    })
+}
+
+/// A wrapper type that acts like TransferStack by wrapping a normal Stack in a CS mutex
+#[cfg(not(target_has_atomic = "ptr"))]
+struct MutexTransferStack<T: Linked<cordyceps::stack::Links<T>>> {
+    inner: critical_section::Mutex<core::cell::UnsafeCell<cordyceps::Stack<T>>>,
+}
+
+#[cfg(not(target_has_atomic = "ptr"))]
+impl<T: Linked<cordyceps::stack::Links<T>>> MutexTransferStack<T> {
+    const fn new() -> Self {
+        Self {
+            inner: critical_section::Mutex::new(core::cell::UnsafeCell::new(cordyceps::Stack::new())),
+        }
+    }
+
+    /// Push an item to the transfer stack, returning whether the stack was previously empty
+    fn push_was_empty(&self, item: T::Handle, token: super::state::Token) -> bool {
+        // SAFETY: The critical-section mutex guarantees that there is no *concurrent* access
+        // for the lifetime of the token, but does NOT protect against re-entrant access.
+        // However, we never *return* the reference, nor do we recurse (or call another method
+        // like `take_all`) that could ever allow for re-entrant aliasing. Therefore, the
+        // presence of the critical section is sufficient to guarantee exclusive access to
+        // the `inner` field for the purposes of this function.
+        let inner = unsafe { &mut *self.inner.borrow(token).get() };
+        let is_empty = inner.is_empty();
+        inner.push(item);
+        is_empty
+    }
+
+    fn take_all(&self) -> cordyceps::Stack<T> {
+        critical_section::with(|cs| {
+            // SAFETY: The critical-section mutex guarantees that there is no *concurrent* access
+            // for the lifetime of the token, but does NOT protect against re-entrant access.
+            // However, we never *return* the reference, nor do we recurse (or call another method
+            // like `push_was_empty`) that could ever allow for re-entrant aliasing. Therefore, the
+            // presence of the critical section is sufficient to guarantee exclusive access to
+            // the `inner` field for the purposes of this function.
+            let inner = unsafe { &mut *self.inner.borrow(cs).get() };
+            inner.take_all()
+        })
+    }
+}
diff --git a/embassy-executor/src/raw/run_queue_atomics.rs b/embassy-executor/src/raw/run_queue_atomics.rs
deleted file mode 100644
index ce511d79a..000000000
--- a/embassy-executor/src/raw/run_queue_atomics.rs
+++ /dev/null
@@ -1,88 +0,0 @@
-use core::ptr;
-use core::ptr::NonNull;
-use core::sync::atomic::{AtomicPtr, Ordering};
-
-use super::{TaskHeader, TaskRef};
-use crate::raw::util::SyncUnsafeCell;
-
-pub(crate) struct RunQueueItem {
-    next: SyncUnsafeCell<Option<TaskRef>>,
-}
-
-impl RunQueueItem {
-    pub const fn new() -> Self {
-        Self {
-            next: SyncUnsafeCell::new(None),
-        }
-    }
-}
-
-/// Atomic task queue using a very, very simple lock-free linked-list queue:
-///
-/// To enqueue a task, task.next is set to the old head, and head is atomically set to task.
-///
-/// Dequeuing is done in batches: the queue is emptied by atomically replacing head with
-/// null. Then the batch is iterated following the next pointers until null is reached.
-///
-/// Note that batches will be iterated in the reverse order as they were enqueued. This is OK
-/// for our purposes: it can't create fairness problems since the next batch won't run until the
-/// current batch is completely processed, so even if a task enqueues itself instantly (for example
-/// by waking its own waker) can't prevent other tasks from running.
-pub(crate) struct RunQueue {
-    head: AtomicPtr<TaskHeader>,
-}
-
-impl RunQueue {
-    pub const fn new() -> Self {
-        Self {
-            head: AtomicPtr::new(ptr::null_mut()),
-        }
-    }
-
-    /// Enqueues an item. Returns true if the queue was empty.
-    ///
-    /// # Safety
-    ///
-    /// `item` must NOT be already enqueued in any queue.
-    #[inline(always)]
-    pub(crate) unsafe fn enqueue(&self, task: TaskRef, _: super::state::Token) -> bool {
-        let mut was_empty = false;
-
-        self.head
-            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |prev| {
-                was_empty = prev.is_null();
-                unsafe {
-                    // safety: the pointer is either null or valid
-                    let prev = NonNull::new(prev).map(|ptr| TaskRef::from_ptr(ptr.as_ptr()));
-                    // safety: there are no concurrent accesses to `next`
-                    task.header().run_queue_item.next.set(prev);
-                }
-                Some(task.as_ptr() as *mut _)
-            })
-            .ok();
-
-        was_empty
-    }
-
-    /// 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.
-    pub(crate) fn dequeue_all(&self, on_task: impl Fn(TaskRef)) {
-        // Atomically empty the queue.
-        let ptr = self.head.swap(ptr::null_mut(), Ordering::AcqRel);
-
-        // safety: the pointer is either null or valid
-        let mut next = unsafe { NonNull::new(ptr).map(|ptr| TaskRef::from_ptr(ptr.as_ptr())) };
-
-        // Iterate the linked list of tasks that were previously in the queue.
-        while let Some(task) = next {
-            // If the task re-enqueues itself, the `next` pointer will get overwritten.
-            // Therefore, first read the next pointer, and only then process the task.
-            // safety: there are no concurrent accesses to `next`
-            next = unsafe { task.header().run_queue_item.next.get() };
-
-            task.header().state.run_dequeue();
-            on_task(task);
-        }
-    }
-}
diff --git a/embassy-executor/src/raw/run_queue_critical_section.rs b/embassy-executor/src/raw/run_queue_critical_section.rs
deleted file mode 100644
index 86c4085ed..000000000
--- a/embassy-executor/src/raw/run_queue_critical_section.rs
+++ /dev/null
@@ -1,74 +0,0 @@
-use core::cell::Cell;
-
-use critical_section::{CriticalSection, Mutex};
-
-use super::TaskRef;
-
-pub(crate) struct RunQueueItem {
-    next: Mutex<Cell<Option<TaskRef>>>,
-}
-
-impl RunQueueItem {
-    pub const fn new() -> Self {
-        Self {
-            next: Mutex::new(Cell::new(None)),
-        }
-    }
-}
-
-/// Atomic task queue using a very, very simple lock-free linked-list queue:
-///
-/// To enqueue a task, task.next is set to the old head, and head is atomically set to task.
-///
-/// Dequeuing is done in batches: the queue is emptied by atomically replacing head with
-/// null. Then the batch is iterated following the next pointers until null is reached.
-///
-/// Note that batches will be iterated in the reverse order as they were enqueued. This is OK
-/// for our purposes: it can't create fairness problems since the next batch won't run until the
-/// current batch is completely processed, so even if a task enqueues itself instantly (for example
-/// by waking its own waker) can't prevent other tasks from running.
-pub(crate) struct RunQueue {
-    head: Mutex<Cell<Option<TaskRef>>>,
-}
-
-impl RunQueue {
-    pub const fn new() -> Self {
-        Self {
-            head: Mutex::new(Cell::new(None)),
-        }
-    }
-
-    /// Enqueues an item. Returns true if the queue was empty.
-    ///
-    /// # Safety
-    ///
-    /// `item` must NOT be already enqueued in any queue.
-    #[inline(always)]
-    pub(crate) unsafe fn enqueue(&self, task: TaskRef, cs: CriticalSection<'_>) -> bool {
-        let prev = self.head.borrow(cs).replace(Some(task));
-        task.header().run_queue_item.next.borrow(cs).set(prev);
-
-        prev.is_none()
-    }
-
-    /// 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.
-    pub(crate) fn dequeue_all(&self, on_task: impl Fn(TaskRef)) {
-        // Atomically empty the queue.
-        let mut next = critical_section::with(|cs| self.head.borrow(cs).take());
-
-        // Iterate the linked list of tasks that were previously in the queue.
-        while let Some(task) = next {
-            // If the task re-enqueues itself, the `next` pointer will get overwritten.
-            // Therefore, first read the next pointer, and only then process the task.
-
-            critical_section::with(|cs| {
-                next = task.header().run_queue_item.next.borrow(cs).get();
-                task.header().state.run_dequeue(cs);
-            });
-
-            on_task(task);
-        }
-    }
-}