pure_stage/
time.rs

1// Copyright 2025 PRAGMA
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
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
15use std::{
16    cell::RefCell,
17    fmt::{Display, Formatter},
18    sync::{
19        LazyLock,
20        atomic::{AtomicU64, Ordering},
21    },
22    time::Duration,
23};
24
25use parking_lot::Mutex;
26
27use crate::drop_guard::DropGuard;
28
29/// A simulation clock that is driven explicitly by the simulation.
30pub trait Clock {
31    /// Get the current time.
32    fn now(&self) -> Instant;
33
34    /// Advance the clock to the given time.
35    ///
36    /// This method is expected to panic when attempting to go backwards in time.
37    fn advance_to(&self, instant: Instant);
38}
39
40impl Clock for AtomicU64 {
41    fn now(&self) -> Instant {
42        *EPOCH + Duration::from_nanos(self.load(Ordering::Relaxed))
43    }
44
45    fn advance_to(&self, instant: Instant) {
46        let nanos = instant.saturating_since(*EPOCH).as_nanos();
47        assert!(nanos < u64::MAX as u128, "simulation is not supposed to run for more than 584 years");
48        let nanos = nanos as u64;
49        let old = self.swap(nanos, Ordering::Relaxed);
50        assert!(old <= nanos, "clock is not monotonic");
51    }
52}
53
54impl Clock for Mutex<Instant> {
55    fn now(&self) -> Instant {
56        *self.lock()
57    }
58
59    fn advance_to(&self, instant: Instant) {
60        *self.lock() = instant;
61    }
62}
63
64/// A point in time in the simulation.
65///
66/// Note that this is an opaque type that serialises and prints as a duration since the [`EPOCH`].
67#[derive(Clone, Copy, Eq, PartialOrd, Ord)]
68pub struct Instant(tokio::time::Instant);
69
70thread_local! {
71    static TOLERANCE: RefCell<Duration> = const { RefCell::new(Duration::from_nanos(0)) };
72}
73
74impl PartialEq for Instant {
75    fn eq(&self, other: &Self) -> bool {
76        let tolerance = TOLERANCE.with(|tolerance| *tolerance.borrow());
77        if tolerance.is_zero() {
78            self.0 == other.0
79        } else if self > other {
80            self.0 - other.0 <= tolerance
81        } else {
82            other.0 - self.0 <= tolerance
83        }
84    }
85}
86
87impl std::fmt::Debug for Instant {
88    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
89        f.debug_tuple("Instant").field(&self.saturating_since(*EPOCH)).finish()
90    }
91}
92
93impl Display for Instant {
94    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
95        let duration = self.saturating_since(*EPOCH);
96        write!(f, "{:.6?}", duration)
97    }
98}
99
100impl<'de> serde::Deserialize<'de> for Instant {
101    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
102    where
103        D: serde::Deserializer<'de>,
104    {
105        let duration = Duration::deserialize(deserializer)?;
106        Ok(*EPOCH + duration)
107    }
108}
109
110impl serde::Serialize for Instant {
111    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
112    where
113        S: serde::Serializer,
114    {
115        let duration = self.saturating_since(*EPOCH);
116        duration.serialize(serializer)
117    }
118}
119
120impl Instant {
121    pub fn with_tolerance_for_test(tolerance: Duration) -> DropGuard<Duration, fn(Duration)> {
122        fn restore(tolerance: Duration) {
123            TOLERANCE.with_borrow_mut(|t2| *t2 = tolerance)
124        }
125        TOLERANCE.with_borrow_mut(|t| DropGuard::new(std::mem::replace(t, tolerance), restore as fn(Duration)))
126    }
127
128    pub(crate) fn from_tokio(instant: tokio::time::Instant) -> Self {
129        Self(instant)
130    }
131
132    pub(crate) fn to_tokio(self) -> tokio::time::Instant {
133        self.0
134    }
135
136    pub(crate) fn now() -> Self {
137        Self(tokio::time::Instant::now())
138    }
139
140    pub fn pretty(self, now: Self) -> String {
141        if let Some(duration) = self.checked_since(now) {
142            format!("{:?} in the future", duration)
143        } else if let Some(duration) = now.checked_since(self) {
144            format!("{:?} ago", duration)
145        } else {
146            "(time bug)".to_string()
147        }
148    }
149
150    pub fn saturating_since(&self, other: Self) -> Duration {
151        self.0.duration_since(other.0)
152    }
153
154    pub fn checked_since(&self, other: Self) -> Option<Duration> {
155        self.0.checked_duration_since(other.0)
156    }
157
158    pub fn at_offset(offset: Duration) -> Self {
159        *EPOCH + offset
160    }
161}
162
163impl std::ops::Add<Duration> for Instant {
164    type Output = Instant;
165
166    #[expect(clippy::expect_used)]
167    fn add(self, duration: Duration) -> Self {
168        Instant(
169            self.0.checked_add(duration).expect("simulation is not supposed to run for more than 290 billion years"),
170        )
171    }
172}
173
174impl std::ops::Sub<Duration> for Instant {
175    type Output = Instant;
176
177    #[expect(clippy::expect_used)]
178    fn sub(self, duration: Duration) -> Self {
179        Instant(
180            self.0.checked_sub(duration).expect("simulation is not supposed to run for more than 290 billion years"),
181        )
182    }
183}
184
185/// The concrete value of the epoch is completely opaque and irrelevant, we only persist
186/// durations. The only guarantee needed is that the epoch stays constant for the duration of
187/// the simulation.
188pub static EPOCH: LazyLock<Instant> = LazyLock::new(Instant::now);
189
190#[test]
191fn instant() {
192    let now = Instant::now();
193    let later = now + Duration::from_secs(1);
194
195    assert_eq!(later.checked_since(now).unwrap(), Duration::from_secs(1));
196    assert_eq!(now.checked_since(later), None);
197
198    assert_eq!(later.saturating_since(now), Duration::from_secs(1));
199    assert_eq!(now.saturating_since(later), Duration::from_secs(0));
200
201    assert_eq!(now + Duration::from_secs(1), later);
202    assert_eq!(later - Duration::from_secs(1), now);
203}
pure_stage/time.rs

pure_stage/
time.rs
