cat_gateway/service/utilities/middleware/
metrics_updater.rs

1//! A middleware module wrapping around `prometheus::Registry`.
2//! It acts as an endpoint to export metrics data to the Prometheus service.
3//! For every request to this endpoint, it will call the `updater` function to update
4//! metrics to the latest before sending to the service.
5
6use poem::{
7    http::{Method, StatusCode},
8    Endpoint, Request, Response, Result,
9};
10use prometheus::{Encoder, Registry, TextEncoder};
11
12use crate::metrics::{init_prometheus, metrics_updater_fn};
13
14/// A Middleware wrapping the Prometheus registry to report as metrics.
15///
16/// The middleware is originally from `poem::endpoint::PrometheusExporter`.
17pub struct MetricsUpdaterMiddleware {
18    /// The Prometheus registry.
19    registry: Registry,
20}
21
22impl MetricsUpdaterMiddleware {
23    /// Create a `PrometheusExporter` endpoint.
24    pub fn new() -> Self {
25        let registry = init_prometheus();
26        Self { registry }
27    }
28}
29
30impl Endpoint for MetricsUpdaterMiddleware {
31    type Output = Response;
32
33    async fn call(
34        &self,
35        req: Request,
36    ) -> Result<Self::Output> {
37        if req.method() != Method::GET {
38            return Ok(StatusCode::METHOD_NOT_ALLOWED.into());
39        }
40
41        metrics_updater_fn();
42
43        let encoder = TextEncoder::new();
44        let metric_families = self.registry.gather();
45        let mut result = Vec::new();
46        match encoder.encode(&metric_families, &mut result) {
47            Ok(()) => {
48                Ok(Response::builder()
49                    .content_type(encoder.format_type())
50                    .body(result))
51            },
52            Err(_) => Err(StatusCode::INTERNAL_SERVER_ERROR.into()),
53        }
54    }
55}
cat_gateway/service/utilities/middleware/metrics_updater.rs

cat_gateway/service/utilities/middleware/
metrics_updater.rs
