Codex Handbook
shell-escalation/src/unix/stopwatch.rs 237 lines
use std::future::Future;use std::sync::Arc;use std::time::Duration;use std::time::Instant;use tokio::sync::Mutex;use tokio::sync::Notify;use tokio_util::sync::CancellationToken;#[derive(Clone, Debug)]pub struct Stopwatch {    limit: Option<Duration>,    inner: Arc<Mutex<StopwatchState>>,    notify: Arc<Notify>,}#[derive(Debug)]struct StopwatchState {    elapsed: Duration,    running_since: Option<Instant>,    active_pauses: u32,}impl Stopwatch {    pub fn new(limit: Duration) -> Self {        Self {            inner: Arc::new(Mutex::new(StopwatchState {                elapsed: Duration::ZERO,                running_since: Some(Instant::now()),                active_pauses: 0,            })),            notify: Arc::new(Notify::new()),            limit: Some(limit),        }    }    pub fn unlimited() -> Self {        Self {            inner: Arc::new(Mutex::new(StopwatchState {                elapsed: Duration::ZERO,                running_since: Some(Instant::now()),                active_pauses: 0,            })),            notify: Arc::new(Notify::new()),            limit: None,        }    }    pub fn cancellation_token(&self) -> CancellationToken {        let token = CancellationToken::new();        let Some(limit) = self.limit else {            return token;        };        let cancel = token.clone();        let inner = Arc::clone(&self.inner);        let notify = Arc::clone(&self.notify);        tokio::spawn(async move {            loop {                let (remaining, running) = {                    let guard = inner.lock().await;                    let elapsed = guard.elapsed                        + guard                            .running_since                            .map(|since| since.elapsed())                            .unwrap_or_default();                    if elapsed >= limit {                        break;                    }                    (limit - elapsed, guard.running_since.is_some())                };                if !running {                    notify.notified().await;                    continue;                }                let sleep = tokio::time::sleep(remaining);                tokio::pin!(sleep);                tokio::select! {                    _ = &mut sleep => {                        break;                    }                    _ = notify.notified() => {                        continue;                    }                }            }            cancel.cancel();        });        token    }    /// Runs `fut`, pausing the stopwatch while the future is pending. The clock    /// resumes automatically when the future completes. Nested/overlapping    /// calls are reference-counted so the stopwatch only resumes when every    /// pause is lifted.    pub async fn pause_for<F, T>(&self, fut: F) -> T    where        F: Future<Output = T>,    {        self.pause().await;        let result = fut.await;        self.resume().await;        result    }    async fn pause(&self) {        let mut guard = self.inner.lock().await;        guard.active_pauses += 1;        if guard.active_pauses == 1            && let Some(since) = guard.running_since.take()        {            guard.elapsed += since.elapsed();            self.notify.notify_waiters();        }    }    async fn resume(&self) {        let mut guard = self.inner.lock().await;        if guard.active_pauses == 0 {            return;        }        guard.active_pauses -= 1;        if guard.active_pauses == 0 && guard.running_since.is_none() {            guard.running_since = Some(Instant::now());            self.notify.notify_waiters();        }    }}#[cfg(test)]mod tests {    use super::Stopwatch;    use tokio::time::Duration;    use tokio::time::Instant;    use tokio::time::sleep;    use tokio::time::timeout;    #[tokio::test]    async fn cancellation_receiver_fires_after_limit() {        let stopwatch = Stopwatch::new(Duration::from_millis(50));        let token = stopwatch.cancellation_token();        let start = Instant::now();        token.cancelled().await;        assert!(start.elapsed() >= Duration::from_millis(50));    }    #[tokio::test]    async fn pause_prevents_timeout_until_resumed() {        let stopwatch = Stopwatch::new(Duration::from_millis(50));        let token = stopwatch.cancellation_token();        let pause_handle = tokio::spawn({            let stopwatch = stopwatch.clone();            async move {                stopwatch                    .pause_for(async {                        sleep(Duration::from_millis(100)).await;                    })                    .await;            }        });        assert!(            timeout(Duration::from_millis(30), token.cancelled())                .await                .is_err()        );        pause_handle.await.expect("pause task should finish");        token.cancelled().await;    }    #[tokio::test]    async fn overlapping_pauses_only_resume_once() {        let stopwatch = Stopwatch::new(Duration::from_millis(50));        let token = stopwatch.cancellation_token();        // First pause.        let pause1 = {            let stopwatch = stopwatch.clone();            tokio::spawn(async move {                stopwatch                    .pause_for(async {                        sleep(Duration::from_millis(80)).await;                    })                    .await;            })        };        // Overlapping pause that ends sooner.        let pause2 = {            let stopwatch = stopwatch.clone();            tokio::spawn(async move {                stopwatch                    .pause_for(async {                        sleep(Duration::from_millis(30)).await;                    })                    .await;            })        };        // While both pauses are active, the cancellation should not fire.        assert!(            timeout(Duration::from_millis(40), token.cancelled())                .await                .is_err()        );        pause2.await.expect("short pause should complete");        // Still paused because the long pause is active.        assert!(            timeout(Duration::from_millis(30), token.cancelled())                .await                .is_err()        );        pause1.await.expect("long pause should complete");        // Now the stopwatch should resume and hit the limit shortly after.        token.cancelled().await;    }    #[tokio::test]    async fn unlimited_stopwatch_never_cancels() {        let stopwatch = Stopwatch::unlimited();        let token = stopwatch.cancellation_token();        assert!(            timeout(Duration::from_millis(30), token.cancelled())                .await                .is_err()        );    }}