trillium_prometheus/
lib.rs

1//! This is a small utility crate that provides a Prometheus metrics endpoint as a Trillium handler.
2//! It responds to GET requests to "/metrics" with metrics from the provided registry, using
3//! text-format encoding.
4//!
5//! # Example:
6//!
7//! ```
8//! # let stopper = trillium_smol::Stopper::new();
9//! # stopper.stop();
10//! let registry = prometheus::Registry::new();
11//! let handler = trillium_prometheus::text_format_handler(registry);
12//! trillium_smol::config()
13//!     .with_host("0.0.0.0")
14//!     .with_port(9464)
15//! #   .with_stopper(stopper)
16//!     .run(handler);
17//! ```
18use prometheus::{Encoder, Registry, TextEncoder};
19use tracing::error;
20use trillium::{KnownHeaderName, Status};
21use trillium_router::Router;
22
23/// Creates a handler that responds to GET requests for "/metrics".
24pub fn text_format_handler(registry: Registry) -> Router {
25    Router::new().get("metrics", move |conn: trillium::Conn| {
26        let registry = registry.clone();
27        async move {
28            let mut buffer = Vec::new();
29            let encoder = TextEncoder::new();
30            match encoder.encode(&registry.gather(), &mut buffer) {
31                Ok(()) => conn
32                    .with_response_header(
33                        KnownHeaderName::ContentType,
34                        encoder.format_type().to_owned(),
35                    )
36                    .ok(buffer),
37                Err(error) => {
38                    error!(%error, "Failed to encode Prometheus metrics");
39                    conn.with_status(Status::InternalServerError)
40                }
41            }
42        }
43    })
44}
45
46#[cfg(test)]
47mod tests {
48    use prometheus::{IntGauge, Registry};
49    use trillium_testing::{assert_response, prelude::get};
50
51    use crate::text_format_handler;
52
53    #[test]
54    fn text_format_encode() {
55        let registry = Registry::new();
56        let gauge = IntGauge::new("my_gauge", "Test fixture").unwrap();
57        gauge.set(5);
58        registry.register(Box::new(gauge)).unwrap();
59
60        let handler = text_format_handler(registry);
61        assert_response!(
62            get("metrics").on(&handler),
63            200,
64            "# HELP my_gauge Test fixture\n# TYPE my_gauge gauge\nmy_gauge 5"
65        );
66    }
67}