linera_service/
prometheus_server.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::fmt::Debug;
5
6use axum::{http::StatusCode, response::IntoResponse, routing::get, Router};
7use tokio::net::ToSocketAddrs;
8use tokio_util::sync::CancellationToken;
9use tracing::info;
10
11pub fn start_metrics(
12    address: impl ToSocketAddrs + Debug + Send + 'static,
13    shutdown_signal: CancellationToken,
14) {
15    info!("Starting to serve metrics on {:?}", address);
16    let prometheus_router = Router::new().route("/metrics", get(serve_metrics));
17
18    tokio::spawn(async move {
19        if let Err(e) = axum::serve(
20            tokio::net::TcpListener::bind(address).await.unwrap(),
21            prometheus_router,
22        )
23        .with_graceful_shutdown(shutdown_signal.cancelled_owned())
24        .await
25        {
26            panic!("Error serving metrics: {}", e);
27        }
28    });
29}
30
31async fn serve_metrics() -> Result<String, AxumError> {
32    let metric_families = prometheus::gather();
33    Ok(prometheus::TextEncoder::new()
34        .encode_to_string(&metric_families)
35        .map_err(anyhow::Error::from)?)
36}
37
38struct AxumError(anyhow::Error);
39
40impl IntoResponse for AxumError {
41    fn into_response(self) -> axum::response::Response {
42        (
43            StatusCode::INTERNAL_SERVER_ERROR,
44            format!("Something went wrong: {}", self.0),
45        )
46            .into_response()
47    }
48}
49
50impl<E> From<E> for AxumError
51where
52    E: Into<anyhow::Error>,
53{
54    fn from(err: E) -> Self {
55        Self(err.into())
56    }
57}