feat: prevent too large results (>500Mb) from OOMing database

This commit is contained in:
Ruben Fiszel
2025-07-30 12:55:39 +00:00
parent 14aa6d9681
commit 4b9683f146
6 changed files with 102 additions and 15 deletions
+1
View File
@@ -14993,6 +14993,7 @@ dependencies = [
"bytes",
"chrono",
"const_format",
"flume",
"futures",
"lazy_static",
"object_store",
+40
View File
@@ -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()
+1
View File
@@ -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
+48 -7
View File
@@ -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<Vec<String>>;
fn size(&self) -> usize;
}
#[derive(serde::Deserialize)]
@@ -588,6 +589,10 @@ impl ValidableJson for WrappedError {
fn wm_labels(&self) -> Option<Vec<String>> {
None
}
fn size(&self) -> usize {
0
}
}
impl ValidableJson for Box<RawValue> {
@@ -600,6 +605,10 @@ impl ValidableJson for Box<RawValue> {
.ok()
.map(|r| r.wm_labels)
}
fn size(&self) -> usize {
self.get().len()
}
}
impl<T: ValidableJson> ValidableJson for Arc<T> {
@@ -610,6 +619,10 @@ impl<T: ValidableJson> ValidableJson for Arc<T> {
fn wm_labels(&self) -> Option<Vec<String>> {
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<T: ValidableJson> ValidableJson for Json<T> {
@@ -632,6 +649,10 @@ impl<T: ValidableJson> ValidableJson for Json<T> {
fn wm_labels(&self) -> Option<Vec<String>> {
self.0.wm_labels()
}
fn size(&self) -> usize {
self.0.size()
}
}
pub async fn register_metric<T, F, F2, R>(
@@ -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<String> = 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<T: Serialize + Send + Sync + ValidableJson>(
@@ -753,15 +775,34 @@ pub async fn add_completed_job<T: Serialize + Send + Sync + ValidableJson>(
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<T: Serialize + Send + Sync + ValidableJson>(
.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:#?}"
+1 -1
View File
@@ -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()}),
}
}
@@ -644,13 +644,14 @@ async fn handle_non_flow_job_error(
job: &MiniPulledJob,
mem_peak: i32,
canceled_by: Option<CanceledBy>,
err: Value,
err_string: String,
err_json: Value,
worker_name: &str,
) -> Result<WrappedError, Error> {
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,