Add health probes runner (#17)

HealthProbes manager keyed by probe id, spawning one task per
probe. Sync diffs old vs new spec signature so unchanged probes
keep running. HTTP via reqwest with timeout + optional status/body
expectations; TCP via tokio::net::TcpStream::connect with timeout.
Report only on state transitions (or first sample after sync).

Author: sppidy

src/health.rs

@@ -0,0 +1,164 @@
+//! Per-agent health probes — HTTP and TCP. Each probe runs on its own
+//! tokio task; reports flow back through the shared message channel
+//! when the state transitions (or on first sample after a sync).
+
+use shared::{HealthProbeKind, HealthProbeResult, HealthProbeSpec, HealthProbeState, Message};
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::sync::Mutex;
+use tokio::task::JoinHandle;
+use tokio::time::Instant;
+
+#[derive(Default, Clone)]
+pub struct HealthProbes {
+    inner: Arc<Mutex<Inner>>,
+}
+
+#[derive(Default)]
+struct Inner {
+    /// id → (signature, JoinHandle). Signature is the serde_json
+    /// representation of the spec, so we can detect when *anything*
+    /// changed and respawn.
+    tasks: HashMap<String, (String, JoinHandle<()>)>,
+}
+
+impl HealthProbes {
+    /// Apply a new probe set. Spawns/aborts tasks so the running set
+    /// matches `probes`.
+    pub async fn sync(&self, probes: Vec<HealthProbeSpec>, tx: tokio::sync::mpsc::UnboundedSender<Message>) {
+        let mut inner = self.inner.lock().await;
+        let mut keep: HashMap<String, ()> = HashMap::with_capacity(probes.len());
+        for p in probes {
+            let sig = match serde_json::to_string(&p) {
+                Ok(s) => s,
+                Err(_) => continue,
+            };
+            keep.insert(p.id.clone(), ());
+            if let Some((existing_sig, _)) = inner.tasks.get(&p.id) {
+                if existing_sig == &sig {
+                    continue;
+                }
+                // Spec changed — abort old and respawn below.
+                if let Some((_, h)) = inner.tasks.remove(&p.id) {
+                    h.abort();
+                }
+            }
+            let tx_clone = tx.clone();
+            let id_for_task = p.id.clone();
+            let id_for_map = p.id.clone();
+            let handle = tokio::spawn(async move {
+                run_probe(p, tx_clone).await;
+                tracing_drop(&id_for_task);
+            });
+            inner.tasks.insert(id_for_map, (sig, handle));
+        }
+        // Abort tasks not present in the new set.
+        let to_remove: Vec<String> = inner
+            .tasks
+            .keys()
+            .filter(|k| !keep.contains_key(*k))
+            .cloned()
+            .collect();
+        for id in to_remove {
+            if let Some((_, h)) = inner.tasks.remove(&id) {
+                h.abort();
+            }
+        }
+    }
+}
+
+fn tracing_drop(_id: &str) {}
+
+fn now_unix() -> i64 {
+    std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .map(|d| d.as_secs() as i64)
+        .unwrap_or(0)
+}
+
+async fn run_probe(spec: HealthProbeSpec, tx: tokio::sync::mpsc::UnboundedSender<Message>) {
+    let interval = Duration::from_secs(spec.interval_secs.max(1) as u64);
+    let timeout = Duration::from_secs(spec.timeout_secs.max(1) as u64);
+    let mut last_state: Option<HealthProbeState> = None;
+    // A short initial delay so a brand-new probe doesn't fire all at once
+    // alongside every other probe on the host.
+    let jitter = Duration::from_millis((spec.id.bytes().fold(0u64, |a, b| a.wrapping_add(b as u64)) % 5_000) as u64);
+    tokio::time::sleep(jitter).await;
+    loop {
+        let started = Instant::now();
+        let (state, detail) = match spec.kind {
+            HealthProbeKind::Http => probe_http(&spec, timeout).await,
+            HealthProbeKind::Tcp => probe_tcp(&spec, timeout).await,
+        };
+        let latency_ms = started.elapsed().as_millis().min(u32::MAX as u128) as u32;
+        if last_state != Some(state) {
+            // Send a report whenever state flips (or on first sample).
+            let _ = tx.send(Message::HealthProbeReport {
+                results: vec![HealthProbeResult {
+                    id: spec.id.clone(),
+                    state,
+                    latency_ms,
+                    detail: detail.clone(),
+                    at: now_unix(),
+                }],
+            });
+            last_state = Some(state);
+        }
+        tokio::time::sleep(interval).await;
+    }
+}
+
+async fn probe_http(
+    spec: &HealthProbeSpec,
+    timeout: Duration,
+) -> (HealthProbeState, String) {
+    let client = match reqwest::Client::builder().timeout(timeout).build() {
+        Ok(c) => c,
+        Err(e) => return (HealthProbeState::Red, format!("client build: {e}")),
+    };
+    let resp = match client.get(&spec.target).send().await {
+        Ok(r) => r,
+        Err(e) => return (HealthProbeState::Red, format!("connect: {e}")),
+    };
+    let status = resp.status().as_u16();
+    let status_ok = match spec.expect_status {
+        Some(want) => status == want,
+        None => resp.status().is_success(),
+    };
+    if !status_ok {
+        return (
+            HealthProbeState::Red,
+            format!("unexpected status {status}"),
+        );
+    }
+    if let Some(want_body) = &spec.expect_body {
+        let body = match resp.text().await {
+            Ok(t) => t,
+            Err(e) => return (HealthProbeState::Red, format!("read body: {e}")),
+        };
+        if !body.contains(want_body) {
+            return (
+                HealthProbeState::Red,
+                format!("body missing {want_body:?} (status {status})"),
+            );
+        }
+    }
+    (HealthProbeState::Green, format!("ok ({status})"))
+}
+
+async fn probe_tcp(
+    spec: &HealthProbeSpec,
+    timeout: Duration,
+) -> (HealthProbeState, String) {
+    let target = spec.target.clone();
+    let connect = tokio::net::TcpStream::connect(target);
+    match tokio::time::timeout(timeout, connect).await {
+        Ok(Ok(_)) => (HealthProbeState::Green, "ok".to_string()),
+        Ok(Err(e)) => (HealthProbeState::Red, format!("connect: {e}")),
+        Err(_) => (
+            HealthProbeState::Red,
+            format!("timeout after {}s", timeout.as_secs()),
+        ),
+    }
+}

src/main.rs

@@ -1,6 +1,7 @@
 mod apt;
 mod deploy;
 mod docker;
+mod health;
 mod journal;
 mod logs;
 mod stats;
@@ -132,6 +133,7 @@ async fn main() {
     let mut term_session: Option<terminal::TerminalSession> = None;
     let log_streams = logs::LogStreams::default();
     let journal_streams = journal::JournalStreams::default();
+    let health_probes = health::HealthProbes::default();
 
     // Watchdog: if the WebSocket goes silent for 75s the connection is
     // probably dead at the TCP layer (Cloudflare or the kernel may drop
@@ -284,6 +286,13 @@ async fn main() {
                                     streams.stop(&unit).await;
                                 });
                             }
+                            Message::HealthProbeSyncRequest { probes } => {
+                                let probes_mgr = health_probes.clone();
+                                let tx_clone = tx.clone();
+                                tokio::spawn(async move {
+                                    probes_mgr.sync(probes, tx_clone).await;
+                                });
+                            }
                             Message::SwarmServiceInspectRequest { name } => {
                                 let tx_clone = tx.clone();
                                 tokio::spawn(async move {