athena_rs 3.18.0

//! Daemon
//! Houses background tasks that run alongside the HTTP server.

mod runtime_registry;

use actix_web::rt::time::{Instant as RtInstant, interval_at};
use actix_web::web::Data;
use std::env::var;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::Interval;
use tracing::{debug, info, warn};
use uuid::Uuid;

/// ANSI codes for terminal styling (only when stdout is a TTY; otherwise no-op if stripped).
const B: &str = "\x1b[1m"; // bold
const Y: &str = "\x1b[33m"; // yellow
const R: &str = "\x1b[31m"; // red
const _G: &str = "\x1b[32m"; // green
const Z: &str = "\x1b[0m"; // reset

use crate::AppState;
use crate::config_validation::{RuntimeEnvSettings, runtime_env_settings};
use crate::data::client_connections::{insert_connection_snapshot, prune_connection_snapshots};
use crate::drivers::postgresql::sqlx_driver::{
    ClientConnectionTarget, PostgresClientRegistry, RegisteredClient,
};
use crate::features::connection_pooler::ConnectionPoolSnapshot;
use crate::utils::best_effort_pg_backoff::{
    BestEffortPgBackoffCause, best_effort_pg_write_backoff_active,
    maybe_activate_best_effort_pg_write_backoff,
};

pub use runtime_registry::{resolve_process_daemon_id, spawn_runtime_registry_heartbeat};

const CLIENT_CONNECTIONS_BACKOFF_DOMAIN: &str = "client_connections_pg_logging";

/// Spawn a background task that periodically records pool occupancy into the
/// `client_connections` table inside the logging database.
///
/// The monitor is intentionally lightweight: it samples counts from sqlx pool
/// metadata and writes a single row per client per interval. When the logging
/// client is misconfigured or unavailable, the monitor logs a warning and exits
/// quietly without impacting the HTTP server.
///
/// **Important:** The spawned task uses `PgPool` (sqlx); it must run on the
/// Actix/Tokio runtime. Do not run pool or sqlx operations from `spawn_blocking`
/// or other threads that lack a Tokio context, or you may see "this functionality
/// requires a Tokio context" panics.
pub fn spawn_connection_monitor(app_state: Data<AppState>) {
    let Some(logging_client_name) = app_state.logging_client_name.clone() else {
        debug!("No logging client configured; skipping connection monitor");
        return;
    };

    let Some(logging_pool) = app_state.pg_registry.get_pool(&logging_client_name) else {
        warn!(
            client = %logging_client_name,
            "Logging client pool unavailable; skipping connection monitor"
        );
        return;
    };

    let registry: Arc<PostgresClientRegistry> = app_state.pg_registry.clone();
    let runtime_env: &RuntimeEnvSettings = runtime_env_settings();
    let host: String = var("HOSTNAME").unwrap_or_else(|_| "unknown".to_string());
    let instance_id: Uuid = Uuid::new_v4();
    let interval_secs: u64 = runtime_env.pool_monitor_interval_secs;
    let retention_hours: i64 = runtime_env.pool_monitor_retention_hours;

    info!(
        client = %logging_client_name,
        interval_secs,
        retention_hours,
        host,
        instance_id = %instance_id,
        "Starting connection pool monitor"
    );

    actix_web::rt::spawn(async move {
        let start: RtInstant = RtInstant::now();
        let mut ticker: Interval = interval_at(start, Duration::from_secs(interval_secs));

        loop {
            ticker.tick().await;

            let snapshots: Vec<ConnectionPoolSnapshot> = registry.pool_snapshots();
            for snapshot in &snapshots {
                let max: u32 = snapshot.max_connections.max(1);
                let active_pct: u64 =
                    (snapshot.active_connections as u64).saturating_mul(100) / u64::from(max);
                let saturated: bool = snapshot.idle_connections == 0
                    && snapshot.active_connections > 0
                    && snapshot.pool_size >= snapshot.max_connections;
                if active_pct >= 85 || saturated {
                    warn!(
                        target: "athena::pool",
                        client_name = %snapshot.client_name,
                        active_connections = snapshot.active_connections,
                        idle_connections = snapshot.idle_connections,
                        pool_size = snapshot.pool_size,
                        max_connections = snapshot.max_connections,
                        active_pct,
                        saturated,
                        "athena engine connection pool under pressure"
                    );
                }
            }
            if best_effort_pg_write_backoff_active(CLIENT_CONNECTIONS_BACKOFF_DOMAIN) {
                continue;
            }
            let total: usize = snapshots.len();
            let mut failed: i32 = 0;
            let mut last_err: Option<sqlx::Error> = None;
            for snapshot in &snapshots {
                if let Err(err) =
                    insert_connection_snapshot(&logging_pool, snapshot, &host, instance_id).await
                {
                    failed += 1;
                    if last_err.is_none() {
                        last_err = Some(err);
                    }
                }
            }
            if failed > 0 {
                let err = last_err.unwrap();
                let err_message = err.to_string();
                match maybe_activate_best_effort_pg_write_backoff(
                    CLIENT_CONNECTIONS_BACKOFF_DOMAIN,
                    "client_connections_snapshot",
                    &err_message,
                ) {
                    Some(BestEffortPgBackoffCause::StorageExhaustion) => {
                        warn!(
                            failed,
                            total,
                            error = %err,
                            "{}⚠{} {}Failed to write client_connections snapshot(s){} ({}failed{}={}, total={}); {}logging DB storage is full, pausing snapshot writes{}",
                            Y, Z, B, Z, B, Z, failed, total, R, Z
                        );
                    }
                    Some(BestEffortPgBackoffCause::PoolTimeout) => {
                        warn!(
                            failed,
                            total,
                            error = %err,
                            "{}⚠{} {}Failed to write client_connections snapshot(s){} ({}failed{}={}, total={}); {}logging DB pool is saturated, pausing snapshot writes{}",
                            Y, Z, B, Z, B, Z, failed, total, R, Z
                        );
                    }
                    None => {
                        warn!(
                            failed,
                            total,
                            error = %err,
                            "{}⚠{} {}Failed to write client_connections snapshot(s){} ({}failed{}={}, total={}); {}logging DB may be unreachable{}",
                            Y, Z, B, Z, B, Z, failed, total, R, Z
                        );
                    }
                }
            }

            // Opportunistically prune old snapshots to keep the table bounded.
            if let Err(err) = prune_connection_snapshots(&logging_pool, retention_hours).await {
                let err_message = err.to_string();
                if maybe_activate_best_effort_pg_write_backoff(
                    CLIENT_CONNECTIONS_BACKOFF_DOMAIN,
                    "client_connections_prune",
                    &err_message,
                )
                .is_none()
                {
                    debug!(error = %err, "Failed to prune old client_connections rows");
                }
            }
        }
    });
}

