diff --git a/backend/Cargo.toml b/backend/Cargo.toml index 53210ac9d0..3ea692746a 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -433,7 +433,9 @@ tower = "^0" tower-http = { version = "^0.6", features = ["trace", "cors", "catch-panic"] } tower-cookies = "^0.11" serde = "^1" -serde_json = { version = "^1", features = ["preserve_order", "raw_value"] } +# 1.0.151 introduced RawValue::from_string_unchecked, which the SQL executors use +# to avoid re-parsing every collected row. +serde_json = { version = "^1.0.151", features = ["preserve_order", "raw_value"] } serde_yml = "0.0.12" uuid = { version = "^1", features = ["serde", "v4", "js"] } thiserror = "^2" diff --git a/backend/windmill-duckdb-ffi-internal/Cargo.toml b/backend/windmill-duckdb-ffi-internal/Cargo.toml index 6151ff3615..73e132a2d7 100644 --- a/backend/windmill-duckdb-ffi-internal/Cargo.toml +++ b/backend/windmill-duckdb-ffi-internal/Cargo.toml @@ -21,7 +21,9 @@ duckdb = { git = "https://github.com/windmill-labs/duckdb-rs", rev = "7190adfcf5 regex = "1" rust_decimal = "1.37.2" serde = { version = "1.0", features = ["derive"] } -serde_json = { version = "^1", features = ["preserve_order", "raw_value"] } +# 1.0.151 introduced RawValue::from_string_unchecked, which row collection uses +# to avoid re-parsing every row. +serde_json = { version = "^1.0.151", features = ["preserve_order", "raw_value"] } [lib] crate-type = ["cdylib"] diff --git a/backend/windmill-duckdb-ffi-internal/src/lib.rs b/backend/windmill-duckdb-ffi-internal/src/lib.rs index 289fbb6c08..69820b285f 100644 --- a/backend/windmill-duckdb-ffi-internal/src/lib.rs +++ b/backend/windmill-duckdb-ffi-internal/src/lib.rs @@ -909,11 +909,6 @@ fn interpolate_named_args<'a>( /// the callers pass a `serde_json::Map` of `Value`s: serializing one cannot fail /// for any reason except the budget. A caller passing a type with a fallible /// `Serialize` would have its error silently retold as a size limit. -/// -/// The bytes come back through `RawValue::from_string`, which re-parses to -/// validate what `serde_json` just wrote. `to_raw_value` skipped that, so this -/// costs one extra pass per row — the price of getting the bytes through a writer -/// we can bound, since the unchecked constructor is not public. fn to_raw_value_within(value: &T, budget: usize) -> Option> { struct Budgeted { buf: Vec, @@ -921,6 +916,12 @@ fn to_raw_value_within(value: &T, budget: usize) -> Option< } impl std::io::Write for Budgeted { fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.write_all(bytes)?; + Ok(bytes.len()) + } + // `Vec` overrides this too: the default implementation loops over + // `write`, and serde_json emits a great many small pieces per row. + fn write_all(&mut self, bytes: &[u8]) -> std::io::Result<()> { if bytes.len() > self.left { return Err(std::io::Error::new( std::io::ErrorKind::WriteZero, @@ -929,7 +930,7 @@ fn to_raw_value_within(value: &T, budget: usize) -> Option< } self.left -= bytes.len(); self.buf.extend_from_slice(bytes); - Ok(bytes.len()) + Ok(()) } fn flush(&mut self) -> std::io::Result<()> { Ok(()) @@ -938,9 +939,15 @@ fn to_raw_value_within(value: &T, budget: usize) -> Option< let mut writer = Budgeted { buf: Vec::new(), left: budget }; serde_json::to_writer(&mut writer, value).ok()?; - String::from_utf8(writer.buf) - .ok() - .and_then(|s| RawValue::from_string(s).ok()) + let json = String::from_utf8(writer.buf).ok()?; + // SAFETY: `to_writer` returned `Ok`, so `json` holds one complete, well-formed + // JSON value with no surrounding whitespace. Running out of budget is the only + // way a partial write happens, and it takes the `?` above instead of reaching + // here. The safe constructor re-parses every row to learn the same thing, which + // measured ~1.8x the cost of serializing it in the first place; serde_json + // itself builds a `RawValue` this way in `to_raw_value`, and `debug_assert!`s + // the invariant by re-parsing in debug builds. + Some(unsafe { RawValue::from_string_unchecked(json) }) } /// How much a row may still expand to, carried through every value it contains. diff --git a/backend/windmill-worker/src/pg_executor.rs b/backend/windmill-worker/src/pg_executor.rs index 0135cf4fcc..1413924bf3 100644 --- a/backend/windmill-worker/src/pg_executor.rs +++ b/backend/windmill-worker/src/pg_executor.rs @@ -46,7 +46,7 @@ use crate::handle_child::run_future_with_polling_update_job_poller; use crate::sanitized_sql_params::sanitize_and_interpolate_unsafe_sql_args; use crate::sql_s3_input::fetch_s3object_as_json_text; use crate::sql_utils::remove_comments; -use crate::MAX_RESULT_SIZE; +use crate::{max_sql_result_size, sql_result_too_large_error, to_raw_value_within}; use bytes::Buf; use lazy_static::lazy_static; use windmill_common::client::AuthedClient; @@ -579,47 +579,64 @@ fn do_postgresql_inner<'a>( rows.boxed() }; - let rows = rows.try_collect::>().await.map_err(to_anyhow)?; + // The stream is consumed one row at a time so the cap below can still + // refuse. Collecting it into a `Vec` first holds every wire buffer + // and every converted row at once, and the worker is already past the + // budget by the time the first check gets to run. + futures::pin_mut!(rows); + let max_result_size = max_sql_result_size(); + let mut envelope = raw_output + .then(|| crate::pg_raw_output::RawOutputEnvelopeBuilder::new(max_result_size)); + let mut column_names: Option> = None; - if let Some(column_order) = column_order { - *column_order = Some( - rows.first() - .map(|x| { - x.columns() - .iter() - .map(|x| x.name().to_string()) - .collect::>() - }) - .unwrap_or_default(), - ); + while let Some(row) = rows.try_next().await.map_err(to_anyhow)? { + if column_names.is_none() { + column_names = Some( + row.columns() + .iter() + .map(|x| x.name().to_string()) + .collect::>(), + ); + } + + if let Some(envelope) = envelope.as_mut() { + envelope.push(row, &format_state, siz)?; + continue; + } + + let v = postgres_row_to_json_value_with_state(row, &format_state)?; + // Serialized under what is left of the budget: escaping can expand + // a row that fit in memory past what remains, and a check placed + // after the write happens once the allocation already did. + let raw = to_raw_value_within( + &v, + max_result_size.saturating_sub(siz.load(Ordering::Relaxed)), + ) + .ok_or_else(|| sql_result_too_large_error(max_result_size))?; + // Both are proxies for what the row costs the worker, and neither + // dominates: the value tree is wider than its JSON for small + // scalars, narrower once escaping expands the text. + siz.fetch_add(sizeof_val(&v).max(raw.get().len()), Ordering::Relaxed); + if siz.load(Ordering::Relaxed) > max_result_size { + return Err(sql_result_too_large_error(max_result_size)); + } + res.push(raw); } - if raw_output { - let envelope = crate::pg_raw_output::build_envelope(rows, &format_state, siz)?; - res.push(to_raw_value(&envelope)); - } else { - for row in rows.into_iter() { - let r = postgres_row_to_json_value_with_state(row, &format_state); - if let Ok(v) = r.as_ref() { - let size = sizeof_val(v); - siz.fetch_add(size, Ordering::Relaxed); - } - if *CLOUD_HOSTED { - let siz = siz.load(Ordering::Relaxed); - if siz > MAX_RESULT_SIZE * 4 { - return Err(Error::ExecutionErr(format!( - "Query result too large for cloud (size = {} > {})", - siz, - MAX_RESULT_SIZE * 4, - ))); - } - } - if let Ok(v) = r { - res.push(to_raw_value(&v)); - } else { - return Err(to_anyhow(r.err().unwrap()).into()); - } - } + if let Some(column_order) = column_order { + // A statement that returned no rows reports no columns. + *column_order = Some(column_names.unwrap_or_default()); + } + + if let Some(envelope) = envelope { + // The envelope is budgeted against the whole cap: its rows were + // already charged as text on the way in, and this is that same + // text serialized, so charging it against the remainder would + // reject a result that passed every row-level check. + res.push( + to_raw_value_within(&envelope.finish(), max_result_size) + .ok_or_else(|| sql_result_too_large_error(max_result_size))?, + ); } } diff --git a/backend/windmill-worker/src/pg_raw_output.rs b/backend/windmill-worker/src/pg_raw_output.rs index 85f8b6d1aa..5e558f65b2 100644 --- a/backend/windmill-worker/src/pg_raw_output.rs +++ b/backend/windmill-worker/src/pg_raw_output.rs @@ -16,11 +16,10 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use serde::Serialize; use serde_json::value::RawValue; use tokio_postgres::Row; -use windmill_common::error::{self, to_anyhow, Error}; -use windmill_common::worker::CLOUD_HOSTED; +use windmill_common::error::{self, to_anyhow}; use crate::pg_executor::{postgres_row_to_row_data_with_state, JSONValue, ResultFormatState}; -use crate::MAX_RESULT_SIZE; +use crate::sql_result_too_large_error; #[derive(Serialize, Clone, Debug)] pub struct RawOutputColumn { @@ -44,59 +43,79 @@ impl RawOutputEnvelope { } } -/// Build a `RawOutputEnvelope` for one statement's worth of rows. Reuses the -/// existing JSON-cell formatter (numeric precision warning, interval/timetz -/// coercion, JSON columns…) and only post-processes its output into the text -/// form Postgres's wire protocol expects. -pub fn build_envelope( - rows: Vec, - format_state: &ResultFormatState, - siz: &AtomicUsize, -) -> error::Result { - let columns: Vec = rows - .first() - .map(|r| { - r.columns() +/// What one converted row costs the worker. Containers count, not just the text +/// they hold: a row of NULLs still allocates a `Vec>` and one +/// `Option` per column, so charging only for string bytes would let an +/// all-NULL or all-empty result grow without ever moving the total. +fn text_row_size(text_row: &[Option]) -> usize { + std::mem::size_of::>>() + + text_row.len() * std::mem::size_of::>() + + text_row + .iter() + .filter_map(|v| v.as_deref().map(str::len)) + .sum::() +} + +/// Builds a `RawOutputEnvelope` one row at a time, charging each against the +/// caller's running total as it goes. Reuses the existing JSON-cell formatter +/// (numeric precision warning, interval/timetz coercion, JSON columns…) and only +/// post-processes its output into the text form Postgres's wire protocol expects. +/// +/// Taking rows one at a time is what lets the caller drop each wire buffer as its +/// text form is appended, so the whole result is never resident twice. +pub struct RawOutputEnvelopeBuilder { + columns: Vec, + seen_first_row: bool, + rows: Vec>>, + max_result_size: usize, +} + +impl RawOutputEnvelopeBuilder { + /// Takes the caller's budget rather than reading it again: the envelope is + /// serialized under the same cap once it is finished, and the two have to be + /// the same number for that to mean anything. + pub fn new(max_result_size: usize) -> Self { + Self { columns: Vec::new(), seen_first_row: false, rows: Vec::new(), max_result_size } + } + + pub fn push( + &mut self, + row: Row, + format_state: &ResultFormatState, + siz: &AtomicUsize, + ) -> error::Result<()> { + if !self.seen_first_row { + self.seen_first_row = true; + self.columns = row + .columns() .iter() .map(|c| RawOutputColumn { name: c.name().to_string(), oid: c.type_().oid(), type_name: c.type_().name().to_string(), }) - .collect() - }) - .unwrap_or_default(); + .collect(); + } - let mut text_rows: Vec>> = Vec::with_capacity(rows.len()); - for row in rows { let row_data = postgres_row_to_row_data_with_state(row, format_state).map_err(to_anyhow)?; - let text_row: Vec> = columns + let text_row: Vec> = self + .columns .iter() .map(|c| { json_value_to_pg_text(row_data.get(&c.name).cloned().unwrap_or(JSONValue::Null)) }) .collect(); - siz.fetch_add( - text_row - .iter() - .filter_map(|v| v.as_deref().map(str::len)) - .sum(), - Ordering::Relaxed, - ); - if *CLOUD_HOSTED { - let total = siz.load(Ordering::Relaxed); - if total > MAX_RESULT_SIZE * 4 { - return Err(Error::ExecutionErr(format!( - "Query result too large for cloud (size = {} > {})", - total, - MAX_RESULT_SIZE * 4, - ))); - } + siz.fetch_add(text_row_size(&text_row), Ordering::Relaxed); + if siz.load(Ordering::Relaxed) > self.max_result_size { + return Err(sql_result_too_large_error(self.max_result_size)); } - text_rows.push(text_row); + self.rows.push(text_row); + Ok(()) } - Ok(RawOutputEnvelope { columns, rows: text_rows }) + pub fn finish(self) -> RawOutputEnvelope { + RawOutputEnvelope { columns: self.columns, rows: self.rows } + } } /// Extract the single envelope `Box` from the per-statement results @@ -194,6 +213,20 @@ mod tests { ); } + /// The size cap is the only thing standing between a large raw_output result + /// and the OOM killer, so a row has to cost something even when every cell is + /// NULL — otherwise millions of them accumulate against a total that never + /// moves. + #[test] + fn null_and_empty_cells_still_count_toward_the_cap() { + assert!(text_row_size(&[None, None, None]) > 0); + assert!(text_row_size(&[Some(String::new())]) > 0); + assert!(text_row_size(&[None, None]) > text_row_size(&[None])); + assert!( + text_row_size(&[Some("abcde".to_string())]) > text_row_size(&[Some(String::new())]) + ); + } + #[test] fn extract_envelope_falls_back_to_empty_when_no_statement_produced_one() { let raw = extract_envelope_or_empty(vec![]); diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index bf8a3a5bf8..4c2c2c64f6 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/backend/windmill-worker/src/worker.rs @@ -1270,19 +1270,18 @@ pub const MAX_RESULT_SIZE: usize = 1024 * 1024 * 2; // 2MB // costs several times the JSON they serialize to — a separately allocated value // per row, then a contiguous buffer holding all of them — and the worker still // needs the rest of its budget for what it already has resident. -#[cfg(feature = "duckdb")] const SQL_RESULT_SIZE_FRACTION: f64 = 0.15; // Under this a result cannot threaten a worker of any size, so capping it would // only reject work that would have succeeded. -#[cfg(feature = "duckdb")] const MIN_MAX_SQL_RESULT_SIZE: usize = 8 * 1024 * 1024; /// `"512"`, `"512MB"`, `"2GiB"`, `"1.5GB"` -> bytes. Suffixes are case-insensitive /// and binary, so `MB` and `MiB` both mean 1024². /// -/// Fractions are accepted because `format_byte_size` emits them, and the limit it -/// renders into an error is meant to be usable as a setting verbatim. -#[cfg(feature = "duckdb")] +/// Fractions have to be accepted even though `format_byte_size` never emits one: +/// the duckdb error is rendered by the FFI crate's own copy of that helper, which +/// rounds to a fraction above 1 GiB, and every limit an error quotes is meant to +/// be usable as a setting verbatim. fn parse_byte_size(v: &str) -> Option { let upper = v.trim().to_ascii_uppercase(); // Longest-first: `GB` would otherwise swallow `GIB`, and `B` every other suffix. @@ -1306,9 +1305,8 @@ fn parse_byte_size(v: &str) -> Option { (bytes <= usize::MAX as f64).then(|| bytes as usize) } -#[cfg(feature = "duckdb")] lazy_static::lazy_static! { - /// Bytes one duckdb job may collect before the executor gives up — the budget + /// Bytes one SQL job may collect before its executor gives up — the budget /// spans every query block in the job, since what the worker cannot survive is /// the total it ends up holding. Nothing else bounds it at collection time: /// every row is accumulated before anything can stream the result out, so an @@ -1322,11 +1320,9 @@ lazy_static::lazy_static! { /// with no cgroup reading to scale from both mean no cap. /// /// It bounds what is *collected*, not the process: the collected rows are - /// still live while `serde_json::to_string` builds a second whole copy, and - /// the worker parses that copy back for every strategy but - /// `AllStatementsAllRows`, so peak sits near twice the cap. The derived - /// default leaves room for that — it is a fraction of the worker's budget, - /// not the whole of it. + /// still live while a second whole copy is serialized out of them, so peak + /// sits near twice the cap. The derived default leaves room for that — it is + /// a fraction of the worker's budget, not the whole of it. pub(crate) static ref MAX_SQL_RESULT_SIZE: usize = { let explicit = std::env::var("MAX_SQL_RESULT_SIZE").ok().and_then(|v| { let parsed = parse_byte_size(&v); @@ -1357,6 +1353,113 @@ lazy_static::lazy_static! { }; } +/// The limit postgres collection is bounded by: the cloud product cap where one +/// applies, and otherwise the worker-survival cap, which is the only thing worth +/// enforcing on a deployment that has no product limit to answer to. +/// +/// Duckdb deliberately does not come through here — its cap is a survival limit +/// only, so it reads `MAX_SQL_RESULT_SIZE` directly and is never narrowed to the +/// cloud product limit. +pub(crate) fn max_sql_result_size() -> usize { + if *CLOUD_HOSTED { + MAX_RESULT_SIZE * 4 + } else { + *MAX_SQL_RESULT_SIZE + } +} + +/// Renders a byte count so the figure in the error names the limit exactly and +/// can be set as `MAX_SQL_RESULT_SIZE` verbatim. +/// +/// A unit is only used when it divides the count evenly. Rounding to the nearest +/// MB reads better but names a threshold nobody configured — a 1.5MB limit shown +/// as `1MB` both misreports it and lowers it if pasted back — and the +/// memory-derived default is rarely a whole number of MB, which is exactly the +/// case an operator is most likely to copy. +fn format_byte_size(bytes: usize) -> String { + [("GB", 1usize << 30), ("MB", 1 << 20), ("KB", 1 << 10)] + .into_iter() + .find(|(_, unit)| bytes >= *unit && bytes % unit == 0) + .map(|(suffix, unit)| format!("{}{suffix}", bytes / unit)) + .unwrap_or_else(|| format!("{bytes}B")) +} + +/// Serializes `value` to JSON, refusing to allocate more than `budget` bytes. +/// +/// A running total kept over values in memory does not see what serializing them +/// costs: JSON escaping expands text on the way out — one control character +/// becomes the six-byte escape `\u0001` — so a value that fit the budget unescaped +/// can still allocate several times it while being written, long past the point +/// where a check between rows could help. Bounding the writer is what keeps that +/// expansion inside the budget rather than inside the cgroup. +/// +/// `None` means the output did not fit. Serializing a `serde_json::Value` cannot +/// fail for any other reason, which is what makes that reading unambiguous. +pub(crate) fn to_raw_value_within( + value: &T, + budget: usize, +) -> Option> { + struct Budgeted { + buf: Vec, + left: usize, + } + impl std::io::Write for Budgeted { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.write_all(bytes)?; + Ok(bytes.len()) + } + // `Vec` overrides this too: the default implementation loops over + // `write`, and serde_json emits a great many small pieces per row. + fn write_all(&mut self, bytes: &[u8]) -> std::io::Result<()> { + if bytes.len() > self.left { + return Err(std::io::Error::new( + std::io::ErrorKind::WriteZero, + "result over budget", + )); + } + self.left -= bytes.len(); + self.buf.extend_from_slice(bytes); + Ok(()) + } + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + let mut writer = Budgeted { buf: Vec::new(), left: budget }; + serde_json::to_writer(&mut writer, value).ok()?; + let json = String::from_utf8(writer.buf).ok()?; + // SAFETY: `to_writer` returned `Ok`, so `json` holds one complete, well-formed + // JSON value with no surrounding whitespace. Running out of budget is the only + // way a partial write happens, and it takes the `?` above instead of reaching + // here. The safe constructor re-parses every row to learn the same thing, which + // measured ~1.8x the cost of serializing it in the first place; serde_json + // itself builds a `RawValue` this way in `to_raw_value`, and `debug_assert!`s + // the invariant by re-parsing in debug builds. + Some(unsafe { serde_json::value::RawValue::from_string_unchecked(json) }) +} + +/// Wording shared by the SQL executors that collect rows in the worker process. +/// `MAX_SQL_RESULT_SIZE` is not settable on cloud, so only mention it off-cloud. +/// +/// Only the limit is quoted: collection stops on the row that crosses it, so the +/// running total is the threshold plus one row, not the size of the result. +pub(crate) fn sql_result_too_large_error(limit: usize) -> Error { + // Each branch is a whole sentence: splicing a prefix in leaves the cloud + // message starting mid-sentence, since there is no prefix to splice there. + let remedy = if *CLOUD_HOSTED { + "Return fewer rows" + } else { + "Raise MAX_SQL_RESULT_SIZE, or return fewer rows" + }; + Error::ExecutionErr(format!( + "Query result too large: collecting it passed the {} limit. {remedy} — \ + aggregate, add a LIMIT, or write the rows out from the query instead of \ + returning them.", + format_byte_size(limit), + )) +} + #[derive(Clone)] pub struct SameWorkerSender(pub Sender, pub Arc); @@ -4960,7 +5063,7 @@ pub async fn write_module_files( Ok(()) } -#[cfg(all(test, feature = "duckdb"))] +#[cfg(test)] mod byte_size_tests { use super::*; @@ -4991,6 +5094,45 @@ mod byte_size_tests { ); } } + + /// Pins the property the budgeted writer exists for: a value is charged by + /// its size in memory, and escaping makes the serialized form diverge from + /// that by up to 6x. + #[test] + fn escaping_cannot_outgrow_the_budget() { + // One control character in, six bytes of `\u0001` out. + let value = serde_json::json!({ "a": "\u{1}".repeat(1000) }); + let serialized_len = serde_json::to_string(&value).unwrap().len(); + assert!( + serialized_len > 6000, + "expected escaping to expand: {serialized_len}" + ); + + // A budget that the unescaped bytes would clear, and the escaped ones cannot. + assert!(to_raw_value_within(&value, 2000).is_none()); + + // The `RawValue` is built without re-parsing, so nothing but this check + // stands between a serializer change and a malformed value being handed + // out as valid JSON. Escaped text is the case most likely to expose it. + let fitted = to_raw_value_within(&value, serialized_len).expect("fits its own length"); + assert_eq!(fitted.get(), serde_json::to_string(&value).unwrap()); + } + + /// The error quotes the limit so it can be set verbatim, which only holds if + /// the rendered figure means the same number of bytes. Rounding to a unit + /// that does not divide it evenly parses fine and still names a different + /// limit, so parseability alone is not the property worth pinning. + #[test] + fn every_rendered_limit_round_trips_exactly() { + for bytes in [512, 8 << 20, 307 << 20, 2 << 30, 2_576_980_377, 16 << 30] { + let rendered = format_byte_size(bytes); + assert_eq!( + parse_byte_size(&rendered), + Some(bytes), + "format_byte_size({bytes}) = {rendered:?}, which names a different limit" + ); + } + } } #[cfg(test)]