From 15d214e2372308fa1d12b576a675c9e2cbf6cde1 Mon Sep 17 00:00:00 2001 From: Ellie Huxtable Date: Thu, 16 Nov 2023 23:18:13 +0000 Subject: [PATCH] feat: add metrics server and http metrics (#1394) * feat: add metrics server and http metrics * setup metrics * update default config * fix tests --- Cargo.lock | 130 ++++++++++++++++++++++++++++++++++- atuin-server/Cargo.toml | 2 + atuin-server/server.toml | 5 ++ atuin-server/src/lib.rs | 27 +++++++- atuin-server/src/metrics.rs | 52 ++++++++++++++ atuin-server/src/router.rs | 4 +- atuin-server/src/settings.rs | 21 ++++++ atuin/src/command/server.rs | 9 ++- atuin/tests/sync.rs | 1 + 9 files changed, 247 insertions(+), 4 deletions(-) create mode 100644 atuin-server/src/metrics.rs diff --git a/Cargo.lock b/Cargo.lock index 0ffee9a88fa..1d04bdacf2c 100644 --- a/Cargo.lock +++ b/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" diff --git a/atuin-server/Cargo.toml b/atuin-server/Cargo.toml index 95bd3e0939f..e5390b00b69 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 3aed7f9d8b3..b2468ddbd8f 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 6778b099fb6..2d2a9c7855a 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( 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 00000000000..2e3e6894ae2 --- /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, 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(req: Request, next: Next) -> impl IntoResponse { + let start = Instant::now(); + + let path = if let Some(matched_path) = req.extensions().get::() { + 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 e1220e56b12..90e726d36a7 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(database: DB, settings: 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 744f4ec234e..d6f1867c654 100644 --- a/atuin-server/src/settings.rs +++ b/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 { pub host: String, @@ -18,6 +35,7 @@ pub struct Settings { pub page_size: i64, pub register_webhook_url: Option, pub register_webhook_username: String, + pub metrics: Metrics, #[serde(flatten)] pub db_settings: DbSettings, @@ -46,6 +64,9 @@ impl Settings { .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("_") diff --git a/atuin/src/command/server.rs b/atuin/src/command/server.rs index bfecdd751c0..4bcf19db9a9 100644 --- a/atuin/src/command/server.rs +++ b/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::(settings, &host, port).await } Self::DefaultConfig => { diff --git a/atuin/tests/sync.rs b/atuin/tests/sync.rs index 6dbc724424c..765b9cb8d64 100644 --- a/atuin/tests/sync.rs +++ b/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();