later 0.0.49

//! Process-local Prometheus metrics for job servers.

use crate::{
    job_timing,
    models::{AmqpCommand, Job, Stage},
};
use prometheus::{
    Encoder, GaugeVec, HistogramOpts, HistogramVec, IntCounterVec, IntGaugeVec, Opts, Registry,
    TextEncoder,
};
use std::{sync::OnceLock, time::Instant};

static METRICS: OnceLock<Result<Metrics, String>> = OnceLock::new();

struct Metrics {
    registry: Registry,
    job_transitions: IntCounterVec,
    handler_duration: HistogramVec,
    job_duration: HistogramVec,
    job_wait_duration: HistogramVec,
    commands: IntCounterVec,
    workers: IntGaugeVec,
    queue_jobs: IntGaugeVec,
    partition_claims: IntCounterVec,
    partition_heads_claimed: IntCounterVec,
    partition_lease_renewals: IntCounterVec,
    partitions_active: IntGaugeVec,
    partition_workers_live: IntGaugeVec,
    partition_backlog: IntGaugeVec,
    partition_oldest_head_age_seconds: GaugeVec,
    retained_log_consumer_lag: IntGaugeVec,
    partition_queue_depth: IntGaugeVec,
}

impl Metrics {
    fn new() -> Result<Self, prometheus::Error> {
        let registry = Registry::new();
        let job_transitions = IntCounterVec::new(
            Opts::new(
                "later_job_transitions_total",
                "Durable job stage transitions in this process.",
            ),
            &["queue", "job_type", "stage"],
        )?;
        let handler_duration = HistogramVec::new(
            HistogramOpts::new(
                "later_job_handler_duration_seconds",
                "Time each worker spends in application job handlers.",
            )
            .buckets(vec![
                0.000_1, 0.000_25, 0.000_5, 0.001, 0.002_5, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25,
                0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0,
            ]),
            &["queue", "worker_id", "job_type", "outcome"],
        )?;
        let duration_buckets = vec![
            0.001, 0.002_5, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0,
            60.0, 300.0, 900.0, 3_600.0, 21_600.0, 86_400.0,
        ];
        let job_duration = HistogramVec::new(
            HistogramOpts::new(
                "later_job_duration_seconds",
                "End-to-end time from job creation to its terminal stage.",
            )
            .buckets(duration_buckets.clone()),
            &["queue", "job_type", "outcome"],
        )?;
        let job_wait_duration = HistogramVec::new(
            HistogramOpts::new(
                "later_job_wait_duration_seconds",
                "Time completed jobs spent outside application handlers.",
            )
            .buckets(duration_buckets),
            &["queue", "job_type", "outcome"],
        )?;
        let commands = IntCounterVec::new(
            Opts::new(
                "later_worker_commands_total",
                "Delivery commands handled by each worker in this process.",
            ),
            &["queue", "worker_id", "command", "outcome"],
        )?;
        let workers = IntGaugeVec::new(
            Opts::new(
                "later_workers",
                "Each worker task in this process by current state.",
            ),
            &["queue", "worker_id", "state"],
        )?;
        let queue_jobs = IntGaugeVec::new(
            Opts::new(
                "later_queue_jobs",
                "Durable jobs in each active queue state.",
            ),
            &["queue", "state"],
        )?;
        let partition_claims = IntCounterVec::new(
            Opts::new(
                "later_partition_claims_total",
                "Partition head claim attempts by this process, whether or not one was found.",
            ),
            &["queue", "outcome"],
        )?;
        let partition_heads_claimed = IntCounterVec::new(
            Opts::new(
                "later_partition_heads_claimed_total",
                "Partition heads claimed by this process, by topic.",
            ),
            &["queue", "topic"],
        )?;
        let partition_lease_renewals = IntCounterVec::new(
            Opts::new(
                "later_partition_lease_renewals_total",
                "Partition lease renewal attempts by this process, by topic and outcome.",
            ),
            &["queue", "topic", "outcome"],
        )?;
        let partitions_active = IntGaugeVec::new(
            Opts::new(
                "later_partitions_active",
                "Partition head handlers currently running in this process, by topic.",
            ),
            &["queue", "topic"],
        )?;
        let partition_workers_live = IntGaugeVec::new(
            Opts::new(
                "later_partition_workers_live",
                "Live partition worker count for this namespace, as last observed by this process.",
            ),
            &["queue"],
        )?;
        let partition_backlog = IntGaugeVec::new(
            Opts::new(
                "later_partition_backlog",
                "Ready, unleased partition heads by topic, as last observed by this process.",
            ),
            &["queue", "topic"],
        )?;
        let partition_oldest_head_age_seconds = GaugeVec::new(
            Opts::new(
                "later_partition_oldest_head_age_seconds",
                "Age of the longest-waiting ready partition head by topic, or 0 when none are \
                 ready, as last observed by this process.",
            ),
            &["queue", "topic"],
        )?;
        let retained_log_consumer_lag = IntGaugeVec::new(
            Opts::new(
                "later_retained_log_consumer_lag",
                "Retained-log records appended but not yet committed by one consumer group, by \
                 topic and group, as last observed by this process. Only set when Config::retained_log_lag \
                 is configured.",
            ),
            &["queue", "topic", "group"],
        )?;
        let partition_queue_depth = IntGaugeVec::new(
            Opts::new(
                "later_partition_queue_depth",
                "Not-yet-completed jobs currently enqueued to one topic partition (including \
                 the one at the head), as last observed by this process. The job-queue analog \
                 of later_retained_log_consumer_lag - always available for every configured \
                 topic, no separate opt-in needed.",
            ),
            &["queue", "topic", "partition"],
        )?;

        registry.register(Box::new(job_transitions.clone()))?;
        registry.register(Box::new(handler_duration.clone()))?;
        registry.register(Box::new(job_duration.clone()))?;
        registry.register(Box::new(job_wait_duration.clone()))?;
        registry.register(Box::new(commands.clone()))?;
        registry.register(Box::new(workers.clone()))?;
        registry.register(Box::new(queue_jobs.clone()))?;
        registry.register(Box::new(partition_claims.clone()))?;
        registry.register(Box::new(partition_heads_claimed.clone()))?;
        registry.register(Box::new(partition_lease_renewals.clone()))?;
        registry.register(Box::new(partitions_active.clone()))?;
        registry.register(Box::new(partition_workers_live.clone()))?;
        registry.register(Box::new(partition_backlog.clone()))?;
        registry.register(Box::new(partition_oldest_head_age_seconds.clone()))?;
        registry.register(Box::new(retained_log_consumer_lag.clone()))?;
        registry.register(Box::new(partition_queue_depth.clone()))?;

        Ok(Self {
            registry,
            job_transitions,
            handler_duration,
            job_duration,
            job_wait_duration,
            commands,
            workers,
            queue_jobs,
            partition_claims,
            partition_heads_claimed,
            partition_lease_renewals,
            partitions_active,
            partition_workers_live,
            partition_backlog,
            partition_oldest_head_age_seconds,
            retained_log_consumer_lag,
            partition_queue_depth,
        })
    }
}

