neiam /dms

//! Gauges describing the archive itself, alongside the HTTP and process
//! metrics.
//!
//! These answer questions the request metrics cannot: how much has been
//! stored, and — the reason this module exists — whether the processing
//! pipeline behind an upload is actually finishing. They come from counting
//! rows, so they are sampled on a timer rather than updated on every write; a
//! gauge that every code path touching a table has to keep correct is a gauge
//! that drifts.
//!
//! **Every replica reports the same numbers.** They describe the database, not
//! the pod, so a dashboard must aggregate them with `max` and never `sum`;
//! summing counts the same uploads once per web pod.
//!
//! The job gauges are the point. An upload here fans out into up to five
//! independent jobs — thumbnail, ocr, describe, extract_tags, similarity — run
//! by four separate worker Deployments. When one of those stalls, the site
//! stays up, uploads keep succeeding, and the only visible symptom is that
//! thumbnails quietly stop appearing. [`sample`] breaks the job table down by
//! `kind` and `status` so the stalled one names itself.
use std::sync::Arc;
use std::time::Duration;

use chrono::Utc;
use sea_orm::{
    ColumnTrait, ConnectionTrait, DatabaseBackend, DatabaseConnection, EntityTrait, PaginatorTrait,
    QueryFilter, Statement,
};

use crate::entities::{
    ai_descriptions, api_tokens, ocr_results, tags, upload_jobs, uploads, user_sessions, users,
};
use crate::jobs::{
    STATUS_COMPLETE, STATUS_FAILED, STATUS_PENDING, STATUS_PROCESSING, UPLOAD_STATUS_COMPLETE,
    UPLOAD_STATUS_PARTIAL, UPLOAD_STATUS_PROCESSING,
};

/// How often the counts are refreshed.
const SAMPLE_EVERY: Duration = Duration::from_secs(60);

/// How many ticks between refreshes of the one query that scans a big table.
///
/// Summing `file_size_bytes` across every upload cannot use an index, so
/// Postgres reads the heap. Once a minute forever is not a reasonable price
/// for a number that is read as a slow-moving line; ten minutes is well below
/// the resolution anyone looks at it.
const SLOW_EVERY: u64 = 10;

/// How long a job sits in `processing` before it is called stuck.
///
/// Thumbnailing is seconds and a `describe` against the VLM is at worst a
/// couple of minutes, so an hour is an order of magnitude past anything
/// healthy. This is the signal for a worker killed mid-job: the row is claimed
/// and no longer queued, so nothing will ever pick it up again.
const STUCK_AFTER: chrono::TimeDelta = chrono::TimeDelta::hours(1);

/// The job kinds, in the order `queue::JobKind` declares them.
///
/// Listed rather than discovered so a kind that has never run still reports
/// zero instead of dropping its series — a missing line and a line at zero
/// mean very different things on the pipeline panel.
const KINDS: [&str; 5] = ["thumbnail", "ocr", "describe", "extract_tags", "similarity"];

const JOB_STATUSES: [&str; 4] = [STATUS_PENDING, STATUS_PROCESSING, STATUS_COMPLETE, STATUS_FAILED];

const UPLOAD_STATUSES: [&str; 3] =
    [UPLOAD_STATUS_PROCESSING, UPLOAD_STATUS_COMPLETE, UPLOAD_STATUS_PARTIAL];

/// Register the metric names and what they mean.
///
/// Called at boot so the names exist in `/metrics` from the first scrape.
/// Without this a fresh pod exposes nothing until the first sample lands, and
/// a panel over a just-restarted deployment reads as an outage.
pub fn describe() {
    metrics::describe_gauge!(
        "dms_uploads",
        "Uploads held, by processing state. `partial` means every job finished but at least \
         one failed — the file is stored and usable, it just has no thumbnail or no OCR text"
    );
    metrics::describe_gauge!(
        "dms_uploads_bytes",
        "Total stored bytes across every upload. Refreshed every ten minutes rather than every \
         minute; see SLOW_EVERY"
    );
    metrics::describe_gauge!(
        "dms_jobs",
        "Processing jobs by kind and status. This is the pipeline: five kinds fan out from one \
         upload, and a worker Deployment stalling shows up as its kind stuck in `pending`"
    );
    metrics::describe_gauge!(
        "dms_jobs_stuck",
        "Jobs claimed as `processing` more than an hour ago. The worker that took them is gone \
         and the row is no longer queued, so nothing will finish them without intervention"
    );
    metrics::describe_gauge!(
        "dms_jobs_exhausted",
        "Failed jobs that have used every retry (attempts >= WORKER_MAX_ATTEMPTS). These will \
         not be retried again and are what the /failed page exists to show"
    );
    metrics::describe_gauge!("dms_ocr_results", "Uploads with extracted text");
    metrics::describe_gauge!("dms_ai_descriptions", "Uploads the VLM has described");
    metrics::describe_gauge!("dms_tags", "Distinct tags across the archive");
    metrics::describe_gauge!("dms_users", "Accounts that exist");
    metrics::describe_gauge!("dms_sessions_active", "Sessions not yet expired");
    metrics::describe_gauge!("dms_api_tokens", "API tokens issued");
}

