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

* feat: add metrics server and http metrics

* setup metrics

* update default config

* fix tests

Cargo.lock

@@ -256,6 +256,8 @@ dependencies = [
  "eyre",
  "fs-err",
  "http",
+ "metrics",
+ "metrics-exporter-prometheus",
  "rand 0.8.5",
  "reqwest",
  "semver",
@@ -701,6 +703,19 @@ version = "2.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9cace84e55f07e7301bae1c519df89cdad8cc3cd868413d3fdbdeca9ff3db484"
 
+[[package]]
+name = "crossbeam-epoch"
+version = "0.9.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ae211234986c545741a7dc064309f67ee1e5ad243d0e48335adc0484d960bcc7"
+dependencies = [
+ "autocfg",
+ "cfg-if",
+ "crossbeam-utils",
+ "memoffset 0.9.0",
+ "scopeguard",
+]
+
 [[package]]
 name = "crossbeam-queue"
 version = "0.3.8"
@@ -1321,6 +1336,15 @@ version = "0.12.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
 
+[[package]]
+name = "hashbrown"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "33ff8ae62cd3a9102e5637afc8452c55acf3844001bd5374e0b0bd7b6616c038"
+dependencies = [
+ "ahash",
+]
+
 [[package]]
 name = "hashbrown"
 version = "0.14.2"
@@ -1700,6 +1724,15 @@ dependencies = [
  "hashbrown 0.14.2",
 ]
 
+[[package]]
+name = "mach2"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6d0d1830bcd151a6fc4aea1369af235b36c1528fe976b8ff678683c9995eade8"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "malloc_buf"
 version = "0.0.6"
@@ -1749,6 +1782,70 @@ dependencies = [
  "autocfg",
 ]
 
+[[package]]
+name = "memoffset"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5a634b1c61a95585bd15607c6ab0c4e5b226e695ff2800ba0cdccddf208c406c"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "metrics"
+version = "0.21.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fde3af1a009ed76a778cb84fdef9e7dbbdf5775ae3e4cc1f434a6a307f6f76c5"
+dependencies = [
+ "ahash",
+ "metrics-macros",
+ "portable-atomic",
+]
+
+[[package]]
+name = "metrics-exporter-prometheus"
+version = "0.12.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8a4964177ddfdab1e3a2b37aec7cf320e14169abb0ed73999f558136409178d5"
+dependencies = [
+ "base64 0.21.5",
+ "hyper",
+ "indexmap 1.9.3",
+ "ipnet",
+ "metrics",
+ "metrics-util",
+ "quanta",
+ "thiserror",
+ "tokio",
+ "tracing",
+]
+
+[[package]]
+name = "metrics-macros"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ddece26afd34c31585c74a4db0630c376df271c285d682d1e55012197830b6df"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.38",
+]
+
+[[package]]
+name = "metrics-util"
+version = "0.15.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4de2ed6e491ed114b40b732e4d1659a9d53992ebd87490c44a6ffe23739d973e"
+dependencies = [
+ "crossbeam-epoch",
+ "crossbeam-utils",
+ "hashbrown 0.13.1",
+ "metrics",
+ "num_cpus",
+ "quanta",
+ "sketches-ddsketch",
+]
+
 [[package]]
 name = "mime"
 version = "0.3.17"
@@ -1797,7 +1894,7 @@ dependencies = [
  "bitflags 1.3.2",
  "cfg-if",
  "libc",
- "memoffset",
+ "memoffset 0.6.5",
 ]
 
 [[package]]
@@ -2216,6 +2313,22 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "quanta"
+version = "0.11.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a17e662a7a8291a865152364c20c7abc5e60486ab2001e8ec10b24862de0b9ab"
+dependencies = [
+ "crossbeam-utils",
+ "libc",
+ "mach2",
+ "once_cell",
+ "raw-cpuid",
+ "wasi 0.11.0+wasi-snapshot-preview1",
+ "web-sys",
+ "winapi",
+]
+
 [[package]]
 name = "quote"
 version = "1.0.33"
@@ -2314,6 +2427,15 @@ dependencies = [
  "unicode-width",
 ]
 
+[[package]]
+name = "raw-cpuid"
+version = "10.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6c297679cb867470fa8c9f67dbba74a78d78e3e98d7cf2b08d6d71540f797332"
+dependencies = [
+ "bitflags 1.3.2",
+]
+
 [[package]]
 name = "redox_syscall"
 version = "0.2.16"
@@ -2883,6 +3005,12 @@ dependencies = [
  "rand_core 0.6.4",
 ]
 
+[[package]]
+name = "sketches-ddsketch"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "68a406c1882ed7f29cd5e248c9848a80e7cb6ae0fea82346d2746f2f941c07e1"
+
 [[package]]
 name = "slab"
 version = "0.4.9"

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();