diff --git a/doc/user/content/sql/create-source/webhook.md b/doc/user/content/sql/create-source/webhook.md index 20533b61bd50e..09b119dc20a50 100644 --- a/doc/user/content/sql/create-source/webhook.md +++ b/doc/user/content/sql/create-source/webhook.md @@ -316,11 +316,15 @@ SELECT COUNT(body) FROM webhook_source_ndjson; Webhook sources apply the following limits to received requests: -* The maximum size of the request body is **`2MB`**. Requests larger than this +* The maximum size of the request body is **`5MB`**. Requests larger than this will fail with `413 Payload Too Large`. * The maximum number of concurrent requests across **all** webhook sources is **500**. Trying to connect when the server is at capacity will fail with `429 Too Many Requests`. +* A `CHECK` expression may use at most **`20MB`** of temporary memory while + validating a single request. A `CHECK` that needs more, for example one that + builds a large string out of the request body, will fail with + `400 Bad Request`. * Requests that contain a header name specified more than once will be rejected with `401 Unauthorized`. diff --git a/misc/python/materialize/mzcompose/__init__.py b/misc/python/materialize/mzcompose/__init__.py index 18afa831d4b44..baefa66a07b7e 100644 --- a/misc/python/materialize/mzcompose/__init__.py +++ b/misc/python/materialize/mzcompose/__init__.py @@ -767,6 +767,7 @@ def get_default_system_parameters( "mcp_request_timeout", "user_id_pool_batch_size", "webhook_max_request_size_bytes", + "webhook_validation_memory_budget_bytes", "cluster_controller_tick_interval", "default_cluster_reconfiguration_timeout", "read_then_write_max_dependencies", diff --git a/misc/python/materialize/parallel_workload/action.py b/misc/python/materialize/parallel_workload/action.py index 03f25fd615ca1..60bb6bad07fec 100644 --- a/misc/python/materialize/parallel_workload/action.py +++ b/misc/python/materialize/parallel_workload/action.py @@ -3057,6 +3057,15 @@ def __init__( "5242880", "10485760", ] + # The CHECK expressions this workload generates allocate at most a few + # bytes of temporary storage, well under the 1 MiB floor here, so none of + # these values can make one fail. + self.flags_with_values["webhook_validation_memory_budget_bytes"] = [ + # 1 MiB, 20 MiB (default), 100 MiB + "1048576", + "20971520", + "104857600", + ] self.flags_with_values["aws_prefetch_sts_connect_timeout"] = [ "'3100ms'", "'30s'", diff --git a/src/adapter-types/src/dyncfgs.rs b/src/adapter-types/src/dyncfgs.rs index 327af7a7c464e..5d7f2e3661b2f 100644 --- a/src/adapter-types/src/dyncfgs.rs +++ b/src/adapter-types/src/dyncfgs.rs @@ -276,6 +276,23 @@ pub const WEBHOOK_MAX_REQUEST_SIZE_BYTES: Config = Config::new( "The maximum size in bytes of a webhook request body, measured after decompression.", ); +/// Maximum temporary storage a webhook `CHECK` expression may allocate while +/// validating one request. A `CHECK` that exceeds it fails the request with HTTP +/// 400 rather than holding the memory. +/// +/// This exists because a `CHECK` can allocate a multiple of the request body, +/// and `environmentd` evaluates one per in-flight request: without a bound +/// proportionate to the request, a bounded amount of network input becomes an +/// unbounded amount of heap on a process shared by every connection. The default +/// is 4x `WEBHOOK_MAX_REQUEST_SIZE_BYTES`, far above what a realistic `CHECK` +/// (an HMAC, a `decode`, a `concat` with a secret) needs, and well below the +/// 100 MiB per-call ceiling that applies in a cluster. +pub const WEBHOOK_VALIDATION_MEMORY_BUDGET_BYTES: Config = Config::new( + "webhook_validation_memory_budget_bytes", + 20 * 1024 * 1024, + "The maximum bytes of temporary storage a webhook CHECK expression may allocate while validating one request.", +); + /// Number of user IDs to pre-allocate in a batch. Pre-allocating IDs avoids /// a persist write + oracle call per DDL statement. pub const USER_ID_POOL_BATCH_SIZE: Config = Config::new( @@ -445,6 +462,7 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet { .add(&MCP_MAX_RESPONSE_SIZE) .add(&MCP_REQUEST_TIMEOUT) .add(&WEBHOOK_MAX_REQUEST_SIZE_BYTES) + .add(&WEBHOOK_VALIDATION_MEMORY_BUDGET_BYTES) .add(&USER_ID_POOL_BATCH_SIZE) .add(&GROUP_COMMIT_MAX_ATTEMPTS) .add(&CONSOLE_OIDC_CLIENT_ID) diff --git a/src/adapter/src/error.rs b/src/adapter/src/error.rs index ac348a8dd26d3..c362362c9af8b 100644 --- a/src/adapter/src/error.rs +++ b/src/adapter/src/error.rs @@ -474,6 +474,7 @@ fn eval_error_code(err: &EvalError) -> SqlState { EvalError::StringValueTooLong { .. } => SqlState::STRING_DATA_RIGHT_TRUNCATION, EvalError::LikePatternTooLong | EvalError::LengthTooLarge + | EvalError::TempStorageBudgetExceeded | EvalError::NullCharacterNotPermitted | EvalError::MaxArraySizeExceeded(_) | EvalError::LetRecLimitExceeded(_) => SqlState::PROGRAM_LIMIT_EXCEEDED, diff --git a/src/adapter/src/webhook.rs b/src/adapter/src/webhook.rs index ce641983ba9c5..a86a169efde64 100644 --- a/src/adapter/src/webhook.rs +++ b/src/adapter/src/webhook.rs @@ -76,11 +76,17 @@ impl AppendWebhookValidator { } } + /// Runs the validation expression against one request. + /// + /// `memory_budget` caps the temporary storage the expression may allocate. The expression is + /// user-authored and runs in `environmentd`, so without a cap proportionate to the request one + /// `CHECK` can turn a bounded body into an unbounded amount of heap on a shared process. pub async fn eval( self, body: bytes::Bytes, headers: Arc>, received_at: DateTime, + memory_budget: usize, ) -> Result { let AppendWebhookValidator { validation, @@ -125,9 +131,7 @@ impl AppendWebhookValidator { // work. let validate = move || { // Gather our Datums for evaluation - // - // TODO(parkmycar): Re-use the RowArena when we implement rate limiting. - let temp_storage = RowArena::default(); + let temp_storage = RowArena::with_budget(memory_budget); let mut datums = Vec::with_capacity( body_columns.len() + header_columns.len() + secret_contents.len(), ); diff --git a/src/environmentd/src/http/webhook.rs b/src/environmentd/src/http/webhook.rs index 94a24f506043e..7058e5f52b344 100644 --- a/src/environmentd/src/http/webhook.rs +++ b/src/environmentd/src/http/webhook.rs @@ -26,7 +26,9 @@ use axum::extract::{Path, State}; use axum::response::IntoResponse; use bytes::Bytes; use http::StatusCode; -use mz_adapter_types::dyncfgs::WEBHOOK_MAX_REQUEST_SIZE_BYTES; +use mz_adapter_types::dyncfgs::{ + WEBHOOK_MAX_REQUEST_SIZE_BYTES, WEBHOOK_VALIDATION_MEMORY_BUDGET_BYTES, +}; use thiserror::Error; use crate::http::WebhookState; @@ -42,6 +44,7 @@ pub async fn handle_webhook( body: Body, ) -> impl IntoResponse { let max_request_size = WEBHOOK_MAX_REQUEST_SIZE_BYTES.get(&dyncfgs); + let validation_memory_budget = WEBHOOK_VALIDATION_MEMORY_BUDGET_BYTES.get(&dyncfgs); let body = axum::body::to_bytes(body, max_request_size) .await .map_err(|err| { @@ -88,6 +91,7 @@ pub async fn handle_webhook( &name, &body, &headers, + validation_memory_budget, ) .await; @@ -114,6 +118,7 @@ async fn append_webhook( name: &str, body: &Bytes, headers: &Arc>, + validation_memory_budget: usize, ) -> Result<(), AppendWebhookError> { // Shenanigans to get the types working for the async retry. let (database, schema, name) = (database.to_string(), schema.to_string(), name.to_string()); @@ -164,7 +169,12 @@ async fn append_webhook( // If this source requires validation, then validate! if let Some(validator) = validator { let valid = validator - .eval(Bytes::clone(body), Arc::clone(headers), received_at) + .eval( + Bytes::clone(body), + Arc::clone(headers), + received_at, + validation_memory_budget, + ) .await?; if !valid { return Err(AppendWebhookError::ValidationFailed); diff --git a/src/environmentd/tests/server.rs b/src/environmentd/tests/server.rs index 7e3c0dd58f669..882713185d29d 100644 --- a/src/environmentd/tests/server.rs +++ b/src/environmentd/tests/server.rs @@ -3193,6 +3193,108 @@ fn webhook_max_request_size() { .expect("2 KiB body rejected after lowering the limit to 1 KiB"); } +/// A `CHECK` expression that allocates a multiple of the request body must be refused rather than +/// allowed to hold the memory (SQL-431). +/// +/// `environmentd` evaluates one `CHECK` per in-flight request, so before the budget existed 150 +/// concurrent 5 MB posts to a source checking `length(repeat(body, 20)) >= 0` took the process from +/// 355 MiB to 8.7 GiB, growing with the number of clients. Every request returned 200. Nothing +/// refused the work, so the memory was just held. +#[mz_ore::test] +#[cfg_attr(miri, ignore)] // too slow +#[allow(clippy::disallowed_methods)] +fn webhook_validation_memory_budget() { + let server = test_util::TestHarness::default() + .unsafe_mode() + .start_blocking(); + + let mut mz_client = server + .pg_config_internal() + .user(&SYSTEM_USER.name) + .connect(postgres::NoTls) + .unwrap(); + + let mut client = server.connect(postgres::NoTls).unwrap(); + + client + .execute( + "CREATE CLUSTER webhook_cluster (SIZE 'scale=1,workers=1');", + &[], + ) + .expect("failed to create cluster"); + client + .execute( + "CREATE SOURCE webhook_amplify IN CLUSTER webhook_cluster \ + FROM WEBHOOK BODY FORMAT TEXT \ + CHECK (WITH (BODY) length(repeat(body, 20)) >= 0)", + &[], + ) + .expect("failed to create source"); + + let http_client = reqwest::Client::new(); + let webhook_url = format!( + "http://{}/api/webhook/materialize/public/webhook_amplify", + server.http_local_addr(), + ); + + let post = |len: usize| { + let body = vec![b'a'; len]; + server.runtime().block_on(async { + http_client + .post(&webhook_url) + .body(body) + .send() + .await + .expect("request failed") + .status() + }) + }; + + // The default budget is 20 MiB. A 2 MiB body amplifies to 40 MiB, so the `CHECK` is refused + // with 400 rather than allocating. Note that nothing else rejects this: 40 MiB is comfortably + // under the 100 MiB per-call ceiling that applies in a cluster, and 2 MiB is under the 5 MiB + // request-size limit. + assert_eq!(post(2 * 1024 * 1024).as_u16(), 400); + + // A body whose amplified size fits the budget still succeeds through the same source, so the + // rejection above is the budget and not the `CHECK` or the source being broken. + assert!(post(512 * 1024).is_success()); + + // Raising the budget above the amplified size accepts the body that was just rejected. This is + // what pins the test to the budget: with enforcement removed the 2 MiB case would already have + // succeeded and the first assertion would fail. + mz_client + .batch_execute("ALTER SYSTEM SET webhook_validation_memory_budget_bytes = 67108864") + .unwrap(); + // The dyncfg propagates to the shared persist ConfigSet asynchronously, so retry briefly. + Retry::default() + .max_duration(std::time::Duration::from_secs(30)) + .retry(|_| { + if post(2 * 1024 * 1024).is_success() { + Ok(()) + } else { + Err(()) + } + }) + .expect("2 MiB body accepted after raising the budget to 64 MiB"); + + // And lowering it below what the smaller body needs rejects that too, confirming the knob is + // live in both directions rather than the first result being a fixed threshold. + mz_client + .batch_execute("ALTER SYSTEM SET webhook_validation_memory_budget_bytes = 1024") + .unwrap(); + Retry::default() + .max_duration(std::time::Duration::from_secs(30)) + .retry(|_| { + if post(512 * 1024).as_u16() == 400 { + Ok(()) + } else { + Err(()) + } + }) + .expect("512 KiB body rejected after lowering the budget to 1 KiB"); +} + #[mz_ore::test] #[cfg_attr(miri, ignore)] // too slow #[allow(clippy::disallowed_methods)] diff --git a/src/expr/src/scalar.proto b/src/expr/src/scalar.proto index c016fa370beb9..d3a03324ef9d3 100644 --- a/src/expr/src/scalar.proto +++ b/src/expr/src/scalar.proto @@ -165,5 +165,6 @@ message ProtoEvalError { string invalid_catalog_json = 81; string redact_error = 82; google.protobuf.Empty negative_rows_from_subquery = 83; + google.protobuf.Empty temp_storage_budget_exceeded = 84; } } diff --git a/src/expr/src/scalar.rs b/src/expr/src/scalar.rs index a7f37a0d3629e..288736f4c6dc8 100644 --- a/src/expr/src/scalar.rs +++ b/src/expr/src/scalar.rs @@ -1157,6 +1157,22 @@ impl MirScalarExpr { } } +/// Fails once `temp_storage` holds more than the budget it was built with. +/// +/// Checked after each function call rather than inside the arena, because `RowArena`'s pushes are +/// infallible: refusing one would hand back a truncated value. Between calls is the innermost point +/// that can return an error, so a budgeted arena reaches at most its budget plus whatever the call +/// that crossed it allocated. Functions that can predict their own size cut that overshoot by +/// consulting [`crate::func::max_string_func_result_bytes`] first. +/// +/// An unbudgeted arena, which is every arena in a dataflow, costs one branch on a `None`. +fn check_temp_storage_budget(temp_storage: &RowArena) -> Result<(), EvalError> { + if temp_storage.over_budget() { + return Err(EvalError::TempStorageBudgetExceeded); + } + Ok(()) +} + impl Eval for MirScalarExpr { fn eval<'a>( &'a self, @@ -1176,13 +1192,19 @@ impl Eval for MirScalarExpr { format!("cannot evaluate unmaterializable function: {:?}", x).into(), )), MirScalarExpr::CallUnary { func, expr } => { - func.eval(datums, temp_storage, expr.as_ref()) + let datum = func.eval(datums, temp_storage, expr.as_ref())?; + check_temp_storage_budget(temp_storage)?; + Ok(datum) } MirScalarExpr::CallBinary { func, expr1, expr2 } => { - func.eval(datums, temp_storage, &[expr1.as_ref(), expr2.as_ref()]) + let datum = func.eval(datums, temp_storage, &[expr1.as_ref(), expr2.as_ref()])?; + check_temp_storage_budget(temp_storage)?; + Ok(datum) } MirScalarExpr::CallVariadic { func, exprs } => { - func.eval(datums, temp_storage, exprs.as_slice()) + let datum = func.eval(datums, temp_storage, exprs.as_slice())?; + check_temp_storage_budget(temp_storage)?; + Ok(datum) } MirScalarExpr::If { cond, then, els } => match cond.eval(datums, temp_storage)? { Datum::True => then.eval(datums, temp_storage), @@ -1827,6 +1849,10 @@ pub enum EvalError { // printer. IfNullError(Box), LengthTooLarge, + // A budgeted `RowArena` (`mz_repr::RowArena::with_budget`) exceeded its budget while an + // expression was being evaluated. Only a budgeted arena raises this, so it never arises in a + // dataflow, only on the webhook `CHECK` path that runs user expressions in `environmentd`. + TempStorageBudgetExceeded, AclArrayNullElement, MzAclArrayNullElement, PrettyError(Box), @@ -2048,6 +2074,9 @@ impl fmt::Display for EvalError { } EvalError::IfNullError(s) => f.write_str(s), EvalError::LengthTooLarge => write!(f, "requested length too large"), + EvalError::TempStorageBudgetExceeded => { + write!(f, "expression exceeded its temporary storage limit") + } EvalError::AclArrayNullElement => write!(f, "ACL arrays must not contain null values"), EvalError::MzAclArrayNullElement => { write!(f, "MZ_ACL arrays must not contain null values") @@ -2310,6 +2339,7 @@ impl RustType for EvalError { }), EvalError::IfNullError(s) => IfNullError(s.into_proto()), EvalError::LengthTooLarge => LengthTooLarge(()), + EvalError::TempStorageBudgetExceeded => TempStorageBudgetExceeded(()), EvalError::AclArrayNullElement => AclArrayNullElement(()), EvalError::MzAclArrayNullElement => MzAclArrayNullElement(()), EvalError::InvalidIanaTimezoneId(s) => InvalidIanaTimezoneId(s.into_proto()), @@ -2438,6 +2468,7 @@ impl RustType for EvalError { }), IfNullError(v) => Ok(EvalError::IfNullError(v.into())), LengthTooLarge(()) => Ok(EvalError::LengthTooLarge), + TempStorageBudgetExceeded(()) => Ok(EvalError::TempStorageBudgetExceeded), AclArrayNullElement(()) => Ok(EvalError::AclArrayNullElement), MzAclArrayNullElement(()) => Ok(EvalError::MzAclArrayNullElement), InvalidIanaTimezoneId(s) => Ok(EvalError::InvalidIanaTimezoneId(s.into())), @@ -2480,6 +2511,129 @@ mod tests { ); } + /// A budgeted arena must stop an amplifying expression that the per-call constant allows + /// (SQL-431). `repeat(body, 20)` on a 1 MiB body is 20 MiB: far under + /// `MAX_STRING_FUNC_RESULT_BYTES`, so nothing rejects it without a budget, and it is exactly + /// the shape that made a webhook `CHECK` turn a bounded request into unbounded heap. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] // multi-MB allocations; the small-size UB coverage is in `row.rs` + fn test_repeat_respects_arena_budget() { + use crate::scalar::func::RepeatString; + + let body = "a".repeat(1024 * 1024); + let expr = MirScalarExpr::column(0).call_binary( + MirScalarExpr::literal_ok(Datum::Int32(20), ReprScalarType::Int32), + RepeatString, + ); + let datums = [Datum::String(&body)]; + + // Unbudgeted: allowed, and the arena really does hold the 20 MiB. + let arena = RowArena::new(); + let datum = expr + .eval(&datums, &arena) + .expect("under the 100 MiB ceiling"); + assert_eq!(datum.unwrap_str().len(), 20 * 1024 * 1024); + assert!(arena.allocated_bytes() >= 20 * 1024 * 1024); + + // Budgeted below the result: rejected, and the pre-check means the arena never grew, i.e. + // the bytes were never allocated rather than allocated and then complained about. + let arena = RowArena::with_budget(4 * 1024 * 1024); + assert_eq!( + expr.eval(&datums, &arena), + Err(EvalError::LengthTooLarge), + "an over-budget result must be refused" + ); + assert_eq!(arena.allocated_bytes(), 0); + + // Budgeted above the result: unaffected. + let arena = RowArena::with_budget(64 * 1024 * 1024); + let datum = expr.eval(&datums, &arena).expect("within budget"); + assert_eq!(datum.unwrap_str().len(), 20 * 1024 * 1024); + } + + /// The budget also has to catch a function with no size pre-check of its own, which allocates + /// straight into the arena. Enforcement for those is the evaluator's post-call check. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] // multi-MB allocations; the small-size UB coverage is in `row.rs` + fn test_arena_built_result_respects_budget() { + use crate::scalar::func::variadic::StringToArray; + + let body = "a".repeat(256 * 1024); + let expr = MirScalarExpr::call_variadic( + StringToArray, + vec![ + MirScalarExpr::column(0), + MirScalarExpr::literal_ok(Datum::String("a"), ReprScalarType::String), + ], + ); + let datums = [Datum::String(&body)]; + + let arena = RowArena::new(); + expr.eval(&datums, &arena).expect("no ceiling applies"); + let unbudgeted = arena.allocated_bytes(); + assert!(unbudgeted > 0); + + let arena = RowArena::with_budget(unbudgeted / 2); + assert_eq!( + expr.eval(&datums, &arena), + Err(EvalError::TempStorageBudgetExceeded), + "an over-budget arena-built result must be refused" + ); + } + + /// `array_fill` sizes its result from a parameter rather than its input, so a budgeted arena + /// has to refuse it in its own pre-check the way the string amplifiers do (SQL-431). Its result + /// stays under the `array_fill` size ceiling, so without the budget-aware pre-check the only + /// thing that would catch it is the evaluator's post-call check, after the spike has happened. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] // multi-MB allocations; the small-size UB coverage is in `row.rs` + fn test_array_fill_respects_arena_budget() { + use crate::scalar::func::variadic::ArrayFill; + use mz_repr::adt::array::ArrayDimension; + + // array_fill(1, ARRAY[fill_count]) builds a one-dimensional int array of `fill_count` ones. + let fill_count: usize = 512 * 1024; + let dims_storage = RowArena::new(); + let dims = dims_storage + .try_make_datum(|packer| { + packer.try_push_array( + &[ArrayDimension { + lower_bound: 1, + length: 1, + }], + [Datum::Int32(i32::try_from(fill_count).unwrap())], + ) + }) + .unwrap(); + let expr = MirScalarExpr::call_variadic( + ArrayFill { + elem_type: mz_repr::SqlScalarType::Int32, + }, + vec![MirScalarExpr::column(0), MirScalarExpr::column(1)], + ); + let datums = [Datum::Int32(1), dims]; + + // Unbudgeted: allowed, and the arena really holds the packed array. + let arena = RowArena::new(); + expr.eval(&datums, &arena) + .expect("under the array-size ceiling"); + assert!(arena.allocated_bytes() > 0); + + // Budgeted below what the call needs: refused by the pre-check, so nothing was allocated. + // The intermediate `Vec` alone is 512 Ki elements, well over this 1 MiB budget. + let arena = RowArena::with_budget(1024 * 1024); + assert_eq!( + expr.eval(&datums, &arena), + Err(EvalError::TempStorageBudgetExceeded), + "an over-budget array_fill must be refused before it allocates" + ); + assert_eq!(arena.allocated_bytes(), 0); + + // Budgeted well above both the packed result and the intermediate: unaffected. + let arena = RowArena::with_budget(256 * 1024 * 1024); + expr.eval(&datums, &arena).expect("within budget"); + } + #[mz_ore::test] #[cfg_attr(miri, ignore)] // error: unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux` fn test_reduce() { diff --git a/src/expr/src/scalar/func.rs b/src/expr/src/scalar/func.rs index a69fc02aa4086..30fc4b127dbf0 100644 --- a/src/expr/src/scalar/func.rs +++ b/src/expr/src/scalar/func.rs @@ -171,6 +171,22 @@ func_name! { /// function where it applies. pub const MAX_STRING_FUNC_RESULT_BYTES: usize = 1024 * 1024 * 100; +/// The largest result a string function may build into `temp_storage`. +/// +/// [`MAX_STRING_FUNC_RESULT_BYTES`] unless the arena carries a tighter budget, which is how an +/// expression evaluated in `environmentd` on behalf of a request (a webhook `CHECK`) is held to a +/// size proportionate to that request rather than to the constant, which is sized for a cluster. +/// +/// A function that can predict its result size must consult this *before* building the result: the +/// arena's own budget is only observable after the bytes exist, which for an amplifying function is +/// exactly too late. +pub fn max_string_func_result_bytes(temp_storage: &RowArena) -> usize { + std::cmp::min( + MAX_STRING_FUNC_RESULT_BYTES, + temp_storage.budget_remaining(), + ) +} + pub fn jsonb_stringify<'a>(a: Datum<'a>, temp_storage: &'a RowArena) -> Option<&'a str> { match a { Datum::JsonNull => None, @@ -469,10 +485,10 @@ fn encode(bytes: &[u8], format: &str) -> Result { } #[sqlfunc] -fn decode(string: &str, format: &str) -> Result, EvalError> { +fn decode(string: &str, format: &str, temp_storage: &RowArena) -> Result, EvalError> { let format = encoding::lookup_format(format)?; let out = format.decode(string)?; - if out.len() > MAX_STRING_FUNC_RESULT_BYTES { + if out.len() > max_string_func_result_bytes(temp_storage) { Err(EvalError::LengthTooLarge) } else { Ok(out) @@ -2521,8 +2537,8 @@ fn starts_with(a: &str, b: &str) -> bool { // 'A' < 'AA' but 'AZ' > 'AAZ'.) is_monotone = (false, true), )] -fn text_concat_binary(a: &str, b: &str) -> Result { - if a.len() + b.len() > MAX_STRING_FUNC_RESULT_BYTES { +fn text_concat_binary(a: &str, b: &str, temp_storage: &RowArena) -> Result { + if a.len() + b.len() > max_string_func_result_bytes(temp_storage) { return Err(EvalError::LengthTooLarge); } let mut buf = String::with_capacity(a.len() + b.len()); @@ -2639,9 +2655,9 @@ pub fn build_regex(needle: &str, flags: &str) -> Result { } #[sqlfunc(sqlname = "repeat")] -fn repeat_string(string: &str, count: i32) -> Result { +fn repeat_string(string: &str, count: i32, temp_storage: &RowArena) -> Result { let len = usize::try_from(count).unwrap_or(0); - if (len * string.len()) > MAX_STRING_FUNC_RESULT_BYTES { + if len.saturating_mul(string.len()) > max_string_func_result_bytes(temp_storage) { return Err(EvalError::LengthTooLarge); } Ok(string.repeat(len)) diff --git a/src/expr/src/scalar/func/variadic.rs b/src/expr/src/scalar/func/variadic.rs index f32c739067be1..6ea2908e1d966 100644 --- a/src/expr/src/scalar/func/variadic.rs +++ b/src/expr/src/scalar/func/variadic.rs @@ -40,7 +40,7 @@ use mz_repr::{ use serde::{Deserialize, Serialize}; use crate::func::{ - CaseLiteral, MAX_STRING_FUNC_RESULT_BYTES, array_create_scalar, build_regex, date_bin, + CaseLiteral, array_create_scalar, build_regex, date_bin, max_string_func_result_bytes, parse_timezone, regexp_match_static, regexp_replace_parse_flags, regexp_split_to_array_re, stringify_datum, timezone_time, }; @@ -323,6 +323,19 @@ fn array_fill<'a>( return Err(EvalError::MaxArraySizeExceeded(MAX_SIZE)); } + // The packed array lands in `temp_storage`, and building it first allocates a transient + // `Vec` of `fill_count` elements. Unlike the string amplifiers, `array_fill` sizes its + // result from a parameter rather than its input, so without this check a single call could + // spike far past a budgeted arena before the evaluator's post-call check ever runs. Refuse both + // allocations up front instead. Without a budget `budget_remaining` is `usize::MAX` and both + // comparisons fold away, so this is free in a dataflow. + let budget_remaining = temp_storage.budget_remaining(); + let packed_size = mz_repr::datum_size(&fill).saturating_mul(fill_count); + let build_size = fill_count.saturating_mul(std::mem::size_of::>()); + if packed_size > budget_remaining || build_size > budget_remaining { + return Err(EvalError::TempStorageBudgetExceeded); + } + let array_dimensions = if fill_count == 0 { vec![ArrayDimension { lower_bound: 1, @@ -1196,7 +1209,12 @@ impl LazyVariadicFunc for Or { } #[sqlfunc(sqlname = "lpad")] -fn pad_leading(string: &str, raw_len: i32, pad: OptionalArg<&str>) -> Result { +fn pad_leading( + string: &str, + raw_len: i32, + pad: OptionalArg<&str>, + temp_storage: &RowArena, +) -> Result { let len = match usize::try_from(raw_len) { Ok(len) => len, Err(_) => { @@ -1205,7 +1223,7 @@ fn pad_leading(string: &str, raw_len: i32, pad: OptionalArg<&str>) -> Result MAX_STRING_FUNC_RESULT_BYTES { + if len > max_string_func_result_bytes(temp_storage) { return Err(EvalError::LengthTooLarge); } @@ -1271,18 +1289,19 @@ fn regexp_replace<'a>( } #[sqlfunc] -fn replace(text: &str, from: &str, to: &str) -> Result { +fn replace(text: &str, from: &str, to: &str, temp_storage: &RowArena) -> Result { // As a compromise to avoid always nearly duplicating the work of replace by doing size estimation, // we first check if it's possible for the fully replaced string to exceed the limit by assuming that // every possible substring is replaced. // // If that estimate exceeds the limit, we then do a more precise (and expensive) estimate by counting // the actual number of replacements that would occur, and using that to calculate the final size. + let max_result_bytes = max_string_func_result_bytes(temp_storage); let possible_size = text.len() * to.len(); - if possible_size > MAX_STRING_FUNC_RESULT_BYTES { + if possible_size > max_result_bytes { let replacement_count = text.matches(from).count(); let estimated_size = text.len() + replacement_count * (to.len().saturating_sub(from.len())); - if estimated_size > MAX_STRING_FUNC_RESULT_BYTES { + if estimated_size > max_result_bytes { return Err(EvalError::LengthTooLarge); } } @@ -1456,12 +1475,13 @@ fn split_part<'a>(string: &'a str, delimiter: &str, field: i32) -> Result<&'a st } #[sqlfunc(is_associative = true)] -fn concat(strs: Variadic>) -> Result { +fn concat(strs: Variadic>, temp_storage: &RowArena) -> Result { + let max_result_bytes = max_string_func_result_bytes(temp_storage); let mut total_size = 0; for s in &strs { if let Some(s) = s { total_size += s.len(); - if total_size > MAX_STRING_FUNC_RESULT_BYTES { + if total_size > max_result_bytes { return Err(EvalError::LengthTooLarge); } } @@ -1476,13 +1496,18 @@ fn concat(strs: Variadic>) -> Result { } #[sqlfunc] -fn concat_ws(ws: &str, rest: Variadic>) -> Result { +fn concat_ws( + ws: &str, + rest: Variadic>, + temp_storage: &RowArena, +) -> Result { + let max_result_bytes = max_string_func_result_bytes(temp_storage); let mut total_size = 0; for s in &rest { if let Some(s) = s { total_size += s.len(); total_size += ws.len(); - if total_size > MAX_STRING_FUNC_RESULT_BYTES { + if total_size > max_result_bytes { return Err(EvalError::LengthTooLarge); } } diff --git a/src/repr/src/row.rs b/src/repr/src/row.rs index bfd0e9e4a0e3c..94e8c6a100241 100644 --- a/src/repr/src/row.rs +++ b/src/repr/src/row.rs @@ -837,6 +837,17 @@ pub struct RowArena { // writer's lifetime. That keeps nested writers sound: a writer obtained while another is live // finds the slot empty and allocates its own buffer instead of double-borrowing. scratch: RefCell>>, + // Optional ceiling on the bytes this arena will hold, and a running total of what it holds. + // `None` is unbounded, which is what every arena in a dataflow uses. A budget is for evaluating + // a user-authored expression in a shared process, where that expression's memory use has to be + // bounded (see `mz_adapter::webhook`). + // + // NOTE: exceeding the budget does not make a push fail. The pushes are infallible, and a + // refused push would hand back a truncated value, i.e. a corrupt datum. The budget is instead a + // *reported* condition: `over_budget` is polled by whoever is able to return an error, which + // for scalar expressions is the evaluator between calls. + budget: Option, + allocated: Cell, } // DatumList and DatumDict defined here rather than near Datum because we need private access to the unsafe data field @@ -3078,6 +3089,45 @@ impl RowArena { RowArena { inner: RefCell::new(vec![]), scratch: RefCell::new(None), + budget: None, + allocated: Cell::new(0), + } + } + + /// Creates a `RowArena` that reports itself [`RowArena::over_budget`] once it holds more than + /// `budget` bytes. + /// + /// The budget is advisory to the arena itself: pushes still succeed, because handing back a + /// truncated value would corrupt the datum. It is the caller's job to poll `over_budget` at a + /// point where it can fail, so the bytes an arena actually reaches is `budget` plus whatever the + /// operation in flight at the time added. + pub fn with_budget(budget: usize) -> Self { + RowArena { + budget: Some(budget), + ..RowArena::new() + } + } + + /// Bytes this arena currently holds. + pub fn allocated_bytes(&self) -> usize { + self.allocated.get() + } + + /// Whether this arena holds more than its budget. Always false without one. + pub fn over_budget(&self) -> bool { + self.budget + .is_some_and(|budget| self.allocated.get() > budget) + } + + /// Bytes this arena can still take before it is [`RowArena::over_budget`], or `usize::MAX` + /// without a budget. + /// + /// Intended for an operation that can predict its own size and would rather fail than build a + /// value it is about to be told is too big. + pub fn budget_remaining(&self) -> usize { + match self.budget { + None => usize::MAX, + Some(budget) => budget.saturating_sub(self.allocated.get()), } } @@ -3090,7 +3140,7 @@ impl RowArena { } RowArena { inner: RefCell::new(inner), - scratch: RefCell::new(None), + ..RowArena::new() } } @@ -3148,6 +3198,7 @@ impl RowArena { let region = inner.last_mut().expect("region present"); let start = region.len(); region.extend_from_slice(bytes); + self.allocated.set(self.allocated.get() + need); let copied = ®ion[start..]; unsafe { // This is safe because: @@ -3163,11 +3214,54 @@ impl RowArena { } } - /// Copies `string` into the arena and returns a reference valid for its lifetime. + /// Moves `bytes` into the arena and returns a reference valid for its lifetime. + /// + /// Prefer this to [`RowArena::push_bytes`] whenever the bytes are already owned: when they do + /// not fit the active region, their allocation is adopted as a region instead of a fresh region + /// being allocated and copied into, which for a large value halves the peak. + pub fn push_owned_bytes<'a>(&'a self, bytes: Vec) -> &'a [u8] { + let need = bytes.len(); + if need == 0 { + return &[]; + } + + let mut inner = self.inner.borrow_mut(); + let has_room = inner + .last() + .map_or(false, |region| region.capacity() - region.len() >= need); + if has_room { + // Copying into a region that is already paid for beats giving these bytes one of their + // own. Adopting unconditionally would turn every small string into its own allocation + // and defeat the bump allocator. + drop(inner); + return self.push_bytes(&bytes[..]); + } + + // `push_bytes` would allocate a fresh region here and copy into it, so adopt the caller's + // allocation as that region. Sound for the same reasons as `push_bytes`: the reference + // points into a heap buffer the arena now owns for `'a`, and the buffer is never resized + // while it holds data. + // + // Inserted *below* the active region rather than appended, because `push_bytes` sizes a new + // region as twice the last one's capacity: leaving a large adopted buffer on top would make + // the next push allocate twice its size. + self.allocated.set(self.allocated.get() + need); + let idx = inner.len().saturating_sub(1); + inner.insert(idx, bytes); + if inner.len() == 1 { + // There was no active region to insert below, so keep an empty one on top for the same + // reason. `Vec::new` does not allocate. + inner.push(Vec::new()); + } + let adopted = &inner[idx][..]; + unsafe { transmute::<&[u8], &'a [u8]>(adopted) } + } + + /// Moves `string` into the arena and returns a reference valid for its lifetime. pub fn push_string<'a>(&'a self, string: String) -> &'a str { - let copied = self.push_bytes(string.as_bytes()); + let copied = self.push_owned_bytes(string.into_bytes()); unsafe { - // This is safe because we just copied the bytes of a valid `String`. + // This is safe because we just moved in the bytes of a valid `String`. std::str::from_utf8_unchecked(copied) } } @@ -3294,6 +3388,7 @@ impl RowArena { inner.truncate(1); inner[0].clear(); } + self.allocated.set(0); } } @@ -3596,6 +3691,63 @@ mod tests { assert_eq!(arena.push_bytes(Vec::::new()), empty); } + #[mz_ore::test] + fn miri_test_arena_adopts_owned_bytes_and_keeps_references() { + // `push_owned_bytes` adopts a buffer too large for the active region instead of copying it, + // which puts a region the arena never wrote into in the middle of the stack. References + // handed out before and after that must all stay valid. + let arena = RowArena::new(); + let before = arena.push_bytes(vec![1u8; 8]); + let adopted = arena.push_owned_bytes(vec![2u8; 64 * 1024]); + let after = arena.push_bytes(vec![3u8; 8]); + // A small buffer fits the active region, so it is copied rather than given a region. + let small = arena.push_owned_bytes(vec![4u8; 4]); + + assert_eq!(before, &[1u8; 8]); + assert_eq!(adopted, &vec![2u8; 64 * 1024][..]); + assert_eq!(after, &[3u8; 8]); + assert_eq!(small, &[4u8; 4]); + + let empty: &[u8] = &[]; + assert_eq!(arena.push_owned_bytes(vec![]), empty); + } + + #[mz_ore::test] + fn miri_test_arena_budget() { + // Without a budget nothing is ever over it, however much is pushed. + let arena = RowArena::new(); + let _ = arena.push_bytes(vec![0u8; 1024]); + assert!(!arena.over_budget()); + assert_eq!(arena.budget_remaining(), usize::MAX); + + let arena = RowArena::with_budget(100); + assert!(!arena.over_budget()); + assert_eq!(arena.budget_remaining(), 100); + + // Staying within the budget leaves it satisfied, and the remaining count tracks what a + // caller that predicts its own size would consult. + let _ = arena.push_bytes(vec![0u8; 60]); + assert!(!arena.over_budget()); + assert_eq!(arena.budget_remaining(), 40); + assert_eq!(arena.allocated_bytes(), 60); + + // Crossing it reports, rather than refusing the push: a truncated push would corrupt the + // datum, so the value is intact and it is the caller's job to fail. + let pushed = arena.push_bytes(vec![7u8; 80]); + assert_eq!(pushed, &[7u8; 80]); + assert!(arena.over_budget()); + assert_eq!(arena.budget_remaining(), 0); + + // An adopted buffer counts against the budget too, or adoption would be a way around it. + let mut arena = RowArena::with_budget(100); + let _ = arena.push_owned_bytes(vec![0u8; 101]); + assert!(arena.over_budget()); + + arena.clear(); + assert!(!arena.over_budget()); + assert_eq!(arena.allocated_bytes(), 0); + } + #[mz_ore::test] fn miri_test_arena_writer() { use std::io::Write; diff --git a/src/repr/src/scalar.rs b/src/repr/src/scalar.rs index 4275ec05b3e16..f62069e26a17c 100644 --- a/src/repr/src/scalar.rs +++ b/src/repr/src/scalar.rs @@ -2775,7 +2775,7 @@ impl<'a, E> OutputDatumType<'a, E> for Vec { } fn into_result(self, temp_storage: &'a RowArena) -> Result, E> { - Ok(Datum::Bytes(temp_storage.push_bytes(self))) + Ok(Datum::Bytes(temp_storage.push_owned_bytes(self))) } } diff --git a/src/storage-types/src/errors.rs b/src/storage-types/src/errors.rs index 66d5dffa4a51e..2441141e13e74 100644 --- a/src/storage-types/src/errors.rs +++ b/src/storage-types/src/errors.rs @@ -668,6 +668,7 @@ mod columnation { | e @ EvalError::InvalidTimezoneInterval | e @ EvalError::InvalidTimezoneConversion | e @ EvalError::LengthTooLarge + | e @ EvalError::TempStorageBudgetExceeded | e @ EvalError::AclArrayNullElement | e @ EvalError::MzAclArrayNullElement => e.clone(), EvalError::Unsupported {