Merge #959

959: Generic, executor-agnostic queue implementation r=ivmarkov a=ivmarkov

Hopefully relatively well documented.

Implementation relies on a fixed-size `SortedLinkedList` from `heapless`. (By default, for up to 128 timer schedules, but we can lower this number to - say - 64.)

As discussed earlier, on queue overflow, the `WakerRegistration` approach is utilized, whereas the waker that is ordered first in the queue is awoken to make room for the incoming one (which might be the waker that would be awoken after all!). Wakers are compared with `Waker::will_wake`, so the queue should actually not fill up that easily, if at all.

I've left provisions for the user to manually instantiate the queue using a dedicated macro - `generic_queue!` so that users willing to adjust the queue size, or users (like me) who have to use the queue in a complex "on-top-of-RTOS-but-the-timer-driver-calling-back-from-ISR" scenario can customize the mutex that protects the queue.

The one thing I'm not completely happy with is the need to call `{ embassy_time::queue::initialize() }` early on before any futures using embassy-time are polled, which is currently on the shoulders of the user. I'm open to any ideas where we can get rid of this and do it on the first call to `_embassy_time_schedule_wake`, without introducing very complex combinations of critical sections, atomics and whatnot.




Co-authored-by: ivmarkov <[email protected]>
Co-authored-by: Dario Nieuwenhuis <[email protected]>

1 files changed, 449 insertions, 0 deletions

