expr: bound the memory a webhook CHECK may allocate (SQL-431)

[mtabebe]

Problem:
A webhook CHECK is a user defined expression that environmentd evaluates once per in-flight request, and nothing bounds what it could allocate.

The webhook path caps request count (500) and body size (5 MiB), but the evaluation's footprint is per-request and was bounded only by MAX_STRING_FUNC_RESULT_BYTES, a 100 MiB per-call ceiling sized for a cluster.

That ceiling multiplies by request concurrency, so a CHECK doing length(repeat(body, 20)) >= 0 turned 150 concurrent 5 MB posts into 8.7 GiB of RSS, growing with the number of clients.

Solution:

• Give RowArena an optional budget and build the validation arena with one, defaulting to 20 MiB via a new dyncfg, webhook_validation_memory_budget_bytes.
• Functions that can predict their result size consult max_string_func_result_bytes, which narrows the constant to the arena's remaining budget, so an over-budget result is never allocated at all.
• Everything else is caught by a post-call check in the evaluator, which covers functions with no ceiling of their own.
• Stop copying an owned result into the arena. String and Vec outputs went through push_string/push_bytes, which copied and dropped the original.

An unbudgeted arena, which is every arena in a dataflow, is unaffected

Testing:

• New integration test webhook_validation_memory_budget: a 2 MiB body that amplifies past the 20 MiB default is refused with 400, a smaller body through the same source succeeds, and raising then lowering the dyncfg flips the result both ways
• New expr unit tests cover the pre-check and post-call paths
• New unit tests cover the arena bookkeeping and buffer adoption

[antiguru]

I think this looks good, thank you!

One caveat: If we ever use the budget in compute dataflows, we need to ensure that it never changes for the lifetime of a cluster replica. Changing it mid-run would make the evaluation non-deterministic, in which case we would produce non-accumulating diffs. It might be worth flagging this in the implementation so we don't accidentally run into this problem in the future.

[def-]

Two failing extra tests:

diff --git a/src/expr/src/scalar.rs b/src/expr/src/scalar.rs
index 288736f4c6..e5ed44bdaa 100644
--- a/src/expr/src/scalar.rs
+++ b/src/expr/src/scalar.rs
@@ -2634,6 +2634,82 @@ mod tests {
         expr.eval(&datums, &arena).expect("within budget");
     }

