halide-cache-server: monitoring

halide-cache-server/Cargo.toml

@@ -7,12 +7,13 @@ description = "HTTP server sharing halide-cache entries between machines"
 [dependencies]
 lager = { path = "../lager" }
 anyhow = "1.0.102"
-axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "json"] }
+axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "json", "query"] }
 bytesize = "2.7.0"
 clap = { version = "4.5.53", features = ["derive"] }
 futures-util = { version = "0.3", default-features = false }
 serde = { version = "1", features = ["derive"] }
 tokio = { version = "1.53.1", features = ["rt-multi-thread", "macros", "net", "fs", "signal", "time"] }
 tokio-util = { version = "0.7", features = ["io", "io-util"] }
 tracing = "0.1"
+rustix = { version = "1.1.5", features = ["fs"] }
 tracing-subscriber = { version = "0.3.23", features = ["env-filter"] }

halide-cache-server/src/blobs.rs

@@ -1,6 +1,6 @@
 use crate::AppState;
-use crate::extract::BlobAddress;
-use crate::metrics::Event;
+use crate::extract::{BlobAddress, ClientId};
+use crate::metrics::{Event, Request as Timed};
 use axum::{
     body::Body,
     extract::{Request, State},
@@ -10,6 +10,7 @@ use axum::{
 use futures_util::TryStreamExt;
 use lager::Address;
 use std::sync::Arc;
+use std::time::Instant;
 use tokio_util::io::{ReaderStream, StreamReader, SyncIoBridge};
 use tracing::{error, info, warn};
 
@@ -51,10 +52,16 @@ fn blob_headers(len: u64) -> HeaderMap {
 pub async fn get(
     State(state): State<Arc<AppState>>,
     BlobAddress(address): BlobAddress,
+    ClientId(client): ClientId,
 ) -> Response {
-    match open_blob(&state, address).await {
+    let started = Instant::now();
+    let found = open_blob(&state, address).await;
+    if found.is_ok() {
+        state.metrics.record_latency(Timed::Get, started.elapsed());
+    }
+    match found {
         Ok(Some(blob)) => {
-            state.metrics.record(Event::Hit, blob.len);
+            state.metrics.record(Event::Hit, &client, blob.len);
             info!(%address, len = blob.len, "hit");
             let stream = ReaderStream::new(tokio::fs::File::from_std(blob.file));
             (
@@ -65,7 +72,7 @@ pub async fn get(
                 .into_response()
         }
         Ok(None) => {
-            state.metrics.record(Event::Miss, 0);
+            state.metrics.record(Event::Miss, &client, 0);
             info!(%address, "miss");
             StatusCode::NOT_FOUND.into_response()
         }
@@ -76,8 +83,10 @@ pub async fn get(
 pub async fn put(
     State(state): State<Arc<AppState>>,
     BlobAddress(address): BlobAddress,
+    ClientId(client): ClientId,
     request: Request,
 ) -> Response {
+    let started = Instant::now();
     let body = request
         .into_body()
         .into_data_stream()
@@ -100,11 +109,13 @@ pub async fn put(
             } else {
                 (Event::Duplicate, StatusCode::OK, "replaced existing")
             };
-            state.metrics.record(event, len);
+            state.metrics.record(event, &client, len);
+            state.metrics.record_latency(Timed::Put, started.elapsed());
             info!(%address, len, what);
             status.into_response()
         }
         Ok(Err(lager::Error::Io(e))) if e.kind() == std::io::ErrorKind::FileTooLarge => {
+            state.metrics.record(Event::Rejected, &client, 0);
             warn!(%address, "upload rejected: larger than --max-blob-size");
             (
                 StatusCode::PAYLOAD_TOO_LARGE,
@@ -113,6 +124,7 @@ pub async fn put(
                 .into_response()
         }
         Ok(Err(e)) => {
+            state.metrics.record(Event::Failed, &client, 0);
             warn!(%address, err = %e, "upload failed");
             (StatusCode::BAD_REQUEST, "upload failed\n").into_response()
         }