Skip to main content

microclaw_engine/
supervision.rs

1//! Supervised background loops.
2//!
3//! A plain `tokio::spawn`ed loop that panics dies silently — the scheduler,
4//! reflector, or heartbeat simply stops until the process restarts, with
5//! nothing but one log line at panic time. Loops spawned through
6//! [`spawn_supervised`] instead restart with capped exponential backoff, and
7//! every restart is counted in a process-wide registry that `insights`
8//! surfaces — "scheduler restarted 14× today" becomes visible instead of a
9//! mystery.
10
11use 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
19/// Backoff cap between restarts of a crash-looping task.
20const 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
35/// Snapshot of `(task_name, restart_count)` sorted by name. Empty when no
36/// supervised task has ever panicked.
37pub 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
44/// Spawn a background loop that survives its own panics. `make` builds a
45/// fresh instance of the loop future; on panic the task is rebuilt and
46/// restarted after a capped exponential backoff. A future that returns
47/// normally ends the supervision (loops that finish on purpose stay
48/// finished).
49pub 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                // Third attempt completes normally — supervision ends.
110            }
111        });
112        // Two panics with 5ms/10ms backoff: give it ample time.
113        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}