at-jet 0.7.2

//! Prometheus metrics exporter setup
//!
//! Provides unified metrics initialization and a Prometheus scrape endpoint.
//!
//! # Example
//!
//! ```ignore
//! use at_jet::metrics::{init_metrics, metrics_router, MetricsConfig};
//! use std::sync::Arc;
//!
//! let config = MetricsConfig {
//!     enabled: true,
//!     ..Default::default()
//! };
//!
//! if let Some(guard) = init_metrics(&config) {
//!     let router = metrics_router(Arc::new(guard), "/metrics");
//!     // merge into your JetServer
//! }
//! ```

use {axum::{Router,
            extract::State,
            response::IntoResponse,
            routing::get},
     metrics::gauge,
     metrics_exporter_prometheus::{Matcher,
                                   PrometheusBuilder,
                                   PrometheusHandle},
     std::sync::Arc,
     tracing::{error,
               info}};

/// Metrics configuration
#[derive(Debug, Clone)]
pub struct MetricsConfig {
  /// Enable metrics collection
  pub enabled:        bool,
  /// Histogram buckets for HTTP request duration (in seconds)
  pub http_buckets:   Vec<f64>,
  /// Application-specific histogram bucket overrides.
  ///
  /// Each entry is `(metric_suffix, buckets)` and is matched via
  /// `Matcher::Suffix`. For example:
  ///
  /// ```ignore
  /// vec![
  ///   ("db_query_duration_seconds".into(), vec![0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1.0]),
  ///   ("cache_operation_duration_seconds".into(), vec![0.0005, 0.001, 0.005, 0.01, 0.05]),
  /// ]
  /// ```
  pub custom_buckets: Vec<(String, Vec<f64>)>,
}

impl Default for MetricsConfig {
  fn default() -> Self {
    Self {
      enabled:        true,
      http_buckets:   vec![0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0],
      custom_buckets: Vec::new(),
    }
  }
}

/// Guard that holds the Prometheus handle.
///
/// Returned from [`init_metrics`] for use with [`metrics_router`].
pub struct MetricsGuard {
  handle: PrometheusHandle,
}

impl MetricsGuard {
  /// Get the Prometheus handle for rendering metrics
  pub fn handle(&self) -> &PrometheusHandle {
    &self.handle
  }

  /// Render current metrics as Prometheus text format
  pub fn render(&self) -> String {
    self.handle.render()
  }
}

/// Initialize Prometheus metrics exporter.
///
/// Returns a [`MetricsGuard`] if metrics are enabled, `None` otherwise.
/// The guard contains the [`PrometheusHandle`] needed for serving metrics.
pub fn init_metrics(config: &MetricsConfig) -> Option<MetricsGuard> {
  if !config.enabled {
    info!("Metrics collection disabled");
    return None;
  }

  // Build the Prometheus exporter with custom histogram buckets
  let mut builder = PrometheusBuilder::new()
    .set_buckets_for_metric(
      Matcher::Suffix("http_request_duration_seconds".to_string()),
      &config.http_buckets,
    )
    .ok()?;

  // Apply application-specific bucket overrides
  for (suffix, buckets) in &config.custom_buckets {
    builder = builder
      .set_buckets_for_metric(Matcher::Suffix(suffix.clone()), buckets)
      .ok()?;
  }

  match builder.install_recorder() {
    | Ok(handle) => {
      info!("Metrics exporter initialized");
      Some(MetricsGuard { handle })
    }
    | Err(e) => {
      error!(error = %e, "Failed to install Prometheus recorder");
      None
    }
  }
}

/// Create an Axum router that serves Prometheus metrics at the given path.
///
/// Merge this router into your service's existing HTTP server so that
/// metrics are served on the same port as the service itself.
///
/// # Arguments
///
/// * `guard` - [`MetricsGuard`] from [`init_metrics`], wrapped in `Arc` for sharing
/// * `path` - The endpoint path, e.g. `"/metrics"` or `"/my-service/metrics"`
pub fn metrics_router(guard: Arc<MetricsGuard>, path: &str) -> Router {
  Router::new().route(path, get(metrics_handler)).with_state(guard)
}

async fn metrics_handler(State(guard): State<Arc<MetricsGuard>>) -> impl IntoResponse {
  (
    [("content-type", "text/plain; version=0.0.4; charset=utf-8")],
    guard.render(),
  )
}

/// Record application info as a gauge metric.
///
/// Creates a gauge `app_info` with value 1 and labels for service name and version.
/// This is useful for tracking which versions are deployed.
///
/// # Example
///
/// ```ignore
/// use at_jet::metrics::record_app_info;
///
/// record_app_info("my-service", "1.0.0");
/// ```
pub fn record_app_info(service: &'static str, version: &'static str) {
  gauge!("app_info", "service" => service, "version" => version).set(1.0);
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
  use super::*;

  #[test]
  fn test_default_config() {
    let config = MetricsConfig::default();
    assert!(config.enabled);
    assert!(!config.http_buckets.is_empty());
    assert!(config.custom_buckets.is_empty());
  }

  #[test]
  fn test_disabled_metrics() {
    let config = MetricsConfig {
      enabled: false,
      ..Default::default()
    };
    let guard = init_metrics(&config);
    assert!(guard.is_none());
  }
}