1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

use libprometheus::{Encoder, Registry, TextEncoder};

use crate::{
    http::{Method, StatusCode},
    Endpoint, IntoEndpoint, Request, Response, Result,
};

/// An endpoint that exports metrics for Prometheus.
///
/// # Example
///
/// ```
/// use libprometheus::Registry;
/// use poem::{endpoint::PrometheusExporter, Route};
///
/// let registry = Registry::new();
/// let app = Route::new().nest("/metrics", PrometheusExporter::new(registry));
/// ```
#[cfg_attr(docsrs, doc(cfg(feature = "prometheus")))]
pub struct PrometheusExporter {
    registry: Registry,
}

impl PrometheusExporter {
    /// Create a `PrometheusExporter` endpoint.
    pub fn new(registry: Registry) -> Self {
        Self { registry }
    }
}

impl IntoEndpoint for PrometheusExporter {
    type Endpoint = PrometheusExporterEndpoint;

    fn into_endpoint(self) -> Self::Endpoint {
        PrometheusExporterEndpoint {
            registry: self.registry.clone(),
        }
    }
}

#[doc(hidden)]
pub struct PrometheusExporterEndpoint {
    registry: Registry,
}

impl Endpoint for PrometheusExporterEndpoint {
    type Output = Response;

    async fn call(&self, req: Request) -> Result<Self::Output> {
        if req.method() != Method::GET {
            return Ok(StatusCode::METHOD_NOT_ALLOWED.into());
        }

        let encoder = TextEncoder::new();
        let metric_families = self.registry.gather();
        let mut result = Vec::new();
        match encoder.encode(&metric_families, &mut result) {
            Ok(()) => Ok(Response::builder().content_type("text/plain").body(result)),
            Err(_) => Err(StatusCode::INTERNAL_SERVER_ERROR.into()),
        }
    }
}