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

* feat: add metrics server and http metrics

* setup metrics

* update default config

* fix tests

6 files changed, 109 insertions, 2 deletions

diff --git a/atuin-server/Cargo.toml b/atuin-server/Cargo.toml
index 95bd3e09..e5390b00 100644
--- a/atuin-server/Cargo.toml
+++ b/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"
diff --git a/atuin-server/server.toml b/atuin-server/server.toml
index 3aed7f9d..b2468ddb 100644
--- a/atuin-server/server.toml
+++ b/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
diff --git a/atuin-server/src/lib.rs b/atuin-server/src/lib.rs
index 6778b099..2d2a9c78 100644
--- a/atuin-server/src/lib.rs
+++ b/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(())
+}
diff --git a/atuin-server/src/metrics.rs b/atuin-server/src/metrics.rs
new file mode 100644
index 00000000..2e3e6894
--- /dev/null
+++ b/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
+}
diff --git a/atuin-server/src/router.rs b/atuin-server/src/router.rs
index e1220e56..90e726d3 100644
--- a/atuin-server/src/router.rs
+++ b/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)),
     )
 }
diff --git a/atuin-server/src/settings.rs b/atuin-server/src/settings.rs
index 744f4ec2..d6f1867c 100644
--- a/atuin-server/src/settings.rs
+++ b/atuin-server/src/settings.rs
@@ -8,6 +8,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,
     pub port: u16,
@@ -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("_")