compute: bound a stashed peek answer across the workers that produced it

`max_result_size` is a bound on the answer a client receives, and a
stashed answer escaped it. Each worker decides on its own whether its
share is large enough to stash, so an oversized result arrives as several
stashed responses each well under the ceiling, and the controller measured
only the rows a response carries inline. The batches were invisible to it,
which left the effective bound at the ceiling times the worker count.

Carry the size of the stashed rows in the response, measured the way an
inline answer measures its own, and accumulate that in the controller's
per-peek total. The encoded size the response already carried is what the
batches cost in blob storage, which is a different number and not the one
this bounds.

The gap was hidden by the adapter's per-chunk check on read-back, which
rejects a chunk larger than the ceiling and so happened to catch this at
the default read batch size. `max_result_size.slt` now covers the case
with that batch size lowered, so the chunk check cannot be what rejects
the result, and with a stash threshold low enough that every worker's
share reaches the stash. Without this change the query returns roughly
1.1 MiB under a 1 MiB cap.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: antiguru

src/compute-client/src/protocol/response.rs

@@ -217,6 +217,18 @@ impl PeekResponse {
             Self::Error(_) | Self::Canceled => 0,
         }
     }
+
+    /// The size of the rows this response answers with, stashed rows included.
+    ///
+    /// This is what `max_result_size` bounds, so a stashed answer has to be measured the way an
+    /// inline one is. Reading [`PeekResponse::inline_byte_len`] instead would leave a stashed
+    /// answer bounded per worker only, since the batches a worker writes are invisible to it.
+    pub fn answer_byte_len(&self) -> usize {
+        match self {
+            Self::Rows(_) | Self::Error(_) | Self::Canceled => self.inline_byte_len(),
+            Self::Stashed(stashed) => stashed.answer_byte_len(),
+        }
+    }
 }
 
 /// The error of an unsuccessful peek.
@@ -280,7 +292,15 @@ pub struct StashedPeekResponse {
     /// This does _NOT_ include rows in `inline_rows`.
     pub num_rows_batches: u64,
     /// The sum of the encoded sizes of all batches in this response.
+    ///
+    /// What the batches cost in blob storage, which is not what `max_result_size` bounds. That is
+    /// [`StashedPeekResponse::stashed_byte_len`], the size of the same rows as an answer carries
+    /// them.
     pub encoded_size_bytes: usize,
+    /// The size of the stashed rows, measured as an inline answer measures its own.
+    ///
+    /// Does _NOT_ include `inline_rows`.
+    pub stashed_byte_len: usize,
     /// [RelationDesc] for the rows in these stashed batches of results.
     pub relation_desc: RelationDesc,
     /// The [ShardId] under which result batches have been stashed.
@@ -311,6 +331,13 @@ impl StashedPeekResponse {
 
         self.encoded_size_bytes + inline_size
     }
+
+    /// The size of the rows in this result, measured as an inline answer measures its own.
+    pub fn answer_byte_len(&self) -> usize {
+        let inline_size: usize = self.inline_rows.iter().map(|r| r.byte_len()).sum();
+
+        self.stashed_byte_len.saturating_add(inline_size)
+    }
 }
 
 /// Various responses that can be communicated after a COPY TO command.

src/compute-client/src/service.rs

