expose stats for prometheus to scrape

1 files changed, 28 insertions, 7 deletions

diff --git a/src/api/stats.rs b/src/api/stats.rs
index 3822b5f..b39e054 100644
--- a/src/api/stats.rs
+++ b/src/api/stats.rs
@@ -1,6 +1,14 @@
-//! /admin/stats: how this node is doing, in numbers.
+//! /admin/stats: how this node is doing, in numbers. JSON, or the Prometheus text
+//! format when the client asks for it (`prometheus`).
 
-use axum::{Json, extract::State};
+mod prometheus;
+
+use axum::{
+    Json,
+    extract::State,
+    http::{HeaderMap, header},
+    response::{IntoResponse, Response},
+};
 use serde::Serialize;
 
 use super::peers::EpochInfo;
@@ -107,7 +115,20 @@ struct StabilityStats {
 }
 
 /// `GET /admin/stats`.
-pub(crate) async fn get_stats(State(state): State<AppState>) -> Json<StatsResponse> {
+pub(crate) async fn get_stats(State(state): State<AppState>, headers: HeaderMap) -> Response {
+    let stats = collect(&state).await;
+    let accept = headers
+        .get(header::ACCEPT)
+        .and_then(|value| value.to_str().ok());
+    if prometheus::wanted(accept) {
+        let body = prometheus::render(&state, &stats).await;
+        ([(header::CONTENT_TYPE, prometheus::CONTENT_TYPE)], body).into_response()
+    } else {
+        Json(stats).into_response()
+    }
+}
+
+async fn collect(state: &AppState) -> StatsResponse {
     let node = state.stats.snapshot();
 
     let (heads, tier) = {
@@ -117,7 +138,7 @@ pub(crate) async fn get_stats(State(state): State<AppState>) -> Json<StatsRespon
     };
     let saved = state.stats.saved_size();
 
-    let (connected_fraction, observed) = stability::measured(&state);
+    let (connected_fraction, observed) = stability::measured(state);
 
     let mut incidents = Incidents::default();
     let peers: Vec<_> = state.peer_states.read().await.values().cloned().collect();
@@ -125,7 +146,7 @@ pub(crate) async fn get_stats(State(state): State<AppState>) -> Json<StatsRespon
         incidents.add(&peer_state.lock().await.incidents);
     }
 
-    Json(StatsResponse {
+    StatsResponse {
         uptime_secs: node.uptime.as_secs(),
         ready: state.is_ready().await,
         epoch: state.local_epoch().into(),
@@ -143,7 +164,7 @@ pub(crate) async fn get_stats(State(state): State<AppState>) -> Json<StatsRespon
         },
         topology: TopologyStats {
             k: state.neighbours_k,
-            view: topology::summary(&state).map(|s| TopologyView {
+            view: topology::summary(state).map(|s| TopologyView {
                 members: s.members,
                 reachable: s.reachable,
                 stable: s.stable,
@@ -170,5 +191,5 @@ pub(crate) async fn get_stats(State(state): State<AppState>) -> Json<StatsRespon
                 .map(|(le_ms, count)| PropagationBucket { le_ms, count })
                 .collect(),
         },
-    })
+    }
 }