Added basic prometheus metrics

src/main.rs

@@ -1,11 +1,18 @@
-use std::net::SocketAddr;
+use std::{future::ready, net::SocketAddr};
 
-use axum::{extract::Extension, http::StatusCode, routing::get_service, Router};
+use axum::{
+    extract::Extension,
+    http::StatusCode,
+    middleware,
+    routing::{get, get_service},
+    Router,
+};
 use tower_http::{auth::RequireAuthorizationLayer, services::ServeDir, trace::TraceLayer};
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
 
 mod api;
 mod matrix;
+mod metrics;
 
 const STATIC_DIR_NAME: &str = "static";
 
@@ -27,6 +34,8 @@ async fn main()
 
     std::fs::create_dir_all(&static_dir);
 
+    let recorder_handle = metrics::setup_metrics_recorder();
+
     let app = Router::new()
         // Handle Matrix .well-known files
         .nest("/", matrix::router())
@@ -35,6 +44,7 @@ async fn main()
             "/api",
             api::router().layer(RequireAuthorizationLayer::bearer(&api_key)),
         )
+        .route("/metrics", get(move || ready(recorder_handle.render())))
         // The fallback option is to serve the actual static files
         .fallback(get_service(ServeDir::new(static_dir)).handle_error(
             |error: std::io::Error| async move {
@@ -44,8 +54,9 @@ async fn main()
                 )
             },
         ))
-        .layer(TraceLayer::new_for_http())
-        .layer(Extension(data_dir));
+        .layer(middleware::from_fn(metrics::track_metrics))
+        .layer(Extension(data_dir))
+        .layer(TraceLayer::new_for_http());
 
     let addr = SocketAddr::from(([0, 0, 0, 0], 3000));
     tracing::debug!("listening on {}", addr);

src/metrics.rs

@@ -0,0 +1,48 @@
+use std::time::Instant;
+
+use axum::{extract::MatchedPath, http::Request, middleware::Next, response::IntoResponse};
+use metrics_exporter_prometheus::{Matcher, PrometheusBuilder, PrometheusHandle};
+
+pub fn setup_metrics_recorder() -> PrometheusHandle
+{
+    const EXPONENTIAL_SECONDS: &[f64] = &[
+        0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0,
+    ];
+
+    PrometheusBuilder::new()
+        .set_buckets_for_metric(
+            Matcher::Full("http_requests_duration_seconds".to_string()),
+            EXPONENTIAL_SECONDS,
+        )
+        .unwrap()
+        .install_recorder()
+        .unwrap()
+}
+
+// Taken from the examples section
+pub async fn track_metrics<B>(req: Request<B>, next: Next<B>) -> impl IntoResponse
+{
+    let start = Instant::now();
+    let path = if let Some(matched_path) = req.extensions().get::<MatchedPath>() {
+        matched_path.as_str().to_owned()
+    } else {
+        req.uri().path().to_owned()
+    };
+    let method = req.method().clone();
+
+    let response = next.run(req).await;
+
+    let latency = start.elapsed().as_secs_f64();
+    let status = response.status().as_u16().to_string();
+
+    let labels = [
+        ("method", method.to_string()),
+        ("path", path),
+        ("status", status),
+    ];
+
+    metrics::increment_counter!("http_requests_total", &labels);
+    metrics::histogram!("http_requests_duration_seconds", latency, &labels);
+
+    response
+}