opencode_ctl/
clock.rs

1//! Time seam: `SystemClock` for production, `TestClock` for deterministic tests.
2
3use std::collections::BTreeMap;
4use std::future::Future;
5use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
6use std::sync::{Arc, Mutex};
7use std::time::Duration;
8
9pub trait Clock: Send + Sync + 'static {
10    fn now_ms(&self) -> i64;
11    fn sleep(&self, d: Duration) -> impl Future<Output = ()> + Send;
12}
13
14pub struct SystemClock;
15
16impl Clock for SystemClock {
17    fn now_ms(&self) -> i64 {
18        std::time::SystemTime::now()
19            .duration_since(std::time::UNIX_EPOCH)
20            .map(|d| d.as_millis() as i64)
21            .unwrap_or(0)
22    }
23
24    fn sleep(&self, d: Duration) -> impl Future<Output = ()> + Send {
25        tokio::time::sleep(d)
26    }
27}
28
29struct Inner {
30    now: AtomicI64,
31    notify: tokio::sync::Notify,
32    deadlines: Mutex<BTreeMap<u64, i64>>,
33    next_id: AtomicU64,
34}
35
36#[derive(Clone)]
37pub struct TestClock {
38    inner: Arc<Inner>,
39}
40
41struct DeadlineGuard {
42    inner: Arc<Inner>,
43    id: u64,
44}
45
46impl Drop for DeadlineGuard {
47    fn drop(&mut self) {
48        self.inner.deadlines.lock().unwrap().remove(&self.id);
49    }
50}
51
52impl TestClock {
53    pub fn new(start_ms: i64) -> Self {
54        Self {
55            inner: Arc::new(Inner {
56                now: AtomicI64::new(start_ms),
57                notify: tokio::sync::Notify::new(),
58                deadlines: Mutex::new(BTreeMap::new()),
59                next_id: AtomicU64::new(0),
60            }),
61        }
62    }
63
64    pub fn advance(&self, ms: i64) {
65        self.inner.now.fetch_add(ms, Ordering::SeqCst);
66        self.inner.notify.notify_waiters();
67    }
68
69    pub fn set(&self, ms: i64) {
70        self.inner.now.store(ms, Ordering::SeqCst);
71        self.inner.notify.notify_waiters();
72    }
73
74    pub fn pending_sleepers(&self) -> usize {
75        self.inner.deadlines.lock().unwrap().len()
76    }
77
78    pub fn next_deadline(&self) -> Option<i64> {
79        self.inner.deadlines.lock().unwrap().values().min().copied()
80    }
81
82    /// Drive virtual time: once every runnable task has parked on a sleep
83    /// (the sleeper set is stable for 32 consecutive yields), jump to the
84    /// earliest deadline. Never returns; use inside `tokio::select!` on a
85    /// current-thread runtime.
86    pub async fn run_until_idle(&self) {
87        const STABLE: u32 = 32;
88        const MAX_YIELDS: u64 = 100_000;
89        let mut total: u64 = 0;
90        loop {
91            let mut last = (self.pending_sleepers(), self.next_deadline());
92            let mut stable = 0u32;
93            loop {
94                tokio::task::yield_now().await;
95                total += 1;
96                let cur = (self.pending_sleepers(), self.next_deadline());
97                if cur == last {
98                    stable += 1;
99                } else {
100                    stable = 0;
101                    last = cur;
102                }
103                if stable >= STABLE && cur.0 >= 1 {
104                    break;
105                }
106                if total > MAX_YIELDS {
107                    panic!("TestClock::run_until_idle: no quiescence after {MAX_YIELDS} yields (sleepers, next_deadline) = {cur:?}");
108                }
109            }
110            if let Some(d) = self.next_deadline() {
111                self.set(d);
112            }
113        }
114    }
115}
116
117impl Clock for TestClock {
118    fn now_ms(&self) -> i64 {
119        self.inner.now.load(Ordering::SeqCst)
120    }
121
122    fn sleep(&self, d: Duration) -> impl Future<Output = ()> + Send {
123        // The deadline is fixed when `sleep` is CALLED (not on first poll), so a
124        // future handed to `spawn` and then outrun by `advance` still fires.
125        let inner = self.inner.clone();
126        let deadline = inner.now.load(Ordering::SeqCst) + d.as_millis() as i64;
127        let id = inner.next_id.fetch_add(1, Ordering::SeqCst);
128        inner.deadlines.lock().unwrap().insert(id, deadline);
129        let guard = DeadlineGuard {
130            inner: inner.clone(),
131            id,
132        };
133        async move {
134            let _guard = guard;
135            loop {
136                if inner.now.load(Ordering::SeqCst) >= deadline {
137                    return;
138                }
139                let n = inner.notify.notified();
140                tokio::pin!(n);
141                n.as_mut().enable();
142                if inner.now.load(Ordering::SeqCst) >= deadline {
143                    return;
144                }
145                n.await;
146            }
147        }
148    }
149}
opencode_ctl/clock.rs

opencode_ctl/
clock.rs