fn metrics() -> anyhow::Result<&'static Metrics> {
    match METRICS.get_or_init(|| Metrics::new().map_err(|error| error.to_string())) {
        Ok(metrics) => Ok(metrics),
        Err(error) => Err(anyhow::anyhow!(error.clone())),
    }
}

pub(crate) fn encode() -> anyhow::Result<String> {
    let metrics = metrics()?;
    let mut buffer = Vec::new();
    TextEncoder::new().encode(&metrics.registry.gather(), &mut buffer)?;
    Ok(String::from_utf8(buffer)?)
}

/// Every registered metric's current value, as plain JSON rather than the
/// Prometheus text exposition format - used by the dashboard's own
/// Monitoring tab, which polls this in-process (no scraper, no separate
/// time-series store) and keeps its own rolling history client-side. Shaped
/// generically (name/help/type/labelled samples) so a chart on the
/// dashboard side never has to know about a specific metric ahead of time.
pub(crate) fn snapshot() -> anyhow::Result<serde_json::Value> {
    let metrics = metrics()?;
    Ok(serde_json::Value::Array(
        metrics.registry.gather().iter().map(family_json).collect(),
    ))
}

fn family_json(family: &prometheus::proto::MetricFamily) -> serde_json::Value {
    use prometheus::proto::MetricType;
    let metric_type = match family.get_field_type() {
        MetricType::COUNTER => "counter",
        MetricType::GAUGE => "gauge",
        MetricType::HISTOGRAM => "histogram",
        MetricType::SUMMARY => "summary",
        MetricType::UNTYPED => "untyped",
    };
    serde_json::json!({
        "name": family.get_name(),
        "help": family.get_help(),
        "type": metric_type,
        "metrics": family.get_metric().iter().map(metric_json).collect::<Vec<_>>(),
    })
}

