1use std::future::Future;
6use std::pin::Pin;
7use std::sync::{Arc, Mutex};
8use std::task::{Context, Poll, Waker};
9
10use rustev_core::seams::BoxFuture;
11
12pub type Instant = u64;
14
15pub trait Clock: Send + Sync + 'static {
16 fn now(&self) -> Instant;
17 fn sleep_until(&self, at: Instant) -> BoxFuture<'static, ()>;
19}
20
21pub struct TokioClock {
23 origin: tokio::time::Instant,
24}
25
26impl TokioClock {
27 pub fn new() -> Self {
28 TokioClock {
29 origin: tokio::time::Instant::now(),
30 }
31 }
32}
33
34impl Default for TokioClock {
35 fn default() -> Self {
36 Self::new()
37 }
38}
39
40impl Clock for TokioClock {
41 fn now(&self) -> Instant {
42 u64::try_from(self.origin.elapsed().as_millis()).unwrap_or(u64::MAX)
43 }
44
45 fn sleep_until(&self, at: Instant) -> BoxFuture<'static, ()> {
46 let deadline = self
47 .origin
48 .checked_add(std::time::Duration::from_millis(at))
49 .unwrap_or_else(|| self.origin + std::time::Duration::from_secs(86_400 * 365 * 30));
50 Box::pin(tokio::time::sleep_until(deadline))
51 }
52}
53
54#[derive(Clone, Default)]
56pub struct ManualClock(Arc<Mutex<ManualState>>);
57
58#[derive(Default)]
59struct ManualState {
60 now: Instant,
61 next_id: u64,
62 sleepers: Vec<(u64, Instant, Waker)>,
63}
64
65impl ManualClock {
66 pub fn new() -> Self {
67 Self::default()
68 }
69
70 pub fn advance(&self, ms: u64) {
72 let due: Vec<Waker> = {
73 let mut s = self.0.lock().unwrap_or_else(|e| e.into_inner());
74 s.now = s.now.saturating_add(ms);
75 let now = s.now;
76 let (due, keep): (Vec<_>, Vec<_>) = std::mem::take(&mut s.sleepers)
77 .into_iter()
78 .partition(|(_, at, _)| *at <= now);
79 s.sleepers = keep;
80 due.into_iter().map(|(_, _, w)| w).collect()
81 };
82 for w in due {
83 w.wake();
84 }
85 }
86
87 pub fn pending_delays(&self) -> Vec<u64> {
91 let s = self.0.lock().unwrap_or_else(|e| e.into_inner());
92 let mut v: Vec<u64> = s
93 .sleepers
94 .iter()
95 .map(|(_, at, _)| at.saturating_sub(s.now))
96 .collect();
97 v.sort_unstable();
98 v
99 }
100
101 pub fn sleepers(&self) -> usize {
103 self.0
104 .lock()
105 .unwrap_or_else(|e| e.into_inner())
106 .sleepers
107 .len()
108 }
109}
110
111impl Clock for ManualClock {
112 fn now(&self) -> Instant {
113 self.0.lock().unwrap_or_else(|e| e.into_inner()).now
114 }
115
116 fn sleep_until(&self, at: Instant) -> BoxFuture<'static, ()> {
117 let id = {
118 let mut s = self.0.lock().unwrap_or_else(|e| e.into_inner());
119 s.next_id += 1;
120 s.next_id
121 };
122 Box::pin(ManualSleep {
123 clock: self.clone(),
124 id,
125 at,
126 })
127 }
128}
129
130struct ManualSleep {
131 clock: ManualClock,
132 id: u64,
133 at: Instant,
134}
135
136impl Future for ManualSleep {
137 type Output = ();
138
139 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
140 let mut s = self.clock.0.lock().unwrap_or_else(|e| e.into_inner());
141 if s.now >= self.at {
142 return Poll::Ready(());
143 }
144 match s.sleepers.iter_mut().find(|(id, _, _)| *id == self.id) {
145 Some(entry) => entry.2 = cx.waker().clone(),
146 None => s.sleepers.push((self.id, self.at, cx.waker().clone())),
147 }
148 Poll::Pending
149 }
150}
151
152impl Drop for ManualSleep {
153 fn drop(&mut self) {
154 let mut s = self.clock.0.lock().unwrap_or_else(|e| e.into_inner());
155 s.sleepers.retain(|(id, _, _)| *id != self.id);
156 }
157}