diff --git a/Cargo.lock b/Cargo.lock index c9ef544..e2b332d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1051,6 +1051,7 @@ dependencies = [ "lru", "once_cell", "parking_lot", + "prometheus", "reqwest", "rustls", "serde", @@ -1329,6 +1330,43 @@ dependencies = [ "unicode-xid", ] +[[package]] +name = "procfs" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab8809e0c18450a2db0f236d2a44ec0b4c1412d0eb936233579f0990faa5d5cd" +dependencies = [ + "bitflags", + "byteorder", + "flate2", + "hex", + "lazy_static", + "libc", +] + +[[package]] +name = "prometheus" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5986aa8d62380092d2f50f8b1cdba9cb9b6731ffd4b25b51fd126b6c3e05b99c" +dependencies = [ + "cfg-if", + "fnv", + "lazy_static", + "libc", + "memchr", + "parking_lot", + "procfs", + "protobuf", + "thiserror", +] + +[[package]] +name = "protobuf" +version = "2.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45604fc7a88158e7d514d8e22e14ac746081e7a70d7690074dd0029ee37458d6" + [[package]] name = "quote" version = "1.0.9" diff --git a/Cargo.toml b/Cargo.toml index 59f91a4..68492a9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,6 +30,7 @@ log = "0.4" lfu_cache = "1" lru = "0.6" parking_lot = "0.11" +prometheus = { version = "0.12", features = [ "process" ] } reqwest = { version = "0.11", default_features = false, features = [ "json", "stream", "rustls-tls" ] } rustls = "0.19" serde = "1" diff --git a/src/main.rs b/src/main.rs index 9f9feb3..50dc3d6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -34,6 +34,7 @@ use crate::state::DynamicServerCert; mod cache; mod config; +mod metrics; mod ping; mod routes; mod state; @@ -102,6 +103,8 @@ async fn main() -> Result<(), Box> { ENCRYPTION_KEY.set(gen_key()).unwrap(); } + // HTTP Server init + let server = ServerState::init(&client_secret, &cli_args).await?; let data_0 = Arc::new(RwLockServerState(RwLock::new(server))); let data_1 = Arc::clone(&data_0); @@ -161,6 +164,7 @@ async fn main() -> Result<(), Box> { HttpServer::new(move || { App::new() .service(routes::token_data) + .service(routes::metrics) .service(routes::token_data_saver) .route("{tail:.*}", web::get().to(routes::default)) .app_data(Data::from(Arc::clone(&data_1))) diff --git a/src/metrics.rs b/src/metrics.rs new file mode 100644 index 0000000..e07d2ee --- /dev/null +++ b/src/metrics.rs @@ -0,0 +1,36 @@ +use once_cell::sync::Lazy; +use prometheus::{register_int_counter, IntCounter}; + +pub static CACHE_HIT_COUNTER: Lazy = + Lazy::new(|| register_int_counter!("cache.hit", "The number of cache hits").unwrap()); + +pub static CACHE_MISS_COUNTER: Lazy = + Lazy::new(|| register_int_counter!("cache.miss", "The number of cache misses").unwrap()); + +pub static REQUESTS_TOTAL_COUNTER: Lazy = Lazy::new(|| { + register_int_counter!("requests.total", "The total number of requests served.").unwrap() +}); + +pub static REQUESTS_DATA_COUNTER: Lazy = Lazy::new(|| { + register_int_counter!( + "requests.data", + "The number of requests served from the /data endpoint." + ) + .unwrap() +}); + +pub static REQUESTS_DATA_SAVER_COUNTER: Lazy = Lazy::new(|| { + register_int_counter!( + "requests.data-saver", + "The number of requests served from the /data-saver endpoint." + ) + .unwrap() +}); + +pub static REQUESTS_OTHER_COUNTER: Lazy = Lazy::new(|| { + register_int_counter!( + "requests.other", + "The total number of request not served by primary endpoints." + ) + .unwrap() +}); diff --git a/src/routes.rs b/src/routes.rs index 56f6895..46926a2 100644 --- a/src/routes.rs +++ b/src/routes.rs @@ -13,6 +13,7 @@ use chrono::{DateTime, Utc}; use futures::{Stream, TryStreamExt}; use log::{debug, error, info, warn}; use once_cell::sync::Lazy; +use prometheus::{Encoder, TextEncoder}; use reqwest::{Client, StatusCode}; use serde::Deserialize; use sodiumoxide::crypto::box_::{open_precomputed, Nonce, PrecomputedKey, NONCEBYTES}; @@ -21,6 +22,10 @@ use thiserror::Error; use crate::cache::{Cache, CacheKey, ImageMetadata, UpstreamError}; use crate::client_api_version; use crate::config::{SEND_SERVER_VERSION, VALIDATE_TOKENS}; +use crate::metrics::{ + CACHE_HIT_COUNTER, CACHE_MISS_COUNTER, REQUESTS_DATA_COUNTER, REQUESTS_DATA_SAVER_COUNTER, + REQUESTS_OTHER_COUNTER, REQUESTS_TOTAL_COUNTER, +}; use crate::state::RwLockServerState; pub const BASE64_CONFIG: base64::Config = base64::Config::new(base64::CharacterSet::UrlSafe, false); @@ -46,7 +51,10 @@ impl Responder for ServerResponse { fn respond_to(self, req: &HttpRequest) -> HttpResponse { match self { Self::TokenValidationError(e) => e.respond_to(req), - Self::HttpResponse(resp) => resp.respond_to(req), + Self::HttpResponse(resp) => { + REQUESTS_TOTAL_COUNTER.inc(); + resp.respond_to(req) + } } } } @@ -58,6 +66,7 @@ async fn token_data( cache: Data, path: Path<(String, String, String)>, ) -> impl Responder { + REQUESTS_DATA_COUNTER.inc(); let (token, chapter_hash, file_name) = path.into_inner(); if VALIDATE_TOKENS.load(Ordering::Acquire) { if let Err(e) = validate_token(&state.0.read().precomputed_key, token, &chapter_hash) { @@ -74,6 +83,7 @@ async fn token_data_saver( cache: Data, path: Path<(String, String, String)>, ) -> impl Responder { + REQUESTS_DATA_SAVER_COUNTER.inc(); let (token, chapter_hash, file_name) = path.into_inner(); if VALIDATE_TOKENS.load(Ordering::Acquire) { if let Err(e) = validate_token(&state.0.read().precomputed_key, token, &chapter_hash) { @@ -86,6 +96,7 @@ async fn token_data_saver( #[allow(clippy::future_not_send)] pub async fn default(state: Data, req: HttpRequest) -> impl Responder { + REQUESTS_OTHER_COUNTER.inc(); let path = &format!( "{}{}", state.0.read().image_server, @@ -109,6 +120,17 @@ pub async fn default(state: Data, req: HttpRequest) -> impl R ServerResponse::HttpResponse(resp_builder.body(resp.bytes().await.unwrap_or_default())) } +#[allow(clippy::future_not_send)] +#[get("/metrics")] +pub async fn metrics() -> impl Responder { + let metric_families = prometheus::gather(); + let mut buffer = Vec::new(); + TextEncoder::new() + .encode(&metric_families, &mut buffer) + .unwrap(); + String::from_utf8(buffer).unwrap() +} + #[derive(Error, Debug)] enum TokenValidationError { #[error("Failed to decode base64 token.")] @@ -198,6 +220,7 @@ async fn fetch_image( match cache.get(&key).await { Some(Ok((image, metadata))) => { + CACHE_HIT_COUNTER.inc(); return construct_response(image, &metadata); } Some(Err(_)) => { @@ -206,6 +229,8 @@ async fn fetch_image( _ => (), } + CACHE_MISS_COUNTER.inc(); + // It's important to not get a write lock before this request, else we're // holding the read lock until the await resolves.