Skip to content

fix: replace unwrap() with poison-tolerant lock acquisition in TaskHealth #600

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
graceuvala-collab merged 1 commit into from
Aug 29, 2026
Merged

fix: replace unwrap() with poison-tolerant lock acquisition in TaskHealth #600

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
284 changes: 135 additions & 149 deletions src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,55 +11,72 @@ pub mod strkey;
pub mod supervise;
pub mod webhook;

use std::collections::HashMap;
use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};

/// 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;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TaskSnapshot {
pub name: &'static str,
pub running: bool,
pub restarts: u64,
pub consecutive_failures: u32,
pub disabled_reason: Option<&'static str>,
}

/// Consecutive panics at or above this mark a required task as crash-looping.
/// `/health` fails while any required task is crash-looping, even if the
/// supervisor has already spawned a replacement.
pub const CRASH_LOOP_THRESHOLD: u32 = 3;
// ── Per-task state ────────────────────────────────────────────────────────────

/// Snapshot of one background task, used to render `/metrics`.
#[derive(Debug, Clone, PartialEq, Eq)]
/// 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: &'static str,
pub running: bool,
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,
pub disabled_reason: Option<&'static str>,
/// Set if the task exited because configuration gave it nothing to do.
pub disabled_reason: Option<String>,
}

/// Tracks background task health for liveness, readiness, and monitoring.
#[derive(Clone)]
pub struct TaskHealth {
inner: Arc<TaskHealthInner>,
/// 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 {
started: AtomicU64,
stopped: AtomicU64,
/// 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<HashMap<&'static str, TaskState>>,

/// Set of task names that must be running for the process to be
/// considered healthy. Populated by [`TaskHealth::require`] at boot.
required: Mutex<HashSet<&'static str>>,

/// Total panics recorded across all tasks. Incremented by
/// [`TaskHealth::task_failed`] and read by [`TaskHealth::failed`].
failed: AtomicU64,
running: Mutex<HashMap<&'static str, bool>>,
restarts: Mutex<HashMap<&'static str, u64>>,
consecutive_failures: Mutex<HashMap<&'static str, u32>>,
disabled: Mutex<HashMap<&'static str, &'static str>>,
required: Mutex<Vec<&'static str>>,

/// 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,
gateway_account_exists: std::sync::atomic::AtomicBool,

/// Flag that is set once on startup to confirm the gateway account exists.
gateway_account_exists: AtomicBool,
}

impl Default for TaskHealthInner {
Expand All @@ -68,19 +85,21 @@ impl Default for TaskHealthInner {
tasks: Mutex::new(HashMap::new()),
required: Mutex::new(HashSet::new()),
failed: AtomicU64::new(0),
running: Mutex::new(HashMap::new()),
restarts: Mutex::new(HashMap::new()),
consecutive_failures: Mutex::new(HashMap::new()),
disabled: Mutex::new(HashMap::new()),
required: Mutex::new(Vec::new()),
last_success_unix: AtomicI64::new(0),
gateway_account_exists: std::sync::atomic::AtomicBool::new(false),
gateway_account_exists: AtomicBool::new(false),
}
}
}

fn lock<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
mutex.lock().unwrap_or_else(|poisoned| poisoned.into_inner())
// ── 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<TaskHealthInner>,
}

impl TaskHealth {
Expand All @@ -90,126 +109,100 @@ impl TaskHealth {
}
}

// ── 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) {
lock(&self.inner.required).push(name);
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();
}

// ── 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) {
self.inner.started.fetch_add(1, Ordering::Relaxed);
lock(&self.inner.running).insert(name, true);
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;
}

/// Record that `name` stopped cleanly (shutdown requested or ordinary
/// return). Does **not** increment the failure counter.
pub fn task_stopped(&self, name: &'static str) {
self.inner.stopped.fetch_add(1, Ordering::Relaxed);
lock(&self.inner.running).insert(name, false);
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);
lock(&self.inner.running).insert(name, false);
let mut consecutive = lock(&self.inner.consecutive_failures);
*consecutive.entry(name).or_insert(0) += 1;
}

