Skip to main content

backbone_observability/
metrics.rs

1//! Metrics collection with Prometheus export
2//!
3//! Provides production-ready metrics collection with Prometheus HTTP endpoint.
4
5use std::net::SocketAddr;
6use std::sync::OnceLock;
7use thiserror::Error;
8use crate::config::MetricsConfig;
9use tokio::sync::oneshot;
10
11// Re-export http types used in middleware
12pub use http::{Method, StatusCode, Response, Request};
13
14// Prometheus registry (available when feature is enabled)
15#[cfg(feature = "prometheus-metrics")]
16pub use prometheus::Registry;
17
18/// Metrics initialization error
19#[derive(Debug, Error)]
20pub enum MetricsError {
21    /// Failed to create the metrics exporter
22    #[error("Failed to create exporter: {0}")]
23    ExporterFailed(String),
24
25    /// Failed to setup metrics (registration, configuration)
26    #[error("Metrics setup error: {0}")]
27    SetupError(String),
28
29    /// Invalid configuration provided
30    #[error("Invalid configuration: {0}")]
31    InvalidConfig(String),
32
33    /// IO error during metrics operation
34    #[error("IO error: {0}")]
35    IoError(#[from] std::io::Error),
36}
37
38/// Metrics handle - keeps the metrics server alive
39pub struct MetricsHandle {
40    /// Shutdown channel for the metrics server
41    shutdown_tx: Option<oneshot::Sender<()>>,
42    /// The address where metrics are served
43    pub metrics_addr: SocketAddr,
44}
45
46impl MetricsHandle {
47    /// Shutdown the metrics server gracefully
48    pub async fn shutdown(mut self) {
49        if let Some(tx) = self.shutdown_tx.take() {
50            let _ = tx.send(());
51        }
52    }
53}
54
55/// Initialize metrics collection with Prometheus exporter
56///
57/// # Arguments
58/// * `config` - Metrics configuration
59///
60/// # Returns
61/// Returns a handle that keeps the metrics server running
62///
63/// # Example
64/// ```no_run
65/// use backbone_observability::{MetricsConfig, MetricsExporterType, init_metrics};
66///
67/// # #[tokio::main]
68/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
69/// let config = MetricsConfig {
70///     enabled: true,
71///     exporter: MetricsExporterType::Prometheus,
72///     port: 9090,
73/// };
74///
75/// let metrics_handle = init_metrics(&config)?;
76///
77/// // Metrics are now available at http://localhost:9090/metrics
78///
79/// // Keep the handle alive while your app runs
80/// // tokio::select! {
81/// //     _ = shutdown_signal() => {}
82/// //     _ = metrics_handle.shutdown() => {}
83/// // }
84/// # Ok(())
85/// # }
86/// ```
87pub fn init_metrics(config: &MetricsConfig) -> Result<MetricsHandle, MetricsError> {
88    const DEFAULT_METRICS_ADDR: &str = "0.0.0.0:0";
89
90    if !config.enabled {
91        return Ok(MetricsHandle {
92            shutdown_tx: None,
93            metrics_addr: DEFAULT_METRICS_ADDR
94                .parse()
95                .map_err(|e| MetricsError::InvalidConfig(format!("Invalid default metrics address: {}", e)))?,
96        });
97    }
98
99    match config.exporter {
100        crate::config::MetricsExporterType::Prometheus => {
101            init_prometheus(config)
102        }
103        crate::config::MetricsExporterType::Stdout => {
104            tracing::info!("Metrics exporter: stdout (logging via tracing)");
105            Ok(MetricsHandle {
106                shutdown_tx: None,
107                metrics_addr: "0.0.0.0:0".parse().unwrap(),
108            })
109        }
110    }
111}
112
113/// Initialize Prometheus exporter with HTTP server
114#[cfg(feature = "prometheus-metrics")]
115fn init_prometheus(config: &MetricsConfig) -> Result<MetricsHandle, MetricsError> {
116    use std::net::SocketAddr;
117
118    // Serve the GLOBAL default registry. The HTTP observability middleware
119    // (`ObservabilityLayer` → `record_http_metrics`) registers its counters /
120    // histograms on `prometheus::default_registry()` via `prometheus::register`,
121    // so the metrics server MUST serve that same registry or the middleware's
122    // observations land in a registry nothing scrapes (every HTTP metric stays
123    // flat at zero). Registry is Clone (cheap — inner Arc), so clone the static.
124    let registry = prometheus::default_registry().clone();
125
126    // Create default metrics
127    setup_default_metrics(&registry)?;
128
129    // Get the port to bind to
130    let port = config.port;
131
132    // Create shutdown channel
133    let (shutdown_tx, shutdown_rx) = oneshot::channel();
134
135    // Spawn the metrics server task
136    tokio::spawn(async move {
137        metrics_server_task(registry, port, shutdown_rx).await;
138    });
139
140    // Return immediately with the handle
141    // The actual address will be set by the server task
142    Ok(MetricsHandle {
143        shutdown_tx: Some(shutdown_tx),
144        metrics_addr: SocketAddr::from(([127, 0, 0, 1], port)),
145    })
146}
147
148/// Metrics server task
149#[cfg(feature = "prometheus-metrics")]
150async fn metrics_server_task(
151    registry: Registry,
152    port: u16,
153    shutdown_rx: oneshot::Receiver<()>,
154) {
155    use axum::{routing::get, Router};
156    use tokio::net::TcpListener;
157
158    // Bind to the configured port
159    let addr = SocketAddr::from(([0, 0, 0, 0], port));
160
161    match TcpListener::bind(addr).await {
162        Ok(listener) => {
163            let actual_addr = listener.local_addr().unwrap_or(addr);
164
165            tracing::info!("Prometheus metrics server listening on http://{}", actual_addr);
166            tracing::info!("Metrics endpoint: http://{}/metrics", actual_addr);
167
168            // Build the metrics router
169            let app = Router::new()
170                .route("/metrics", get(metrics_handler))
171                .route("/health", get(health_handler))
172                .with_state(registry);
173
174            // Run the server
175            let _ = axum::serve(listener, app)
176                .with_graceful_shutdown(async {
177                    shutdown_rx.await.ok();
178                    tracing::info!("Metrics server shutting down gracefully");
179                })
180                .await
181                .map_err(|e| tracing::error!("Metrics server error: {}", e));
182        }
183        Err(e) => {
184            tracing::error!("Failed to bind metrics server: {}", e);
185        }
186    }
187}
188
189/// Setup default Prometheus metrics
190///
191/// Registers default metrics to the provided registry for HTTP request tracking,
192/// database query monitoring, and connection state.
193#[cfg(feature = "prometheus-metrics")]
194fn setup_default_metrics(registry: &Registry) -> Result<(), MetricsError> {
195    use prometheus::{
196        IntCounterVec, Histogram, IntGaugeVec,
197        Opts, HistogramOpts,
198    };
199
200    // NOTE: the LIVE HTTP metrics (`http_requests_total` counter and
201    // `http_request_duration_seconds` histogram) are NOT registered here.
202    // They are owned by the observability middleware (`record_http_metrics`),
203    // which registers + observes them on `prometheus::default_registry()`.
204    // Registering dead copies here would (a) never be observed and (b) conflict
205    // with the middleware's live registrations on the now-shared global registry.
206    // This function registers only the db / connection metrics below.
207
208    // Database query counter
209    let _db_queries_total = IntCounterVec::new(
210        Opts::new("db_queries_total", "Total number of database queries"),
211        &["operation", "table"]
212    )
213    .map_err(|e| MetricsError::SetupError(format!("Failed to create db_queries_total: {}", e)))?;
214    registry.register(Box::new(_db_queries_total.clone()))
215        .map_err(|e| MetricsError::SetupError(format!("Failed to register db_queries_total: {}", e)))?;
216
217    // Database query duration histogram
218    let _db_query_duration_seconds = Histogram::with_opts(
219        HistogramOpts::new("db_query_duration_seconds", "Database query latency in seconds")
220            .buckets(vec![0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0])
221            .into(),
222    )
223    .map_err(|e| MetricsError::SetupError(format!("Failed to create db_query_duration_seconds: {}", e)))?;
224    registry.register(Box::new(_db_query_duration_seconds.clone()))
225        .map_err(|e| MetricsError::SetupError(format!("Failed to register db_query_duration_seconds: {}", e)))?;
226
227    // Active connections gauge
228    let _http_active_connections = IntGaugeVec::new(
229        Opts::new("http_active_connections", "Number of active HTTP connections"),
230        &["state"]
231    )
232    .map_err(|e| MetricsError::SetupError(format!("Failed to create http_active_connections: {}", e)))?;
233    registry.register(Box::new(_http_active_connections.clone()))
234        .map_err(|e| MetricsError::SetupError(format!("Failed to register http_active_connections: {}", e)))?;
235
236    tracing::debug!("Default Prometheus metrics registered successfully");
237
238    Ok(())
239}
240
241/// Metrics endpoint handler
242#[cfg(feature = "prometheus-metrics")]
243async fn metrics_handler(
244    axum::extract::State(registry): axum::extract::State<Registry>,
245) -> Result<axum::response::Response, String> {
246    use prometheus::{Encoder, TextEncoder};
247    use http::header;
248
249    let metrics_families = registry.gather();
250    let encoder = TextEncoder::new();
251    let mut buffer = Vec::new();
252
253    encoder
254        .encode(&metrics_families, &mut buffer)
255        .map_err(|e| format!("Failed to encode metrics: {}", e))?;
256
257    Ok(axum::response::Response::builder()
258        .status(axum::http::StatusCode::OK)
259        .header(header::CONTENT_TYPE, encoder.format_type())
260        .body(axum::body::Body::from(buffer))
261        .unwrap())
262}
263
264/// Health check endpoint for metrics server
265#[cfg(feature = "prometheus-metrics")]
266async fn health_handler() -> &'static str {
267    "OK"
268}
269
270/// Stub for when prometheus-metrics feature is not enabled
271#[cfg(not(feature = "prometheus-metrics"))]
272fn init_prometheus(_config: &MetricsConfig) -> Result<MetricsHandle, MetricsError> {
273    tracing::warn!("Prometheus metrics requested but 'prometheus-metrics' feature is not enabled");
274    tracing::warn!("Enable with: cargo build --features prometheus-metrics");
275    Ok(MetricsHandle {
276        shutdown_tx: None,
277        metrics_addr: "0.0.0.0:0".parse().unwrap(),
278    })
279}
280
281// ==============================================================================
282// DATABASE POOL METRICS
283// ==============================================================================
284
285/// Database connection pool statistics
286///
287/// Callers (e.g., backbone-orm) extract pool stats from their connection pool
288/// implementation (e.g., `sqlx::PgPool`) and pass them here for Prometheus recording.
289/// This avoids a hard dependency on sqlx in the observability crate.
290#[derive(Debug, Clone, Copy)]
291pub struct PoolStats {
292    /// Total number of connections in the pool (active + idle)
293    pub total: u32,
294    /// Number of connections currently in use
295    pub active: u32,
296    /// Number of idle connections available
297    pub idle: u32,
298    /// Maximum number of connections configured for the pool
299    pub max: u32,
300}
301
302/// Global DB pool gauges (lazy-initialized on first use)
303#[cfg(feature = "prometheus-metrics")]
304static DB_POOL_TOTAL: OnceLock<prometheus::IntGaugeVec> = OnceLock::new();
305#[cfg(feature = "prometheus-metrics")]
306static DB_POOL_ACTIVE: OnceLock<prometheus::IntGaugeVec> = OnceLock::new();
307#[cfg(feature = "prometheus-metrics")]
308static DB_POOL_IDLE: OnceLock<prometheus::IntGaugeVec> = OnceLock::new();
309#[cfg(feature = "prometheus-metrics")]
310static DB_POOL_MAX: OnceLock<prometheus::IntGaugeVec> = OnceLock::new();
311
312/// Record database pool metrics for the default pool
313///
314/// Updates Prometheus gauges for database connection pool statistics.
315/// Should be called periodically (e.g., every 15 seconds) by the ORM layer.
316#[cfg(feature = "prometheus-metrics")]
317pub fn record_pool_metrics(stats: &PoolStats) {
318    record_pool_metrics_named("default", stats);
319}
320
321/// Record database pool metrics for a named pool
322///
323/// Allows tracking metrics for multiple database pools (e.g., "primary", "replica").
324#[cfg(feature = "prometheus-metrics")]
325pub fn record_pool_metrics_named(pool_name: &str, stats: &PoolStats) {
326    use prometheus::{IntGaugeVec, Opts};
327
328    let total = DB_POOL_TOTAL.get_or_init(|| {
329        let metric = IntGaugeVec::new(
330            Opts::new(common::DB_POOL_CONNECTIONS_TOTAL, "Total connections in database pool"),
331            &[common::LABEL_POOL],
332        ).expect("Failed to create db_pool_connections_total metric");
333        let _ = prometheus::register(Box::new(metric.clone()));
334        metric
335    });
336
337    let active = DB_POOL_ACTIVE.get_or_init(|| {
338        let metric = IntGaugeVec::new(
339            Opts::new(common::DB_POOL_CONNECTIONS_ACTIVE, "Active connections in database pool"),
340            &[common::LABEL_POOL],
341        ).expect("Failed to create db_pool_connections_active metric");
342        let _ = prometheus::register(Box::new(metric.clone()));
343        metric
344    });
345
346    let idle = DB_POOL_IDLE.get_or_init(|| {
347        let metric = IntGaugeVec::new(
348            Opts::new(common::DB_POOL_CONNECTIONS_IDLE, "Idle connections in database pool"),
349            &[common::LABEL_POOL],
350        ).expect("Failed to create db_pool_connections_idle metric");
351        let _ = prometheus::register(Box::new(metric.clone()));
352        metric
353    });
354
355    let max = DB_POOL_MAX.get_or_init(|| {
356        let metric = IntGaugeVec::new(
357            Opts::new(common::DB_POOL_MAX_CONNECTIONS, "Max configured connections for database pool"),
358            &[common::LABEL_POOL],
359        ).expect("Failed to create db_pool_max_connections metric");
360        let _ = prometheus::register(Box::new(metric.clone()));
361        metric
362    });
363
364    total.with_label_values(&[pool_name]).set(stats.total as i64);
365    active.with_label_values(&[pool_name]).set(stats.active as i64);
366    idle.with_label_values(&[pool_name]).set(stats.idle as i64);
367    max.with_label_values(&[pool_name]).set(stats.max as i64);
368}
369
370/// No-op stub when prometheus-metrics feature is disabled
371#[cfg(not(feature = "prometheus-metrics"))]
372pub fn record_pool_metrics(_stats: &PoolStats) {}
373
374/// No-op stub when prometheus-metrics feature is disabled
375#[cfg(not(feature = "prometheus-metrics"))]
376pub fn record_pool_metrics_named(_pool_name: &str, _stats: &PoolStats) {}
377
378// Common metrics definitions
379///
380/// Common metric names and label names used throughout the observability system.
381pub mod common {
382    /// Request counter metric name
383    pub const HTTP_REQUESTS_TOTAL: &str = "http_requests_total";
384
385    /// Request duration metric name
386    pub const HTTP_REQUEST_DURATION_SECONDS: &str = "http_request_duration_seconds";
387
388    /// Active connections metric name
389    pub const HTTP_ACTIVE_CONNECTIONS: &str = "http_active_connections";
390
391    /// Database query counter metric name
392    pub const DB_QUERIES_TOTAL: &str = "db_queries_total";
393
394    /// Database query duration metric name
395    pub const DB_QUERY_DURATION_SECONDS: &str = "db_query_duration_seconds";
396
397    /// Metric labels for HTTP requests
398    pub const LABEL_METHOD: &str = "method";
399    /// Label for the HTTP method (GET, POST, etc.)
400    pub const LABEL_PATH: &str = "path";
401    /// Label for the HTTP request path
402    pub const LABEL_STATUS: &str = "status";
403    /// Label for the HTTP status code
404
405    /// Metric labels for DB queries
406    pub const LABEL_OPERATION: &str = "operation";
407    /// Label for the database operation (SELECT, INSERT, etc.)
408    pub const LABEL_TABLE: &str = "table";
409
410    /// Metric labels for connection state
411    pub const LABEL_STATE: &str = "state";
412
413    /// Database pool total connections metric name
414    pub const DB_POOL_CONNECTIONS_TOTAL: &str = "db_pool_connections_total";
415    /// Database pool active connections metric name
416    pub const DB_POOL_CONNECTIONS_ACTIVE: &str = "db_pool_connections_active";
417    /// Database pool idle connections metric name
418    pub const DB_POOL_CONNECTIONS_IDLE: &str = "db_pool_connections_idle";
419    /// Database pool max connections metric name
420    pub const DB_POOL_MAX_CONNECTIONS: &str = "db_pool_max_connections";
421    /// Label for pool name (supports multiple pools)
422    pub const LABEL_POOL: &str = "pool";
423}