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        while let Err(e) =
117            self.now
118                .compare_exchange_weak(prev, next, Ordering::Release, Ordering::Relaxed)
119        {
120            prev = e;
121            next = prev + by;
122        }
123    }
124}
125
126impl PartialEq for FakeRelativeClock {
127    fn eq(&self, other: &Self) -> bool {
128        self.now.load(Ordering::Relaxed) == other.now.load(Ordering::Relaxed)
129    }
130}
131
132impl Clock for FakeRelativeClock {
133    type Instant = Nanos;
134
135    fn now(&self) -> Self::Instant {
136        self.now.load(Ordering::Relaxed).into()
137    }
138
139    fn sleep(&self, duration: Duration) -> impl Future<Output = ()> + Send + '_ {
140        self.advance(duration);
141        std::future::ready(())
142    }
143}
144
145/// The monotonic clock implemented by [`Instant`].
146#[derive(Clone, Debug, Default)]
147pub struct MonotonicClock;
148
149impl Add<Nanos> for Instant {
150    type Output = Self;
151
152    fn add(self, other: Nanos) -> Self {
153        let other: Duration = other.into();
154        self + other
155    }
156}
157
158impl Reference for Instant {
159    fn duration_since(&self, earlier: Self) -> Nanos {
160        if earlier < *self {
161            (*self - earlier).into()
162        } else {
163            Nanos::from(Duration::new(0, 0))
164        }
165    }
166
167    fn saturating_sub(&self, duration: Nanos) -> Self {
168        self.checked_sub(duration.into()).unwrap_or(*self)
169    }
170}
171
172impl Clock for MonotonicClock {
173    type Instant = Instant;
174
175    fn now(&self) -> Self::Instant {
176        Instant::now()
177    }
178
179    async fn sleep(&self, duration: Duration) {
180        crate::dst::time::sleep(duration).await;
181    }
182}
183
184#[cfg(test)]
185mod tests {
186    use std::{sync::Arc, thread, time::Duration};
187
188    use rstest::rstest;
189
190    use super::*;
191
192    #[rstest]
193    fn fake_clock_parallel_advances() {
194        let clock = Arc::new(FakeRelativeClock::default());
195        let threads = std::iter::repeat_n((), 10)
196            .map(|()| {
197                let clock = Arc::clone(&clock);
198
199                thread::spawn(move || {
200                    for _ in 0..1_000_000 {
201                        let now = clock.now();
202                        clock.advance(Duration::from_nanos(1));
203                        assert!(clock.now() > now);
204                    }
205                })
206            })
207            .collect::<Vec<_>>();
208
209        for t in threads {
210            t.join().unwrap();
211        }
212
213        assert_eq!(clock.now(), Nanos::new(10_000_000));
214    }
215
216    #[rstest]
217    fn duration_addition_coverage() {
218        let d = Duration::from_secs(1);
219        let one_ns = Nanos::from(1);
220        assert_eq!(d + one_ns, Duration::new(1, 1));
221    }
222
223    #[rstest]
224    #[case(12, 5, 7)]
225    #[case(12, 12, 0)]
226    #[case(5, 12, 0)]
227    fn duration_since_saturates(#[case] now: u64, #[case] earlier: u64, #[case] expected: u64) {
228        assert_eq!(
229            Reference::duration_since(&Duration::from_nanos(now), Duration::from_nanos(earlier)),
230            Nanos::new(expected)
231        );
232    }
233
234    #[rstest]
235    #[case(12, 5, 7)]
236    #[case(12, 12, 0)]
237    #[case(5, 12, 5)]
238    fn duration_subtraction_preserves_reference_on_underflow(
239        #[case] now: u64,
240        #[case] subtract: u64,
241        #[case] expected: u64,
242    ) {
243        assert_eq!(
244            Reference::saturating_sub(&Duration::from_nanos(now), Nanos::new(subtract)),
245            Duration::from_nanos(expected)
246        );
247    }
248
249    #[rstest]
250    #[tokio::test]
251    async fn fake_sleep_advances_shared_clock() {
252        let clock = FakeRelativeClock::default();
253        let clone = clock.clone();
254        clock.advance(Duration::from_nanos(13));
255
256        clone.sleep(Duration::from_nanos(29)).await;
257
258        assert_eq!(clock.now(), Nanos::new(42));
259        assert_eq!(clone.now(), Nanos::new(42));
260        assert_eq!(clock, clone);
261        assert_ne!(clock, FakeRelativeClock::default());
262    }
263
264    #[rstest]
265    #[should_panic(expected = "Cannot represent durations greater than 584 years")]
266    fn fake_clock_rejects_unrepresentable_duration() {
267        FakeRelativeClock::default().advance(Duration::MAX);
268    }
269
270    #[rstest]
271    #[cfg_attr(not(all(feature = "simulation", madsim)), tokio::test)]
272    #[cfg_attr(all(feature = "simulation", madsim), madsim::test)]
273    async fn instant_reference_arithmetic_preserves_exact_offsets() {
274        let start = Instant::now();
275        let later = start + Nanos::new(37);
276
277        assert_eq!(later, start + Duration::from_nanos(37));
278        assert_eq!(Reference::duration_since(&later, start), Nanos::new(37));
279        assert_eq!(Reference::duration_since(&start, later), Nanos::new(0));
280        assert_eq!(Reference::duration_since(&start, start), Nanos::new(0));
281        assert_eq!(Reference::saturating_sub(&later, Nanos::new(37)), start);
282    }
283
284    #[rstest]
285    #[cfg(not(all(feature = "simulation", madsim)))]
286    #[tokio::test(start_paused = true)]
287    async fn monotonic_sleep_advances_clock_by_requested_duration() {
288        let clock = MonotonicClock;
289        let start = clock.now();
290
291        clock.sleep(Duration::from_millis(37)).await;
292
293        assert_eq!(clock.now() - start, Duration::from_millis(37));
294    }
295
296    // Under madsim, `MonotonicClock::sleep` runs on the virtual clock with
297    // sub-ms scheduling epsilon. If the cfg gate fell through to real tokio,
298    // `sleep` would block on the OS scheduler with ~5-15ms of jitter and the
299    // tight upper bound would fail.
300    #[cfg(all(feature = "simulation", madsim))]
301    #[madsim::test]
302    async fn test_monotonic_clock_sleep_uses_virtual_time() {
303        let clock = MonotonicClock;
304        let start = Instant::now();
305        clock.sleep(Duration::from_millis(100)).await;
306        let elapsed = start.elapsed();
307        assert!(elapsed >= Duration::from_millis(100));
308        assert!(
309            elapsed < Duration::from_millis(101),
310            "virtual sleep showed real-tokio jitter: {elapsed:?}"
311        );
312    }
313}
nautilus_network/ratelimiter/clock.rs

nautilus_network/ratelimiter/
clock.rs