+    /// A budget has to bound what a single call allocates, not only what is observable between
+    /// calls (SQL-431).
+    ///
+    /// The evaluator polls the budget only after `func.eval` has built its result and moved it into
+    /// the arena, and erroring then does not give the bytes back: they stay resident until the arena
+    /// drops, which for a webhook `CHECK` is the end of the request. So the property is arena
+    /// residency, not the returned `Result`. Nor is the overshoot a constant. It scales with the
+    /// body and with a multiplier the `CHECK` author picks at DDL time.
+    #[mz_ore::test]
+    #[cfg_attr(miri, ignore)] // multi-MB allocations; the small-size UB coverage is in `row.rs`
+    fn test_single_call_respects_arena_budget() {
+        use crate::scalar::func::variadic::{ArrayCreate, PadLeading, Translate};
+
+        // Scaled down from the shipped 5 MiB body and 20 MiB budget. Every amplifier here is linear
+        // in the body, so the ratios hold at any scale.
+        const BODY_BYTES: usize = 1024 * 1024;
+        const BUDGET: usize = 2 * 1024 * 1024;
+        const WIDE: &str = "\u{1F4A5}"; // one character, four bytes
+
+        let str_lit = |s| MirScalarExpr::literal_ok(Datum::String(s), ReprScalarType::String);
+        // `ARRAY[body, ...]` is not a string function, so no ceiling of its own applies and the
+        // multiplier is just how many times the author wrote `body`.
+        let array_of = |n| {
+            let elem_type = mz_repr::SqlScalarType::String;
+            let refs = vec![MirScalarExpr::column(0); n];
+            MirScalarExpr::call_variadic(ArrayCreate { elem_type }, refs)
+        };
+        let cases = [
+            ("ARRAY[body x4]", array_of(4)),
+            ("ARRAY[body x16]", array_of(16)),
+            // `lpad`'s pre-check is budget-aware but compares `len`, a character count, against a
+            // budget in bytes, so a 4-byte pad passes a check for exactly the budget then writes 4x.
+            (
+                "lpad(body, BUDGET, wide)",
+                MirScalarExpr::call_variadic(
+                    PadLeading,
+                    vec![
+                        MirScalarExpr::column(0),
+                        MirScalarExpr::literal_ok(
+                            Datum::Int32(i32::try_from(BUDGET).unwrap()),
+                            ReprScalarType::Int32,
+                        ),
+                        str_lit(WIDE),
+                    ],
+                ),
+            ),
+            // `translate` has no pre-check at all, and widening each body byte is a 4x amplifier
+            // that needs no length argument to drive it.
+            (
+                "translate(body, 'a', wide)",
+                MirScalarExpr::call_variadic(
+                    Translate,
+                    vec![MirScalarExpr::column(0), str_lit("a"), str_lit(WIDE)],
+                ),
+            ),
+        ];
+
+        let body = "a".repeat(BODY_BYTES);
+        let datums = [Datum::String(&body)];
+        let mut over = Vec::new();
+        for (name, expr) in cases {
+            let arena = RowArena::with_budget(BUDGET);
+            let _ = expr.eval(&datums, &arena); // refused or not is beside the point
+            let held = arena.allocated_bytes();
+            if held > BUDGET {
+                let ratio = held as f64 / BUDGET as f64;
+                over.push(format!("  {name}: held {held} bytes, {ratio:.1}x"));
+            }
+        }
+        assert!(
+            over.is_empty(),
+            "a single call left a {BUDGET} byte arena holding more:\n{}",
+            over.join("\n"),
+        );
+    }
+
     #[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/repr/src/row.rs b/src/repr/src/row.rs
index 94e8c6a100..b04ae011d1 100644
--- a/src/repr/src/row.rs
+++ b/src/repr/src/row.rs
@@ -3712,6 +3712,53 @@ mod tests {
         assert_eq!(arena.push_owned_bytes(vec![]), empty);
     }

