llm_optimizer_processor/metrics/
server.rs

1//! HTTP server for Prometheus metrics endpoint
2//!
3//! This module provides an HTTP server that exposes metrics at /metrics
4//! along with health check and readiness probe endpoints.
5
6use super::registry::MetricsRegistry;
7use super::MetricsError;
8use axum::{
9    extract::State,
10    http::StatusCode,
11    response::{IntoResponse, Response},
12    routing::get,
13    Json, Router,
14};
15use serde::{Deserialize, Serialize};
16use std::net::SocketAddr;
17use std::sync::Arc;
18use std::time::{Duration, Instant};
19use tokio::net::TcpListener;
20use tracing::{debug, error, info};
21
22/// Configuration for the metrics server
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct MetricsServerConfig {
25    /// Bind address for the metrics server
26    pub bind_address: String,
27
28    /// Port for the metrics server
29    pub port: u16,
30}
31
32impl Default for MetricsServerConfig {
33    fn default() -> Self {
34        Self {
35            bind_address: "0.0.0.0".to_string(),
36            port: 9090,
37        }
38    }
39}
40
41impl MetricsServerConfig {
42    /// Create a new metrics server config
43    pub fn new(bind_address: impl Into<String>, port: u16) -> Self {
44        Self {
45            bind_address: bind_address.into(),
46            port,
47        }
48    }
49
50    /// Get the socket address
51    pub fn socket_addr(&self) -> Result<SocketAddr, MetricsError> {
52        format!("{}:{}", self.bind_address, self.port)
53            .parse()
54            .map_err(|e| MetricsError::BindError {
55                address: format!("{}:{}", self.bind_address, self.port),
56                source: std::io::Error::new(std::io::ErrorKind::InvalidInput, e),
57            })
58    }
59}
60
61/// Health status response
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct HealthStatus {
64    pub status: String,
65    pub version: String,
66    pub uptime_seconds: u64,
67}
68
69impl HealthStatus {
70    /// Create a healthy status
71    pub fn healthy() -> Self {
72        Self {
73            status: "healthy".to_string(),
74            version: env!("CARGO_PKG_VERSION").to_string(),
75            uptime_seconds: 0,
76        }
77    }
78
79    /// Create a healthy status with uptime
80    pub fn healthy_with_uptime(uptime: Duration) -> Self {
81        Self {
82            status: "healthy".to_string(),
83            version: env!("CARGO_PKG_VERSION").to_string(),
84            uptime_seconds: uptime.as_secs(),
85        }
86    }
87}
88
89/// Readiness status response
90#[derive(Debug, Clone, Serialize, Deserialize)]
91pub struct ReadinessStatus {
92    pub ready: bool,
93    pub checks: Vec<ReadinessCheck>,
94}
95
96#[derive(Debug, Clone, Serialize, Deserialize)]
97pub struct ReadinessCheck {
98    pub name: String,
99    pub ready: bool,
100    pub message: Option<String>,
101}
102
103impl ReadinessStatus {
104    /// Create a ready status
105    pub fn ready() -> Self {
106        Self {
107            ready: true,
108            checks: vec![ReadinessCheck {
109                name: "processor".to_string(),
110                ready: true,
111                message: None,
112            }],
113        }
114    }
115
116    /// Create a not ready status
117    pub fn not_ready(message: impl Into<String>) -> Self {
118        Self {
119            ready: false,
120            checks: vec![ReadinessCheck {
121                name: "processor".to_string(),
122                ready: false,
123                message: Some(message.into()),
124            }],
125        }
126    }
127}
128
129/// Shared state for the metrics server
130#[derive(Clone)]
131struct ServerState {
132    registry: Arc<MetricsRegistry>,
133    start_time: Instant,
134}
135
136/// HTTP server for exposing Prometheus metrics
137pub struct MetricsServer {
138    config: MetricsServerConfig,
139    registry: Arc<MetricsRegistry>,
140}
141
142impl MetricsServer {
143    /// Create a new metrics server
144    pub fn new(config: MetricsServerConfig) -> Self {
145        Self {
146            config,
147            registry: MetricsRegistry::global(),
148        }
149    }
150
151    /// Create a metrics server with custom registry
152    pub fn with_registry(config: MetricsServerConfig, registry: Arc<MetricsRegistry>) -> Self {
153        Self { config, registry }
154    }
155
156    /// Start the metrics server
157    pub async fn start(self) -> Result<(), MetricsError> {
158        let addr = self.config.socket_addr()?;
159
160        info!(
161            "Starting metrics server on http://{}:{}",
162            self.config.bind_address, self.config.port
163        );
164
165        let state = ServerState {
166            registry: self.registry,
167            start_time: Instant::now(),
168        };
169
170        let app = Router::new()
171            .route("/metrics", get(metrics_handler))
172            .route("/health", get(health_handler))
173            .route("/ready", get(ready_handler))
174            .with_state(state);
175
176        let listener = TcpListener::bind(addr)
177            .await
178            .map_err(|e| MetricsError::BindError {
179                address: format!("{}:{}", self.config.bind_address, self.config.port),
180                source: e,
181            })?;
182
183        info!(
184            "Metrics server listening on http://{}",
185            listener.local_addr().map_err(|e| MetricsError::BindError {
186                address: format!("{}:{}", self.config.bind_address, self.config.port),
187                source: e,
188            })?
189        );
190
191        axum::serve(listener, app)
192            .await
193            .map_err(|e| MetricsError::ServerStartError(e.to_string()))?;
194
195        Ok(())
196    }
197
198    /// Get the metrics server configuration
199    pub fn config(&self) -> &MetricsServerConfig {
200        &self.config
201    }
202}
203
204/// Handler for /metrics endpoint
205async fn metrics_handler(State(state): State<ServerState>) -> Response {
206    debug!("Metrics endpoint called");
207
208    match state.registry.encode() {
209        Ok(metrics) => {
210            debug!("Successfully encoded metrics");
211            (
212                StatusCode::OK,
213                [("content-type", "text/plain; version=0.0.4")],
214                metrics,
215            )
216                .into_response()
217        }
218        Err(e) => {
219            error!("Failed to encode metrics: {}", e);
220            (
221                StatusCode::INTERNAL_SERVER_ERROR,
222                format!("Failed to encode metrics: {}", e),
223            )
224                .into_response()
225        }
226    }
227}
228
229/// Handler for /health endpoint
230async fn health_handler(State(state): State<ServerState>) -> Json<HealthStatus> {
231    debug!("Health check endpoint called");
232    let uptime = state.start_time.elapsed();
233    Json(HealthStatus::healthy_with_uptime(uptime))
234}
235
236/// Handler for /ready endpoint
237async fn ready_handler() -> Response {
238    debug!("Readiness probe endpoint called");
239    let status = ReadinessStatus::ready();
240
241    if status.ready {
242        (StatusCode::OK, Json(status)).into_response()
243    } else {
244        (StatusCode::SERVICE_UNAVAILABLE, Json(status)).into_response()
245    }
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251
252    #[test]
253    fn test_config_default() {
254        let config = MetricsServerConfig::default();
255        assert_eq!(config.bind_address, "0.0.0.0");
256        assert_eq!(config.port, 9090);
257    }
258
259    #[test]
260    fn test_config_custom() {
261        let config = MetricsServerConfig::new("127.0.0.1", 8080);
262        assert_eq!(config.bind_address, "127.0.0.1");
263        assert_eq!(config.port, 8080);
264    }
265
266    #[test]
267    fn test_socket_addr() {
268        let config = MetricsServerConfig::new("127.0.0.1", 8080);
269        let addr = config.socket_addr().unwrap();
270        assert_eq!(addr.port(), 8080);
271    }
272
273    #[test]
274    fn test_health_status() {
275        let status = HealthStatus::healthy();
276        assert_eq!(status.status, "healthy");
277        assert_eq!(status.version, env!("CARGO_PKG_VERSION"));
278    }
279
280    #[test]
281    fn test_readiness_status() {
282        let ready = ReadinessStatus::ready();
283        assert!(ready.ready);
284        assert_eq!(ready.checks.len(), 1);
285
286        let not_ready = ReadinessStatus::not_ready("Not initialized");
287        assert!(!not_ready.ready);
288        assert!(not_ready.checks[0].message.is_some());
289    }
290
291    #[tokio::test]
292    async fn test_metrics_server_creation() {
293        let config = MetricsServerConfig::new("127.0.0.1", 9091);
294        let server = MetricsServer::new(config);
295        assert_eq!(server.config().port, 9091);
296    }
297}
llm_optimizer_processor/metrics/server.rs

llm_optimizer_processor/metrics/
server.rs
