AkurAI-Monitor / src/main.rs

Latest change 811ef9b7c8ba017fa242f008efa90faab07a0454 - Adopt AkurAI-Framework UI layer and add framework-version tracking by Ólafur Búi Ólafsson

mod alert;
mod auth;
mod collectors;
mod config;
mod db;
mod dns_metrics;
mod framework_version;
mod health;
mod journald;
mod management;
mod pibun_push;
mod routes;
mod schema;
mod stream;
mod tailer;
mod version_poller;

use axum::{
    Router,
    http::{HeaderValue, header},
    middleware,
    response::Response,
    routing::{get, post},
};
use tower_http::services::{ServeDir, ServeFile};

// Force browsers to revalidate static assets so a redeploy is picked up
// immediately instead of serving a stale cached app.js / style.css.
async fn add_no_cache(mut response: Response) -> Response {
    response
        .headers_mut()
        .insert(header::CACHE_CONTROL, HeaderValue::from_static("no-cache"));
    response
}

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("rust_monitor=info")),
        )
        .init();

    tracing::info!("rust-monitor starting up");

    let cfg = config::get();

    // Initialize DB (triggers LazyLock)
    db::with_db(|_| {});
    management::ensure_configured_channels().expect("failed to configure alert channels");
    // Self-register this monitor in the applications table so version tracking
    // knows our own framework/version info.
    db::with_db(|conn| {
        let now = chrono::Utc::now().timestamp();
        let fw = framework_version::FRAMEWORK_VERSION;
        let ver = env!("CARGO_PKG_VERSION");
        conn.execute(
            "INSERT INTO applications (host_id, name, service_name, health_url, enabled, \
             app_version, framework, framework_version, versions_reported_at, created_at, updated_at) \
             VALUES ((SELECT id FROM hosts WHERE name = 'local'), 'rust-monitor', 'rust-monitor', \
             NULL, 1, ?1, 'AkurAI-Framework', ?2, ?3, ?3, ?3) \
             ON CONFLICT(host_id, name) DO UPDATE SET \
             app_version = excluded.app_version, \
             framework = excluded.framework, \
             framework_version = excluded.framework_version, \
             versions_reported_at = excluded.versions_reported_at, \
             updated_at = excluded.updated_at",
            rusqlite::params![ver, fw, now],
        )
        .ok();
    });

    // Initialize the live SSE broadcast channel
    stream::init();

    // Follow journald for configured systemd units (other apps log here)
    journald::spawn(cfg.journal_units.clone());

    // Spawn metric collector task.
    // Samples on a fast "live" cadence and broadcasts every snapshot to SSE
    // subscribers, but only persists to the DB on the configured interval
    // (keeping history — and chart queries — at the intended resolution).
    let interval = cfg.interval_secs;
    let live_secs = interval.clamp(1, 5);
    tokio::spawn(async move {
        let mut ticker = tokio::time::interval(std::time::Duration::from_secs(live_secs));
        let mut last_persist: i64 = 0;
        loop {
            ticker.tick().await;
            let mut metrics = tokio::task::spawn_blocking(collectors::collect_all)
                .await
                .unwrap_or_default();
            // Append the latest DNS window (derived from akurai-dns logs) so
            // DNS badges stay live between persist intervals.
            metrics.extend(dns_metrics::latest());

            let now = chrono::Utc::now().timestamp();

            // Push live snapshot to any connected dashboards
            stream::publish(
                "status",
                serde_json::json!({ "metrics": metrics, "ts": now }).to_string(),
            );

            // Persist at the configured interval only
            if now - last_persist >= interval as i64 {
                last_persist = now;
                // Drain the DNS aggregator exactly once per interval so counts
                // and qps reflect the true window, then persist alongside the
                // system metrics (collect_all values are still current).
                let dns = dns_metrics::drain(interval);
                let mut metrics = metrics;
                metrics.retain(|m| !m.name.starts_with("dns."));
                metrics.extend(dns);
                let count = metrics.len();
                db::with_db(|conn| {
                    let tx = conn.unchecked_transaction().ok();
                    for m in &metrics {
                        conn.execute(
                            "INSERT INTO metrics (host_id, name, value, ts)
                             VALUES ((SELECT id FROM hosts WHERE name = 'local'), ?1, ?2, ?3)",
                            rusqlite::params![m.name, m.value, now],
                        )
                        .ok();
                    }
                    if let Some(tx) = tx {
                        tx.commit().ok();
                    }
                });
                tracing::debug!(count, "collected metrics");
            }
        }
    });

    // Spawn log tailer task — tails on the fast cadence, persists every batch,
    // and broadcasts new lines live to connected log views.
    // Register tailed log files for health tracking.
    for f in &cfg.log_files {
        health::register_source(f, "tailer");
    }
    let log_files = cfg.log_files.clone();
    tokio::spawn(async move {
        let mut log_tailer = tailer::LogTailer::new(&log_files);
        let mut ticker = tokio::time::interval(std::time::Duration::from_secs(live_secs));
        loop {
            ticker.tick().await;
            let entries = log_tailer.read_new_lines();
            if entries.is_empty() {
                continue;
            }

            let now = chrono::Utc::now().timestamp();
            let count = entries.len();
            db::with_db(|conn| {
                let tx = conn.unchecked_transaction().ok();
                for (source, line) in &entries {
                    conn.execute(
                        "INSERT OR IGNORE INTO logs (host_id, source, line, origin, ts)
                         VALUES ((SELECT id FROM hosts WHERE name = 'local'), ?1, ?2, 'server', ?3)",
                        rusqlite::params![source, line, now],
                    )
                    .ok();
                }
                if let Some(tx) = tx {
                    tx.commit().ok();
                }
            });

            let logs: Vec<_> = entries
                .iter()
                .map(|(source, line)| serde_json::json!({ "source": source, "line": line, "origin": "server", "ts": now }))
                .collect();
            stream::publish("log", serde_json::json!({ "logs": logs }).to_string());

            tracing::debug!(count, "tailed log lines");
        }
    });

    // Spawn alert engine task
    let alert_log = cfg.alert_log.clone();
    tokio::spawn(async move {
        let mut ticker = tokio::time::interval(std::time::Duration::from_secs(60));
        loop {
            ticker.tick().await;
            alert::evaluate_alerts(&alert_log).await;
        }
    });

    // Spawn retention cleanup task
    let retention_days = cfg.retention_days;
    let log_retention_days = cfg.log_retention_days;
    let log_debug_retention_days = cfg.log_debug_retention_days;
    let log_max_rows = cfg.log_max_rows;
    tokio::spawn(async move {
        let mut ticker = tokio::time::interval(std::time::Duration::from_secs(3600));
        loop {
            ticker.tick().await;
            let now = chrono::Utc::now().timestamp();
            let metric_cutoff = now - (retention_days * 86400);
            let log_cutoff = now - (log_retention_days * 86400);
            let debug_cutoff = now - (log_debug_retention_days * 86400);

            let mut deleted = [0usize; 3];
            for (index, (sql, cutoff)) in [
                (
                    "DELETE FROM metrics WHERE id IN
                     (SELECT id FROM metrics WHERE ts < ?1 LIMIT 5000)",
                    metric_cutoff,
                ),
                (
                    "DELETE FROM logs WHERE id IN
                     (SELECT id FROM logs WHERE ts < ?1 LIMIT 5000)",
                    log_cutoff,
                ),
                // trace/debug are high-volume and low-value once aged; purge sooner.
                (
                    "DELETE FROM logs WHERE id IN
                     (SELECT id FROM logs
                      WHERE ts < ?1 AND level IN ('trace', 'debug') LIMIT 5000)",
                    debug_cutoff,
                ),
            ]
            .into_iter()
            .enumerate()
            {
                loop {
                    let batch = db::with_db(|conn| {
                        conn.execute(sql, rusqlite::params![cutoff]).unwrap_or(0)
                    });
                    deleted[index] += batch;
                    if batch < 5000 {
                        break;
                    }
                    tokio::task::yield_now().await;
                }
            }

            // Hard cap total log rows to bound SQLite growth under log floods.
            let mut capped = 0usize;
            loop {
                let batch = db::with_db(|conn| {
                    conn.execute(
                        "DELETE FROM logs WHERE id IN
                         (SELECT id FROM logs ORDER BY id DESC LIMIT 5000 OFFSET ?1)",
                        rusqlite::params![log_max_rows],
                    )
                    .unwrap_or(0)
                });
                capped += batch;
                if batch < 5000 {
                    break;
                }
                tokio::task::yield_now().await;
            }

            if deleted != [0, 0, 0] || capped != 0 {
                tracing::info!(
                    metrics_deleted = deleted[0],
                    logs_deleted = deleted[1] + deleted[2],
                    logs_capped = capped,
                    "retention cleanup"
                );
            }
        }
    });

    // Spawn version poller — fetches /api/health from every enabled application
    // and records version / framework info.
    version_poller::spawn();

    // Build the router
    let index_path = format!("{}/index.html", cfg.static_dir);
    let static_service = ServeDir::new(&cfg.static_dir)
        .append_index_html_on_directories(true)
        .fallback(ServeFile::new(&index_path));

    let app = Router::new()
        .route("/api/health", get(routes::health))
        .route("/api/status", get(routes::status))
        .route("/api/metrics", get(routes::metrics))
        .route("/api/logs", get(routes::logs))
        .route("/api/alerts", get(routes::alerts))
        .route("/api/ingestion/health", get(routes::ingestion_health))
        .route("/api/stream", get(stream::sse_handler))
        .route("/api/ingest", post(routes::ingest))
        .route(
            "/api/manage/{resource}",
            get(management::list).post(management::create),
        )
        .route(
            "/api/manage/{resource}/{id}",
            get(management::get)
                .put(management::update)
                .delete(management::delete),
        )
        .route(
            "/api/manage/channels/{id}/targets",
            get(management::channel_targets),
        )
        .route(
            "/api/manage/channels/{id}/test",
            post(management::test_channel),
        )
        .route("/auth/callback", get(auth::auth_callback))
        .route("/auth/logout", get(auth::auth_logout))
        .fallback_service(static_service)
        .layer(middleware::from_fn(auth::auth_middleware))
        .layer(middleware::map_response(add_no_cache));

    let listener = tokio::net::TcpListener::bind(&cfg.listen_addr)
        .await
        .unwrap();
    tracing::info!("listening on {}", cfg.listen_addr);
    axum::serve(listener, app).await.unwrap();
}