6 files changed, 532 insertions, 11 deletions

diff --git a/embassy-time/src/driver.rs b/embassy-time/src/driver.rs
index 79ae14b91..5c2ad3b23 100644
--- a/embassy-time/src/driver.rs
+++ b/embassy-time/src/driver.rs
@@ -105,20 +105,21 @@ pub trait Driver: Send + Sync + 'static {
     /// Sets an alarm at the given timestamp. When the current timestamp reaches the alarm
     /// timestamp, the provided callback function will be called.
     ///
-    /// If `timestamp` is already in the past, the alarm callback must be immediately fired.
-    /// In this case, it is allowed (but not mandatory) to call the alarm callback synchronously from `set_alarm`.
+    /// The `Driver` implementation should guarantee that the alarm callback is never called synchronously from `set_alarm`.
+    /// Rather - if `timestamp` is already in the past - `false` should be returned and alarm should not be set,
+    /// or alternatively, the driver should return `true` and arrange to call the alarm callback as soon as possible, but not synchronously.
     ///
     /// When callback is called, it is guaranteed that now() will return a value greater or equal than timestamp.
     ///
     /// Only one alarm can be active at a time for each AlarmHandle. This overwrites any previously-set alarm if any.
-    fn set_alarm(&self, alarm: AlarmHandle, timestamp: u64);
+    fn set_alarm(&self, alarm: AlarmHandle, timestamp: u64) -> bool;
 }
 
 extern "Rust" {
     fn _embassy_time_now() -> u64;
     fn _embassy_time_allocate_alarm() -> Option<AlarmHandle>;
     fn _embassy_time_set_alarm_callback(alarm: AlarmHandle, callback: fn(*mut ()), ctx: *mut ());
-    fn _embassy_time_set_alarm(alarm: AlarmHandle, timestamp: u64);
+    fn _embassy_time_set_alarm(alarm: AlarmHandle, timestamp: u64) -> bool;
 }
 
 /// See [`Driver::now`]
@@ -139,7 +140,7 @@ pub fn set_alarm_callback(alarm: AlarmHandle, callback: fn(*mut ()), ctx: *mut (
 }
 
 /// See [`Driver::set_alarm`]
-pub fn set_alarm(alarm: AlarmHandle, timestamp: u64) {
+pub fn set_alarm(alarm: AlarmHandle, timestamp: u64) -> bool {
     unsafe { _embassy_time_set_alarm(alarm, timestamp) }
 }
 
@@ -167,7 +168,7 @@ macro_rules! time_driver_impl {
         }
 
         #[no_mangle]
-        fn _embassy_time_set_alarm(alarm: $crate::driver::AlarmHandle, timestamp: u64) {
+        fn _embassy_time_set_alarm(alarm: $crate::driver::AlarmHandle, timestamp: u64) -> bool {
             <$t as $crate::driver::Driver>::set_alarm(&$name, alarm, timestamp)
         }
     };
diff --git a/embassy-time/src/driver_std.rs b/embassy-time/src/driver_std.rs
index 2ddb2e604..fc7fd1979 100644
--- a/embassy-time/src/driver_std.rs
+++ b/embassy-time/src/driver_std.rs
@@ -127,12 +127,14 @@ impl Driver for TimeDriver {
         alarm.ctx = ctx;
     }
 
-    fn set_alarm(&self, alarm: AlarmHandle, timestamp: u64) {
+    fn set_alarm(&self, alarm: AlarmHandle, timestamp: u64) -> bool {
         self.init();
         let mut alarms = unsafe { self.alarms.as_ref() }.lock().unwrap();
         let alarm = &mut alarms[alarm.id() as usize];
         alarm.timestamp = timestamp;
         unsafe { self.signaler.as_ref() }.signal();
+
+        true
     }
 }
 
diff --git a/embassy-time/src/driver_wasm.rs b/embassy-time/src/driver_wasm.rs
index e4497e6a2..63d049897 100644
--- a/embassy-time/src/driver_wasm.rs
+++ b/embassy-time/src/driver_wasm.rs
@@ -90,15 +90,23 @@ impl Driver for TimeDriver {
         }));
     }
 
