feat: add metrics server and http metrics (#1394)

* feat: add metrics server and http metrics

* setup metrics

* update default config

* fix tests

atuin-server/Cargo.toml

@@ -33,3 +33,5 @@ tower-http = { version = "0.4", features = ["trace"] }
 reqwest = { workspace = true }
 argon2 = "0.5.0"
 semver = { workspace = true }
+metrics-exporter-prometheus = "0.12.1"
+metrics = "0.21.1"

atuin-server/server.toml

@@ -22,3 +22,8 @@
 
 ## Default page size for requests
 # page_size = 1100
+
+# [metrics]
+# enable = false
+# host = 127.0.0.1
+# port = 9001

atuin-server/src/lib.rs

@@ -3,16 +3,20 @@
 use std::{future::Future, net::TcpListener};
 
 use atuin_server_database::Database;
+use axum::Router;
 use axum::Server;
 use eyre::{Context, Result};
 
 mod handlers;
+mod metrics;
 mod router;
-mod settings;
 mod utils;
 
 pub use settings::example_config;
 pub use settings::Settings;
+
+pub mod settings;
+
 use tokio::signal;
 
 #[cfg(target_family = "unix")]
@@ -70,3 +74,24 @@ pub async fn launch_with_listener<Db: Database>(
 
     Ok(())
 }
+
+// The separate listener means it's much easier to ensure metrics are not accidentally exposed to
+// the public.
+pub async fn launch_metrics_server(host: String, port: u16) -> Result<()> {
+    let listener = TcpListener::bind((host, port)).context("failed to bind metrics tcp")?;
+
+    let recorder_handle = metrics::setup_metrics_recorder();
+
+    let router = Router::new().route(
+        "/metrics",
+        axum::routing::get(move || std::future::ready(recorder_handle.render())),
+    );
+
+    Server::from_tcp(listener)
+        .context("could not launch server")?
+        .serve(router.into_make_service())
+        .with_graceful_shutdown(shutdown_signal())
+        .await?;
+
+    Ok(())
+}

atuin-server/src/metrics.rs

@@ -0,0 +1,52 @@
+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()
+}
+
+/// Middleware to record some common HTTP metrics
+/// Generic over B to allow for arbitrary body types (eg Vec<u8>, Streams, a deserialized thing, etc)
+/// Someday tower-http might provide a metrics middleware: https://github.com/tower-rs/tower-http/issues/57
+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();
+
+    // Run the rest of the request handling first, so we can measure it and get response
+    // codes.
+    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
+}

atuin-server/src/router.rs

@@ -16,6 +16,7 @@ use tower_http::trace::TraceLayer;
 use super::handlers;
 use crate::{
     handlers::{ErrorResponseStatus, RespExt},
+    metrics,
     settings::Settings,
 };
 use atuin_server_database::{models::User, Database, DbError};
@@ -124,6 +125,7 @@ pub fn router<DB: Database>(database: DB, settings: Settings<DB::Settings>) -> R
     .layer(
         ServiceBuilder::new()
             .layer(axum::middleware::from_fn(clacks_overhead))
-            .layer(TraceLayer::new_for_http()),
+            .layer(TraceLayer::new_for_http())
+            .layer(axum::middleware::from_fn(metrics::track_metrics)),
     )
 }

atuin-server/src/settings.rs

@@ -7,6 +7,23 @@ use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
 static EXAMPLE_CONFIG: &str = include_str!("../server.toml");
 
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct Metrics {
+    pub enable: bool,
+    pub host: String,
+    pub port: u16,
+}
+
+impl Default for Metrics {
+    fn default() -> Self {
+        Self {
+            enable: false,
+            host: String::from("127.0.0.1"),
+            port: 9001,
+        }
+    }
+}
+
 #[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct Settings<DbSettings> {
     pub host: String,
@@ -18,6 +35,7 @@ pub struct Settings<DbSettings> {
     pub page_size: i64,
     pub register_webhook_url: Option<String>,
     pub register_webhook_username: String,
+    pub metrics: Metrics,
 
     #[serde(flatten)]
     pub db_settings: DbSettings,
@@ -46,6 +64,9 @@ impl<DbSettings: DeserializeOwned> Settings<DbSettings> {
             .set_default("path", "")?
             .set_default("register_webhook_username", "")?
             .set_default("page_size", 1100)?
+            .set_default("metrics.enable", false)?
+            .set_default("metrics.host", "127.0.0.1")?
+            .set_default("metrics.port", 9001)?
             .add_source(
                 Environment::with_prefix("atuin")
                     .prefix_separator("_")

atuin/src/command/server.rs

@@ -4,7 +4,7 @@ use tracing_subscriber::{fmt, prelude::*, EnvFilter};
 use clap::Parser;
 use eyre::{Context, Result};
 
-use atuin_server::{example_config, launch, Settings};
+use atuin_server::{example_config, launch, launch_metrics_server, Settings};
 
 #[derive(Parser, Debug)]
 #[clap(infer_subcommands = true)]
@@ -40,6 +40,13 @@ impl Cmd {
                 let host = host.as_ref().unwrap_or(&settings.host).clone();
                 let port = port.unwrap_or(settings.port);
 
+                if settings.metrics.enable {
+                    tokio::spawn(launch_metrics_server(
+                        settings.metrics.host.clone(),
+                        settings.metrics.port,
+                    ));
+                }
+
                 launch::<Postgres>(settings, &host, port).await
             }
             Self::DefaultConfig => {

atuin/tests/sync.rs

@@ -37,6 +37,7 @@ async fn start_server(path: &str) -> (String, oneshot::Sender<()>, JoinHandle<()
         register_webhook_url: None,
         register_webhook_username: String::new(),
         db_settings: PostgresSettings { db_uri },
+        metrics: atuin_server::settings::Metrics::default(),
     };
 
     let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();