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(&self, req: Request) -> Result<Self::Output> {
34        if req.method() != Method::GET {
35            return Ok(StatusCode::METHOD_NOT_ALLOWED.into());
36        }
37
38        metrics_updater_fn().await;
39
40        let encoder = TextEncoder::new();
41        let metric_families = self.registry.gather();
42        let mut result = Vec::new();
43        match encoder.encode(&metric_families, &mut result) {
44            Ok(()) => {
45                Ok(Response::builder()
46                    .content_type(encoder.format_type())
47                    .body(result))
48            },
49            Err(_) => Err(StatusCode::INTERNAL_SERVER_ERROR.into()),
50        }
51    }
52}
cat_gateway/service/utilities/middleware/metrics_updater.rs

cat_gateway/service/utilities/middleware/
metrics_updater.rs