fn metric_json(metric: &prometheus::proto::Metric) -> serde_json::Value {
    let labels = serde_json::Value::Object(
        metric
            .get_label()
            .iter()
            .map(|pair| {
                (
                    pair.get_name().to_string(),
                    serde_json::Value::String(pair.get_value().to_string()),
                )
            })
            .collect(),
    );
    if metric.has_counter() {
        serde_json::json!({ "labels": labels, "value": metric.get_counter().get_value() })
    } else if metric.has_gauge() {
        serde_json::json!({ "labels": labels, "value": metric.get_gauge().get_value() })
    } else if metric.has_histogram() {
        let histogram = metric.get_histogram();
        serde_json::json!({
            "labels": labels,
            "histogram": {
                "sample_count": histogram.get_sample_count(),
                "sample_sum": histogram.get_sample_sum(),
            },
        })
    } else {
        serde_json::json!({ "labels": labels })
    }
}

pub(crate) fn record_job_transition(queue: &str, job: &Job) {
    let stage = match &job.stage {
        Stage::Delayed(_) => "delayed",
        Stage::Waiting(_) => "waiting",
        Stage::Enqueued(_) => "enqueued",
        Stage::Running(_) => "running",
        Stage::Requeued(_) => "requeued",
        Stage::Success(_) => "success",
        Stage::Failed(_) => "failed",
    };
    if let Ok(metrics) = metrics() {
        metrics
            .job_transitions
            .with_label_values(&[queue, &job.payload_type, stage])
            .inc();

        let outcome = match &job.stage {
            Stage::Success(_) => Some("success"),
            Stage::Failed(_) => Some("failed"),
            Stage::Delayed(_)
            | Stage::Waiting(_)
            | Stage::Enqueued(_)
            | Stage::Running(_)
            | Stage::Requeued(_) => None,
        };
        if let Some(outcome) = outcome {
            let timing =
                job_timing::calculate(&job.previous_stages, &job.stage, chrono::Utc::now());
            metrics
                .job_duration
                .with_label_values(&[queue, &job.payload_type, outcome])
                .observe(timing.total_ms as f64 / 1_000.0);
            metrics
                .job_wait_duration
                .with_label_values(&[queue, &job.payload_type, outcome])
                .observe(timing.waiting_ms as f64 / 1_000.0);
        }
    }
}

pub(crate) fn command_name(command: &AmqpCommand) -> &'static str {
    match command {
        AmqpCommand::ExecuteJob(_) => "execute_job",
        AmqpCommand::PartitionReady { .. } => "partition_ready",
    }
}

pub(crate) fn record_command(queue: &str, worker_id: &str, command: &str, outcome: &str) {
    if let Ok(metrics) = metrics() {
        metrics
            .commands
            .with_label_values(&[queue, worker_id, command, outcome])
            .inc();
    }
}

pub(crate) fn set_queue_depth(queue: &str, state: &str, count: usize) {
    if let Ok(metrics) = metrics() {
        let count = i64::try_from(count).unwrap_or(i64::MAX);
        metrics
            .queue_jobs
            .with_label_values(&[queue, state])
            .set(count);
    }
}

/// Records one `claim_partition_head` call. `topic` is `Some` only when a
/// head was actually claimed, since the topic isn't known otherwise.
pub(crate) fn record_partition_claim(queue: &str, topic: Option<&str>) {
    if let Ok(metrics) = metrics() {
        let outcome = if topic.is_some() { "claimed" } else { "none" };
        metrics
            .partition_claims
            .with_label_values(&[queue, outcome])
            .inc();
        if let Some(topic) = topic {
            metrics
                .partition_heads_claimed
                .with_label_values(&[queue, topic])
                .inc();
        }
    }
}