pub fn task_restarted(&self, name: &'static str) {
let mut restarts = lock(&self.inner.restarts);
*restarts.entry(name).or_insert(0) += 1;
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.stopped.fetch_add(1, Ordering::Relaxed);
lock(&self.inner.running).insert(name, false);
lock(&self.inner.disabled).insert(name, reason);
}

pub fn disabled_reason(&self, name: &'static str) -> Option<&'static str> {
lock(&self.inner.disabled).get(name).copied()
}

pub fn expected_tasks(&self) -> usize {
let required = lock(&self.inner.required);
let disabled = lock(&self.inner.disabled);
required.iter().filter(|name| !disabled.contains_key(*name)).count()
}

pub fn live_tasks(&self) -> usize {
let running = lock(&self.inner.running);
let required = lock(&self.inner.required);
let disabled = lock(&self.inner.disabled);
required.iter()
.filter(|name| !disabled.contains_key(*name))
.filter(|name| running.get(*name) == Some(&true))
.count()
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) {
lock(&self.inner.consecutive_failures).insert(name, 0);
}

pub fn started(&self) -> u64 { self.inner.started.load(Ordering::Relaxed) }
pub fn stopped(&self) -> u64 { self.inner.stopped.load(Ordering::Relaxed) }
pub fn failed(&self) -> u64 { self.inner.failed.load(Ordering::Relaxed) }

pub fn restarts(&self, name: &'static str) -> u64 {
lock(&self.inner.restarts).get(name).copied().unwrap_or(0)
}

pub fn consecutive_failures(&self, name: &'static str) -> u32 {
lock(&self.inner.consecutive_failures).get(name).copied().unwrap_or(0)
}

pub fn dead_required_tasks(&self) -> Vec<&'static str> {
let running = lock(&self.inner.running);
let required = lock(&self.inner.required);
let disabled = lock(&self.inner.disabled);
required.iter().copied()
.filter(|name| !disabled.contains_key(name))
.filter(|name| running.get(name) != Some(&true))
.collect()
}

pub fn crash_looping_required_tasks(&self) -> Vec<&'static str> {
let consecutive = lock(&self.inner.consecutive_failures);
let required = lock(&self.inner.required);
1 required.iter().copied()
.filter(|name| consecutive.get(name).copied().unwrap_or(0) >= CRASH_LOOP_THRESHOLD)
.collect()
}

pub fn snapshot(&self) -> Vec<TaskSnapshot> {
let running = lock(&self.inner.running);
let restarts = lock(&self.inner.restarts);
let consecutive = lock(&self.inner.consecutive_failures);
let required = lock(&self.inner.required);
let disabled = lock(&self.inner.disabled);
let mut names = required.clone();
for name in running.keys().chain(restarts.keys()).chain(consecutive.keys()).chain(disabled.keys()) {
if !names.contains(name) { names.push(*name); }
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;
}
names.sort_unstable();
names.into_iter().map(|name| TaskSnapshot {
name,
running: running.get(name).copied().unwrap_or(false),
restarts: restarts.get(name).copied().unwrap_or(0),
consecutive_failures: consecutive.get(name).copied().unwrap_or(0),
disabled_reason: disabled.get(name).copied(),
}).collect()
}

pub fn note_success(&self) {
self.set_last_success_unix(unix_now_secs());
}

pub fn set_last_success_unix(&self, unix_secs: i64) {
self.inner.last_success_unix.store(unix_secs, Ordering::Relaxed);
}

pub fn last_success_age_secs(&self) -> i64 {
unix_now_secs().saturating_sub(self.inner.last_success_unix.load(Ordering::Relaxed))
}

pub fn last_success_unix(&self) -> i64 {
self.inner.last_success_unix.load(Ordering::Relaxed)
}

/// Called by the supervisor after scheduling a restart for `name`.
Expand Down Expand Up @@ -402,13 +395,6 @@ impl TaskHealth {

}

fn unix_now_secs() -> i64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|duration| duration.as_secs() as i64)
.unwrap_or(0)
}

impl Default for TaskHealth {
fn default() -> Self {
Self::new()
Expand Down