diff --git a/embassy-time/src/queue_generic.rs b/embassy-time/src/queue_generic.rs
new file mode 100644
index 000000000..20ae7e6cc
--- /dev/null
+++ b/embassy-time/src/queue_generic.rs
@@ -0,0 +1,449 @@
+use core::cell::RefCell;
+use core::cmp::{min, Ordering};
+use core::task::Waker;
+
+use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
+use embassy_sync::blocking_mutex::Mutex;
+use heapless::Vec;
+
+use crate::driver::{allocate_alarm, set_alarm, set_alarm_callback, AlarmHandle};
+use crate::queue::TimerQueue;
+use crate::Instant;
+
+#[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 = 32;
+#[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;
+
+#[derive(Debug)]
+struct Timer {
+    at: Instant,
+    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)
+    }
+}
+
+struct InnerQueue {
+    queue: Vec<Timer, QUEUE_SIZE>,
+    alarm: AlarmHandle,
+}
+
+impl InnerQueue {
+    fn schedule_wake(&mut self, at: Instant, waker: &Waker) {
+        self.queue
+            .iter_mut()
+            .find(|timer| timer.waker.will_wake(waker))
+            .map(|mut timer| {
+                timer.at = min(timer.at, at);
+            })
+            .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();
+                }
+            });
+
+        // Don't wait for the alarm callback to trigger and directly
+        // dispatch all timers that are already due
+        //
+        // Then update the alarm if necessary
+        self.dispatch();
+    }
+
+    fn dispatch(&mut self) {
+        loop {
+            let now = Instant::now();
+
+            let mut next_alarm = Instant::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;
+                }
+            }
+
+            if self.update_alarm(next_alarm) {
+                break;
+            }
+        }
+    }
+
+    fn update_alarm(&mut self, next_alarm: Instant) -> bool {
+        if next_alarm == Instant::MAX {
+            true
+        } else {
+            set_alarm(self.alarm, next_alarm.as_ticks())
+        }
+    }
+
+    fn handle_alarm(&mut self) {
+        self.dispatch();
+    }
+}
+
+struct Queue {
+    inner: Mutex<CriticalSectionRawMutex, RefCell<Option<InnerQueue>>>,
+}
+
+impl Queue {
+    const fn new() -> Self {
+        Self {
+            inner: Mutex::new(RefCell::new(None)),
+        }
+    }
+
+    fn schedule_wake(&'static self, at: Instant, waker: &Waker) {
+        self.inner.lock(|inner| {
+            let mut inner = inner.borrow_mut();
+
+            if inner.is_none() {}
+
+            inner
+                .get_or_insert_with(|| {
+                    let handle = unsafe { allocate_alarm() }.unwrap();
+                    set_alarm_callback(handle, Self::handle_alarm_callback, self as *const _ as _);
+                    InnerQueue {
+                        queue: Vec::new(),
+                        alarm: handle,
+                    }
+                })
+                .schedule_wake(at, waker)
+        });
+    }
+
+    fn handle_alarm(&self) {
+        self.inner
+            .lock(|inner| inner.borrow_mut().as_mut().unwrap().handle_alarm());
+    }
+
+    fn handle_alarm_callback(ctx: *mut ()) {
+        unsafe { (ctx as *const Self).as_ref().unwrap() }.handle_alarm();
+    }
+}
+
+impl TimerQueue for Queue {
+    fn schedule_wake(&'static self, at: Instant, waker: &Waker) {
+        Queue::schedule_wake(self, at, waker);
+    }
+}
+
+crate::timer_queue_impl!(static QUEUE: Queue = Queue::new());
+
+#[cfg(test)]
+mod tests {
+    use core::cell::Cell;
+    use core::task::{RawWaker, RawWakerVTable, Waker};
+    use std::rc::Rc;
+    use std::sync::Mutex;
+
+    use serial_test::serial;
+
+    use super::InnerQueue;
+    use crate::driver::{AlarmHandle, Driver};
+    use crate::queue_generic::QUEUE;
+    use crate::Instant;
+
+    struct InnerTestDriver {
+        now: u64,
+        alarm: u64,
+        callback: fn(*mut ()),
+        ctx: *mut (),
+    }
+
+    impl InnerTestDriver {
+        const fn new() -> Self {
+            Self {
+                now: 0,
+                alarm: u64::MAX,
+                callback: Self::noop,
+                ctx: core::ptr::null_mut(),
+            }
+        }
+
+        fn noop(_ctx: *mut ()) {}
+    }
+
+    unsafe impl Send for InnerTestDriver {}
+
+    struct TestDriver(Mutex<InnerTestDriver>);
+
+    impl TestDriver {
+        const fn new() -> Self {
+            Self(Mutex::new(InnerTestDriver::new()))
+        }
+
+        fn reset(&self) {
+            *self.0.lock().unwrap() = InnerTestDriver::new();
+        }
+
+        fn set_now(&self, now: u64) {
+            let notify = {
+                let mut inner = self.0.lock().unwrap();
+
+                if inner.now < now {
+                    inner.now = now;
+
+                    if inner.alarm <= now {
+                        inner.alarm = u64::MAX;
+
+                        Some((inner.callback, inner.ctx))
+                    } else {
+                        None
+                    }
+                } else {
+                    panic!("Going back in time?");
+                }
+            };
+
+            if let Some((callback, ctx)) = notify {
+                (callback)(ctx);
+            }
+        }
+    }
+
+    impl Driver for TestDriver {
+        fn now(&self) -> u64 {
+            self.0.lock().unwrap().now
+        }
+
+        unsafe fn allocate_alarm(&self) -> Option<AlarmHandle> {
+            Some(AlarmHandle::new(0))
+        }
+
+        fn set_alarm_callback(&self, _alarm: AlarmHandle, callback: fn(*mut ()), ctx: *mut ()) {
+            let mut inner = self.0.lock().unwrap();
+
+            inner.callback = callback;
+            inner.ctx = ctx;
+        }
+
+        fn set_alarm(&self, _alarm: AlarmHandle, timestamp: u64) -> bool {
+            let mut inner = self.0.lock().unwrap();
+
+            if timestamp <= inner.now {
+                false
+            } else {
+                inner.alarm = timestamp;
+                true
+            }
+        }
+    }
+
+    struct TestWaker {
+        pub awoken: Rc<Cell<bool>>,
+        pub waker: Waker,
+    }
+
+    impl TestWaker {
+        fn new() -> Self {
+            let flag = Rc::new(Cell::new(false));
+
+            const VTABLE: RawWakerVTable = RawWakerVTable::new(
+                |data: *const ()| {
+                    unsafe {
+                        Rc::increment_strong_count(data as *const Cell<bool>);
+                    }
+
+                    RawWaker::new(data as _, &VTABLE)
+                },
+                |data: *const ()| unsafe {
+                    let data = data as *const Cell<bool>;
+                    data.as_ref().unwrap().set(true);
+                    Rc::decrement_strong_count(data);
+                },
+                |data: *const ()| unsafe {
+                    (data as *const Cell<bool>).as_ref().unwrap().set(true);
+                },
+                |data: *const ()| unsafe {
+                    Rc::decrement_strong_count(data);
+                },
+            );
+
+            let raw = RawWaker::new(Rc::into_raw(flag.clone()) as _, &VTABLE);
+
+            Self {
+                awoken: flag.clone(),
+                waker: unsafe { Waker::from_raw(raw) },
+            }
+        }
+    }
+
+    crate::time_driver_impl!(static DRIVER: TestDriver = TestDriver::new());
+
+    fn setup() {
+        DRIVER.reset();
+
+        QUEUE.inner.lock(|inner| {
+            *inner.borrow_mut() = InnerQueue::new();
+        });
+    }
+
+    fn queue_len() -> usize {
+        QUEUE.inner.lock(|inner| inner.borrow().queue.iter().count())
+    }
+
+    #[test]
+    #[serial]
+    fn test_schedule() {
+        setup();
+
+        assert_eq!(queue_len(), 0);
+
+        let waker = TestWaker::new();
+
+        QUEUE.schedule_wake(Instant::from_secs(1), &waker.waker);
+
+        assert!(!waker.awoken.get());
+        assert_eq!(queue_len(), 1);
+    }
+
+    #[test]
+    #[serial]
+    fn test_schedule_same() {
+        setup();
+
+        let waker = TestWaker::new();
+
+        QUEUE.schedule_wake(Instant::from_secs(1), &waker.waker);
+
+        assert_eq!(queue_len(), 1);
+
+        QUEUE.schedule_wake(Instant::from_secs(1), &waker.waker);
+
+        assert_eq!(queue_len(), 1);
+
+        QUEUE.schedule_wake(Instant::from_secs(100), &waker.waker);
+
+        assert_eq!(queue_len(), 1);
+
+        let waker2 = TestWaker::new();
+
+        QUEUE.schedule_wake(Instant::from_secs(100), &waker2.waker);
+
+        assert_eq!(queue_len(), 2);
+    }
+
+    #[test]
+    #[serial]
+    fn test_trigger() {
+        setup();
+
+        let waker = TestWaker::new();
+
+        QUEUE.schedule_wake(Instant::from_secs(100), &waker.waker);
+
+        assert!(!waker.awoken.get());
+
+        DRIVER.set_now(Instant::from_secs(99).as_ticks());
+
+        assert!(!waker.awoken.get());
+
+        assert_eq!(queue_len(), 1);
+
+        DRIVER.set_now(Instant::from_secs(100).as_ticks());
+
+        assert!(waker.awoken.get());
+
+        assert_eq!(queue_len(), 0);
+    }
+
+    #[test]
+    #[serial]
+    fn test_immediate_trigger() {
+        setup();
+
+        let waker = TestWaker::new();
+
+        QUEUE.schedule_wake(Instant::from_secs(100), &waker.waker);
+
+        DRIVER.set_now(Instant::from_secs(50).as_ticks());
+
+        let waker2 = TestWaker::new();
+
+        QUEUE.schedule_wake(Instant::from_secs(40), &waker2.waker);
+
+        assert!(!waker.awoken.get());
+        assert!(waker2.awoken.get());
+        assert_eq!(queue_len(), 1);
+    }
+
+    #[test]
+    #[serial]
+    fn test_queue_overflow() {
+        setup();
+
+        for i in 1..super::QUEUE_SIZE {
+            let waker = TestWaker::new();
+
+            QUEUE.schedule_wake(Instant::from_secs(310), &waker.waker);
+
+            assert_eq!(queue_len(), i);
+            assert!(!waker.awoken.get());
+        }
+
+        let first_waker = TestWaker::new();
+
+        QUEUE.schedule_wake(Instant::from_secs(300), &first_waker.waker);
+
+        assert_eq!(queue_len(), super::QUEUE_SIZE);
+        assert!(!first_waker.awoken.get());
+
+        let second_waker = TestWaker::new();
+
+        QUEUE.schedule_wake(Instant::from_secs(305), &second_waker.waker);
+
+        assert_eq!(queue_len(), super::QUEUE_SIZE);
+        assert!(first_waker.awoken.get());
+
+        QUEUE.schedule_wake(Instant::from_secs(320), &TestWaker::new().waker);
+        assert_eq!(queue_len(), super::QUEUE_SIZE);
+        assert!(second_waker.awoken.get());
+    }
+}