microclaw_engine/
supervision.rs1use std::collections::HashMap;
12use std::future::Future;
13use std::panic::AssertUnwindSafe;
14use std::sync::{Mutex, OnceLock};
15
16use futures_util::FutureExt;
17use tracing::error;
18
19const MAX_BACKOFF_MS: u64 = 300_000;
21
22static RESTARTS: OnceLock<Mutex<HashMap<&'static str, u64>>> = OnceLock::new();
23
24fn registry() -> &'static Mutex<HashMap<&'static str, u64>> {
25 RESTARTS.get_or_init(|| Mutex::new(HashMap::new()))
26}
27
28fn record_restart(task_name: &'static str) -> u64 {
29 let mut counts = registry().lock().expect("restart registry poisoned");
30 let count = counts.entry(task_name).or_insert(0);
31 *count += 1;
32 *count
33}
34
35pub fn restart_counts() -> Vec<(&'static str, u64)> {
38 let counts = registry().lock().expect("restart registry poisoned");
39 let mut out: Vec<_> = counts.iter().map(|(k, v)| (*k, *v)).collect();
40 out.sort_by_key(|(name, _)| *name);
41 out
42}
43
44pub fn spawn_supervised<F, Fut>(task_name: &'static str, make: F)
50where
51 F: Fn() -> Fut + Send + 'static,
52 Fut: Future<Output = ()> + Send + 'static,
53{
54 spawn_supervised_with_backoff(task_name, 1_000, make);
55}
56
57fn spawn_supervised_with_backoff<F, Fut>(task_name: &'static str, initial_backoff_ms: u64, make: F)
58where
59 F: Fn() -> Fut + Send + 'static,
60 Fut: Future<Output = ()> + Send + 'static,
61{
62 tokio::spawn(async move {
63 let mut backoff_ms = initial_backoff_ms.max(1);
64 loop {
65 match AssertUnwindSafe(make()).catch_unwind().await {
66 Ok(()) => break,
67 Err(payload) => {
68 let restarts = record_restart(task_name);
69 error!(
70 "Background task '{task_name}' panicked (restart #{restarts}); \
71 restarting in {backoff_ms}ms. reason={}",
72 panic_message(&*payload)
73 );
74 tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)).await;
75 backoff_ms = (backoff_ms.saturating_mul(2)).min(MAX_BACKOFF_MS);
76 }
77 }
78 }
79 });
80}
81
82fn panic_message(payload: &(dyn std::any::Any + Send)) -> String {
83 if let Some(s) = payload.downcast_ref::<&str>() {
84 (*s).to_string()
85 } else if let Some(s) = payload.downcast_ref::<String>() {
86 s.clone()
87 } else {
88 "non-string panic payload".to_string()
89 }
90}
91
92#[cfg(test)]
93mod tests {
94 use super::*;
95 use std::sync::atomic::{AtomicU64, Ordering};
96 use std::sync::Arc;
97
98 #[tokio::test]
99 async fn panicking_task_restarts_and_is_counted() {
100 let attempts = Arc::new(AtomicU64::new(0));
101 let attempts_in = attempts.clone();
102 spawn_supervised_with_backoff("test-crashy-loop", 5, move || {
103 let attempts = attempts_in.clone();
104 async move {
105 let n = attempts.fetch_add(1, Ordering::SeqCst);
106 if n < 2 {
107 panic!("boom {n}");
108 }
109 }
111 });
112 tokio::time::timeout(std::time::Duration::from_secs(5), async {
114 while attempts.load(Ordering::SeqCst) < 3 {
115 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
116 }
117 })
118 .await
119 .expect("supervised task did not restart in time");
120
121 let counts = restart_counts();
122 let entry = counts.iter().find(|(name, _)| *name == "test-crashy-loop");
123 assert_eq!(entry.map(|(_, c)| *c), Some(2));
124 }
125
126 #[tokio::test]
127 async fn completing_task_is_not_counted() {
128 spawn_supervised_with_backoff("test-clean-loop", 5, || async {});
129 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
130 assert!(restart_counts()
131 .iter()
132 .all(|(name, _)| *name != "test-clean-loop"));
133 }
134}