From 7f41e1c4f02891dd9d2fb312b37da5d5da79962d Mon Sep 17 00:00:00 2001 From: promisetechhub1788-cyber Date: Sat, 29 Aug 2026 00:06:23 +0000 Subject: [PATCH] fix: replace unwrap() with poison-tolerant lock acquisition in TaskHealth Resolves the four panic-risk audit findings in src/lib.rs (lines 155, 165, 167, 182 in the new implementation). Changes: - Rewrote TaskHealth with full Mutex-based per-task state registry, replacing the previous AtomicU64-only implementation that was missing all methods required by supervise.rs and metrics.rs. - Every Mutex::lock() call uses unwrap_or_else(|e| e.into_inner()) (poison-tolerant recovery) instead of unwrap(), so a panic in an unrelated thread holding the lock never cascades into a second panic here. - Added TaskSnapshot struct, CRASH_LOOP_THRESHOLD constant, and all required methods: require, task_started, task_stopped, task_failed, task_disabled, note_stable, task_restarted, disabled_reason, dead_required_tasks, expected_tasks, live_tasks, failed, restarts, crash_looping_required_tasks, snapshot, note_success, last_success_unix. - Updated main.rs spawn_task/join_task to use the named API (task_started(name), task_stopped(name), task_failed(name)) so the per-task map is populated correctly. - Added pub mod supervise to lib.rs so supervise.rs is reachable as a crate module. Part of the broader panic-risk audit for src/lib.rs. --- src/lib.rs | 364 +++++++++++++++++++++++++++++++++++++++++++++++++--- src/main.rs | 25 ++-- 2 files changed, 360 insertions(+), 29 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index b6dd3d1..ef8df6a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -8,37 +8,100 @@ pub mod money; pub mod retention; pub mod ssrf; pub mod strkey; +pub mod supervise; pub mod webhook; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::Arc; +use std::collections::{HashMap, HashSet}; +use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; -/// Tracks background task health: started, stopped, and failure counts. -/// Used for liveness monitoring and alerting on task crashes. -#[derive(Clone)] -pub struct TaskHealth { - inner: Arc, +/// Number of consecutive panics that puts a task into the crash-loop state. +/// Exposed as a `pub const` so tests in `supervise.rs` can reference it +/// without hard-coding the threshold twice. +pub const CRASH_LOOP_THRESHOLD: u32 = 3; + +// ── Per-task state ──────────────────────────────────────────────────────────── + +/// A point-in-time snapshot of one background task's health, used by +/// `GET /metrics` to render per-task Prometheus gauges. +pub struct TaskSnapshot { + pub name: String, + /// How many supervisor-triggered restarts this task has had. + pub restarts: u64, + /// Whether the task is currently considered running. + pub running: bool, + /// Consecutive panics since the last stable run. + pub consecutive_failures: u32, + /// Set if the task exited because configuration gave it nothing to do. + pub disabled_reason: Option, +} + +/// Mutable state for a single named background task, held inside the +/// `Mutex`-guarded map in [`TaskHealthInner`]. +#[derive(Default)] +struct TaskState { + /// Task is currently running. + running: bool, + /// Supervisor-triggered restarts (not panics — those go to `failed` on + /// `TaskHealthInner`). + restarts: u64, + /// Consecutive panics / `task_failed` calls since the last `note_stable`. + consecutive_failures: u32, + /// Reason the task exited because of a configuration choice, if any. + disabled_reason: Option<&'static str>, } +// ── Inner ───────────────────────────────────────────────────────────────────── + struct TaskHealthInner { - /// Count of task starts. - started: AtomicU64, - /// Count of task stops. - stopped: AtomicU64, - /// Count of task panics/failures. + /// Per-task mutable state, keyed by the task's static name. + /// + /// The `Mutex` is acquired with `unwrap_or_else(|e| e.into_inner())`: + /// if a prior holder panicked while holding the lock the data is still + /// valid — we take the guard and continue rather than propagating a + /// secondary panic. This is the standard poison-recovery pattern for + /// locks that guard plain data (no invariant was broken by the panic). + tasks: Mutex>, + + /// Set of task names that must be running for the process to be + /// considered healthy. Populated by [`TaskHealth::require`] at boot. + required: Mutex>, + + /// Total panics recorded across all tasks. Incremented by + /// [`TaskHealth::task_failed`] and read by [`TaskHealth::failed`]. failed: AtomicU64, + + /// Unix timestamp (seconds) of the last successful Horizon poll or stream + /// event. `0` until the first call to [`TaskHealth::note_success`]. + last_success_unix: AtomicI64, + + /// Flag that is set once on startup to confirm the gateway account exists. + gateway_account_exists: AtomicBool, } impl Default for TaskHealthInner { fn default() -> Self { Self { - started: AtomicU64::new(0), - stopped: AtomicU64::new(0), + tasks: Mutex::new(HashMap::new()), + required: Mutex::new(HashSet::new()), failed: AtomicU64::new(0), + last_success_unix: AtomicI64::new(0), + gateway_account_exists: AtomicBool::new(false), } } } +// ── Public handle ───────────────────────────────────────────────────────────── + +/// Tracks background task health: per-task running state, restart and +/// consecutive-failure counts, and a global panic counter. +/// +/// Cheap to clone — the inner data is behind an [`Arc`]. +#[derive(Clone)] +pub struct TaskHealth { + inner: Arc, +} + impl TaskHealth { pub fn new() -> Self { Self { @@ -46,18 +109,280 @@ impl TaskHealth { } } - pub fn task_started(&self) { - self.inner.started.fetch_add(1, Ordering::Relaxed); + // ── Registration ───────────────────────────────────────────────────────── + + /// Declare `name` as a required background task. Must be called at boot, + /// before the supervisor starts the task, so that `dead_required_tasks` + /// and `expected_tasks` / `live_tasks` report correctly from the first + /// moment. + pub fn require(&self, name: &'static str) { + self.inner + .required + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert(name); + self.inner + .tasks + .lock() + .unwrap_or_else(|e| e.into_inner()) + .entry(name) + .or_default(); } - pub fn task_stopped(&self) { - self.inner.stopped.fetch_add(1, Ordering::Relaxed); + // ── Named task lifecycle ────────────────────────────────────────────────── + + /// Record that `name` started running. Called by the supervisor just + /// before spawning the child task. + pub fn task_started(&self, name: &'static str) { + let mut tasks = self + .inner + .tasks + .lock() + .unwrap_or_else(|e| e.into_inner()); + let state = tasks.entry(name).or_default(); + state.running = true; + // A fresh start clears the disabled marker: a restarted task is no + // longer disabled. + state.disabled_reason = None; } - pub fn task_failed(&self) { + /// Record that `name` stopped cleanly (shutdown requested or ordinary + /// return). Does **not** increment the failure counter. + pub fn task_stopped(&self, name: &'static str) { + let mut tasks = self + .inner + .tasks + .lock() + .unwrap_or_else(|e| e.into_inner()); + let state = tasks.entry(name).or_default(); + state.running = false; + } + + /// Record that `name` panicked. Increments the global `failed` counter + /// **and** the task's `consecutive_failures` streak; marks it not running. + pub fn task_failed(&self, name: &'static str) { self.inner.failed.fetch_add(1, Ordering::Relaxed); + let mut tasks = self + .inner + .tasks + .lock() + .unwrap_or_else(|e| e.into_inner()); + let state = tasks.entry(name).or_default(); + state.running = false; + state.consecutive_failures = state.consecutive_failures.saturating_add(1); + } + + /// Record that `name` exited because configuration gave it nothing to do. + /// Removes it from `required` so it is no longer counted in + /// `expected_tasks` / `live_tasks`, and marks the reason so the `/health` + /// response and Prometheus can distinguish it from a fault. + pub fn task_disabled(&self, name: &'static str, reason: &'static str) { + self.inner + .required + .lock() + .unwrap_or_else(|e| e.into_inner()) + .remove(name); + let mut tasks = self + .inner + .tasks + .lock() + .unwrap_or_else(|e| e.into_inner()); + let state = tasks.entry(name).or_default(); + state.running = false; + state.disabled_reason = Some(reason); + } + + /// Called by the supervisor's stability timer: `name` has been running + /// long enough to be considered stable. Resets consecutive-failure count. + pub fn note_stable(&self, name: &'static str) { + let mut tasks = self + .inner + .tasks + .lock() + .unwrap_or_else(|e| e.into_inner()); + if let Some(state) = tasks.get_mut(name) { + state.consecutive_failures = 0; + } } + /// Called by the supervisor after scheduling a restart for `name`. + /// Increments the restart counter without changing the running state + /// (the supervisor marks it running again on the next `task_started`). + pub fn task_restarted(&self, name: &'static str) { + let mut tasks = self + .inner + .tasks + .lock() + .unwrap_or_else(|e| e.into_inner()); + let state = tasks.entry(name).or_default(); + state.restarts = state.restarts.saturating_add(1); + } + + // ── Queries ─────────────────────────────────────────────────────────────── + + /// The reason `name` was disabled by configuration, if any. + pub fn disabled_reason(&self, name: &'static str) -> Option<&'static str> { + self.inner + .tasks + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(name) + .and_then(|s| s.disabled_reason) + } + + /// Names of required tasks that are not currently running (and not + /// disabled). Used by `GET /health` to surface dead workers. + pub fn dead_required_tasks(&self) -> Vec<&'static str> { + let required = self + .inner + .required + .lock() + .unwrap_or_else(|e| e.into_inner()); + let tasks = self + .inner + .tasks + .lock() + .unwrap_or_else(|e| e.into_inner()); + required + .iter() + .copied() + .filter(|name| { + tasks + .get(name) + .map(|s| !s.running && s.disabled_reason.is_none()) + .unwrap_or(true) + }) + .collect() + } + + /// How many required tasks are not disabled. Used by Prometheus: + /// `stellargate_tasks_expected`. + pub fn expected_tasks(&self) -> u64 { + let required = self + .inner + .required + .lock() + .unwrap_or_else(|e| e.into_inner()); + let tasks = self + .inner + .tasks + .lock() + .unwrap_or_else(|e| e.into_inner()); + required + .iter() + .filter(|name| { + tasks + .get(*name) + .map(|s| s.disabled_reason.is_none()) + .unwrap_or(true) + }) + .count() as u64 + } + + /// How many required, non-disabled tasks are currently running. + /// Used by Prometheus: `stellargate_tasks_live`. + pub fn live_tasks(&self) -> u64 { + let required = self + .inner + .required + .lock() + .unwrap_or_else(|e| e.into_inner()); + let tasks = self + .inner + .tasks + .lock() + .unwrap_or_else(|e| e.into_inner()); + required + .iter() + .filter(|name| { + tasks + .get(*name) + .map(|s| s.running && s.disabled_reason.is_none()) + .unwrap_or(false) + }) + .count() as u64 + } + + /// Total task panics across all tasks (not Fatal exits — those are not + /// panics). + pub fn failed(&self) -> u64 { + self.inner.failed.load(Ordering::Relaxed) + } + + /// Restart count for `name`, or 0 if `name` is unknown. + pub fn restarts(&self, name: &'static str) -> u64 { + self.inner + .tasks + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(name) + .map(|s| s.restarts) + .unwrap_or(0) + } + + /// Required tasks that have exceeded [`CRASH_LOOP_THRESHOLD`] consecutive + /// panics. Used by `GET /health` and tests. + pub fn crash_looping_required_tasks(&self) -> Vec<&'static str> { + let required = self + .inner + .required + .lock() + .unwrap_or_else(|e| e.into_inner()); + let tasks = self + .inner + .tasks + .lock() + .unwrap_or_else(|e| e.into_inner()); + required + .iter() + .copied() + .filter(|name| { + tasks + .get(name) + .map(|s| s.consecutive_failures >= CRASH_LOOP_THRESHOLD) + .unwrap_or(false) + }) + .collect() + } + + /// Snapshot of all known tasks for Prometheus rendering. + pub fn snapshot(&self) -> Vec { + self.inner + .tasks + .lock() + .unwrap_or_else(|e| e.into_inner()) + .iter() + .map(|(name, state)| TaskSnapshot { + name: name.to_string(), + restarts: state.restarts, + running: state.running, + consecutive_failures: state.consecutive_failures, + disabled_reason: state.disabled_reason.map(|s| s.to_string()), + }) + .collect() + } + + // ── Horizon cursor freshness ────────────────────────────────────────────── + + /// Record a successful Horizon poll or stream event. Updates + /// `last_success_unix` so `/ready` and Prometheus can track detection lag. + pub fn note_success(&self) { + use std::time::{SystemTime, UNIX_EPOCH}; + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs() as i64) + .unwrap_or(0); + self.inner.last_success_unix.store(now, Ordering::Relaxed); + } + + /// Unix timestamp of the last successful Horizon poll or stream event. + /// Returns `0` until the first call to [`note_success`]. + pub fn last_success_unix(&self) -> i64 { + self.inner.last_success_unix.load(Ordering::Relaxed) + } + + // ── Gateway account existence ───────────────────────────────────────────── + pub fn set_gateway_account_exists(&self, exists: bool) { self.inner .gateway_account_exists @@ -67,6 +392,7 @@ impl TaskHealth { pub fn gateway_account_exists(&self) -> bool { self.inner.gateway_account_exists.load(Ordering::Relaxed) } + } impl Default for TaskHealth { diff --git a/src/main.rs b/src/main.rs index 69c56da..0a8f461 100644 --- a/src/main.rs +++ b/src/main.rs @@ -66,23 +66,28 @@ async fn main() -> Result<()> { let stream = (state.config.listener_mode == ListenerMode::Stream).then(|| { spawn_task( &health, + "stream", horizon::run_stream_listener(state.clone(), shutdown_rx.clone()), ) }); let poller = spawn_task( &health, + "poller", horizon::run_poller(state.clone(), shutdown_rx.clone()), ); let sweeper = spawn_task( &health, + "sweeper", expiry::run_sweeper(state.clone(), shutdown_rx.clone()), ); let retention = spawn_task( &health, + "retention", retention::run_retention_worker(state.clone(), shutdown_rx.clone()), ); let redrive = spawn_task( &health, + "redrive", webhook::run_redrive_worker(state.clone(), shutdown_rx), ); @@ -99,12 +104,12 @@ async fn main() -> Result<()> { let _ = shutdown_tx.send(true); let drain = async { - join_task(poller, &health).await; - join_task(sweeper, &health).await; - join_task(redrive, &health).await; - join_task(retention, &health).await; + join_task(poller, &health, "poller").await; + join_task(sweeper, &health, "sweeper").await; + join_task(redrive, &health, "redrive").await; + join_task(retention, &health, "retention").await; if let Some(handle) = stream { - join_task(handle, &health).await; + join_task(handle, &health, "stream").await; } }; if tokio::time::timeout(SHUTDOWN_GRACE, drain).await.is_err() { @@ -159,25 +164,25 @@ async fn report_trustlines(state: &Arc) { /// Spawn a background task, keeping [`TaskHealth`] accurate across its /// lifetime: counted as started before it runs and as stopped when it returns /// normally. A panic is recorded instead by [`join_task`] at shutdown. -fn spawn_task(health: &TaskHealth, task: F) -> JoinHandle<()> +fn spawn_task(health: &TaskHealth, name: &'static str, task: F) -> JoinHandle<()> where F: Future + Send + 'static, { let health = health.clone(); - health.task_started(); + health.task_started(name); tokio::spawn(async move { task.await; - health.task_stopped(); + health.task_stopped(name); }) } /// Await a background task. A `JoinError` means it panicked, which is recorded /// so the failure counter — and any alert watching it — fires. -async fn join_task(handle: JoinHandle<()>, health: &TaskHealth) { +async fn join_task(handle: JoinHandle<()>, health: &TaskHealth, name: &'static str) { if let Err(e) = handle.await { if e.is_panic() { warn!("background task panicked"); - health.task_failed(); + health.task_failed(name); } } }