nautilus_network/ratelimiter/
clock.rs1use std::{
26 fmt::Debug,
27 ops::Add,
28 prelude::v1::*,
29 sync::{
30 Arc,
31 atomic::{AtomicU64, Ordering},
32 },
33 time::{Duration, Instant},
34};
35
36use super::nanos::Nanos;
37
38pub trait Reference:
40 Sized + Add<Nanos, Output = Self> + PartialEq + Eq + Ord + Copy + Clone + Send + Sync + Debug
41{
42 fn duration_since(&self, earlier: Self) -> Nanos;
47
48 #[must_use]
52 fn saturating_sub(&self, duration: Nanos) -> Self;
53}
54
55pub trait Clock: Clone {
57 type Instant: Reference;
59
60 fn now(&self) -> Self::Instant;
62}
63
64impl Reference for Duration {
65 fn duration_since(&self, earlier: Self) -> Nanos {
67 self.checked_sub(earlier)
68 .unwrap_or_else(|| Self::new(0, 0))
69 .into()
70 }
71
72 fn saturating_sub(&self, duration: Nanos) -> Self {
74 self.checked_sub(duration.into()).unwrap_or(*self)
75 }
76}
77
78impl Add<Nanos> for Duration {
79 type Output = Self;
80
81 fn add(self, other: Nanos) -> Self {
82 let other: Self = other.into();
83 self + other
84 }
85}
86
87#[derive(Debug, Clone, Default)]
96pub struct FakeRelativeClock {
97 now: Arc<AtomicU64>,
98}
99
100impl FakeRelativeClock {
101 pub fn advance(&self, by: Duration) {
107 let by: u64 = by
108 .as_nanos()
109 .try_into()
110 .expect("Cannot represent durations greater than 584 years");
111
112 let mut prev = self.now.load(Ordering::Acquire);
113 let mut next = prev + by;
114 while let Err(e) =
115 self.now
116 .compare_exchange_weak(prev, next, Ordering::Release, Ordering::Relaxed)
117 {
118 prev = e;
119 next = prev + by;
120 }
121 }
122}
123
124impl PartialEq for FakeRelativeClock {
125 fn eq(&self, other: &Self) -> bool {
126 self.now.load(Ordering::Relaxed) == other.now.load(Ordering::Relaxed)
127 }
128}
129
130impl Clock for FakeRelativeClock {
131 type Instant = Nanos;
132
133 fn now(&self) -> Self::Instant {
134 self.now.load(Ordering::Relaxed).into()
135 }
136}
137
138#[derive(Clone, Debug, Default)]
140pub struct MonotonicClock;
141
142impl Add<Nanos> for Instant {
143 type Output = Self;
144
145 fn add(self, other: Nanos) -> Self {
146 let other: Duration = other.into();
147 self + other
148 }
149}
150
151impl Reference for Instant {
152 fn duration_since(&self, earlier: Self) -> Nanos {
153 if earlier < *self {
154 (*self - earlier).into()
155 } else {
156 Nanos::from(Duration::new(0, 0))
157 }
158 }
159
160 fn saturating_sub(&self, duration: Nanos) -> Self {
161 self.checked_sub(duration.into()).unwrap_or(*self)
162 }
163}
164
165impl Clock for MonotonicClock {
166 type Instant = Instant;
167
168 fn now(&self) -> Self::Instant {
169 Instant::now()
170 }
171}
172
173#[cfg(test)]
174mod test {
175 use std::{sync::Arc, thread, time::Duration};
176
177 use rstest::rstest;
178
179 use super::*;
180
181 #[rstest]
182 fn fake_clock_parallel_advances() {
183 let clock = Arc::new(FakeRelativeClock::default());
184 let threads = std::iter::repeat_n((), 10)
185 .map(move |()| {
186 let clock = Arc::clone(&clock);
187 thread::spawn(move || {
188 for _ in 0..1_000_000 {
189 let now = clock.now();
190 clock.advance(Duration::from_nanos(1));
191 assert!(clock.now() > now);
192 }
193 })
194 })
195 .collect::<Vec<_>>();
196 for t in threads {
197 t.join().unwrap();
198 }
199 }
200
201 #[rstest]
202 fn duration_addition_coverage() {
203 let d = Duration::from_secs(1);
204 let one_ns = Nanos::from(1);
205 assert!(d + one_ns > d);
206 }
207}