Merge pull request #3593 from bugadani/refactor

Rework time-driver contract.

3 files changed, 268 insertions, 43 deletions

diff --git a/embassy-time-queue-driver/src/lib.rs b/embassy-time-queue-driver/src/lib.rs
index 50736e8c7..97c81a124 100644
--- a/embassy-time-queue-driver/src/lib.rs
+++ b/embassy-time-queue-driver/src/lib.rs
@@ -2,38 +2,25 @@
 #![doc = include_str!("../README.md")]
 #![warn(missing_docs)]
 
-//! ## Implementing a timer queue
+//! This crate is an implementation detail of `embassy-time-driver`.
 //!
-//! - Define a struct `MyTimerQueue`
-//! - Implement [`TimerQueue`] for it
-//! - Register it as the global timer queue with [`timer_queue_impl`](crate::timer_queue_impl).
+//! As a HAL user, you should only depend on this crate if your application does not use
+//! `embassy-executor` and your HAL does not configure a generic queue by itself.
 //!
-//! ## Example
-//!
-//! ```
-//! use core::task::Waker;
-//!
-//! use embassy_time::Instant;
-//! use embassy_time::queue::{TimerQueue};
-//!
-//! struct MyTimerQueue{}; // not public!
-//!
-//! impl TimerQueue for MyTimerQueue {
-//!     fn schedule_wake(&'static self, at: u64, waker: &Waker) {
-//!         todo!()
-//!     }
-//! }
-//!
-//! embassy_time_queue_driver::timer_queue_impl!(static QUEUE: MyTimerQueue = MyTimerQueue{});
-//! ```
+//! As a HAL implementer, you need to depend on this crate if you want to implement a time driver,
+//! but how you should do so is documented in `embassy-time-driver`.
+
 use core::task::Waker;
 
