From 4b9683f1462e9c8a577cc9e65a79fcdcd3894da0 Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Wed, 30 Jul 2025 12:55:39 +0000 Subject: [PATCH] feat: prevent too large results (>500Mb) from OOMing database --- backend/Cargo.lock | 1 + backend/windmill-common/src/error.rs | 40 ++++++++++++++ backend/windmill-indexer/Cargo.toml | 1 + backend/windmill-queue/src/jobs.rs | 55 ++++++++++++++++--- backend/windmill-worker/src/common.rs | 2 +- .../windmill-worker/src/result_processor.rs | 18 +++--- 6 files changed, 102 insertions(+), 15 deletions(-) diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 8125ff421e..603504b328 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -14993,6 +14993,7 @@ dependencies = [ "bytes", "chrono", "const_format", + "flume", "futures", "lazy_static", "object_store", diff --git a/backend/windmill-common/src/error.rs b/backend/windmill-common/src/error.rs index a6d43a5fc4..6418968d28 100644 --- a/backend/windmill-common/src/error.rs +++ b/backend/windmill-common/src/error.rs @@ -40,6 +40,8 @@ pub enum Error { RequireAdmin(String), #[error("{0}")] ExecutionErr(String), + #[error("{0}")] + ResultTooLarge(String), #[error("IoErr: {error:#} @{location:#}")] IoErr { error: io::Error, location: String }, #[error("Utf8Err: {error:#} @{location:#}")] @@ -84,6 +86,44 @@ pub enum Error { Generic(StatusCode, String), } +impl Error { + pub fn name(&self) -> &str { + match self { + Self::ExecutionErr(_) => "ExecutionErr", + Self::ResultTooLarge(_) => "ResultTooLarge", + Self::BadRequest(_) => "BadRequest", + Self::QuotaExceeded(_) => "QuotaExceeded", + Self::InternalErr(_) => "InternalErr", + Self::InternalErrLoc { .. } => "InternalErr", + Self::InternalErrAt(_, _) => "InternalErr", + Self::Anyhow { .. } => "Anyhow", + Self::JsonErr(_) => "JsonErr", + Self::AIError(_) => "AIError", + Self::AlreadyCompleted(_) => "AlreadyCompleted", + Self::FindPythonError(_) => "FindPythonError", + Self::ArgumentErr(_) => "ArgumentErr", + Self::Generic(_, _) => "Generic", + Self::IoErr { .. } => "IoErr", + Self::Utf8Err { .. } => "Utf8Err", + Self::UuidErr { .. } => "UuidErr", + Self::SqlErr { .. } => "SqlErr", + Self::SerdeJson { .. } => "SerdeJson", + Self::HexErr { .. } => "HexErr", + Self::DatabaseMigration(_) => "DatabaseMigration", + Self::ExitStatus(_, _) => "ExitStatus", + Self::ExecutionRawError(_) => "ExecutionRawError", + Self::BadGateway(_) => "BadGateway", + Self::BadConfig(_) => "BadConfig", + Self::ConnectingToDatabase(_) => "ConnectingToDatabase", + Self::NotFound(_) => "NotFound", + Self::NotAuthorized(_) => "NotAuthorized", + Self::MetricNotFound(_) => "MetricNotFound", + Self::PermissionDenied(_) => "PermissionDenied", + _ => "InternalErr", + } + } +} + fn prettify_location(location: &'static Location<'static>) -> String { location .to_string() diff --git a/backend/windmill-indexer/Cargo.toml b/backend/windmill-indexer/Cargo.toml index ea95d6237d..8320be4327 100644 --- a/backend/windmill-indexer/Cargo.toml +++ b/backend/windmill-indexer/Cargo.toml @@ -32,3 +32,4 @@ object_store = { workspace = true, optional = true} tokio-tar.workspace = true lazy_static.workspace = true const_format.workspace = true +flume.workspace = true \ No newline at end of file diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index 8ac7641816..eadd75fe45 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -20,7 +20,7 @@ use reqwest::Client; use serde::Deserialize; use serde::{ser::SerializeMap, Serialize}; use serde_json::{json, value::RawValue}; -use sqlx::PgExecutor; +use sqlx::{Encode, PgExecutor}; use sqlx::{types::Json, Pool, Postgres, Transaction}; use tokio::{sync::RwLock, time::sleep}; use ulid::Ulid; @@ -573,6 +573,7 @@ pub struct WrappedError { pub trait ValidableJson { fn is_valid_json(&self) -> bool; fn wm_labels(&self) -> Option>; + fn size(&self) -> usize; } #[derive(serde::Deserialize)] @@ -588,6 +589,10 @@ impl ValidableJson for WrappedError { fn wm_labels(&self) -> Option> { None } + + fn size(&self) -> usize { + 0 + } } impl ValidableJson for Box { @@ -600,6 +605,10 @@ impl ValidableJson for Box { .ok() .map(|r| r.wm_labels) } + + fn size(&self) -> usize { + self.get().len() + } } impl ValidableJson for Arc { @@ -610,6 +619,10 @@ impl ValidableJson for Arc { fn wm_labels(&self) -> Option> { T::wm_labels(&self) } + + fn size(&self) -> usize { + T::size(&self) + } } impl ValidableJson for serde_json::Value { @@ -622,6 +635,10 @@ impl ValidableJson for serde_json::Value { .ok() .map(|r| r.wm_labels) } + + fn size(&self) -> usize { + self.size_hint() + } } impl ValidableJson for Json { @@ -632,6 +649,10 @@ impl ValidableJson for Json { fn wm_labels(&self) -> Option> { self.0.wm_labels() } + + fn size(&self) -> usize { + self.0.size() + } } pub async fn register_metric( @@ -718,6 +739,7 @@ pub async fn add_completed_job_error( lazy_static::lazy_static! { pub static ref GLOBAL_ERROR_HANDLER_PATH_IN_ADMINS_WORKSPACE: Option = std::env::var("GLOBAL_ERROR_HANDLER_PATH_IN_ADMINS_WORKSPACE").ok(); + pub static ref MAX_RESULT_SIZE: usize = std::env::var("MAX_RESULT_SIZE_MB").unwrap_or("500".to_string()).parse().unwrap_or(500); } pub async fn add_completed_job( @@ -753,15 +775,34 @@ pub async fn add_completed_job( let job_id = queued_job.id; // tracing::error!("1 {:?}", start.elapsed()); - tracing::debug!( - "completed job {} {}", - queued_job.id, - serde_json::to_string(&result).unwrap_or_else(|_| "".to_string()) - ); + // tracing::debug!( + // "completed job {} {}", + // queued_job.id, + // serde_json::to_string(&result).unwrap_or_else(|_| "".to_string()) + // ); let mem_peak = mem_peak; // add_time!(bench, "add_completed_job query START"); + let result_size = result.size() / 1024 / 1024; + if result_size > 2 { + if result_size > *MAX_RESULT_SIZE { + tracing::error!("Result of job {} is too large: {}MB > MAX_RESULT_SIZE={}MB", queued_job.id, result_size, *MAX_RESULT_SIZE); + return Err(Error::ResultTooLarge(format!("Result of job {} is too large: {}MB > MAX_RESULT_SIZE={}MB.\nUse external storages such as the Windmill Object Storage to store large results: https://www.windmill.dev/docs/core_concepts/object_storage_in_windmill", queued_job.id, result_size, *MAX_RESULT_SIZE))); + } + append_logs( + &queued_job.id, + &queued_job.workspace_id, + format!("Warning: Result of job {} is large: {}MB.\nRecommended max size is 2MB.\nPrefer using external storages such as the Windmill Object Storage to store large results: https://www.windmill.dev/docs/core_concepts/object_storage_in_windmill", queued_job.id, result_size), + &db.into(), + ) + .await; + if *CLOUD_HOSTED { + return Err(Error::ResultTooLarge(format!("Result of job {} is too large for multi-tenant cloud: {}MB (max 2MB).\nUse external storages such as the Windmill Object Storage to store large results: https://www.windmill.dev/docs/core_concepts/object_storage_in_windmill", queued_job.id, result_size))); + } else { + tracing::warn!("Result of job {} is larger than 2MB: {}MB. Not recommended.", queued_job.id, result_size); + } + } let _duration = sqlx::query_scalar!( "INSERT INTO v2_job_completed AS cj ( workspace_id @@ -1052,7 +1093,7 @@ pub async fn add_completed_job( .with_max_times(5) .build(), ) - .when(|err| !matches!(err, Error::QuotaExceeded(_))) + .when(|err| !matches!(err, Error::QuotaExceeded(_)) && !matches!(err, Error::ResultTooLarge(_))) .notify(|err, dur| { tracing::error!( "Could not insert completed job, retrying in {dur:#?}, err: {err:#?}" diff --git a/backend/windmill-worker/src/common.rs b/backend/windmill-worker/src/common.rs index 7bfd333668..c29f00c963 100644 --- a/backend/windmill-worker/src/common.rs +++ b/backend/windmill-worker/src/common.rs @@ -524,7 +524,7 @@ pub async fn update_worker_ping_for_failed_init_script( pub fn error_to_value(err: Error) -> serde_json::Value { match err { Error::JsonErr(err) => err, - _ => json!({"message": err.to_string(), "name": "InternalErr"}), + _ => json!({"message": err.to_string(), "name": err.name()}), } } diff --git a/backend/windmill-worker/src/result_processor.rs b/backend/windmill-worker/src/result_processor.rs index 05f9bceb20..86ea8f8eec 100644 --- a/backend/windmill-worker/src/result_processor.rs +++ b/backend/windmill-worker/src/result_processor.rs @@ -644,13 +644,14 @@ async fn handle_non_flow_job_error( job: &MiniPulledJob, mem_peak: i32, canceled_by: Option, - err: Value, + err_string: String, + err_json: Value, worker_name: &str, ) -> Result { append_logs( &job.id, &job.workspace_id, - format!("Unexpected error during job execution:\n{err:#?}"), + format!("Unexpected error during job execution:\n{err_string}"), &db.into(), ) .await; @@ -659,7 +660,7 @@ async fn handle_non_flow_job_error( job, mem_peak, canceled_by, - err, + err_json, worker_name, false, None, @@ -682,7 +683,9 @@ pub async fn handle_job_error( job_completed_tx: JobCompletedSender, #[cfg(feature = "benchmark")] bench: &mut BenchmarkIter, ) { - let err = error_to_value(err); + let err_string = format!("{}: {}", err.name(), err.to_string()); + let err_json = error_to_value(err); + let update_job_future = || async { handle_non_flow_job_error( @@ -690,7 +693,8 @@ pub async fn handle_job_error( job, mem_peak, canceled_by.clone(), - err.clone(), + err_string, + err_json.clone(), worker_name, ) .await @@ -709,8 +713,8 @@ pub async fn handle_job_error( (job.id, Uuid::nil()) }; - let wrapped_error = WrappedError { error: err.clone() }; - tracing::error!(parent_flow = %flow, subflow = %job_status_to_update, "handle job error, updating flow status: {err:?}"); + let wrapped_error = WrappedError { error: err_json.clone() }; + tracing::error!(parent_flow = %flow, subflow = %job_status_to_update, "handle job error, updating flow status: {err_json:?}"); let updated_flow = update_flow_status_after_job_completion( db, client,