@@ -566,12 +566,12 @@ impl PendingSubscribe {
 struct PendingPeek {
     /// The responses merged so far.
     response: PeekResponse,
-    /// Inline result bytes seen so far, across all shards.
+    /// Answer bytes seen so far, across all shards, stashed rows included.
     ///
     /// Tracked separately from `response` because a worker's rows are dropped as soon as any
     /// worker reports an error. Without this the aggregate size check would depend on the order
     /// the responses happen to arrive in.
-    inline_byte_len: usize,
+    answer_byte_len: usize,
     /// The shards that have provided responses.
     ready_shards: BTreeSet<usize>,
 }
@@ -580,7 +580,7 @@ impl PendingPeek {
     fn new() -> Self {
         Self {
             response: PeekResponse::Rows(vec![RowCollection::default()]),
-            inline_byte_len: 0,
+            answer_byte_len: 0,
             ready_shards: BTreeSet::new(),
         }
     }
@@ -589,15 +589,17 @@ impl PendingPeek {
         let first = self.ready_shards.insert(shard_id);
         assert!(first, "duplicate peek response");
 
-        self.inline_byte_len = self
-            .inline_byte_len
-            .saturating_add(response.inline_byte_len());
+        // Stashed rows count here as inline ones do, so a peek whose workers each stash a share
+        // under the ceiling is still bounded by what the client receives in total.
+        self.answer_byte_len = self
+            .answer_byte_len
+            .saturating_add(response.answer_byte_len());
         let current = mem::replace(&mut self.response, PeekResponse::Canceled);
         self.response = merge_peek_responses(current, response);
 
         // Merging eagerly is what keeps the controller's memory bounded, so the size check has to
         // happen on every response rather than once at the end.
-        if self.inline_byte_len > max_result_size.cast_into() {
+        if self.answer_byte_len > max_result_size.cast_into() {
             // NOTE: Tests match on this exact message, so nothing else may produce it.
             let error = PeekError::unstructured(format!(
                 "total result exceeds max size of {}",
@@ -636,6 +638,7 @@ fn merge_peek_responses(resp1: PeekResponse, resp2: PeekResponse) -> PeekRespons
             let StashedPeekResponse {
                 num_rows_batches: num_rows_batches1,
                 encoded_size_bytes: encoded_size_bytes1,
+                stashed_byte_len: stashed_byte_len1,
                 relation_desc: relation_desc1,
                 shard_id: shard_id1,
                 batches: mut batches1,
@@ -644,6 +647,7 @@ fn merge_peek_responses(resp1: PeekResponse, resp2: PeekResponse) -> PeekRespons
             let StashedPeekResponse {
                 num_rows_batches: num_rows_batches2,
                 encoded_size_bytes: encoded_size_bytes2,
+                stashed_byte_len: stashed_byte_len2,
                 relation_desc: relation_desc2,
                 shard_id: shard_id2,
                 batches: mut batches2,
@@ -671,6 +675,7 @@ fn merge_peek_responses(resp1: PeekResponse, resp2: PeekResponse) -> PeekRespons
             Stashed(Box::new(StashedPeekResponse {
                 num_rows_batches: num_rows_batches1 + num_rows_batches2,
                 encoded_size_bytes: encoded_size_bytes1 + encoded_size_bytes2,
+                stashed_byte_len: stashed_byte_len1.saturating_add(stashed_byte_len2),
                 relation_desc: relation_desc1,
                 shard_id: shard_id1,
                 batches: batches1,

src/compute-client/src/service/tests.rs

@@ -11,8 +11,10 @@
 
 use super::*;
 use std::num::NonZeroUsize;
+use std::str::FromStr;
 
-use mz_repr::Row;
+use mz_persist_types::ShardId;
+use mz_repr::{RelationDesc, Row};
 
 #[mz_ore::test]
 fn pending_peek_response_precedence() {
@@ -74,3 +76,44 @@ fn peek_max_size_wins_over_row_iteration_limit_in_every_order() {
         assert_eq!(pending.response, expected, "{permutation:?}");
     }
 }
+
+/// A stashed answer is bounded by `max_result_size` across the workers that produced it, not
+/// per worker.
+///
+/// Each worker decides on its own whether its share is large enough to stash, so a peek whose
+/// answer exceeds the ceiling can arrive as several stashed responses each well under it. Reading
+/// only the rows a response carries inline would leave such an answer bounded at the ceiling times
+/// the worker count.
+#[mz_ore::test]
+fn a_stashed_answer_is_bounded_across_workers() {
+    let stashed = |stashed_byte_len| {
+        PeekResponse::Stashed(Box::new(StashedPeekResponse {
+            num_rows_batches: 1,
+            encoded_size_bytes: 0,
+            stashed_byte_len,
+            relation_desc: RelationDesc::empty(),
+            shard_id: ShardId::from_str("s00000000-0000-0000-0000-000000000000").expect("valid"),
+            batches: Vec::new(),
+            inline_rows: Vec::new(),
+        }))
+    };
+    let max_result_size = 100;
+
+    // One worker's share, under the ceiling on its own.
+    let mut pending = PendingPeek::new();
+    pending.absorb(0, stashed(60), max_result_size);
+    assert!(
+        matches!(pending.response, PeekResponse::Stashed(_)),
+        "a share under the ceiling answers with its batches"
+    );
+
+    // A second share of the same size puts the answer over it.
+    pending.absorb(1, stashed(60), max_result_size);
+    assert_eq!(
+        pending.response,
+        PeekResponse::Error(PeekError::unstructured(format!(
+            "total result exceeds max size of {}",
+            ByteSize::b(max_result_size)
+        ))),
+    );
+}

src/compute/src/compute_state/peek_scan.rs

@@ -51,7 +51,7 @@ const _: () = {
 pub(super) type RowBatch = Vec<(Row, NonZeroI64)>;
 
 /// The byte size of a row's count, as an answer built from a [`RowBatch`] stores it.
-const COUNT_BYTE_SIZE: usize = size_of::<NonZeroUsize>();
+pub(super) const COUNT_BYTE_SIZE: usize = size_of::<NonZeroUsize>();
 
 /// What a walk has spent, in the phases the peek metrics report.
 ///

src/compute/src/compute_state/peek_stash.rs

@@ -28,7 +28,7 @@ use tokio::sync::oneshot;
 use tracing::warn;
 use uuid::Uuid;
 
-use crate::compute_state::peek_scan::RowBatch;
+use crate::compute_state::peek_scan::{COUNT_BYTE_SIZE, RowBatch};
 
 /// Whether a [`StashUpload`] has room for more rows.
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
@@ -95,6 +95,9 @@ pub(super) struct StashUpload {
     /// Rows written so far, counting a row with a diff of `n` as `n` rows, which is how the
     /// finishing counts them.
     num_rows: u64,
+    /// The size of those rows, measured as an inline answer measures its own, so that
+    /// `max_result_size` bounds a stashed answer and an inline one alike.
+    stashed_byte_len: usize,
     /// The runtime an abandoned upload's deletion is spawned on, held rather than taken from the
     /// ambient context because [`StashUpload::abandon`] runs where there may be none.
     runtime: Handle,
@@ -155,6 +158,7 @@ impl StashUpload {
             upper,
             max_rows,
             num_rows: 0,
+            stashed_byte_len: 0,
             runtime: Handle::current(),
         })
     }
@@ -175,6 +179,10 @@ impl StashUpload {
             }
 
             self.num_rows += u64::from(NonZeroU64::try_from(diff).expect("diff fits into u64"));
+            self.stashed_byte_len = self
+                .stashed_byte_len
+                .saturating_add(row.byte_len())
+                .saturating_add(COUNT_BYTE_SIZE);
             let diff: i64 = diff.into();
 
             self.batch_builder
@@ -226,6 +234,7 @@ impl StashUpload {
         let stashed_response = StashedPeekResponse {
             num_rows_batches: self.num_rows,
             encoded_size_bytes: batch.encoded_size_bytes(),
+            stashed_byte_len: self.stashed_byte_len,
             relation_desc: self.relation_desc.clone(),
             shard_id: self.shard_id,
             batches: vec![batch.into_transmittable_batch()],

src/compute/src/compute_state/peek_stash/tests.rs

@@ -459,6 +459,7 @@ async fn upload_with_bounds(
         upper: Antichain::from_elem(upper),
         max_rows: None,
         num_rows: 0,
+        stashed_byte_len: 0,
         runtime: Handle::current(),
     }
 }

test/sqllogictest/max_result_size.slt

@@ -141,6 +141,59 @@ SELECT * FROM t1;
 query error total result exceeds max size of 1\.0 MiB
 SUBSCRIBE (SELECT * FROM t1 GROUP BY a, b)
 
+# A peek whose answer goes to the peek stash is bounded by `max_result_size` the same way an
+# inline one is. Each worker decides on its own whether its share is large enough to stash, so a
+# result over the cap arrives as several stashed responses each well under it, and a controller
+# that measured only the rows a response carries inline would bound such an answer at the cap
+# times the worker count.
+simple conn=mz_system,user=mz_system
+ALTER SYSTEM SET max_result_size TO '1MB';
+----
+COMPLETE 0
+
+simple conn=mz_system,user=mz_system
+ALTER SYSTEM SET enable_compute_peek_response_stash TO 'true';
+----
+COMPLETE 0
+
+# Low enough that every worker's share of the result below reaches the stash.
+simple conn=mz_system,user=mz_system
+ALTER SYSTEM SET compute_peek_response_stash_threshold_bytes TO 1024;
+----
+COMPLETE 0
+
+# Small enough that the per-chunk check the adapter applies on read-back cannot be what
+# rejects the result, so what this pins is the bound across the workers that produced it.
+simple conn=mz_system,user=mz_system
+ALTER SYSTEM SET compute_peek_response_stash_read_batch_size_bytes TO 65536;
+----
+COMPLETE 0
+
+# The index is what makes this a fast-path index peek, which is the path that writes to the stash.
+statement ok
+CREATE INDEX t1_idx ON t1 (a);
+
+query error total result exceeds max size of 1\.0 MiB
+SELECT * FROM t1;
+
+statement ok
+DROP INDEX t1_idx;
+
+simple conn=mz_system,user=mz_system
+ALTER SYSTEM RESET compute_peek_response_stash_threshold_bytes;
+----
+COMPLETE 0
+
+simple conn=mz_system,user=mz_system
+ALTER SYSTEM RESET compute_peek_response_stash_read_batch_size_bytes;
+----
+COMPLETE 0
+
+simple conn=mz_system,user=mz_system
+ALTER SYSTEM SET enable_compute_peek_response_stash TO 'false';
+----
+COMPLETE 0
+
 # Regression: `RowSetFinishing` must not reject a result that fits within
 # `max_result_size`. This folds to a fast-path Constant of 10000 rows with a
 # `byte_len()` of 1100000, which fits under the 1140000 cap. The finishing path