/// Records one partition lease renewal attempt.
pub(crate) fn record_partition_lease_renewal(queue: &str, topic: &str, renewed: bool) {
    if let Ok(metrics) = metrics() {
        let outcome = if renewed { "renewed" } else { "lost" };
        metrics
            .partition_lease_renewals
            .with_label_values(&[queue, topic, outcome])
            .inc();
    }
}

/// Sets the last-observed live partition worker count for `queue`.
pub(crate) fn set_partition_workers_live(queue: &str, count: usize) {
    if let Ok(metrics) = metrics() {
        let count = i64::try_from(count).unwrap_or(i64::MAX);
        metrics
            .partition_workers_live
            .with_label_values(&[queue])
            .set(count);
    }
}

/// Sets the ready-backlog count and oldest-ready-head age for one topic, as
/// last observed by this process. `oldest_head_age_seconds` should be `0.0`
/// when `ready_count` is `0`.
pub(crate) fn set_partition_backlog(
    queue: &str,
    topic: &str,
    ready_count: usize,
    oldest_head_age_seconds: f64,
) {
    if let Ok(metrics) = metrics() {
        let ready_count = i64::try_from(ready_count).unwrap_or(i64::MAX);
        metrics
            .partition_backlog
            .with_label_values(&[queue, topic])
            .set(ready_count);
        metrics
            .partition_oldest_head_age_seconds
            .with_label_values(&[queue, topic])
            .set(oldest_head_age_seconds);
    }
}

/// Sets one topic partition's queue depth, as last observed by this
/// process. `partition` is formatted as a plain integer string.
pub(crate) fn set_partition_queue_depth(queue: &str, topic: &str, partition: &str, depth: i64) {
    if let Ok(metrics) = metrics() {
        metrics
            .partition_queue_depth
            .with_label_values(&[queue, topic, partition])
            .set(depth);
    }
}

/// Sets one consumer group's retained-log lag on one topic, as last
/// observed by this process. `group` is host-supplied (whatever name the
/// application commits offsets under), not a fixed, bounded label the way
/// `topic` is - callers should keep the number of distinct group names
/// small, the same caution that applies to any Prometheus label.
pub(crate) fn set_retained_log_consumer_lag(queue: &str, topic: &str, group: &str, lag: i64) {
    if let Ok(metrics) = metrics() {
        metrics
            .retained_log_consumer_lag
            .with_label_values(&[queue, topic, group])
            .set(lag);
    }
}

/// Tracks one running partition head handler for the lifetime of this guard.
pub(crate) struct PartitionActiveGuard<'a> {
    metrics: Option<&'static Metrics>,
    queue: &'a str,
    topic: &'a str,
}

impl<'a> PartitionActiveGuard<'a> {
    pub(crate) fn start(queue: &'a str, topic: &'a str) -> Self {
        let metrics = metrics().ok();
        if let Some(metrics) = metrics {
            metrics
                .partitions_active
                .with_label_values(&[queue, topic])
                .inc();
        }
        Self {
            metrics,
            queue,
            topic,
        }
    }
}

impl Drop for PartitionActiveGuard<'_> {
    fn drop(&mut self) {
        if let Some(metrics) = self.metrics {
            metrics
                .partitions_active
                .with_label_values(&[self.queue, self.topic])
                .dec();
        }
    }
}

pub(crate) struct WorkerGuard<'a> {
    metrics: Option<&'static Metrics>,
    queue: &'a str,
    worker_id: &'a str,
}

impl<'a> WorkerGuard<'a> {
    pub(crate) fn new(queue: &'a str, worker_id: &'a str) -> Self {
        let metrics = metrics().ok();
        if let Some(metrics) = metrics {
            metrics
                .workers
                .with_label_values(&[queue, worker_id, "active"])
                .inc();
        }
        Self {
            metrics,
            queue,
            worker_id,
        }
    }
}

impl Drop for WorkerGuard<'_> {
    fn drop(&mut self) {
        if let Some(metrics) = self.metrics {
            metrics
                .workers
                .with_label_values(&[self.queue, self.worker_id, "active"])
                .dec();
        }
    }
}

