refactor: time driver APIs (#20)

Signed-off-by: tison <[email protected]>

Author: tisonkun

fastimer-driver/src/heap.rs

@@ -0,0 +1,105 @@
+// Copyright 2024 FastLabs Developers
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::BinaryHeap;
+use std::ops::ControlFlow;
+use std::sync::atomic;
+use std::sync::atomic::AtomicBool;
+use std::sync::Arc;
+use std::time::Duration;
+use std::time::Instant;
+
+use crossbeam_queue::SegQueue;
+use parking::Parker;
+use parking::Unparker;
+
+use crate::TimeContext;
+use crate::TimeDriverShutdown;
+use crate::TimeEntry;
+
+/// Returns a new time driver, its time context and the shutdown handle.
+pub fn binary_heap_driver() -> (BinaryHeapTimeDriver, TimeContext, TimeDriverShutdown) {
+    let (parker, unparker) = parking::pair();
+    let timers = BinaryHeap::new();
+    let inbounds = Arc::new(SegQueue::new());
+    let shutdown = Arc::new(AtomicBool::new(false));
+
+    let driver = BinaryHeapTimeDriver {
+        parker,
+        unparker,
+        timers,
+        inbounds,
+        shutdown,
+    };
+
+    let context = TimeContext {
+        unparker: driver.unparker.clone(),
+        inbounds: driver.inbounds.clone(),
+    };
+
+    let shutdown = TimeDriverShutdown {
+        unparker: driver.unparker.clone(),
+        shutdown: driver.shutdown.clone(),
+    };
+
+    (driver, context, shutdown)
+}
+
+/// A heap-based time driver that drives registered timers.
+#[derive(Debug)]
+pub struct BinaryHeapTimeDriver {
+    parker: Parker,
+    unparker: Unparker,
+    timers: BinaryHeap<TimeEntry>,
+    inbounds: Arc<SegQueue<TimeEntry>>,
+    shutdown: Arc<AtomicBool>,
+}
+
+impl BinaryHeapTimeDriver {
+    /// Drives the timers and returns `true` if the driver has been shut down.
+    pub fn turn(&mut self) -> ControlFlow<()> {
+        if self.shutdown.load(atomic::Ordering::Acquire) {
+            return ControlFlow::Break(());
+        }
+
+        match self.timers.peek() {
+            None => self.parker.park(),
+            Some(entry) => {
+                let delta = entry.when.saturating_duration_since(Instant::now());
+                if delta > Duration::ZERO {
+                    self.parker.park_timeout(delta);
+                }
+            }
+        }
+
+        while let Some(entry) = self.inbounds.pop() {
+            self.timers.push(entry);
+        }
+
+        while let Some(entry) = self.timers.peek() {
+            if entry.when <= Instant::now() {
+                entry.waker.wake();
+                let _ = self.timers.pop();
+            } else {
+                break;
+            }
+        }
+
+        if self.shutdown.load(atomic::Ordering::Acquire) {
+            ControlFlow::Break(())
+        } else {
+            ControlFlow::Continue(())
+        }
+    }
+}

fastimer-driver/src/lib.rs

@@ -18,7 +18,6 @@
 //! Runtime-agnostic time driver for creating delay futures.
 
 use std::cmp;
-use std::collections::BinaryHeap;
 use std::future::Future;
 use std::pin::Pin;
 use std::sync::atomic;
@@ -33,11 +32,10 @@ use atomic_waker::AtomicWaker;
 use crossbeam_queue::SegQueue;
 use fastimer::make_instant_from_now;
 use fastimer::MakeDelay;
-use parking::Parker;
 use parking::Unparker;
 
-#[cfg(test)]
-mod tests;
+mod heap;
+pub use heap::*;
 
 #[derive(Debug)]
 struct TimeEntry {
@@ -55,7 +53,7 @@ impl Eq for TimeEntry {}
 
 impl PartialOrd for TimeEntry {
     fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
-        Some(self.when.cmp(&other.when))
+        Some(self.cmp(other))
     }
 }
 
@@ -96,34 +94,6 @@ impl Drop for Delay {
     }
 }
 
-/// Returns a new time driver, its time context and the shutdown handle.
-pub fn driver() -> (TimeDriver, TimeContext, TimeDriverShutdown) {
-    let (parker, unparker) = parking::pair();
-    let timers = BinaryHeap::new();
-    let inbounds = Arc::new(SegQueue::new());
-    let shutdown = Arc::new(AtomicBool::new(false));
-
-    let driver = TimeDriver {
-        parker,
-        unparker,
-        timers,
-        inbounds,
-        shutdown,
-    };
-
-    let context = TimeContext {
-        unparker: driver.unparker.clone(),
-        inbounds: driver.inbounds.clone(),
-    };
-
-    let shutdown = TimeDriverShutdown {
-        unparker: driver.unparker.clone(),
-        shutdown: driver.shutdown.clone(),
-    };
-
-    (driver, context, shutdown)
-}
-
 /// A time context for creating [`Delay`]s.
 #[derive(Debug, Clone)]
 pub struct TimeContext {
@@ -165,50 +135,6 @@ impl TimeDriverShutdown {
     }
 }
 
-/// A time driver that drives registered timers.
-#[derive(Debug)]
-pub struct TimeDriver {
-    parker: Parker,
-    unparker: Unparker,
-    timers: BinaryHeap<TimeEntry>,
-    inbounds: Arc<SegQueue<TimeEntry>>,
-    shutdown: Arc<AtomicBool>,
-}
-
-impl TimeDriver {
-    /// Drives the timers and returns `true` if the driver has been shut down.
-    pub fn turn(&mut self) -> bool {
-        if self.shutdown.load(atomic::Ordering::Acquire) {
-            return true;
-        }
-
-        match self.timers.peek() {
-            None => self.parker.park(),
-            Some(entry) => {
-                let delta = entry.when.saturating_duration_since(Instant::now());
-                if delta > Duration::ZERO {
-                    self.parker.park_timeout(delta);
-                }
-            }
-        }
-
-        while let Some(entry) = self.inbounds.pop() {
-            self.timers.push(entry);
-        }
-
-        while let Some(entry) = self.timers.peek() {
-            if entry.when <= Instant::now() {
-                entry.waker.wake();
-                let _ = self.timers.pop();
-            } else {
-                break;
-            }
-        }
-
-        self.shutdown.load(atomic::Ordering::Acquire)
-    }
-}
-
 /// A delay implementation that uses the given time context.
 #[derive(Debug, Clone)]
 pub struct MakeFastimerDelay(TimeContext);

fastimer-driver/src/tests.rs renamed to fastimer-driver/tests/integration.rs

@@ -16,8 +16,7 @@ use std::time::Duration;
 use std::time::Instant;
 
 use fastimer::make_instant_from_now;
-
-use crate::driver;
+use fastimer_driver::binary_heap_driver;
 
 fn assert_duration_eq(actual: Duration, expected: Duration) {
     if expected.abs_diff(expected) > Duration::from_millis(5) {
@@ -26,11 +25,11 @@ fn assert_duration_eq(actual: Duration, expected: Duration) {
 }
 
 #[test]
-fn test_simple_driver() {
-    let (mut driver, context, shutdown) = driver();
+fn test_binary_heap_driver() {
+    let (mut driver, context, shutdown) = binary_heap_driver();
     let (tx, rx) = std::sync::mpsc::channel();
     std::thread::spawn(move || loop {
-        if driver.turn() {
+        if driver.turn().is_break() {
             tx.send(()).unwrap();
             break;
         }