nautilus_network/ratelimiter/
clock.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Time sources for rate limiters.
17//!
18//! Custom time sources implement [`Reference`], [`Clock`], and `Add<Nanos>`. This supports
19//! deterministic tests without coupling rate-limiting decisions to wall-clock time.
20
21use std::{
22    fmt::Debug,
23    future::Future,
24    ops::Add,
25    sync::{
26        Arc,
27        atomic::{AtomicU64, Ordering},
28    },
29    time::Duration,
30};
31
32use super::nanos::Nanos;
33use crate::dst::time::Instant;
34
35/// A measurement from a clock.
36pub trait Reference:
37    Sized + Add<Nanos, Output = Self> + PartialEq + Eq + Ord + Copy + Clone + Send + Sync + Debug
38{
39    /// Determines the time that separates two measurements of a
40    /// clock. Implementations of this must perform a saturating
41    /// subtraction - if the `earlier` timestamp should be later,
42    /// `duration_since` must return the zero duration.
43    fn duration_since(&self, earlier: Self) -> Nanos;
44
45    /// Returns a reference point that lies at most `duration` in the
46    /// past from the current reference. If an underflow should occur,
47    /// returns the current reference.
48    #[must_use]
49    fn saturating_sub(&self, duration: Nanos) -> Self;
50}
51
52/// A time source used by rate limiters.
53pub trait Clock: Clone {
54    /// A measurement of a monotonically increasing clock.
55    type Instant: Reference;
56
57    /// Returns a measurement of the clock.
58    fn now(&self) -> Self::Instant;
59
60    /// Waits for `duration` on this clock's time base.
61    ///
62    /// Implementations must advance on the same clock as [`Clock::now`] so
63    /// callers using `sleep` together with `now` observe consistent time
64    /// under both real and simulated runtimes.
65    fn sleep(&self, duration: Duration) -> impl Future<Output = ()> + Send + '_;
66}
67
68impl Reference for Duration {
69    /// The internal duration between this point and another.
70    fn duration_since(&self, earlier: Self) -> Nanos {
71        (*self).saturating_sub(earlier).into()
72    }
73
74    /// The internal duration between this point and another.
75    fn saturating_sub(&self, duration: Nanos) -> Self {
76        self.checked_sub(duration.into()).unwrap_or(*self)
77    }
78}
79
80impl Add<Nanos> for Duration {
81    type Output = Self;
82
83    fn add(self, other: Nanos) -> Self {
84        let other: Self = other.into();
85        self + other
86    }
87}
88
89/// A mock implementation of a clock. All it does is keep track of
90/// what "now" is (relative to some point meaningful to the program),
91/// and returns that.
92///
93/// # Thread Safety
94///
95/// The mock time is represented as an atomic u64 count of nanoseconds, behind an [`Arc`].
96/// Clones of this clock will all show the same time, even if the original advances.
97#[derive(Debug, Clone, Default)]
98pub struct FakeRelativeClock {
99    now: Arc<AtomicU64>,
100}
101
102impl FakeRelativeClock {
103    /// Advances the fake clock by the given amount.
104    ///
105    /// # Panics
106    ///
107    /// Panics if `by` cannot be represented as a `u64` number of nanoseconds (i.e., exceeds 584 years).
108    pub fn advance(&self, by: Duration) {
109        let by: u64 = by
110            .as_nanos()
111            .try_into()
112            .expect("Cannot represent durations greater than 584 years");
113
114        let mut prev = self.now.load(Ordering::Acquire);
115        let mut next = prev + by;
116
117        while let Err(e) =
118            self.now
119                .compare_exchange_weak(prev, next, Ordering::Release, Ordering::Relaxed)
120        {
121            prev = e;
122            next = prev + by;
123        }
124    }
125}
126
127impl PartialEq for FakeRelativeClock {
128    fn eq(&self, other: &Self) -> bool {
129        self.now.load(Ordering::Relaxed) == other.now.load(Ordering::Relaxed)
130    }
131}
132
133impl Clock for FakeRelativeClock {
134    type Instant = Nanos;
135
136    fn now(&self) -> Self::Instant {
137        self.now.load(Ordering::Relaxed).into()
138    }
139
140    fn sleep(&self, duration: Duration) -> impl Future<Output = ()> + Send + '_ {
141        self.advance(duration);
142        std::future::ready(())
143    }
144}
145
146/// The monotonic clock implemented by [`Instant`].
147#[derive(Clone, Debug, Default)]
148pub struct MonotonicClock;
149
150impl Add<Nanos> for Instant {
151    type Output = Self;
152
153    fn add(self, other: Nanos) -> Self {
154        let other: Duration = other.into();
155        self + other
156    }
157}
158
159impl Reference for Instant {
160    fn duration_since(&self, earlier: Self) -> Nanos {
161        if earlier < *self {
162            (*self - earlier).into()
163        } else {
164            Nanos::from(Duration::new(0, 0))
165        }
166    }
167
168    fn saturating_sub(&self, duration: Nanos) -> Self {
169        self.checked_sub(duration.into()).unwrap_or(*self)
170    }
171}
172
173impl Clock for MonotonicClock {
174    type Instant = Instant;
175
176    fn now(&self) -> Self::Instant {
177        Instant::now()
178    }
179
180    async fn sleep(&self, duration: Duration) {
181        crate::dst::time::sleep(duration).await;
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use std::{sync::Arc, thread, time::Duration};
188
189    use rstest::rstest;
190
191    use super::*;
192
193    #[rstest]
194    fn fake_clock_parallel_advances() {
195        let clock = Arc::new(FakeRelativeClock::default());
196        let threads = std::iter::repeat_n((), 10)
197            .map(|()| {
198                let clock = Arc::clone(&clock);
199
200                thread::spawn(move || {
201                    for _ in 0..1_000_000 {
202                        let now = clock.now();
203                        clock.advance(Duration::from_nanos(1));
204                        assert!(clock.now() > now);
205                    }
206                })
207            })
208            .collect::<Vec<_>>();
209
210        for t in threads {
211            t.join().unwrap();
212        }
213
214        assert_eq!(clock.now(), Nanos::new(10_000_000));
215    }
216
217    #[rstest]
218    fn duration_addition_coverage() {
219        let d = Duration::from_secs(1);
220        let one_ns = Nanos::from(1);
221        assert_eq!(d + one_ns, Duration::new(1, 1));
222    }
223
224    #[rstest]
225    #[case(12, 5, 7)]
226    #[case(12, 12, 0)]
227    #[case(5, 12, 0)]
228    fn duration_since_saturates(#[case] now: u64, #[case] earlier: u64, #[case] expected: u64) {
229        assert_eq!(
230            Reference::duration_since(&Duration::from_nanos(now), Duration::from_nanos(earlier)),
231            Nanos::new(expected)
232        );
233    }
234
235    #[rstest]
236    #[case(12, 5, 7)]
237    #[case(12, 12, 0)]
238    #[case(5, 12, 5)]
239    fn duration_subtraction_preserves_reference_on_underflow(
240        #[case] now: u64,
241        #[case] subtract: u64,
242        #[case] expected: u64,
243    ) {
244        assert_eq!(
245            Reference::saturating_sub(&Duration::from_nanos(now), Nanos::new(subtract)),
246            Duration::from_nanos(expected)
247        );
248    }
249
250    #[rstest]
251    #[tokio::test]
252    async fn fake_sleep_advances_shared_clock() {
253        let clock = FakeRelativeClock::default();
254        let clone = clock.clone();
255        clock.advance(Duration::from_nanos(13));
256
257        clone.sleep(Duration::from_nanos(29)).await;
258
259        assert_eq!(clock.now(), Nanos::new(42));
260        assert_eq!(clone.now(), Nanos::new(42));
261        assert_eq!(clock, clone);
262        assert_ne!(clock, FakeRelativeClock::default());
263    }
264
265    #[rstest]
266    #[should_panic(expected = "Cannot represent durations greater than 584 years")]
267    fn fake_clock_rejects_unrepresentable_duration() {
268        FakeRelativeClock::default().advance(Duration::MAX);
269    }
270
271    #[rstest]
272    #[cfg_attr(not(all(feature = "simulation", madsim)), tokio::test)]
273    #[cfg_attr(all(feature = "simulation", madsim), madsim::test)]
274    async fn instant_reference_arithmetic_preserves_exact_offsets() {
275        let start = Instant::now();
276        let later = start + Nanos::new(37);
277
278        assert_eq!(later, start + Duration::from_nanos(37));
279        assert_eq!(Reference::duration_since(&later, start), Nanos::new(37));
280        assert_eq!(Reference::duration_since(&start, later), Nanos::new(0));
281        assert_eq!(Reference::duration_since(&start, start), Nanos::new(0));
282        assert_eq!(Reference::saturating_sub(&later, Nanos::new(37)), start);
283    }
284
285    #[rstest]
286    #[cfg(not(all(feature = "simulation", madsim)))]
287    #[tokio::test(start_paused = true)]
288    async fn monotonic_sleep_advances_clock_by_requested_duration() {
289        let clock = MonotonicClock;
290        let start = clock.now();
291
292        clock.sleep(Duration::from_millis(37)).await;
293
294        assert_eq!(clock.now() - start, Duration::from_millis(37));
295    }
296
297    // Under madsim, `MonotonicClock::sleep` runs on the virtual clock with
298    // sub-ms scheduling epsilon. If the cfg gate fell through to real tokio,
299    // `sleep` would block on the OS scheduler with ~5-15ms of jitter and the
300    // tight upper bound would fail.
301    #[cfg(all(feature = "simulation", madsim))]
302    #[madsim::test]
303    async fn test_monotonic_clock_sleep_uses_virtual_time() {
304        let clock = MonotonicClock;
305        let start = Instant::now();
306        clock.sleep(Duration::from_millis(100)).await;
307        let elapsed = start.elapsed();
308        assert!(elapsed >= Duration::from_millis(100));
309        assert!(
310            elapsed < Duration::from_millis(101),
311            "virtual sleep showed real-tokio jitter: {elapsed:?}"
312        );
313    }
314}
nautilus_network/ratelimiter/clock.rs

nautilus_network/ratelimiter/
clock.rs
