nautilus-network 0.65.0

// -------------------------------------------------------------------------------------------------
//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
//  https://nautechsystems.io
//
//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
//  You may not use this file except in compliance with the License.
//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
// -------------------------------------------------------------------------------------------------

//! Time sources for rate limiters.
//!
//! Custom time sources implement [`Reference`], [`Clock`], and `Add<Nanos>`. This supports
//! deterministic tests without coupling rate-limiting decisions to wall-clock time.

use std::{
    fmt::Debug,
    future::Future,
    ops::Add,
    sync::{
        Arc,
        atomic::{AtomicU64, Ordering},
    },
    time::Duration,
};

use super::nanos::Nanos;
use crate::dst::time::Instant;

/// A measurement from a clock.
pub trait Reference:
    Sized + Add<Nanos, Output = Self> + PartialEq + Eq + Ord + Copy + Clone + Send + Sync + Debug
{
    /// Determines the time that separates two measurements of a
    /// clock. Implementations of this must perform a saturating
    /// subtraction - if the `earlier` timestamp should be later,
    /// `duration_since` must return the zero duration.
    fn duration_since(&self, earlier: Self) -> Nanos;

    /// Returns a reference point that lies at most `duration` in the
    /// past from the current reference. If an underflow should occur,
    /// returns the current reference.
    #[must_use]
    fn saturating_sub(&self, duration: Nanos) -> Self;
}

/// A time source used by rate limiters.
pub trait Clock: Clone {
    /// A measurement of a monotonically increasing clock.
    type Instant: Reference;

    /// Returns a measurement of the clock.
    fn now(&self) -> Self::Instant;

    /// Waits for `duration` on this clock's time base.
    ///
    /// Implementations must advance on the same clock as [`Clock::now`] so
    /// callers using `sleep` together with `now` observe consistent time
    /// under both real and simulated runtimes.
    fn sleep(&self, duration: Duration) -> impl Future<Output = ()> + Send + '_;
}

impl Reference for Duration {
    /// The internal duration between this point and another.
    fn duration_since(&self, earlier: Self) -> Nanos {
        (*self).saturating_sub(earlier).into()
    }

    /// The internal duration between this point and another.
    fn saturating_sub(&self, duration: Nanos) -> Self {
        self.checked_sub(duration.into()).unwrap_or(*self)
    }
}

impl Add<Nanos> for Duration {
    type Output = Self;

    fn add(self, other: Nanos) -> Self {
        let other: Self = other.into();
        self + other
    }
}

/// A mock implementation of a clock. All it does is keep track of
/// what "now" is (relative to some point meaningful to the program),
/// and returns that.
///
/// # Thread Safety
///
/// The mock time is represented as an atomic u64 count of nanoseconds, behind an [`Arc`].
/// Clones of this clock will all show the same time, even if the original advances.
#[derive(Debug, Clone, Default)]
pub struct FakeRelativeClock {
    now: Arc<AtomicU64>,
}

impl FakeRelativeClock {
    /// Advances the fake clock by the given amount.
    ///
    /// # Panics
    ///
    /// Panics if `by` cannot be represented as a `u64` number of nanoseconds (i.e., exceeds 584 years).
    pub fn advance(&self, by: Duration) {
        let by: u64 = by
            .as_nanos()
            .try_into()
            .expect("Cannot represent durations greater than 584 years");

        let mut prev = self.now.load(Ordering::Acquire);
        let mut next = prev + by;

        while let Err(e) =
            self.now
                .compare_exchange_weak(prev, next, Ordering::Release, Ordering::Relaxed)
        {
            prev = e;
            next = prev + by;
        }
    }
}

impl PartialEq for FakeRelativeClock {
    fn eq(&self, other: &Self) -> bool {
        self.now.load(Ordering::Relaxed) == other.now.load(Ordering::Relaxed)
    }
}

impl Clock for FakeRelativeClock {
    type Instant = Nanos;

    fn now(&self) -> Self::Instant {
        self.now.load(Ordering::Relaxed).into()
    }

    fn sleep(&self, duration: Duration) -> impl Future<Output = ()> + Send + '_ {
        self.advance(duration);
        std::future::ready(())
    }
}

/// The monotonic clock implemented by [`Instant`].
#[derive(Clone, Debug, Default)]
pub struct MonotonicClock;

impl Add<Nanos> for Instant {
    type Output = Self;

    fn add(self, other: Nanos) -> Self {
        let other: Duration = other.into();
        self + other
    }
}

impl Reference for Instant {
    fn duration_since(&self, earlier: Self) -> Nanos {
        if earlier < *self {
            (*self - earlier).into()
        } else {
            Nanos::from(Duration::new(0, 0))
        }
    }

    fn saturating_sub(&self, duration: Nanos) -> Self {
        self.checked_sub(duration.into()).unwrap_or(*self)
    }
}

