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//! The time sources contained in this module allow the rate limiter
19//! to be (optionally) independent of std, and additionally
20//! allow mocking the passage of time.
21//!
22//! You can supply a custom time source by implementing both [`Reference`]
23//! and [`Clock`] for your own types, and by implementing `Add<Nanos>` for
24//! your [`Reference`] type:
25use 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
38/// A measurement from a clock.
39pub trait Reference:
40    Sized + Add<Nanos, Output = Self> + PartialEq + Eq + Ord + Copy + Clone + Send + Sync + Debug
41{
42    /// Determines the time that separates two measurements of a
43    /// clock. Implementations of this must perform a saturating
44    /// subtraction - if the `earlier` timestamp should be later,
45    /// `duration_since` must return the zero duration.
46    fn duration_since(&self, earlier: Self) -> Nanos;
47
48    /// Returns a reference point that lies at most `duration` in the
49    /// past from the current reference. If an underflow should occur,
50    /// returns the current reference.
51    #[must_use]
52    fn saturating_sub(&self, duration: Nanos) -> Self;
53}
54
55/// A time source used by rate limiters.
56pub trait Clock: Clone {
57    /// A measurement of a monotonically increasing clock.
58    type Instant: Reference;
59
60    /// Returns a measurement of the clock.
61    fn now(&self) -> Self::Instant;
62}
63
64impl Reference for Duration {
65    /// The internal duration between this point and another.
66    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    /// The internal duration between this point and another.
73    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/// A mock implementation of a clock. All it does is keep track of
88/// what "now" is (relative to some point meaningful to the program),
89/// and returns that.
90///
91/// # Thread Safety
92///
93/// The mock time is represented as an atomic u64 count of nanoseconds, behind an [`Arc`].
94/// Clones of this clock will all show the same time, even if the original advances.
95#[derive(Debug, Clone, Default)]
96pub struct FakeRelativeClock {
97    now: Arc<AtomicU64>,
98}
99
100impl FakeRelativeClock {
101    /// Advances the fake clock by the given amount.
102    ///
103    /// # Panics
104    ///
105    /// Panics if `by` cannot be represented as a `u64` number of nanoseconds (i.e., exceeds 584 years).
106    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/// The monotonic clock implemented by [`Instant`].
139#[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}