Merge pull request #2420 from chrisprice/mock-driver

Extend mock time driver to support alarms

6 files changed, 186 insertions, 117 deletions

diff --git a/.github/ci/test.sh b/.github/ci/test.sh
index 369f6d221..b16b3c20d 100755
--- a/.github/ci/test.sh
+++ b/.github/ci/test.sh
@@ -11,7 +11,7 @@ export CARGO_TARGET_DIR=/ci/cache/target
 cargo test --manifest-path ./embassy-sync/Cargo.toml 
 cargo test --manifest-path ./embassy-embedded-hal/Cargo.toml 
 cargo test --manifest-path ./embassy-hal-internal/Cargo.toml 
-cargo test --manifest-path ./embassy-time/Cargo.toml --features generic-queue
+cargo test --manifest-path ./embassy-time/Cargo.toml --features generic-queue,mock-driver
 
 cargo test --manifest-path ./embassy-boot/boot/Cargo.toml
 cargo test --manifest-path ./embassy-boot/boot/Cargo.toml --features ed25519-dalek
diff --git a/ci.sh b/ci.sh
index 84a724ede..dd028ddda 100755
--- a/ci.sh
+++ b/ci.sh
@@ -34,7 +34,7 @@ cargo batch  \
     --- build --release --manifest-path embassy-executor/Cargo.toml --target riscv32imac-unknown-none-elf --features arch-riscv32,executor-thread \
     --- build --release --manifest-path embassy-executor/Cargo.toml --target riscv32imac-unknown-none-elf --features arch-riscv32,executor-thread,integrated-timers \
     --- build --release --manifest-path embassy-sync/Cargo.toml --target thumbv6m-none-eabi --features defmt \
-    --- build --release --manifest-path embassy-time/Cargo.toml --target thumbv6m-none-eabi --features defmt,defmt-timestamp-uptime,tick-hz-32_768,generic-queue-8 \
+    --- build --release --manifest-path embassy-time/Cargo.toml --target thumbv6m-none-eabi --features defmt,defmt-timestamp-uptime,generic-queue-8,mock-driver \
     --- build --release --manifest-path embassy-net/Cargo.toml --target thumbv7em-none-eabi --features defmt,tcp,udp,dns,proto-ipv4,medium-ethernet \
     --- build --release --manifest-path embassy-net/Cargo.toml --target thumbv7em-none-eabi --features defmt,tcp,udp,dns,proto-ipv4,igmp,medium-ethernet \
     --- build --release --manifest-path embassy-net/Cargo.toml --target thumbv7em-none-eabi --features defmt,tcp,udp,dns,dhcpv4,medium-ethernet \
diff --git a/embassy-time/src/driver.rs b/embassy-time/src/driver.rs
index 2afa454fe..f966386b7 100644
--- a/embassy-time/src/driver.rs
+++ b/embassy-time/src/driver.rs
@@ -42,7 +42,6 @@
 //! use embassy_time::driver::{Driver, AlarmHandle};
 //!
 //! struct MyDriver{} // not public!
-//! embassy_time::time_driver_impl!(static DRIVER: MyDriver = MyDriver{});
 //!
 //! impl Driver for MyDriver {
 //!     fn now(&self) -> u64 {
@@ -59,6 +58,9 @@
 //!     }
 //! }
 //! ```
+//! ```ignore
+//! embassy_time::time_driver_impl!(static DRIVER: MyDriver = MyDriver{});
+//! ```
 
 /// Alarm handle, assigned by the driver.
 #[derive(Clone, Copy)]
diff --git a/embassy-time/src/driver_mock.rs b/embassy-time/src/driver_mock.rs
index c255615c7..7533d51e5 100644
--- a/embassy-time/src/driver_mock.rs
+++ b/embassy-time/src/driver_mock.rs
@@ -1,4 +1,4 @@
-use core::cell::Cell;
+use core::cell::RefCell;
 
 use critical_section::Mutex as CsMutex;
 
@@ -8,7 +8,7 @@ use crate::{Duration, Instant};
 /// A mock driver that can be manually advanced.
 /// This is useful for testing code that works with [`Instant`] and [`Duration`].
 ///
-/// This driver cannot currently be used to test runtime functionality, such as
+/// This driver can also be used to test runtime functionality, such as
 /// timers, delays, etc.
 ///
 /// # Example
@@ -26,43 +26,196 @@ use crate::{Duration, Instant};
 ///     assert_eq!(true, has_a_second_passed(reference));
 /// }
 /// ```
-pub struct MockDriver {
-    now: CsMutex<Cell<Instant>>,
-}
+pub struct MockDriver(CsMutex<RefCell<InnerMockDriver>>);
 
