1use std::num::NonZeroU64;
15
16use kernel::event::record::{FiredAction, WatchdogFired};
17use kernel::{AxCode, AxError, Payload, Retries, Retry, RunId, StallVerdict, TimeMs};
18
19#[derive(Debug, Default)]
22pub struct Watchdog {
23 corrections: u32,
24 provider_failures: u32,
25 streak: u32,
31 retries: Retries,
32 jitter_seed: u64,
35}
36
37#[derive(Debug, Clone, PartialEq, Eq)]
40pub enum Disposal {
41 Proceed,
42 CorrectiveSteer {
43 text: String,
44 },
45 BackOff {
51 until: TimeMs,
52 code: AxCode,
53 subject: String,
54 },
55 Freeze {
56 reason: FreezeReason,
57 },
58}
59
60impl Disposal {
61 fn backing_off(until: TimeMs, failure: &AxError) -> Disposal {
63 Disposal::BackOff {
64 until,
65 code: *failure.code(),
66 subject: failure.subject().to_owned(),
67 }
68 }
69}
70
71#[derive(Debug, Clone, Copy, PartialEq, Eq)]
72pub enum FreezeReason {
73 Stall,
74 ProviderRefused,
75}
76
77impl FreezeReason {
78 fn as_str(self) -> &'static str {
79 match self {
80 FreezeReason::Stall => "stall",
81 FreezeReason::ProviderRefused => "provider_refused",
82 }
83 }
84}
85
86impl Watchdog {
87 #[must_use]
88 pub fn new(retries: Retries, run: RunId) -> Watchdog {
89 Watchdog {
90 retries,
91 jitter_seed: seed_of(run),
92 ..Watchdog::default()
93 }
94 }
95
96 pub fn on_stall(&mut self, verdict: &StallVerdict) -> Disposal {
99 match verdict {
100 StallVerdict::Ok => Disposal::Proceed,
101 StallVerdict::Stall { repeats } => {
102 if self.corrections == 0 {
103 self.corrections = 1;
104 Disposal::CorrectiveSteer {
105 text: format!(
106 "You have repeated the same call {repeats} times in a row. \
107 Change the approach or report why the goal cannot be met."
108 ),
109 }
110 } else {
111 Disposal::Freeze {
112 reason: FreezeReason::Stall,
113 }
114 }
115 }
116 }
117 }
118
119 pub fn on_provider_failure(&mut self, failure: &AxError, now: TimeMs) -> Disposal {
138 self.provider_failures = self.provider_failures.saturating_add(1);
139 self.streak = self.streak.saturating_add(1);
140 let own = self.backoff_ms();
141 let wait = failure.retry_after_ms().map_or(own, |told| own.max(told));
142 let not_before = TimeMs::new(now.value().saturating_add(wait));
143 let refused = Disposal::Freeze {
144 reason: FreezeReason::ProviderRefused,
145 };
146 if failure.retry() == Retry::No {
147 return refused;
148 }
149 match self.retries {
150 Retries::UntilHalted => Disposal::backing_off(not_before, failure),
151 Retries::AtMost(ceiling) if self.streak <= ceiling => {
154 Disposal::backing_off(not_before, failure)
155 }
156 Retries::AtMost(_) => refused,
157 }
158 }
159
160 pub fn on_provider_answered(&mut self) {
162 self.streak = 0;
163 }
164
165 fn backoff_ms(&self) -> u64 {
173 const FIRST_MS: u64 = 500;
174 const CEILING_MS: u64 = 60_000;
175 let doublings = self.streak.saturating_sub(1).min(16);
176 let base = FIRST_MS
177 .checked_shl(doublings)
178 .map_or(CEILING_MS, |wait| wait.min(CEILING_MS));
179 let spread = NonZeroU64::MIN.saturating_add(base >> 1);
180 base.saturating_add(mixed(self.jitter_seed ^ u64::from(self.streak)) % spread)
181 }
182
183 pub fn fired_payload(&self, disposal: &Disposal) -> Result<Payload, AxError> {
190 let action = match disposal {
191 Disposal::Proceed => {
192 return Err(AxError::failure(
193 AxCode::InvalidArgs,
194 "encode watchdog_fired",
195 "Proceed does not fire",
196 )
197 .with_recovery(
198 "report this against runtime::watchdog: `Proceed` is the verdict \
199 that records nothing, and the caller asked it for a payload",
200 ));
201 }
202 Disposal::CorrectiveSteer { text } => FiredAction::Steer { text: text.clone() },
203 Disposal::BackOff {
204 until,
205 code,
206 subject,
207 } => FiredAction::BackOff {
208 until_ms: until.value(),
209 code: code.as_str().to_owned(),
210 subject: subject.clone(),
211 },
212 Disposal::Freeze { reason } => FiredAction::Freeze {
213 reason: reason.as_str().to_owned(),
214 },
215 };
216 Payload::of(&WatchdogFired {
217 action,
218 corrections: self.corrections,
219 provider_failures: self.provider_failures,
220 })
221 }
222}
223
224fn seed_of(run: RunId) -> u64 {
228 run.as_bytes()
229 .iter()
230 .fold(0xcbf2_9ce4_8422_2325, |hash, byte| {
231 (hash ^ u64::from(*byte)).wrapping_mul(0x0000_0100_0000_01b3)
232 })
233}
234
235fn mixed(word: u64) -> u64 {
238 let word = (word ^ (word >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
239 let word = (word ^ (word >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
240 word ^ (word >> 31)
241}
242
243#[cfg(test)]
244#[allow(
245 clippy::unwrap_used,
246 clippy::expect_used,
247 clippy::panic,
248 clippy::indexing_slicing,
249 clippy::wildcard_enum_match_arm,
250 clippy::let_underscore_must_use,
251 clippy::let_underscore_untyped,
252 clippy::arithmetic_side_effects,
253 reason = "test code"
254)]
255mod tests {
256 use super::*;
257 use kernel::ActionFingerprint;
258 use kernel::stall::observe;
259
260 const BASES: [u64; 9] = [
262 500, 1_000, 2_000, 4_000, 8_000, 16_000, 32_000, 60_000, 60_000,
263 ];
264
265 fn run(tail: u8) -> RunId {
268 let mut bytes = [
269 0x01, 0x93, 0x2a, 0x7c, 0x10, 0x00, 0x70, 0x00, 0x80, 0, 0, 0, 0, 0, 0, 0,
270 ];
271 bytes[15] = tail;
272 RunId::from_bytes(bytes)
273 }
274
275 fn waits(dog: &mut Watchdog, count: usize) -> Vec<u64> {
278 (0..count)
279 .map(
280 |_| match dog.on_provider_failure(&provider_error(true), TimeMs::new(0)) {
281 Disposal::BackOff { until, .. } => until.value(),
282 other => panic!("a retriable failure under its ceiling backs off: {other:?}"),
283 },
284 )
285 .collect()
286 }
287
288 fn within_schedule(waits: &[u64]) -> bool {
289 waits
290 .iter()
291 .zip(BASES)
292 .all(|(wait, base)| (base..=base + base / 2).contains(wait))
293 }
294
295 #[test]
299 fn each_run_waits_its_own_jittered_schedule() {
300 let schedule = |tail: u8| waits(&mut Watchdog::new(Retries::UntilHalted, run(tail)), 9);
301 let (first, second) = (schedule(1), schedule(2));
302 assert!(
303 within_schedule(&first) && within_schedule(&second),
304 "{first:?} {second:?}"
305 );
306 assert_eq!(first, schedule(1), "the same run waits the same schedule");
307 assert_ne!(first, second, "two runs do not ask again in step");
308 }
309
310 #[test]
311 fn disposal_is_graded_steer_first_freeze_second() {
312 let mut dog = Watchdog::new(Retries::UntilHalted, run(1));
313 let same = ActionFingerprint::derive(b"exec identical");
314 let sample = vec![same, same, same];
315 let verdict = observe(&sample);
316 let first = dog.on_stall(&verdict);
317 match &first {
318 Disposal::CorrectiveSteer { text } => {
319 assert!(text.contains("repeated the same call 3 times"))
320 }
321 other => panic!("first hit must steer, got {other:?}"),
322 }
323 let second = dog.on_stall(&verdict);
324 assert_eq!(
325 second,
326 Disposal::Freeze {
327 reason: FreezeReason::Stall
328 }
329 );
330 let payload = serde_json::to_value(dog.fired_payload(&second).unwrap()).unwrap();
332 assert_eq!(payload["action"], "freeze");
333 assert_eq!(payload["reason"], "stall");
334 assert_eq!(
337 dog.fired_payload(&Disposal::Proceed)
338 .unwrap_err()
339 .recovery(),
340 "report this against runtime::watchdog: `Proceed` is the verdict that records nothing, and the caller asked it for a payload"
341 );
342 }
343
344 #[test]
345 fn ok_verdicts_never_dispose() {
346 let mut dog = Watchdog::new(Retries::UntilHalted, run(1));
347 assert_eq!(dog.on_stall(&StallVerdict::Ok), Disposal::Proceed);
348 assert_eq!(dog.on_stall(&StallVerdict::Ok), Disposal::Proceed);
349 }
350
351 fn provider_error(retriable: bool) -> AxError {
352 let draft = AxError::failure(AxCode::Provider, "call the model", "the provider said no");
353 let draft = if retriable { draft.retriable() } else { draft };
354 draft.with_recovery("dispatch again, or attach a second endpoint")
355 }
356
357 #[test]
358 fn a_failure_the_provider_will_repeat_stops_after_one() {
359 let mut dog = Watchdog::new(Retries::UntilHalted, run(1));
360 assert_eq!(
361 dog.on_provider_failure(&provider_error(false), TimeMs::new(9_000)),
362 Disposal::Freeze {
363 reason: FreezeReason::ProviderRefused
364 },
365 "a request the provider rejected on its shape buys the same rejection again"
366 );
367 }
368
369 #[test]
370 fn an_answer_ends_the_streak_the_backoff_counts() {
371 let mut dog = Watchdog::new(Retries::AtMost(2), run(1));
372 let before = waits(&mut dog, 2);
373 assert!(within_schedule(&before), "{before:?}");
374 dog.on_provider_answered();
375 assert_eq!(
376 waits(&mut dog, 2),
377 before,
378 "an outage after a recovery starts from the first wait, under a fresh ceiling"
379 );
380 }
381
382 #[test]
383 fn a_retriable_failure_backs_off_and_never_freezes_by_itself() {
384 let mut dog = Watchdog::new(Retries::UntilHalted, run(1));
385 let waits: Vec<u64> = (0..64u64)
386 .map(
387 |_| match dog.on_provider_failure(&provider_error(true), TimeMs::new(1_000)) {
388 Disposal::BackOff { until, .. } => until.value().saturating_sub(1_000),
389 other => {
390 panic!("only Halt stops a run that is waiting out a provider: {other:?}")
391 }
392 },
393 )
394 .collect();
395 assert!(
396 within_schedule(&waits),
397 "each failure in a row doubles the wait, up to a minute: {waits:?}"
398 );
399 assert!(
400 waits[9..]
401 .iter()
402 .all(|wait| (60_000..=90_000).contains(wait))
403 );
404 let payload = serde_json::to_value(
405 dog.fired_payload(&Disposal::BackOff {
406 until: TimeMs::new(1_000),
407 code: AxCode::Provider,
408 subject: "the provider said no".to_owned(),
409 })
410 .unwrap(),
411 )
412 .unwrap();
413 assert_eq!(payload["action"], "back_off");
414 assert_eq!(payload["until_ms"], 1_000);
415 assert_eq!(payload["provider_failures"], 64);
416 assert_eq!(
417 payload["subject"], "the provider said no",
418 "a line saying it backed off says what it backed off from"
419 );
420 }
421
422 fn waits_of_fresh_run() -> Vec<u64> {
423 waits(&mut Watchdog::new(Retries::UntilHalted, run(1)), 2)
424 }
425
426 #[test]
427 fn a_provider_that_names_its_wait_is_not_asked_sooner() {
428 let mut dog = Watchdog::new(Retries::UntilHalted, run(1));
429 let told = |wait_ms: u64| {
430 AxError::failure(AxCode::Provider, "call the model", "answered 429")
431 .retriable_after(wait_ms)
432 .with_recovery("wait, then ask again")
433 };
434 let waits: Vec<u64> = [30_000, 10]
435 .map(
436 |wait_ms| match dog.on_provider_failure(&told(wait_ms), TimeMs::new(1_000)) {
437 Disposal::BackOff { until, .. } => until.value().saturating_sub(1_000),
438 other => {
439 panic!("a provider that says when to come back is asked again: {other:?}")
440 }
441 },
442 )
443 .to_vec();
444 let own = waits_of_fresh_run();
445 assert_eq!(
446 waits,
447 [30_000, own[1]],
448 "the longer of the provider's word and the schedule's own wait"
449 );
450 }
451}