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        #[cfg(not(all(feature = "simulation", madsim)))]
182        tokio::time::sleep(duration).await;
183        #[cfg(all(feature = "simulation", madsim))]
184        madsim::time::sleep(duration).await;
185    }
186}
187
188#[cfg(test)]
189mod test {
190    use std::{sync::Arc, thread, time::Duration};
191
192    use rstest::rstest;
193
194    use super::*;
195
196    #[rstest]
197    fn fake_clock_parallel_advances() {
198        let clock = Arc::new(FakeRelativeClock::default());
199        let threads = std::iter::repeat_n((), 10)
200            .map(move |()| {
201                let clock = Arc::clone(&clock);
202
203                thread::spawn(move || {
204                    for _ in 0..1_000_000 {
205                        let now = clock.now();
206                        clock.advance(Duration::from_nanos(1));
207                        assert!(clock.now() > now);
208                    }
209                })
210            })
211            .collect::<Vec<_>>();
212
213        for t in threads {
214            t.join().unwrap();
215        }
216    }
217
218    #[rstest]
219    fn duration_addition_coverage() {
220        let d = Duration::from_secs(1);
221        let one_ns = Nanos::from(1);
222        assert!(d + one_ns > d);
223    }
224
225    // Under madsim, `MonotonicClock::sleep` runs on the virtual clock with
226    // sub-ms scheduling epsilon. If the cfg gate fell through to real tokio,
227    // `sleep` would block on the OS scheduler with ~5-15ms of jitter and the
228    // tight upper bound would fail.
229    #[cfg(all(feature = "simulation", madsim))]
230    #[madsim::test]
231    async fn test_monotonic_clock_sleep_uses_virtual_time() {
232        let clock = MonotonicClock;
233        let start = Instant::now();
234        clock.sleep(Duration::from_millis(100)).await;
235        let elapsed = start.elapsed();
236        assert!(elapsed >= Duration::from_millis(100));
237        assert!(
238            elapsed < Duration::from_millis(101),
239            "virtual sleep showed real-tokio jitter: {elapsed:?}"
240        );
241    }
242}
nautilus_network/ratelimiter/clock.rs

nautilus_network/ratelimiter/
clock.rs