-crate::time_driver_impl!(static DRIVER: MockDriver = MockDriver {
-    now: CsMutex::new(Cell::new(Instant::from_ticks(0))),
-});
+crate::time_driver_impl!(static DRIVER: MockDriver = MockDriver::new());
 
 impl MockDriver {
+    /// Creates a new mock driver.
+    pub const fn new() -> Self {
+        Self(CsMutex::new(RefCell::new(InnerMockDriver::new())))
+    }
+
     /// Gets a reference to the global mock driver.
     pub fn get() -> &'static MockDriver {
         &DRIVER
     }
 
-    /// Advances the time by the specified [`Duration`].
-    pub fn advance(&self, duration: Duration) {
+    /// Resets the internal state of the mock driver
+    /// This will clear and deallocate all alarms, and reset the current time to 0.
+    pub fn reset(&self) {
         critical_section::with(|cs| {
-            let now = self.now.borrow(cs).get().as_ticks();
-            self.now.borrow(cs).set(Instant::from_ticks(now + duration.as_ticks()));
+            self.0.borrow(cs).replace(InnerMockDriver::new());
         });
     }
+
+    /// Advances the time by the specified [`Duration`].
+    /// Calling any alarm callbacks that are due.
+    pub fn advance(&self, duration: Duration) {
+        let notify = {
+            critical_section::with(|cs| {
+                let mut inner = self.0.borrow_ref_mut(cs);
+
+                inner.now += duration;
+
+                let now = inner.now.as_ticks();
+
+                inner
+                    .alarm
+                    .as_mut()
+                    .filter(|alarm| alarm.timestamp <= now)
+                    .map(|alarm| {
+                        alarm.timestamp = u64::MAX;
+
+                        (alarm.callback, alarm.ctx)
+                    })
+            })
+        };
+
+        if let Some((callback, ctx)) = notify {
+            (callback)(ctx);
+        }
+    }
 }
 
 impl Driver for MockDriver {
     fn now(&self) -> u64 {
-        critical_section::with(|cs| self.now.borrow(cs).get().as_ticks() as u64)
+        critical_section::with(|cs| self.0.borrow_ref(cs).now).as_ticks()
     }
 
     unsafe fn allocate_alarm(&self) -> Option<AlarmHandle> {
-        unimplemented!("MockDriver does not support runtime features that require an executor");
+        critical_section::with(|cs| {
+            let mut inner = self.0.borrow_ref_mut(cs);
+
+            if inner.alarm.is_some() {
+                None
+            } else {
+                inner.alarm.replace(AlarmState::new());
+
+                Some(AlarmHandle::new(0))
+            }
+        })
     }
 
-    fn set_alarm_callback(&self, _alarm: AlarmHandle, _callback: fn(*mut ()), _ctx: *mut ()) {
-        unimplemented!("MockDriver does not support runtime features that require an executor");
+    fn set_alarm_callback(&self, _alarm: AlarmHandle, callback: fn(*mut ()), ctx: *mut ()) {
+        critical_section::with(|cs| {
+            let mut inner = self.0.borrow_ref_mut(cs);
+
+            let Some(alarm) = inner.alarm.as_mut() else {
+                panic!("Alarm not allocated");
+            };
+
+            alarm.callback = callback;
+            alarm.ctx = ctx;
+        });
     }
 
-    fn set_alarm(&self, _alarm: AlarmHandle, _timestamp: u64) -> bool {
-        unimplemented!("MockDriver does not support runtime features that require an executor");
+    fn set_alarm(&self, _alarm: AlarmHandle, timestamp: u64) -> bool {
+        critical_section::with(|cs| {
+            let mut inner = self.0.borrow_ref_mut(cs);
+
+            if timestamp <= inner.now.as_ticks() {
+                false
+            } else {
+                let Some(alarm) = inner.alarm.as_mut() else {
+                    panic!("Alarm not allocated");
+                };
+
+                alarm.timestamp = timestamp;
+                true
+            }
+        })
+    }
+}
+
+struct InnerMockDriver {
+    now: Instant,
+    alarm: Option<AlarmState>,
+}
+
+impl InnerMockDriver {
+    const fn new() -> Self {
+        Self {
+            now: Instant::from_ticks(0),
+            alarm: None,
+        }
+    }
+}
+
+struct AlarmState {
+    timestamp: u64,
+    callback: fn(*mut ()),
+    ctx: *mut (),
+}
+
+impl AlarmState {
+    const fn new() -> Self {
+        Self {
+            timestamp: u64::MAX,
+            callback: Self::noop,
+            ctx: core::ptr::null_mut(),
+        }
+    }
+
+    fn noop(_ctx: *mut ()) {}
+}
+
+unsafe impl Send for AlarmState {}
+
+#[cfg(test)]
+mod tests {
+    use serial_test::serial;
+
+    use super::*;
+
+    fn setup() {
+        DRIVER.reset();
+    }
+
+    #[test]
+    #[serial]
+    fn test_advance() {
+        setup();
+
+        let driver = MockDriver::get();
+        let reference = driver.now();
+        driver.advance(Duration::from_secs(1));
+        assert_eq!(Duration::from_secs(1).as_ticks(), driver.now() - reference);
+    }
+
+    #[test]
+    #[serial]
+    fn test_set_alarm_not_in_future() {
+        setup();
+
+        let driver = MockDriver::get();
+        let alarm = unsafe { AlarmHandle::new(0) };
+        assert_eq!(false, driver.set_alarm(alarm, driver.now()));
+    }
+
+    #[test]
+    #[serial]
+    fn test_alarm() {
+        setup();
+
+        let driver = MockDriver::get();
+        let alarm = unsafe { driver.allocate_alarm() }.expect("No alarms available");
+        static mut CALLBACK_CALLED: bool = false;
+        let ctx = &mut () as *mut ();
+        driver.set_alarm_callback(alarm, |_| unsafe { CALLBACK_CALLED = true }, ctx);
+        driver.set_alarm(alarm, driver.now() + 1);
+        assert_eq!(false, unsafe { CALLBACK_CALLED });
+        driver.advance(Duration::from_secs(1));
+        assert_eq!(true, unsafe { CALLBACK_CALLED });
+    }
+
+    #[test]
+    #[serial]
+    fn test_allocate_alarm() {
+        setup();
+
+        let driver = MockDriver::get();
+        assert!(unsafe { driver.allocate_alarm() }.is_some());
+        assert!(unsafe { driver.allocate_alarm() }.is_none());
     }
 }
diff --git a/embassy-time/src/queue.rs b/embassy-time/src/queue.rs
index c6f8b440a..d65197c54 100644
--- a/embassy-time/src/queue.rs
+++ b/embassy-time/src/queue.rs
@@ -23,7 +23,6 @@
 //! 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) {
@@ -31,6 +30,9 @@
 //!     }
 //! }
 //! ```
+//! ```ignore
+//! embassy_time::timer_queue_impl!(static QUEUE: MyTimerQueue = MyTimerQueue{});
+//! ```
 use core::task::Waker;
 
 use crate::Instant;
diff --git a/embassy-time/src/queue_generic.rs b/embassy-time/src/queue_generic.rs
index 77947ab29..829368ffc 100644
--- a/embassy-time/src/queue_generic.rs
+++ b/embassy-time/src/queue_generic.rs
@@ -175,103 +175,17 @@ impl TimerQueue for Queue {
 crate::timer_queue_impl!(static QUEUE: Queue = Queue::new());
 
 #[cfg(test)]
+#[cfg(feature = "mock-driver")]
 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 crate::driver::{AlarmHandle, Driver};
+    use crate::driver_mock::MockDriver;
     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
-            }
-        }
-    }
+    use crate::{Duration, Instant};
 
     struct TestWaker {
         pub awoken: Rc<Cell<bool>>,
@@ -312,10 +226,8 @@ mod tests {
         }
     }
 
-    crate::time_driver_impl!(static DRIVER: TestDriver = TestDriver::new());
-
     fn setup() {
-        DRIVER.reset();
+        MockDriver::get().reset();
         critical_section::with(|cs| *QUEUE.inner.borrow_ref_mut(cs) = None);
     }
 
@@ -382,13 +294,13 @@ mod tests {
 
         assert!(!waker.awoken.get());
 
-        DRIVER.set_now(Instant::from_secs(99).as_ticks());
+        MockDriver::get().advance(Duration::from_secs(99));
 
         assert!(!waker.awoken.get());
 
         assert_eq!(queue_len(), 1);
 
-        DRIVER.set_now(Instant::from_secs(100).as_ticks());
+        MockDriver::get().advance(Duration::from_secs(1));
 
         assert!(waker.awoken.get());
 
@@ -404,7 +316,7 @@ mod tests {
 
         QUEUE.schedule_wake(Instant::from_secs(100), &waker.waker);
 
-        DRIVER.set_now(Instant::from_secs(50).as_ticks());
+        MockDriver::get().advance(Duration::from_secs(50));
 
         let waker2 = TestWaker::new();