1use std::net::SocketAddr;
6use std::sync::OnceLock;
7use thiserror::Error;
8use crate::config::MetricsConfig;
9use tokio::sync::oneshot;
10
11pub use http::{Method, StatusCode, Response, Request};
13
14#[cfg(feature = "prometheus-metrics")]
16pub use prometheus::Registry;
17
18#[derive(Debug, Error)]
20pub enum MetricsError {
21 #[error("Failed to create exporter: {0}")]
23 ExporterFailed(String),
24
25 #[error("Metrics setup error: {0}")]
27 SetupError(String),
28
29 #[error("Invalid configuration: {0}")]
31 InvalidConfig(String),
32
33 #[error("IO error: {0}")]
35 IoError(#[from] std::io::Error),
36}
37
38pub struct MetricsHandle {
40 shutdown_tx: Option<oneshot::Sender<()>>,
42 pub metrics_addr: SocketAddr,
44}
45
46impl MetricsHandle {
47 pub async fn shutdown(mut self) {
49 if let Some(tx) = self.shutdown_tx.take() {
50 let _ = tx.send(());
51 }
52 }
53}
54
55pub 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#[cfg(feature = "prometheus-metrics")]
115fn init_prometheus(config: &MetricsConfig) -> Result<MetricsHandle, MetricsError> {
116 use std::net::SocketAddr;
117
118 let registry = prometheus::default_registry().clone();
125
126 setup_default_metrics(®istry)?;
128
129 let port = config.port;
131
132 let (shutdown_tx, shutdown_rx) = oneshot::channel();
134
135 tokio::spawn(async move {
137 metrics_server_task(registry, port, shutdown_rx).await;
138 });
139
140 Ok(MetricsHandle {
143 shutdown_tx: Some(shutdown_tx),
144 metrics_addr: SocketAddr::from(([127, 0, 0, 1], port)),
145 })
146}
147
148#[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 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 let app = Router::new()
170 .route("/metrics", get(metrics_handler))
171 .route("/health", get(health_handler))
172 .with_state(registry);
173
174 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#[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 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 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 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#[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#[cfg(feature = "prometheus-metrics")]
266async fn health_handler() -> &'static str {
267 "OK"
268}
269
270#[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#[derive(Debug, Clone, Copy)]
291pub struct PoolStats {
292 pub total: u32,
294 pub active: u32,
296 pub idle: u32,
298 pub max: u32,
300}
301
302#[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#[cfg(feature = "prometheus-metrics")]
317pub fn record_pool_metrics(stats: &PoolStats) {
318 record_pool_metrics_named("default", stats);
319}
320
321#[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#[cfg(not(feature = "prometheus-metrics"))]
372pub fn record_pool_metrics(_stats: &PoolStats) {}
373
374#[cfg(not(feature = "prometheus-metrics"))]
376pub fn record_pool_metrics_named(_pool_name: &str, _stats: &PoolStats) {}
377
378pub mod common {
382 pub const HTTP_REQUESTS_TOTAL: &str = "http_requests_total";
384
385 pub const HTTP_REQUEST_DURATION_SECONDS: &str = "http_request_duration_seconds";
387
388 pub const HTTP_ACTIVE_CONNECTIONS: &str = "http_active_connections";
390
391 pub const DB_QUERIES_TOTAL: &str = "db_queries_total";
393
394 pub const DB_QUERY_DURATION_SECONDS: &str = "db_query_duration_seconds";
396
397 pub const LABEL_METHOD: &str = "method";
399 pub const LABEL_PATH: &str = "path";
401 pub const LABEL_STATUS: &str = "status";
403 pub const LABEL_OPERATION: &str = "operation";
407 pub const LABEL_TABLE: &str = "table";
409
410 pub const LABEL_STATE: &str = "state";
412
413 pub const DB_POOL_CONNECTIONS_TOTAL: &str = "db_pool_connections_total";
415 pub const DB_POOL_CONNECTIONS_ACTIVE: &str = "db_pool_connections_active";
417 pub const DB_POOL_CONNECTIONS_IDLE: &str = "db_pool_connections_idle";
419 pub const DB_POOL_MAX_CONNECTIONS: &str = "db_pool_max_connections";
421 pub const LABEL_POOL: &str = "pool";
423}