/// Sample the counts forever, on a timer.
///
/// `max_attempts` is the app's own `WORKER_MAX_ATTEMPTS`, so the exhausted
/// count agrees with the retry limit the workers actually enforce rather than
/// with a second threshold invented here.
///
/// Failures are logged and skipped rather than retried or fatal: a metrics
/// sampler that can take the process down, or that hammers a struggling
/// database, is worse than one that misses a minute.
pub async fn run(db: Arc<DatabaseConnection>, max_attempts: i32) {
    let mut ticker = tokio::time::interval(SAMPLE_EVERY);
    // The first tick fires immediately, so the numbers are right from boot
    // rather than a minute after it.
    ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
    let mut tick: u64 = 0;
    loop {
        ticker.tick().await;
        // Tick zero includes the slow gauge, so the archive size is present
        // from the first scrape rather than ten minutes into the pod's life.
        let include_slow = tick.is_multiple_of(SLOW_EVERY);
        tick = tick.wrapping_add(1);
        if let Err(e) = sample(db.as_ref(), max_attempts, include_slow).await {
            tracing::warn!(error = %e, "could not sample metrics");
        }
    }
}

async fn sample(
    db: &DatabaseConnection,
    max_attempts: i32,
    include_slow: bool,
) -> Result<(), sea_orm::DbErr> {
    let now = Utc::now();

    for status in UPLOAD_STATUSES {
        let n = uploads::Entity::find()
            .filter(uploads::Column::ProcessingStatus.eq(status))
            .count(db)
            .await?;
        metrics::gauge!("dms_uploads", "status" => status).set(n as f64);
    }

    // The pipeline, broken down the way it actually fails. Twenty cheap
    // counted queries rather than one grouped one, because a kind or status
    // with no rows has to report zero and a GROUP BY simply omits it.
    for kind in KINDS {
        for status in JOB_STATUSES {
            let n = upload_jobs::Entity::find()
                .filter(upload_jobs::Column::Kind.eq(kind))
                .filter(upload_jobs::Column::Status.eq(status))
                .count(db)
                .await?;
            metrics::gauge!("dms_jobs", "kind" => kind, "status" => status).set(n as f64);
        }
    }

    // Deliberately unwindowed: a job wedged since last week is exactly the one
    // worth knowing about, and any recent-only window would hide it.
    metrics::gauge!("dms_jobs_stuck").set(
        upload_jobs::Entity::find()
            .filter(upload_jobs::Column::Status.eq(STATUS_PROCESSING))
            .filter(upload_jobs::Column::StartedAt.lt(now - STUCK_AFTER))
            .count(db)
            .await? as f64,
    );
    metrics::gauge!("dms_jobs_exhausted").set(
        upload_jobs::Entity::find()
            .filter(upload_jobs::Column::Status.eq(STATUS_FAILED))
            .filter(upload_jobs::Column::Attempts.gte(max_attempts))
            .count(db)
            .await? as f64,
    );

    metrics::gauge!("dms_ocr_results").set(ocr_results::Entity::find().count(db).await? as f64);
    metrics::gauge!("dms_ai_descriptions")
        .set(ai_descriptions::Entity::find().count(db).await? as f64);
    metrics::gauge!("dms_tags").set(tags::Entity::find().count(db).await? as f64);

    metrics::gauge!("dms_users").set(users::Entity::find().count(db).await? as f64);
    metrics::gauge!("dms_sessions_active").set(
        user_sessions::Entity::find()
            .filter(user_sessions::Column::ExpiresAt.gt(now))
            .count(db)
            .await? as f64,
    );
    metrics::gauge!("dms_api_tokens").set(api_tokens::Entity::find().count(db).await? as f64);

    if include_slow {
        // Raw SQL rather than sea-orm's `sum()`: Postgres returns SUM(bigint)
        // as numeric, which would need a Decimal round-trip to read back. The
        // cast makes the type unambiguous and the gauge is an f64 anyway.
        let row = db
            .query_one(Statement::from_string(
                DatabaseBackend::Postgres,
                "SELECT COALESCE(SUM(file_size_bytes), 0)::float8 AS total FROM uploads",
            ))
            .await?;
        if let Some(row) = row {
            metrics::gauge!("dms_uploads_bytes").set(row.try_get::<f64>("", "total")?);
        }
    }

    Ok(())
}