-    fn set_alarm(&self, alarm: AlarmHandle, timestamp: u64) {
+    fn set_alarm(&self, alarm: AlarmHandle, timestamp: u64) -> bool {
         self.init();
         let mut alarms = unsafe { self.alarms.as_ref() }.lock().unwrap();
         let alarm = &mut alarms[alarm.id() as usize];
-        let timeout = (timestamp - self.now()) as u32;
         if let Some(token) = alarm.token {
             clearTimeout(token);
         }
-        alarm.token = Some(setTimeout(alarm.closure.as_ref().unwrap(), timeout / 1000));
+
+        let now = self.now();
+        if timestamp <= now {
+            false
+        } else {
+            let timeout = (timestamp - now) as u32;
+            alarm.token = Some(setTimeout(alarm.closure.as_ref().unwrap(), timeout / 1000));
+
+            true
+        }
     }
 }
 
diff --git a/embassy-time/src/lib.rs b/embassy-time/src/lib.rs
index 4edc883fe..586aa28de 100644
--- a/embassy-time/src/lib.rs
+++ b/embassy-time/src/lib.rs
@@ -1,4 +1,4 @@
-#![cfg_attr(not(any(feature = "std", feature = "wasm")), no_std)]
+#![cfg_attr(not(any(feature = "std", feature = "wasm", test)), no_std)]
 #![cfg_attr(feature = "nightly", feature(type_alias_impl_trait))]
 #![doc = include_str!("../README.md")]
 #![allow(clippy::new_without_default)]
@@ -11,6 +11,7 @@ mod delay;
 pub mod driver;
 mod duration;
 mod instant;
+pub mod queue;
 mod tick;
 mod timer;
 
@@ -18,6 +19,8 @@ mod timer;
 mod driver_std;
 #[cfg(feature = "wasm")]
 mod driver_wasm;
+#[cfg(feature = "generic-queue")]
+mod queue_generic;
 
 pub use delay::{block_for, Delay};
 pub use duration::Duration;
diff --git a/embassy-time/src/queue.rs b/embassy-time/src/queue.rs
new file mode 100644
index 000000000..c6f8b440a
--- /dev/null
+++ b/embassy-time/src/queue.rs
@@ -0,0 +1,58 @@
+//! Timer queue implementation
+//!
+//! This module defines the interface a timer queue needs to implement to power the `embassy_time` module.
+//!
+//! # Implementing a timer queue
+//!
+//! - Define a struct `MyTimerQueue`
+//! - Implement [`TimerQueue`] for it
+//! - Register it as the global timer queue with [`timer_queue_impl`](crate::timer_queue_impl).
+//!
+//! # Linkage details
+//!
+//! Check the documentation of the [`driver`](crate::driver) module for more information.
+//!
+//! Similarly to driver, if there is none or multiple timer queues in the crate tree, linking will fail.
+//!
+//! # Example
+//!
+//! ```
+//! use core::task::Waker;
+//!
+//! use embassy_time::Instant;
+//! use embassy_time::queue::{TimerQueue};
+//!
+//! struct MyTimerQueue{}; // not public!
+//! embassy_time::timer_queue_impl!(static QUEUE: MyTimerQueue = MyTimerQueue{});
+//!
+//! impl TimerQueue for MyTimerQueue {
+//!     fn schedule_wake(&'static self, at: Instant, waker: &Waker) {
+//!         todo!()
+//!     }
+//! }
+//! ```
+use core::task::Waker;
+
+use crate::Instant;
+
+/// 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: Instant, waker: &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: $crate::Instant, waker: &core::task::Waker) {
+            <$t as $crate::queue::TimerQueue>::schedule_wake(&$name, at, waker);
+        }
+    };
+}
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());
+    }
+}