pub(crate) struct HandlerTimer<'a> {
    metrics: Option<&'static Metrics>,
    queue: &'a str,
    worker_id: &'a str,
    job_type: &'a str,
    started: Instant,
}

impl<'a> HandlerTimer<'a> {
    pub(crate) fn start(queue: &'a str, worker_id: &'a str, job_type: &'a str) -> Self {
        let metrics = metrics().ok();
        if let Some(metrics) = metrics {
            metrics
                .workers
                .with_label_values(&[queue, worker_id, "busy"])
                .inc();
        }
        Self {
            metrics,
            queue,
            worker_id,
            job_type,
            started: Instant::now(),
        }
    }

    pub(crate) fn finish(self, outcome: HandlerOutcome) {
        if let Some(metrics) = self.metrics {
            metrics
                .handler_duration
                .with_label_values(&[self.queue, self.worker_id, self.job_type, outcome.label()])
                .observe(self.started.elapsed().as_secs_f64());
        }
    }
}

impl Drop for HandlerTimer<'_> {
    fn drop(&mut self) {
        if let Some(metrics) = self.metrics {
            metrics
                .workers
                .with_label_values(&[self.queue, self.worker_id, "busy"])
                .dec();
        }
    }
}

pub(crate) enum HandlerOutcome {
    Success,
    Error,
}

impl HandlerOutcome {
    fn label(&self) -> &'static str {
        match self {
            Self::Success => "success",
            Self::Error => "error",
        }
    }
}

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

    #[test]
    fn queue_depth_is_encoded_with_queue_and_state() -> anyhow::Result<()> {
        set_queue_depth("queue-depth-metric-test", "ready", 7);

        assert!(encode()?.lines().any(|line| {
            line == "later_queue_jobs{queue=\"queue-depth-metric-test\",state=\"ready\"} 7"
        }));
        Ok(())
    }

    #[test]
    fn partition_metrics_use_only_queue_and_topic_as_labels() -> anyhow::Result<()> {
        record_partition_claim("partition-metric-test", Some("orders"));
        record_partition_claim("partition-metric-test", None);
        record_partition_lease_renewal("partition-metric-test", "orders", true);
        record_partition_lease_renewal("partition-metric-test", "orders", false);
        set_partition_workers_live("partition-metric-test", 3);
        set_partition_backlog("partition-metric-test", "orders", 5, 12.5);
        let guard = PartitionActiveGuard::start("partition-metric-test", "orders");

        let body = encode()?;
        assert!(body.lines().any(|line| {
            line == "later_partition_claims_total{outcome=\"claimed\",queue=\"partition-metric-test\"} 1"
        }));
        assert!(body.lines().any(|line| {
            line == "later_partition_claims_total{outcome=\"none\",queue=\"partition-metric-test\"} 1"
        }));
        assert!(body.lines().any(|line| {
            line == "later_partition_heads_claimed_total{queue=\"partition-metric-test\",topic=\"orders\"} 1"
        }));
        assert!(body.lines().any(|line| {
            line == "later_partition_lease_renewals_total{outcome=\"renewed\",queue=\"partition-metric-test\",topic=\"orders\"} 1"
        }));
        assert!(body.lines().any(|line| {
            line == "later_partition_lease_renewals_total{outcome=\"lost\",queue=\"partition-metric-test\",topic=\"orders\"} 1"
        }));
        assert!(body.lines().any(|line| {
            line == "later_partition_workers_live{queue=\"partition-metric-test\"} 3"
        }));
        assert!(body.lines().any(|line| {
            line == "later_partitions_active{queue=\"partition-metric-test\",topic=\"orders\"} 1"
        }));
        assert!(body.lines().any(|line| {
            line == "later_partition_backlog{queue=\"partition-metric-test\",topic=\"orders\"} 5"
        }));
        assert!(body.lines().any(|line| {
            line == "later_partition_oldest_head_age_seconds{queue=\"partition-metric-test\",topic=\"orders\"} 12.5"
        }));

        drop(guard);
        let body = encode()?;
        assert!(body.lines().any(|line| {
            line == "later_partitions_active{queue=\"partition-metric-test\",topic=\"orders\"} 0"
        }));
        Ok(())
    }
}