diff --git a/dwctl/src/api/handlers/files.rs b/dwctl/src/api/handlers/files.rs index 05cc3f081..2d49d10b6 100644 --- a/dwctl/src/api/handlers/files.rs +++ b/dwctl/src/api/handlers/files.rs @@ -53,7 +53,7 @@ use tokio_stream::wrappers::ReceiverStream; use uuid::Uuid; // Note: We import multer directly to get typed error matching for body limit errors. // axum's multipart wraps multer, and checking error variants is more robust than string matching. -use crate::limits::MULTIPART_OVERHEAD; +use crate::limits::{MULTIPART_OVERHEAD, UploadLimiter, UploadPermit}; use axum::extract::rejection::LengthLimitError; /// Check if the current user "owns" a resource, considering org context. @@ -544,7 +544,97 @@ fn map_batch_normalize_error(e: BatchNormalizeError, line: u64) -> FileUploadErr } } +/// Counts refused uploads by why they were refused. File uploads don't pass +/// through the inference stack, so `onwards_rejections_total` never sees them, +/// and `dwctl_http_requests_total` only has their status. +const UPLOAD_REJECTIONS: &str = "dwctl_file_upload_rejections_total"; + +/// Every code [`UPLOAD_REJECTIONS`] counts under: one per [`FileUploadError`] +/// variant, and the refusals `upload_file` makes before it reads the file. +const UPLOAD_REJECTION_CODES: &[&str] = &[ + "not_multipart", + "upload_capacity", + "batch_disabled", + "invalid_multipart", + "stream_interrupted", + "file_too_large", + "too_many_lines", + "invalid_json", + "invalid_utf8", + "invalid_compression", + "unsupported_compression", + "no_file", + "empty_file", + "model_access_denied", + "validation_error", + "unprocessable_line", + "image_transient", + "image_store_failed", + "image_access_unavailable", + "image_unfetchable", + "image_fetch_failed", + "invalid_purpose", +]; + +/// Describes [`UPLOAD_REJECTIONS`] and registers each code at zero, so the +/// first refusal of a kind is a visible step for `increase()`. +pub(crate) fn initialize_upload_rejection_metrics() { + metrics::describe_counter!(UPLOAD_REJECTIONS, "File uploads refused, by the reason they were refused"); + for &code in UPLOAD_REJECTION_CODES { + metrics::counter!(UPLOAD_REJECTIONS, "code" => code).increment(0); + } +} + +/// Counts an upload refused for `code` and returns `error`. +fn refuse_upload(code: &'static str, error: Error) -> Error { + metrics::counter!(UPLOAD_REJECTIONS, "code" => code).increment(1); + error +} + +/// An upload slot, when uploads are limited. A refusal (the queue is full, or +/// no slot came free in time) is counted as `upload_capacity`. +async fn acquire_upload_slot(limiter: Option<&UploadLimiter>) -> Result> { + match limiter { + Some(limiter) => limiter + .acquire() + .await + .map(Some) + .map_err(|error| refuse_upload("upload_capacity", error)), + None => Ok(None), + } +} + impl FileUploadError { + /// The bounded code this refusal is counted under. + fn code(&self) -> &'static str { + match self { + FileUploadError::StreamInterrupted { .. } => "stream_interrupted", + FileUploadError::FileTooLarge { .. } => "file_too_large", + FileUploadError::TooManyRequests { .. } => "too_many_lines", + FileUploadError::InvalidJson { .. } => "invalid_json", + FileUploadError::InvalidUtf8 { .. } => "invalid_utf8", + FileUploadError::InvalidCompression { .. } => "invalid_compression", + FileUploadError::UnsupportedCompression { .. } => "unsupported_compression", + FileUploadError::NoFile => "no_file", + FileUploadError::EmptyFile => "empty_file", + FileUploadError::ModelAccessDenied { .. } => "model_access_denied", + FileUploadError::ValidationError { .. } => "validation_error", + FileUploadError::UnprocessableValidationError { .. } => "unprocessable_line", + FileUploadError::ImageTransient { .. } => "image_transient", + FileUploadError::ImageStoreFailed { .. } => "image_store_failed", + FileUploadError::ImageAccessUnavailable { .. } => "image_access_unavailable", + FileUploadError::ImageUnfetchable { .. } => "image_unfetchable", + FileUploadError::ImageFetchFailed { .. } => "image_fetch_failed", + FileUploadError::InvalidPurpose { .. } => "invalid_purpose", + } + } + + /// Counts the refusal by its code and converts it to its HTTP error. + fn into_refusal(self) -> Error { + let code = self.code(); + refuse_upload(code, self.into_http_error()) + } + /// Convert to the appropriate HTTP error type fn into_http_error(self) -> Error { match self { @@ -705,7 +795,7 @@ fn resolve_upload_stream_result( }; if let Some(upload_err) = upload_err { tracing::warn!("File upload aborted with error: {:?}", upload_err); - return Err(upload_err.into_http_error()); + return Err(upload_err.into_refusal()); } Err(Error::Internal { @@ -1194,18 +1284,15 @@ pub async fn upload_file( // Refuse a body that isn't a multipart form before it takes an upload slot. // A malformed form (a bad boundary, say) is still a 400 below. if !is_multipart_form(request.headers()) { - return Err(Error::UnsupportedMediaType { + let error = Error::UnsupportedMediaType { message: "File uploads must be sent as multipart/form-data, with the file in a `file` field".to_string(), - }); + }; + return Err(refuse_upload("not_multipart", error)); } // Acquire upload permit (if limiter is configured) // The permit is held for the duration of the upload to limit concurrency - let _permit = if let Some(ref limiter) = state.limiters.file_uploads { - Some(limiter.acquire().await?) - } else { - None - }; + let _permit = acquire_upload_slot(state.limiters.file_uploads.as_deref()).await?; let config = state.current_config(); let max_file_size = config.limits.files.max_file_size; @@ -1222,14 +1309,18 @@ pub async fn upload_file( && let Ok(length) = length_str.parse::() && length > max_file_size.saturating_add(MULTIPART_OVERHEAD) { - return Err(Error::PayloadTooLarge { + let error = Error::PayloadTooLarge { message: format!("File exceeds the maximum allowed size of {} bytes", max_file_size), - }); + }; + return Err(refuse_upload("file_too_large", error)); } // Extract multipart from request body - let multipart = Multipart::from_request(request, &state).await.map_err(|e| Error::BadRequest { - message: format!("Invalid multipart request: {}", e), + let multipart = Multipart::from_request(request, &state).await.map_err(|e| { + let error = Error::BadRequest { + message: format!("Invalid multipart request: {}", e), + }; + refuse_upload("invalid_multipart", error) })?; // When in org context, attribute file ownership to the org (not the individual user). @@ -1243,7 +1334,12 @@ pub async fn upload_file( let mut conn = state.db.write().acquire().await.map_err(|e| Error::Internal { operation: format!("get db connection for modality check: {}", e), })?; - crate::api::handlers::modalities::ensure_batch_enabled(&mut conn, target_user_id).await?; + crate::api::handlers::modalities::ensure_batch_enabled(&mut conn, target_user_id) + .await + .map_err(|error| match error { + Error::ModalityDisabled { .. } => refuse_upload("batch_disabled", error), + error => error, + })?; } // Re-use the AppState-bound normaliser singleton (built once at @@ -1377,7 +1473,7 @@ pub async fn upload_file( }; if let Some(upload_err) = upload_err { tracing::warn!("File upload aborted with error: {:?}", upload_err); - return upload_err.into_http_error(); + return upload_err.into_refusal(); } // Otherwise it's a fusillade error @@ -2515,10 +2611,21 @@ mod tests { let truncated_gzip = b"\x1f\x8b\x08\x00\x00\x00\x00\x00\x00\x03\xab\x56".to_vec(); let xz = b"\xfd7zXZ\x00\x00\x04\xe6\xd6\xb4\x46".to_vec(); - for (content, status, expected) in [ - (truncated_gzip, axum::http::StatusCode::BAD_REQUEST, "not valid gzip data"), - (xz, axum::http::StatusCode::UNSUPPORTED_MEDIA_TYPE, "compressed with xz"), + for (content, status, expected, code) in [ + ( + truncated_gzip, + axum::http::StatusCode::BAD_REQUEST, + "not valid gzip data", + "invalid_compression", + ), + ( + xz, + axum::http::StatusCode::UNSUPPORTED_MEDIA_TYPE, + "compressed with xz", + "unsupported_compression", + ), ] { + let before = crate::test::utils::upload_rejections(code); let upload_response = app .post("/ai/v1/files") .add_header(&add_auth_headers(&user)[0].0, &add_auth_headers(&user)[0].1) @@ -2533,6 +2640,7 @@ mod tests { upload_response.assert_status(status); let body = upload_response.text(); assert!(body.contains(expected), "{body}"); + assert_eq!(crate::test::utils::upload_rejections(code), before + 1, "{code} is counted"); } } @@ -2542,6 +2650,8 @@ mod tests { let (app, _bg_services) = create_test_app(pool.clone(), false).await; let user = create_test_user_with_roles(&pool, vec![Role::StandardUser, Role::BatchAPIUser]).await; let line = r#"{"custom_id":"request-1","method":"POST","url":"/v1/chat/completions","body":{"model":"gpt-4"}}"#; + let upload_rejections = crate::test::utils::upload_rejections; + let (not_multipart, invalid_multipart) = (upload_rejections("not_multipart"), upload_rejections("invalid_multipart")); for content_type in [Some("application/jsonl"), Some("application/x-www-form-urlencoded"), None] { let mut request = app @@ -2566,6 +2676,9 @@ mod tests { .bytes(line.as_bytes().to_vec().into()) .await; upload_response.assert_status(axum::http::StatusCode::BAD_REQUEST); + + assert_eq!(upload_rejections("not_multipart"), not_multipart + 3); + assert_eq!(upload_rejections("invalid_multipart"), invalid_multipart + 1); } #[dwctl_test_macros::test] @@ -4187,6 +4300,107 @@ mod tests { } } + /// The `dwctl_file_upload_rejections_total` lines rendered while `f` ran. + fn upload_rejection_lines(f: impl FnOnce()) -> Vec { + let recorder = metrics_exporter_prometheus::PrometheusBuilder::new().build_recorder(); + let handle = recorder.handle(); + metrics::with_local_recorder(&recorder, f); + handle + .render() + .lines() + .filter(|line| line.starts_with("dwctl_file_upload_rejections_total{")) + .map(str::to_string) + .collect() + } + + #[test] + fn test_every_upload_refusal_is_counted_by_its_code() { + use super::FileUploadError::*; + let s = || "x".to_string(); + let refusals = [ + StreamInterrupted { message: s() }, + FileTooLarge { max: 1 }, + TooManyRequests { count: 2, max: 1 }, + InvalidJson { line: 1, error: s() }, + InvalidUtf8 { + line: 1, + byte_offset: 0, + error: s(), + }, + InvalidCompression { + format: "gzip", + message: s(), + }, + UnsupportedCompression { format: "xz" }, + NoFile, + EmptyFile, + ModelAccessDenied { model: s(), line: 1 }, + ValidationError { line: 1, message: s() }, + UnprocessableValidationError { line: 1, message: s() }, + ImageTransient { line: 1, message: s() }, + ImageStoreFailed { line: 1, message: s() }, + ImageAccessUnavailable { line: 1 }, + ImageUnfetchable { line: 1, message: s() }, + ImageFetchFailed { line: 1, message: s() }, + InvalidPurpose { purpose: s() }, + ]; + for refusal in refusals { + let code = refusal.code(); + assert!(super::UPLOAD_REJECTION_CODES.contains(&code), "{code} is registered at startup"); + let lines = upload_rejection_lines(|| { + refusal.into_refusal(); + }); + assert_eq!( + lines, + [format!(r#"dwctl_file_upload_rejections_total{{code="{code}"}} 1"#)], + "{code}" + ); + } + } + + #[test] + fn test_upload_capacity_refusals_are_counted() { + let recorder = metrics_exporter_prometheus::PrometheusBuilder::new().build_recorder(); + let handle = recorder.handle(); + metrics::with_local_recorder(&recorder, || { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + let config = crate::config::FileLimitsConfig { + max_concurrent_uploads: 1, + max_upload_wait_secs: 0, + ..Default::default() + }; + let limiter = crate::limits::UploadLimiter::new(&config).unwrap(); + assert!(super::acquire_upload_slot(None).await.unwrap().is_none(), "unlimited"); + let _held = super::acquire_upload_slot(Some(&limiter)).await.unwrap(); + let Err(refused) = super::acquire_upload_slot(Some(&limiter)).await else { + panic!("the only slot is held"); + }; + assert_eq!(refused.status_code(), axum::http::StatusCode::TOO_MANY_REQUESTS); + }); + }); + let rendered = handle.render(); + assert!( + rendered + .lines() + .any(|line| line == r#"dwctl_file_upload_rejections_total{code="upload_capacity"} 1"#), + "{rendered}" + ); + } + + #[test] + fn test_upload_rejection_codes_start_at_zero() { + let lines = upload_rejection_lines(super::initialize_upload_rejection_metrics); + assert_eq!(lines.len(), super::UPLOAD_REJECTION_CODES.len(), "{lines:?}"); + for code in super::UPLOAD_REJECTION_CODES { + let line = format!(r#"dwctl_file_upload_rejections_total{{code="{code}"}} 0"#); + assert!(lines.contains(&line), "{line} in {lines:?}"); + } + } + /// Test that Content-Length header triggers early rejection for oversized files #[dwctl_test_macros::test] #[test_log::test] diff --git a/dwctl/src/inference/body_limit.rs b/dwctl/src/inference/body_limit.rs index 71b0e9d65..c66dc52b4 100644 --- a/dwctl/src/inference/body_limit.rs +++ b/dwctl/src/inference/body_limit.rs @@ -39,6 +39,46 @@ const ZSTD_WINDOW_LOG_MAX: u32 = 23; static DECODE_SLOTS: LazyLock> = LazyLock::new(|| Arc::new(Semaphore::new(std::thread::available_parallelism().map_or(1, |cpus| cpus.get())))); +/// Seconds a compressed body waited for a decode slot. +pub(crate) const DECODE_SLOT_WAIT_SECONDS: &str = "dwctl_inference_decode_slot_wait_seconds"; + +/// Compressed bodies waiting for a decode slot right now. +const DECODES_WAITING: &str = "dwctl_inference_decodes_waiting"; + +/// Bucket edges for [`DECODE_SLOT_WAIT_SECONDS`]. A free slot is taken at +/// once, so most waits fall in the first bucket; 0.25 is an edge so a +/// p99-above-250ms alert reads an exact bucket. +pub(crate) const DECODE_SLOT_WAIT_BUCKETS: &[f64] = &[0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0]; + +/// Describes the decode-slot metrics, with the waiting gauge at zero so it is +/// scraped before the first wait. +pub(crate) fn describe_decode_slot_metrics() { + metrics::describe_histogram!( + DECODE_SLOT_WAIT_SECONDS, + metrics::Unit::Seconds, + "Time a compressed inference body waited for a decode slot" + ); + metrics::describe_gauge!(DECODES_WAITING, "Compressed inference bodies waiting for a decode slot"); + metrics::gauge!(DECODES_WAITING).set(0.0); +} + +/// One decode in [`DECODES_WAITING`] while it lives. +struct Waiting(metrics::Gauge); + +impl Waiting { + fn new() -> Self { + let gauge = metrics::gauge!(DECODES_WAITING); + gauge.increment(1.0); + Self(gauge) + } +} + +impl Drop for Waiting { + fn drop(&mut self) { + self.0.decrement(1.0); + } +} + /// Install outside all inference middleware, including request logging and /// daemon dispatch bypasses. File uploads use a separate router and limit. /// A zero limit retains the configured unlimited behavior. @@ -283,8 +323,17 @@ async fn decode(slots: &Arc, encoding: Encoding, bytes: Bytes, max_by /// moves into the thread with the work: a request dropped mid-decode (the /// client went away) stops waiting for the result, but the work runs on and /// must keep counting against the bound until it ends. +/// +/// The time to get a slot is recorded in [`DECODE_SLOT_WAIT_SECONDS`], and +/// [`DECODES_WAITING`] counts the decodes still waiting for one: a request +/// dropped while waiting leaves the gauge but records no wait. async fn in_slot(slots: &Arc, work: impl FnOnce() -> T + Send + 'static) -> Option { - let slot = slots.clone().acquire_owned().await.ok()?; + let requested = tokio::time::Instant::now(); + let slot = { + let _waiting = Waiting::new(); + slots.clone().acquire_owned().await.ok()? + }; + metrics::histogram!(DECODE_SLOT_WAIT_SECONDS).record(requested.elapsed().as_secs_f64()); tokio::task::spawn_blocking(move || { let _slot = slot; work() @@ -607,6 +656,70 @@ mod tests { assert_eq!(slots.available_permits(), 1, "the slot is returned after decoding"); } + /// Runs `test` on a current-thread runtime with its clock paused, under a + /// fresh recorder it can render, and returns what that rendered at the end. + fn metered>(test: impl FnOnce(metrics_exporter_prometheus::PrometheusHandle) -> F) -> String { + let recorder = metrics_exporter_prometheus::PrometheusBuilder::new() + .set_buckets_for_metric( + metrics_exporter_prometheus::Matcher::Full(DECODE_SLOT_WAIT_SECONDS.to_string()), + DECODE_SLOT_WAIT_BUCKETS, + ) + .unwrap() + .build_recorder(); + let handle = recorder.handle(); + metrics::with_local_recorder(&recorder, || { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .start_paused(true) + .build() + .unwrap() + .block_on(test(handle.clone())) + }); + handle.render() + } + + fn has_line(rendered: &str, line: &str) -> bool { + rendered.lines().any(|l| l == line) + } + + #[test] + fn records_how_long_a_decode_waited_for_its_slot() { + let rendered = metered(|handle| async move { + let slots = Arc::new(Semaphore::new(1)); + in_slot(&slots, || ()).await.unwrap(); + + let busy = slots.acquire().await.unwrap(); + let mut waiting = Box::pin(in_slot(&slots, || ())); + assert!(futures::poll!(&mut waiting).is_pending(), "every slot is taken"); + let during = handle.render(); + assert!(has_line(&during, "dwctl_inference_decodes_waiting 1"), "{during}"); + tokio::time::advance(std::time::Duration::from_millis(300)).await; + drop(busy); + waiting.await.unwrap(); + }); + let wait = |le: &str, count: u32| format!(r#"dwctl_inference_decode_slot_wait_seconds_bucket{{le="{le}"}} {count}"#); + assert!(has_line(&rendered, &wait("0.001", 1)), "a free slot is taken at once: {rendered}"); + assert!(has_line(&rendered, &wait("0.25", 1)), "{rendered}"); + assert!(has_line(&rendered, &wait("0.5", 2)), "the held slot's wait is recorded: {rendered}"); + assert!(has_line(&rendered, "dwctl_inference_decodes_waiting 0"), "{rendered}"); + } + + #[test] + fn a_request_dropped_while_waiting_leaves_the_queue() { + let rendered = metered(|_| async { + let slots = Arc::new(Semaphore::new(1)); + let _busy = slots.acquire().await.unwrap(); + let mut waiting = Box::pin(in_slot(&slots, || ())); + assert!(futures::poll!(&mut waiting).is_pending()); + drop(waiting); + }); + assert!(has_line(&rendered, "dwctl_inference_decodes_waiting 0"), "{rendered}"); + assert!( + !rendered.contains("dwctl_inference_decode_slot_wait_seconds_count"), + "an abandoned wait records no duration: {rendered}" + ); + } + #[tokio::test] async fn a_dropped_request_keeps_its_slot_until_the_decode_ends() { let slots = Arc::new(Semaphore::new(1)); diff --git a/dwctl/src/lib.rs b/dwctl/src/lib.rs index 562ddae71..e6d7f0f40 100644 --- a/dwctl/src/lib.rs +++ b/dwctl/src/lib.rs @@ -452,10 +452,17 @@ fn get_or_install_prometheus_handle() -> PrometheusHandle { onwards::FIRST_TOKEN_SECONDS_BUCKETS, ) .expect("Failed to set custom buckets for onwards_first_token_seconds") + .set_buckets_for_metric( + Matcher::Full(crate::inference::body_limit::DECODE_SLOT_WAIT_SECONDS.to_string()), + crate::inference::body_limit::DECODE_SLOT_WAIT_BUCKETS, + ) + .expect("Failed to set custom buckets for dwctl_inference_decode_slot_wait_seconds") .install_recorder() .expect("Failed to install Prometheus recorder"); initialize_database_error_metrics(); crate::metrics::describe_cache_info_metrics(); + crate::inference::body_limit::describe_decode_slot_metrics(); + crate::api::handlers::files::initialize_upload_rejection_metrics(); handle }) .clone() diff --git a/dwctl/src/test/modalities.rs b/dwctl/src/test/modalities.rs index 507d76022..277fc7472 100644 --- a/dwctl/src/test/modalities.rs +++ b/dwctl/src/test/modalities.rs @@ -11,7 +11,9 @@ use crate::db::handlers::repository::Repository; use crate::db::handlers::{Organizations, api_keys::ApiKeys}; use crate::db::models::api_keys::{ApiKeyCreateDBRequest, ApiKeyPurpose}; use crate::modalities::Modality; -use crate::test::utils::{add_auth_headers, create_test_app, create_test_org, create_test_user, create_test_user_with_roles}; +use crate::test::utils::{ + add_auth_headers, create_test_app, create_test_org, create_test_user, create_test_user_with_roles, upload_rejections, +}; use crate::types::UserId; use serde_json::json; use sqlx::PgPool; @@ -206,6 +208,7 @@ async fn batch_entry_points_refuse_an_org_with_batch_disabled(pool: PgPool) { set_disabled(&pool, org.id, &[Modality::Batch]).await; // In org context: files and batches are refused with the owner's message. + let refused = upload_rejections("batch_disabled"); let resp = server .post("/ai/v1/files") .add_header(&h[0].0, &h[0].1) @@ -214,6 +217,7 @@ async fn batch_entry_points_refuse_an_org_with_batch_disabled(pool: PgPool) { .multipart(jsonl_upload()) .await; assert_modality_disabled(&resp, "batch API is disabled for this organization"); + assert_eq!(upload_rejections("batch_disabled"), refused + 1, "the refused upload is counted"); let resp = server .post("/ai/v1/batches") diff --git a/dwctl/src/test/utils.rs b/dwctl/src/test/utils.rs index c07b9d926..640afe627 100644 --- a/dwctl/src/test/utils.rs +++ b/dwctl/src/test/utils.rs @@ -858,3 +858,13 @@ pub async fn wait_for_advisory_waiter(pool: &PgPool, pid: i32) { .await .expect("replica did not wait for the catalog lock"); } + +/// The process-wide count of file uploads refused for `code`. +pub fn upload_rejections(code: &str) -> u64 { + let series = format!(r#"dwctl_file_upload_rejections_total{{code="{code}"}} "#); + crate::get_or_install_prometheus_handle() + .render() + .lines() + .find_map(|line| line.strip_prefix(&series)) + .map_or(0, |count| count.parse().unwrap()) +}