diff --git a/src/metrics.rs b/src/metrics.rs index 6af6e81..2e36471 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -102,18 +102,56 @@ impl Metrics { self.prom_request_duration.with_label_values(&[method, &status.to_string()]).observe(duration); self.prom_requests_in_flight.with_label_values(&[method]).dec(); } + + /// Returns a vector of metric families from the Prometheus registry + pub fn get_metrics(&self) -> Vec { + self.registry.gather() + } + + /// Returns an iterator over metrics with helper methods to find specific metrics + pub fn metrics_iter(&self) -> MetricsIterator { + MetricsIterator { + metrics: self.get_metrics() + } + } } -async fn metrics_handler(metrics: Arc) -> std::result::Result, Infallible> { - let encoder = TextEncoder::new(); - let metric_families = metrics.registry.gather(); - let mut buffer = Vec::new(); - encoder.encode(&metric_families, &mut buffer).unwrap(); +/// Helper struct to iterate and find metrics easily +pub struct MetricsIterator { + metrics: Vec +} - Ok(Response::builder() - .header("Content-Type", "text/plain") - .body(Body::from(buffer)) - .unwrap()) +impl MetricsIterator { + /// Find a metric by name + pub fn find_metric(&self, name: &str) -> Option<&prometheus::proto::MetricFamily> { + self.metrics.iter().find(|m| m.get_name() == name) + } + + /// Get all metrics + pub fn all(&self) -> &[prometheus::proto::MetricFamily] { + &self.metrics + } +} + +async fn metrics_handler(req: Request, metrics: Arc) -> std::result::Result, Infallible> { + // Only respond to /metrics path + match req.uri().path() { + "/metrics" => { + let encoder = TextEncoder::new(); + let metric_families = metrics.registry.gather(); + let mut buffer = Vec::new(); + encoder.encode(&metric_families, &mut buffer).unwrap(); + + Ok(Response::builder() + .header("Content-Type", "text/plain") + .body(Body::from(buffer)) + .unwrap()) + } + _ => Ok(Response::builder() + .status(404) + .body(Body::from("Not Found")) + .unwrap()) + } } pub async fn run_metrics_server(metrics: Arc, addr: SocketAddr) -> std::result::Result<(), Box> { @@ -121,8 +159,8 @@ pub async fn run_metrics_server(metrics: Arc, addr: SocketAddr) -> std: let make_svc = make_service_fn(move |_conn| { let metrics = metrics.clone(); async move { - Ok::<_, Infallible>(service_fn(move |_req: Request| { - metrics_handler(metrics.clone()) + Ok::<_, Infallible>(service_fn(move |req: Request| { + metrics_handler(req, metrics.clone()) })) } }); diff --git a/tests/integration_test.rs b/tests/integration_test.rs index 724b244..7dfedb8 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -65,19 +65,22 @@ async fn test_server_basic_functionality() -> Result<(), Box Result<(), Box Result<(), Box