/// Spawn a background worker that retries connecting unavailable active clients.
pub fn spawn_registry_reconnect_worker(app_state: Data<AppState>) {
    let registry: Arc<PostgresClientRegistry> = app_state.pg_registry.clone();
    let interval_secs: u64 = runtime_env_settings().client_reconnect_interval_secs;

    info!(interval_secs, "Starting Postgres client reconnect worker");

    actix_web::rt::spawn(async move {
        let start: RtInstant = RtInstant::now();
        let mut ticker: Interval = interval_at(start, Duration::from_secs(interval_secs));
        loop {
            ticker.tick().await;
            let clients: Vec<RegisteredClient> = registry.list_registered_clients();
            for client in clients {
                if !client.is_active || client.is_frozen || client.pool_connected {
                    continue;
                }
                let target: ClientConnectionTarget = ClientConnectionTarget {
                    client_name: client.client_name.clone(),
                    source: client.source.clone(),
                    description: client.description.clone(),
                    pg_uri: client.pg_uri.clone(),
                    pg_uri_env_var: client.pg_uri_env_var.clone(),
                    config_uri_template: client.config_uri_template.clone(),
                    is_active: client.is_active,
                    is_frozen: client.is_frozen,
                };
                match registry.upsert_client(target).await {
                    Ok(()) => info!(client = %client.client_name, "Reconnected Postgres client"),
                    Err(err) => debug!(
                        client = %client.client_name,
                        error = %err,
                        "Reconnect attempt failed"
                    ),
                }
            }
            registry.sync_connection_status();
        }
    });
}