pubky_homeserver/metrics_server/
app.rs

1use std::net::SocketAddr;
2use std::time::Duration;
3
4use crate::metrics_server::routes::metrics::Metrics;
5use crate::AppContext;
6use crate::AppContextConversionError;
7use axum::routing::get;
8use axum::Router;
9use axum::{extract::State, http::StatusCode, response::IntoResponse};
10use axum_server::Handle;
11use tokio::task::JoinHandle;
12
13fn create_app(metrics: Metrics) -> axum::routing::IntoMakeService<Router> {
14    Router::new()
15        .route("/metrics", get(metrics_handler))
16        .with_state(metrics)
17        .into_make_service()
18}
19
20/// Errors that can occur when building a `MetricsServer`.
21#[derive(thiserror::Error, Debug)]
22pub enum MetricsServerBuildError {
23    /// Failed to create metrics server.
24    #[error("Failed to create metrics server: {0}")]
25    Server(anyhow::Error),
26
27    /// Failed to bootstrap from the data directory.
28    #[error("Failed to bootstrap from the data directory: {0}")]
29    DataDir(AppContextConversionError),
30}
31
32/// Metrics server
33///
34/// This server exposes Prometheus metrics on a separate port.
35/// It should be isolated from the public network and only accessible to monitoring systems.
36///
37/// When dropped, the server will stop.
38pub struct MetricsServer {
39    http_handle: Handle<SocketAddr>,
40    join_handle: JoinHandle<()>,
41    socket: SocketAddr,
42}
43
44impl MetricsServer {
45    /// Run the metrics server.
46    pub async fn start(context: &AppContext) -> Result<Self, MetricsServerBuildError> {
47        let metrics = context.metrics.clone();
48        let socket = context.config_toml.metrics.listen_socket;
49        let app = create_app(metrics);
50        let listener = std::net::TcpListener::bind(socket)
51            .map_err(|e| MetricsServerBuildError::Server(e.into()))?;
52        listener
53            .set_nonblocking(true)
54            .map_err(|e| MetricsServerBuildError::Server(e.into()))?;
55        let socket = listener
56            .local_addr()
57            .map_err(|e| MetricsServerBuildError::Server(e.into()))?;
58        let http_handle = Handle::new();
59        let inner_http_handle = http_handle.clone();
60        let server = axum_server::from_tcp(listener)
61            .map_err(|e| MetricsServerBuildError::Server(e.into()))?;
62        let join_handle = tokio::spawn(async move {
63            server
64                .handle(inner_http_handle)
65                .serve(app)
66                .await
67                .unwrap_or_else(|e| tracing::error!("Metrics server error: {}", e));
68        });
69        Ok(Self {
70            http_handle,
71            socket,
72            join_handle,
73        })
74    }
75
76    /// Get the socket address the metrics server is listening on.
77    pub fn listen_socket(&self) -> SocketAddr {
78        self.socket
79    }
80}
81
82impl Drop for MetricsServer {
83    fn drop(&mut self) {
84        self.http_handle
85            .graceful_shutdown(Some(Duration::from_secs(5)));
86        self.join_handle.abort();
87    }
88}
89
90/// HTTP handler for the /metrics endpoint
91pub async fn metrics_handler(State(metrics): State<Metrics>) -> impl IntoResponse {
92    match metrics.render() {
93        Ok(body) => (StatusCode::OK, body).into_response(),
94        Err(e) => (
95            StatusCode::INTERNAL_SERVER_ERROR,
96            format!("# Failed to render metrics: {}\n", e),
97        )
98            .into_response(),
99    }
100}
pubky_homeserver/metrics_server/app.rs

pubky_homeserver/metrics_server/
app.rs