impl Clock for MonotonicClock {
    type Instant = Instant;

    fn now(&self) -> Self::Instant {
        Instant::now()
    }

    async fn sleep(&self, duration: Duration) {
        crate::dst::time::sleep(duration).await;
    }
}

#[cfg(test)]
mod tests {
    use std::{sync::Arc, thread, time::Duration};

    use rstest::rstest;

    use super::*;

    #[rstest]
    fn fake_clock_parallel_advances() {
        let clock = Arc::new(FakeRelativeClock::default());
        let threads = std::iter::repeat_n((), 10)
            .map(|()| {
                let clock = Arc::clone(&clock);

                thread::spawn(move || {
                    for _ in 0..1_000_000 {
                        let now = clock.now();
                        clock.advance(Duration::from_nanos(1));
                        assert!(clock.now() > now);
                    }
                })
            })
            .collect::<Vec<_>>();

        for t in threads {
            t.join().unwrap();
        }

        assert_eq!(clock.now(), Nanos::new(10_000_000));
    }

    #[rstest]
    fn duration_addition_coverage() {
        let d = Duration::from_secs(1);
        let one_ns = Nanos::from(1);
        assert_eq!(d + one_ns, Duration::new(1, 1));
    }

    #[rstest]
    #[case(12, 5, 7)]
    #[case(12, 12, 0)]
    #[case(5, 12, 0)]
    fn duration_since_saturates(#[case] now: u64, #[case] earlier: u64, #[case] expected: u64) {
        assert_eq!(
            Reference::duration_since(&Duration::from_nanos(now), Duration::from_nanos(earlier)),
            Nanos::new(expected)
        );
    }

    #[rstest]
    #[case(12, 5, 7)]
    #[case(12, 12, 0)]
    #[case(5, 12, 5)]
    fn duration_subtraction_preserves_reference_on_underflow(
        #[case] now: u64,
        #[case] subtract: u64,
        #[case] expected: u64,
    ) {
        assert_eq!(
            Reference::saturating_sub(&Duration::from_nanos(now), Nanos::new(subtract)),
            Duration::from_nanos(expected)
        );
    }

    #[rstest]
    #[tokio::test]
    async fn fake_sleep_advances_shared_clock() {
        let clock = FakeRelativeClock::default();
        let clone = clock.clone();
        clock.advance(Duration::from_nanos(13));

        clone.sleep(Duration::from_nanos(29)).await;

        assert_eq!(clock.now(), Nanos::new(42));
        assert_eq!(clone.now(), Nanos::new(42));
        assert_eq!(clock, clone);
        assert_ne!(clock, FakeRelativeClock::default());
    }

    #[rstest]
    #[should_panic(expected = "Cannot represent durations greater than 584 years")]
    fn fake_clock_rejects_unrepresentable_duration() {
        FakeRelativeClock::default().advance(Duration::MAX);
    }

    #[rstest]
    #[cfg_attr(not(all(feature = "simulation", madsim)), tokio::test)]
    #[cfg_attr(all(feature = "simulation", madsim), madsim::test)]
    async fn instant_reference_arithmetic_preserves_exact_offsets() {
        let start = Instant::now();
        let later = start + Nanos::new(37);

        assert_eq!(later, start + Duration::from_nanos(37));
        assert_eq!(Reference::duration_since(&later, start), Nanos::new(37));
        assert_eq!(Reference::duration_since(&start, later), Nanos::new(0));
        assert_eq!(Reference::duration_since(&start, start), Nanos::new(0));
        assert_eq!(Reference::saturating_sub(&later, Nanos::new(37)), start);
    }

    #[rstest]
    #[cfg(not(all(feature = "simulation", madsim)))]
    #[tokio::test(start_paused = true)]
    async fn monotonic_sleep_advances_clock_by_requested_duration() {
        let clock = MonotonicClock;
        let start = clock.now();

        clock.sleep(Duration::from_millis(37)).await;

        assert_eq!(clock.now() - start, Duration::from_millis(37));
    }

    // Under madsim, `MonotonicClock::sleep` runs on the virtual clock with
    // sub-ms scheduling epsilon. If the cfg gate fell through to real tokio,
    // `sleep` would block on the OS scheduler with ~5-15ms of jitter and the
    // tight upper bound would fail.
    #[cfg(all(feature = "simulation", madsim))]
    #[madsim::test]
    async fn test_monotonic_clock_sleep_uses_virtual_time() {
        let clock = MonotonicClock;
        let start = Instant::now();
        clock.sleep(Duration::from_millis(100)).await;
        let elapsed = start.elapsed();
        assert!(elapsed >= Duration::from_millis(100));
        assert!(
            elapsed < Duration::from_millis(101),
            "virtual sleep showed real-tokio jitter: {elapsed:?}"
        );
    }
}