_diffctx/
deadline.rs

1use std::cell::Cell;
2use std::time::{Duration, Instant};
3
4/// Per-run wall-clock ceiling for the compute phases.
5///
6/// The `timeout` parameter only ever reached git subprocesses
7/// (`set_git_timeout`); the native CLI is protected by a process-level
8/// watchdog that exits 124, but the library path — pyo3, and through it the
9/// MCP server — had no ceiling at all: a 420s timeout was observed to sit
10/// through an 8-minute edge build without firing.
11///
12/// A per-run value rather than a process global (#210): the MCP server runs
13/// overlapping pipelines on worker threads, and a shared atomic meant the
14/// last request to arrive overwrote every in-flight request's ceiling — and
15/// an expired ceiling was never cleared, so unrelated later runs in the same
16/// process (the yaml harness above all) inherited it.
17///
18/// Expiry panics with a recognizable message on purpose: the phase runs deep
19/// inside call chains that do not return `Result`, rayon propagates the
20/// unwind to the caller, and pyo3 surfaces it as a Python exception — an
21/// error after `timeout` seconds, where before there was a hang.
22#[derive(Clone, Copy, Debug)]
23pub struct Deadline {
24    expires_at: Option<Instant>,
25}
26
27impl Deadline {
28    /// Saturating: an absurd caller-supplied timeout must clamp to "no
29    /// ceiling", not wrap behind `now` and fire instantly.
30    pub fn from_timeout_secs(timeout_secs: u64) -> Self {
31        Deadline {
32            expires_at: Instant::now().checked_add(Duration::from_secs(timeout_secs)),
33        }
34    }
35
36    pub fn none() -> Self {
37        Deadline { expires_at: None }
38    }
39
40    pub fn check(&self, phase: &str) {
41        if let Some(expires_at) = self.expires_at {
42            check_expired(Instant::now(), expires_at, phase);
43        }
44    }
45
46    /// Publishes this deadline to the current thread for the guard's
47    /// lifetime, so hot loops deep inside edge builders can poll it via
48    /// `check_current_every` without threading a parameter through 51
49    /// `EdgeBuilder::build` implementations. Each rayon worker publishes its
50    /// own copy, so concurrent runs never see each other's ceiling.
51    pub fn enter(&self) -> ScopedDeadline {
52        let prev = CURRENT.with(|c| c.replace(self.expires_at));
53        ScopedDeadline { prev }
54    }
55}
56
57thread_local! {
58    static CURRENT: Cell<Option<Instant>> = const { Cell::new(None) };
59}
60
61pub struct ScopedDeadline {
62    prev: Option<Instant>,
63}
64
65impl Drop for ScopedDeadline {
66    fn drop(&mut self) {
67        let prev = self.prev;
68        CURRENT.with(|c| c.set(prev));
69    }
70}
71
72/// Cheap intra-loop poll: costs a branch on all but every `every`-th
73/// iteration. For the builder loops whose single invocation can outrun the
74/// whole timeout (the envoy 520-`config` cross product, the sentry-scale
75/// config-key scan), where the between-builders check cannot help.
76pub fn check_current_every(i: usize, every: usize, phase: &str) {
77    if i % every != 0 {
78        return;
79    }
80    if let Some(expires_at) = CURRENT.with(|c| c.get()) {
81        check_expired(Instant::now(), expires_at, phase);
82    }
83}
84
85fn check_expired(now: Instant, expires_at: Instant, phase: &str) {
86    if now > expires_at {
87        panic!("diffctx compute deadline exceeded during {phase}");
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn an_expired_deadline_panics_with_the_phase_name() {
97        let deadline = Deadline::from_timeout_secs(0);
98        std::thread::sleep(Duration::from_millis(5));
99        let err = std::panic::catch_unwind(|| deadline.check("edge construction"))
100            .expect_err("deadline did not fire");
101        let msg = err.downcast_ref::<String>().cloned().unwrap_or_default();
102        assert!(msg.contains("edge construction"), "message was: {msg}");
103    }
104
105    #[test]
106    fn an_unexpired_deadline_does_not_fire() {
107        Deadline::from_timeout_secs(1000).check("edge construction");
108        Deadline::none().check("edge construction");
109    }
110
111    #[test]
112    fn concurrent_deadlines_do_not_affect_each_other() {
113        // The #210 defect: request B's short ceiling used to overwrite
114        // request A's. Per-run values make each check see only its own.
115        let short = Deadline::from_timeout_secs(0);
116        let long = Deadline::from_timeout_secs(1000);
117        std::thread::sleep(Duration::from_millis(5));
118        long.check("edge construction");
119        std::panic::catch_unwind(|| short.check("edge construction"))
120            .expect_err("short deadline did not fire");
121        long.check("edge construction");
122    }
123
124    #[test]
125    fn scoped_deadline_clears_on_drop_and_nests() {
126        let outer = Deadline::from_timeout_secs(1000);
127        let guard = outer.enter();
128        check_current_every(0, 1, "edge construction");
129        {
130            let expired = Deadline::from_timeout_secs(0);
131            let inner = expired.enter();
132            std::thread::sleep(Duration::from_millis(5));
133            std::panic::catch_unwind(|| check_current_every(0, 1, "edge construction"))
134                .expect_err("inner deadline did not fire");
135            drop(inner);
136        }
137        check_current_every(0, 1, "edge construction");
138        drop(guard);
139        // No deadline published: must never fire.
140        check_current_every(0, 1, "edge construction");
141    }
142}
_diffctx/deadline.rs

_diffctx/
deadline.rs
