AkurAI-Monitor / src/health.rs

Latest change d0b8d14afb524ab72435e5abd0bbd6049a205349 - Migrate storage layer from SQLite to MariaDB by Ólafur Búi Ólafsson

use std::collections::HashMap;
use std::sync::{LazyLock, Mutex};

use serde::Serialize;

use mysql::prelude::Queryable;

use crate::db;

/// Per-source ingestion health, shared across journald/tailer tasks and the health endpoint.
#[derive(Debug, Clone, Serialize)]
pub struct SourceState {
    pub kind: String,
    pub last_activity_ts: Option<i64>,
    pub last_error: Option<String>,
    pub last_error_ts: Option<i64>,
    pub error_count: u64,
}

/// In-memory mirror of ingestion_state for fast reads from the health endpoint.
/// Updated by journald and tailer tasks on each batch; flushed to DB for cursor persistence.
static STATE: LazyLock<Mutex<HashMap<String, SourceState>>> =
    LazyLock::new(|| Mutex::new(HashMap::new()));

/// Register a source for tracking (idempotent — call on startup for every configured source).
pub fn register_source(source: &str, kind: &str) {
    let mut state = STATE.lock().expect("ingestion health mutex poisoned");
    state
        .entry(source.to_string())
        .or_insert_with(|| SourceState {
            kind: kind.to_string(),
            last_activity_ts: None,
            last_error: None,
            last_error_ts: None,
            error_count: 0,
        });
}

/// Record successful ingestion activity for a source.
pub fn record_activity(source: &str, ts: i64) {
    let mut state = STATE.lock().expect("ingestion health mutex poisoned");
    if let Some(s) = state.get_mut(source) {
        s.last_activity_ts = Some(ts);
        // Clear transient errors on successful activity.
        s.last_error = None;
        s.last_error_ts = None;
    }
}

/// Record an ingestion error for a source.
pub fn record_error(source: &str, ts: i64, error: &str) {
    let mut state = STATE.lock().expect("ingestion health mutex poisoned");
    if let Some(s) = state.get_mut(source) {
        s.last_error = Some(error.to_string());
        s.last_error_ts = Some(ts);
        s.error_count = s.error_count.saturating_add(1);
    }
}

/// Clear a transient source error after the ingestion path recovers.
pub fn clear_error(source: &str) {
    let mut state = STATE.lock().expect("ingestion health mutex poisoned");
    if let Some(s) = state.get_mut(source) {
        s.last_error = None;
        s.last_error_ts = None;
    }
}

/// Drop an unusable journald cursor so the follower can use bounded replay.
pub fn clear_journal_cursors() {
    db::with_db(|conn| {
        if let Err(error) =
            conn.query_drop("UPDATE ingestion_state SET last_cursor = NULL WHERE kind = 'journald'")
        {
            tracing::error!(%error, "failed to clear journald cursors");
        }
    });
}

/// Load the last persisted cursor for *any* journald source (used when following multiple units).
pub fn load_any_cursor() -> Option<String> {
    db::with_db(|conn| {
        conn.query_first(
            "SELECT last_cursor FROM ingestion_state
             WHERE kind = 'journald' AND last_cursor IS NOT NULL
             ORDER BY last_ts DESC LIMIT 1",
        )
        .ok()
        .flatten()
    })
}

/// Snapshot all registered sources for the health endpoint.
pub fn snapshot() -> Vec<SourceStateEntry> {
    let state = STATE.lock().expect("ingestion health mutex poisoned");
    state
        .iter()
        .map(|(source, s)| SourceStateEntry {
            source: source.clone(),
            kind: s.kind.clone(),
            last_activity_ts: s.last_activity_ts,
            last_error: s.last_error.clone(),
            last_error_ts: s.last_error_ts,
            error_count: s.error_count,
        })
        .collect()
}

#[derive(Debug, Clone, Serialize)]
pub struct SourceStateEntry {
    pub source: String,
    pub kind: String,
    pub last_activity_ts: Option<i64>,
    pub last_error: Option<String>,
    pub last_error_ts: Option<i64>,
    pub error_count: u64,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn register_and_track_activity() {
        let source = "test-unit.service";
        register_source(source, "journald");
        record_activity(source, 1000);

        let snap = snapshot();
        let entry = snap.iter().find(|e| e.source == source).unwrap();
        assert_eq!(entry.kind, "journald");
        assert_eq!(entry.last_activity_ts, Some(1000));
        assert_eq!(entry.error_count, 0);
        assert!(entry.last_error.is_none());
    }

    #[test]
    fn record_error_increments_count() {
        let source = "error-test.service";
        register_source(source, "journald");
        record_error(source, 100, "connection refused");
        record_error(source, 200, "timeout");
        record_activity(source, 300);

        let snap = snapshot();
        let entry = snap.iter().find(|e| e.source == source).unwrap();
        assert_eq!(entry.error_count, 2);
        // Successful activity clears the error fields.
        assert!(entry.last_error.is_none());
    }
}