_diffctx/deadline.rs
1use std::sync::atomic::{AtomicU64, Ordering};
2use std::time::Instant;
3
4use once_cell::sync::Lazy;
5
6/// Wall-clock ceiling for the compute phases, shared across rayon workers.
7///
8/// The `timeout` parameter only ever reached git subprocesses
9/// (`set_git_timeout`); the native CLI is protected by a process-level
10/// watchdog that exits 124, but the library path — pyo3, and through it the
11/// MCP server — had no ceiling at all: a 420s timeout was observed to sit
12/// through an 8-minute edge build without firing. The edge phase checks this
13/// between builders, so the overshoot is bounded by the slowest single
14/// builder rather than unbounded.
15///
16/// Expiry panics with a recognizable message on purpose: the phase runs deep
17/// inside call chains that do not return `Result`, rayon propagates the
18/// unwind to the caller, and pyo3 surfaces it as a Python exception — an
19/// error after `timeout` seconds, where before there was a hang.
20static ANCHOR: Lazy<Instant> = Lazy::new(Instant::now);
21static DEADLINE_MS: AtomicU64 = AtomicU64::new(0);
22
23pub fn set_compute_deadline(timeout_secs: u64) {
24 let now_ms = ANCHOR.elapsed().as_millis() as u64;
25 // +1 keeps a zero timeout distinct from the 0 = "no deadline" sentinel.
26 DEADLINE_MS.store(now_ms + timeout_secs * 1000 + 1, Ordering::Relaxed);
27}
28
29pub fn check_compute_deadline(phase: &str) {
30 let deadline = DEADLINE_MS.load(Ordering::Relaxed);
31 if deadline == 0 {
32 return;
33 }
34 let now_ms = ANCHOR.elapsed().as_millis() as u64;
35 check_expired(now_ms, deadline, phase);
36}
37
38fn check_expired(now_ms: u64, deadline_ms: u64, phase: &str) {
39 if now_ms > deadline_ms {
40 panic!("diffctx compute deadline exceeded during {phase}");
41 }
42}
43
44#[cfg(test)]
45mod tests {
46 use super::*;
47
48 // The expiry check is tested against the pure comparison, not the
49 // process-global atomic: an expired global deadline is visible to every
50 // rayon worker in the test process, so mutating it here made any
51 // concurrently running graph test flakily panic mid-build.
52 #[test]
53 fn an_expired_deadline_panics_with_the_phase_name() {
54 let err = std::panic::catch_unwind(|| check_expired(6, 1, "edge construction"))
55 .expect_err("deadline did not fire");
56 let msg = err.downcast_ref::<String>().cloned().unwrap_or_default();
57 assert!(msg.contains("edge construction"), "message was: {msg}");
58 }
59
60 #[test]
61 fn an_unexpired_deadline_does_not_fire() {
62 check_expired(1, 6, "edge construction");
63 }
64}