+    #[mz_ore::test]
+    fn test_arena_owned_pushes_keep_bump_allocating() {
+        // Adoption never *creates* a region with headroom: it inserts the caller's buffer, whose
+        // capacity equals its length, below whatever is on top. Only `push_bytes` grows the arena
+        // geometrically (`new_cap = max(need, last_cap * 2)`), so once the active region cannot fit
+        // an incoming value it never can again and every later owned push adopts: one retained
+        // allocation and one `Vec<u8>` header per value, rather than `O(log n)` regions. That is the
+        // default path for every `String`- and `Vec<u8>`-returning scalar function, and the arenas
+        // in the MFP and join paths outlive a single row, so the region list grows with the number
+        // of string values in a batch. `RowArena::clear` scans every region, so it degrades too.
+        const VALUES: usize = 500;
+        const VALUE: &str = "0123456789";
+
+        let regions = |arena: &RowArena| arena.inner.borrow().len();
+        let push_all = |arena: &RowArena, owned: bool| {
+            for _ in 0..VALUES {
+                match owned {
+                    true => _ = arena.push_string(VALUE.to_string()),
+                    false => _ = arena.push_bytes(VALUE.as_bytes()),
+                }
+            }
+        };
+
+        // The bump allocator working as intended, as the baseline to hold the owned path to.
+        let copied = RowArena::new();
+        push_all(&copied, false);
+
+        let owned = RowArena::new();
+        push_all(&owned, true);
+
+        // Seeding with ordinary copies first must not change the answer. The arena never recovers,
+        // so this is not just the empty-arena case where the placeholder on top has capacity 0.
+        let seeded = RowArena::new();
+        let _ = seeded.push_bytes(VALUE.as_bytes());
+        push_all(&seeded, true);
+
+        // Held to the copy path rather than an absolute count, so this pins the property (a run of
+        // small owned pushes still ends with a region that has headroom) and leaves the adoption
+        // predicate to the fix.
+        let (copied, owned, seeded) = (regions(&copied), regions(&owned), regions(&seeded));
+        assert!(
+            owned <= copied * 2 && seeded <= copied * 2,
+            "{VALUES} owned pushes left {owned} regions on an empty arena and {seeded} on a seeded \
+             one, against {copied} for the same bytes copied",
+        );
+    }
+
     #[mz_ore::test]
     fn miri_test_arena_budget() {
         // Without a budget nothing is ever over it, however much is pushed.

Running cargo nextest run -p mz-expr -p mz-repr --lib --no-fail-fast test_single_call_respects_arena_budget test_arena_owned_pushes_keep_bump_allocating fails:

        FAIL [   0.004s] mz-repr row::tests::test_arena_owned_pushes_keep_bump_allocating
  stdout ───

    running 1 test
    test row::tests::test_arena_owned_pushes_keep_bump_allocating ... FAILED

    failures:

    failures:
        row::tests::test_arena_owned_pushes_keep_bump_allocating

    test result: FAILED. 0 passed; 1 failed; 0 ignored; 0 measured; 159 filtered out; finished in 0.00s

  stderr ───

    thread 'row::tests::test_arena_owned_pushes_keep_bump_allocating' (571281) panicked at src/repr/src/row.rs:3755:9:
    500 owned pushes left 501 regions on an empty arena and 501 on a seeded one, against 9 for the same bytes copied
    note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

        FAIL [   0.120s] mz-expr scalar::tests::test_single_call_respects_arena_budget
  stdout ───

    running 1 test
    test scalar::tests::test_single_call_respects_arena_budget ... FAILED

    failures:

    failures:
        scalar::tests::test_single_call_respects_arena_budget

    test result: FAILED. 0 passed; 1 failed; 0 ignored; 0 measured; 134 filtered out; finished in 0.11s

  stderr ───

    thread 'scalar::tests::test_single_call_respects_arena_budget' (571282) panicked at src/expr/src/scalar.rs:2706:9:
    a single call left a 2097152 byte arena holding more:
      ARRAY[body x4]: held 4194350 bytes, 2.0x
      ARRAY[body x16]: held 16777322 bytes, 8.0x
      lpad(body, BUDGET, wide): held 5242880 bytes, 2.5x
      translate(body, 'a', wide): held 4194304 bytes, 2.0x
    note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

────────────

The memory cap not holding seems significant. Also not great to have increased memory usage in general, but we can live with that I guess.

The nightly failure is unrelated.

[def-]

I extended the test a bit more:

diff --git a/src/expr/src/scalar.rs b/src/expr/src/scalar.rs
index dcb79da9f8..778d2a8bce 100644
--- a/src/expr/src/scalar.rs
+++ b/src/expr/src/scalar.rs
@@ -2579,6 +2579,21 @@ mod tests {
             Err(EvalError::TempStorageBudgetExceeded),
             "an over-budget arena-built result must be refused"
         );
+
+        // A split collects every chunk into a `Vec<&str>` first, and a fat pointer is 16 bytes
+        // against the 2 an empty chunk packs to. A budget the packed array fits under still has to
+        // refuse the call.
+        let intermediate = (body.len() + 1) * std::mem::size_of::<&str>();
+        let budget = 4 * unbudgeted;
+        assert!(
+            unbudgeted < budget && budget < intermediate,
+            "budget sits between"
+        );
+        let arena = RowArena::with_budget(budget);
+        assert!(
+            expr.eval(&datums, &arena).is_err(),
+            "a split costing {intermediate} bytes to build must be refused by a {budget} byte budget"
+        );
     }

     /// `array_fill` sizes its result from a parameter rather than its input, so a budgeted arena

Fails:

    thread 'scalar::tests::test_arena_built_result_respects_budget' (3020774) panicked at src/expr/src/scalar.rs:2594:9:
    a split costing 4194320 bytes to build must be refused by a 2097264 byte budget

[def-]

no qa complaints