-/// Timer queue
-pub trait TimerQueue {
-    /// Schedules a waker in the queue to be awoken at moment `at`.
-    /// If this moment is in the past, the waker might be awoken immediately.
-    fn schedule_wake(&'static self, at: u64, waker: &Waker);
-}
+#[cfg(feature = "_generic-queue")]
+pub mod queue_generic;
+#[cfg(not(feature = "_generic-queue"))]
+pub mod queue_integrated;
+
+#[cfg(feature = "_generic-queue")]
+pub use queue_generic::Queue;
+#[cfg(not(feature = "_generic-queue"))]
+pub use queue_integrated::Queue;
 
 extern "Rust" {
     fn _embassy_time_schedule_wake(at: u64, waker: &Waker);
@@ -41,20 +28,22 @@ extern "Rust" {
 
 /// Schedule the given waker to be woken at `at`.
 pub fn schedule_wake(at: u64, waker: &Waker) {
-    unsafe { _embassy_time_schedule_wake(at, waker) }
-}
-
-/// Set the TimerQueue implementation.
-///
-/// See the module documentation for an example.
-#[macro_export]
-macro_rules! timer_queue_impl {
-    (static $name:ident: $t: ty = $val:expr) => {
-        static $name: $t = $val;
-
-        #[no_mangle]
-        fn _embassy_time_schedule_wake(at: u64, waker: &core::task::Waker) {
-            <$t as $crate::TimerQueue>::schedule_wake(&$name, at, waker);
+    // This function is not implemented in embassy-time-driver because it needs access to executor
+    // internals. The function updates task state, then delegates to the implementation provided
+    // by the time driver.
+    #[cfg(not(feature = "_generic-queue"))]
+    {
+        use embassy_executor::raw::task_from_waker;
+        use embassy_executor::raw::timer_queue::TimerEnqueueOperation;
+        // The very first thing we must do, before we even access the timer queue, is to
+        // mark the task a TIMER_QUEUED. This ensures that the task that is being scheduled
+        // can not be respawn while we are accessing the timer queue.
+        let task = task_from_waker(waker);
+        if unsafe { task.timer_enqueue() } == TimerEnqueueOperation::Ignore {
+            // We are not allowed to enqueue the task in the timer queue. This is because the
+            // task is not spawned, and so it makes no sense to schedule it.
+            return;
         }
-    };
+    }
+    unsafe { _embassy_time_schedule_wake(at, waker) }
 }
diff --git a/embassy-time-queue-driver/src/queue_generic.rs b/embassy-time-queue-driver/src/queue_generic.rs
new file mode 100644
index 000000000..232035bc6
--- /dev/null
+++ b/embassy-time-queue-driver/src/queue_generic.rs
@@ -0,0 +1,146 @@
+//! Generic timer queue implementations.
+//!
+//! Time queue drivers may use this to simplify their implementation.
+
+use core::cmp::{min, Ordering};
+use core::task::Waker;
+
+use heapless::Vec;
+
+#[derive(Debug)]
+struct Timer {
+    at: u64,
+    waker: Waker,
+}
+
+impl PartialEq for Timer {
+    fn eq(&self, other: &Self) -> bool {
+        self.at == other.at
+    }
+}
+
+impl Eq for Timer {}
+
+impl PartialOrd for Timer {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        self.at.partial_cmp(&other.at)
+    }
+}
+
+impl Ord for Timer {
+    fn cmp(&self, other: &Self) -> Ordering {
+        self.at.cmp(&other.at)
+    }
+}
+
+/// A timer queue with a pre-determined capacity.
+pub struct ConstGenericQueue<const QUEUE_SIZE: usize> {
+    queue: Vec<Timer, QUEUE_SIZE>,
+}
+
+impl<const QUEUE_SIZE: usize> ConstGenericQueue<QUEUE_SIZE> {
+    /// Creates a new timer queue.
+    pub const fn new() -> Self {
+        Self { queue: Vec::new() }
+    }
+
+    /// Schedules a task to run at a specific time, and returns whether any changes were made.
+    ///
+    /// 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, waker: &Waker) -> bool {
+        self.queue
+            .iter_mut()
+            .find(|timer| timer.waker.will_wake(waker))
+            .map(|timer| {
+                if timer.at > at {
+                    timer.at = at;
+                    true
+                } else {
+                    false
+                }
+            })
+            .unwrap_or_else(|| {
+                let mut timer = Timer {
+                    waker: waker.clone(),
+                    at,
+                };
+
+                loop {
+                    match self.queue.push(timer) {
+                        Ok(()) => break,
+                        Err(e) => timer = e,
+                    }
+
+                    self.queue.pop().unwrap().waker.wake();
+                }
+
+                true
+            })
+    }
+
+    /// Dequeues expired timers and returns the next alarm time.
+    pub fn next_expiration(&mut self, now: u64) -> u64 {
+        let mut next_alarm = u64::MAX;
+
+        let mut i = 0;
+        while i < self.queue.len() {
+            let timer = &self.queue[i];
+            if timer.at <= now {
+                let timer = self.queue.swap_remove(i);
+                timer.waker.wake();
+            } else {
+                next_alarm = min(next_alarm, timer.at);
+                i += 1;
+            }
+        }
+
+        next_alarm
+    }
+}
+
+#[cfg(feature = "generic-queue-8")]
+const QUEUE_SIZE: usize = 8;
+#[cfg(feature = "generic-queue-16")]
+const QUEUE_SIZE: usize = 16;
+#[cfg(feature = "generic-queue-32")]
+const QUEUE_SIZE: usize = 32;
+#[cfg(feature = "generic-queue-64")]
+const QUEUE_SIZE: usize = 64;
+#[cfg(feature = "generic-queue-128")]
+const QUEUE_SIZE: usize = 128;
+#[cfg(not(any(
+    feature = "generic-queue-8",
+    feature = "generic-queue-16",
+    feature = "generic-queue-32",
+    feature = "generic-queue-64",
+    feature = "generic-queue-128"
+)))]
+const QUEUE_SIZE: usize = 64;
+
+/// A timer queue with a pre-determined capacity.
+pub struct Queue {
+    queue: ConstGenericQueue<QUEUE_SIZE>,
+}
+
+impl Queue {
+    /// Creates a new timer queue.
+    pub const fn new() -> Self {
+        Self {
+            queue: ConstGenericQueue::new(),
+        }
+    }
+
+    /// Schedules a task to run at a specific time, and returns whether any changes were made.
+    ///
+    /// 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, waker: &Waker) -> bool {
+        self.queue.schedule_wake(at, waker)
+    }
+
+    /// Dequeues expired timers and returns the next alarm time.
+    pub fn next_expiration(&mut self, now: u64) -> u64 {
+        self.queue.next_expiration(now)
+    }
+}
diff --git a/embassy-time-queue-driver/src/queue_integrated.rs b/embassy-time-queue-driver/src/queue_integrated.rs
new file mode 100644
index 000000000..6bb4c0c1a
--- /dev/null
+++ b/embassy-time-queue-driver/src/queue_integrated.rs
@@ -0,0 +1,90 @@
+//! Timer queue operations.
+use core::cell::Cell;
+use core::cmp::min;
+use core::task::Waker;
+
+use embassy_executor::raw::TaskRef;
+
+/// A timer queue, with items integrated into tasks.
+pub struct Queue {
+    head: Cell<Option<TaskRef>>,
+}
+
+impl Queue {
+    /// Creates a new timer queue.
+    pub const fn new() -> Self {
+        Self { head: Cell::new(None) }
+    }
+
+    /// 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, waker: &Waker) -> bool {
+        let task = embassy_executor::raw::task_from_waker(waker);
+        let item = task.timer_queue_item();
+        if item.next.get().is_none() {
+            // If not in the queue, add it and update.
+            let prev = self.head.replace(Some(task));
+            item.next.set(if prev.is_none() {
+                Some(unsafe { TaskRef::dangling() })
+            } else {
+                prev
+            });
+            item.expires_at.set(at);
+            true
+        } else if at <= item.expires_at.get() {
+            // If expiration is sooner than previously set, update.
+            item.expires_at.set(at);
+            true
+        } else {
+            // Task does not need to be updated.
+            false
+        }
+    }
+
+    /// 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;
+
+        self.retain(|p| {
+            let item = p.timer_queue_item();
+            let expires = item.expires_at.get();
+
+            if expires <= now {
+                // Timer expired, process task.
+                embassy_executor::raw::wake_task(p);
+                false
+            } else {
+                // Timer didn't yet expire, or never expires.
+                next_expiration = min(next_expiration, expires);
+                expires != u64::MAX
+            }
+        });
+
+        next_expiration
+    }
+
+    fn retain(&self, mut f: impl FnMut(TaskRef) -> bool) {
+        let mut prev = &self.head;
+        while let Some(p) = prev.get() {
+            if unsafe { p == TaskRef::dangling() } {
+                // prev was the last item, stop
+                break;
+            }
+            let item = p.timer_queue_item();
+            if f(p) {
+                // Skip to next
+                prev = &item.next;
+            } else {
+                // Remove it
+                prev.set(item.next.get());
+                item.next.set(None);
+                unsafe { p.timer_dequeue() };
+            }
+        }
+    }
+}