1  use std::{
 2      convert::Infallible,
 3      pin::Pin,
 4      sync::Arc,
 5      task::{Context, Poll},
 6  };
 7  
 8  use axum::{extract::Request, response::Response};
 9  use http::{StatusCode, header::CONTENT_TYPE};
10  use prometheus_client::{encoding::text, registry::Registry};
11  use tower::Service;
12  use tracing::{instrument, warn};
13  
14  pub const OPENMETRICS_TEXT: &str = "application/openmetrics-text";
15  
16  #[derive(Clone)]
17  pub struct MetricsService {
18      registry: Arc<Registry>,
19  }
20  
21  impl MetricsService {
22      pub fn new(registry: Arc<Registry>) -> Self {
23          Self { registry }
24      }
25  
26      #[instrument(skip_all)]
27      fn encode(&self) -> Result<String, std::fmt::Error> {
28          let mut buf = String::new();
29  
30          text::encode(&mut buf, &self.registry)?;
31  
32          Ok(buf)
33      }
34  }
35  
36  impl<Req> Service<Request<Req>> for MetricsService {
37      type Response = Response;
38  
39      type Error = Infallible;
40  
41      type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
42  
43      fn poll_ready(&mut self, _context: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
44          Poll::Ready(Ok(()))
45      }
46  
47      fn call(&mut self, _request: Request<Req>) -> Self::Future {
48          let response = match self.encode() {
49              Ok(encoded) => Response::builder()
50                  .status(StatusCode::OK)
51                  .header(CONTENT_TYPE, OPENMETRICS_TEXT)
52                  .body(encoded.into())
53                  .unwrap(),
54              Err(error) => {
55                  warn!(?error, "encoding metrics");
56  
57                  Response::builder()
58                      .status(StatusCode::INTERNAL_SERVER_ERROR)
59                      .body(().into())
60                      .unwrap()
61              }
62          };
63  
64          Box::pin(async { Ok(response) })
65      }
66  }