compute: make the peek stash a state transition of one scan

Deletes the second walk a stashed peek used to cost. A peek whose accumulated
rows crossed the stash threshold abandoned its walk, and a separate path re-read
the trace bundle from the beginning to stream rows into persist. The offloaded
driver now writes each full batch as the walk produces it, so a peek walks its
arrangement exactly once whatever its answer turns out to be, and never returns
to the worker to reach the stash.

The upload becomes a handle the driver feeds. `StashUpload` keeps everything
about the persist interaction, the shard derived from the peek uuid, the schemas,
the batch builder and the `max_rows` early exit, and takes rows through `push`
instead of through a channel the worker pumps. `Complete` carries whatever the
scan still holds, and the driver assembles the answer around it: a driver that
never opened an upload answers with a row collection, and one that did finishes
and answers with the stashed handle, whose `inline_rows` carry the rows that
never reached the stash. The hand-back the layer below needed goes away, because
the task now has somewhere to write.

`PendingPeek` ends with fewer index-peek states than it had before this work
started. `PendingPeek::Stash`, `StashingPeek` and its worker-driven pump,
`start_stash_upload`, `PeekStatus::UsePeekStash` and `PEEK_STASH_NUM_BATCHES` are
all gone. `PEEK_STASH_BATCH_SIZE` goes too, against the plan: it counted rows,
the upload cuts a batch on the byte threshold the scan already tracks, and
nothing read it. A live tunable that does nothing is worse than an absent one,
because an operator who reaches for it gets silence rather than an error.

Two consequences worth stating plainly rather than discovering later.

The kill switch no longer means no peek ever leaves the worker. With
`UsePeekStash` gone, a large streamable peek would otherwise have no route to the
stash at all, which is a functional regression rather than a placement change.
The switch therefore gates the budget-based offload only: a scan that suspends
because its prefix is batch-ready is handed to a task whichever way the switch is
set, because that offload is for stashing rather than for latency. With the
switch off, ordinary peeks behave exactly as they do today.

The row iteration limit now follows a peek into the peek stash. #38158 stopped at
the stash because a stashed peek restarted its scan and the restart charged the
same rows twice. Deleting the restart is what this change does, so the count
simply continues because the scan does.

Cancellation deletes what it wrote. A cancelled peek used to leave the parts
already written in blob storage, since `impl Drop for Batch` only logs the
dangling keys and the reader-side delete runs after a successful read that a
cancelled peek never reaches. Cancellation aborts the offloaded task rather than
signalling it, so an await placed after the cancellation check would never run,
and the obligation lives on the upload itself: dropping one spawns a
`Batch::delete()` onto a runtime handle captured when it opened. A guard covers
the window inside `finish`, where persist has taken the builder and no upload
holds the parts any more.

Reclaiming those blobs must not cost more than it saves. A builder whose part
write was in flight when its walk was aborted holds a write persist has already
marked as waited on, and finishing it panics rather than returning; this replica
aborts the process on any uncaught panic, so reclaiming one query's blob storage
could take the replica with it. The panic is caught and the shard logged. Failing
to reclaim is an outcome this path already tolerates for a replica that dies
mid-upload, which makes a leak the right answer there and an abort the wrong one.

`mz_index_peek_stashed_total` counts the walks the stash answered. It is
incremented beside the offloaded substrate counter, so it is a strict subset by
construction, and every inline-driver test asserts it stays at zero.
`mz_stashed_peek_seconds` goes: its only observer was the deleted walk, and a
registered histogram that never observes reports a flat zero, which reads worse
in a graph than an absent series. The duration is deliberately not reconstructed,
because one loop now walks and writes, so the write is not separable from the
walk.

Author: antiguru

Cargo.lock

@@ -1144,6 +1144,10 @@ metrics:
   help: Time in seek_fulfillment method including frontier checks and data collection.
   source: src/compute/src/metrics.rs
   visibility: internal
+- name: mz_index_peek_stashed_total
+  help: The number of index peek walks that answered with a handle to the peek response stash, always a subset of the `offloaded` substrate of `mz_index_peek_walks_total`.
+  source: src/compute/src/metrics.rs
+  visibility: internal
 - name: mz_index_peek_total_seconds_bucket
   help: Time one visit to an index peek spent on the timely worker. A peek whose walk was offloaded contributes only the inline slice that offloaded it, and its time away from the worker is `mz_index_peek_offload_seconds`.
   labels:
@@ -3704,25 +3708,6 @@ metrics:
   - version
   source: src/environmentd/src/environmentd/main.rs
   visibility: internal
-- name: mz_stashed_peek_seconds_bucket
-  help: Time spent reading a peek result and stashing it in the peek result stash (aka. persist blob).
-  labels:
-  - le
-  - worker_id
-  source: src/compute/src/metrics.rs
-  visibility: internal
-- name: mz_stashed_peek_seconds_count
-  help: Time spent reading a peek result and stashing it in the peek result stash (aka. persist blob).
-  labels:
-  - worker_id
-  source: src/compute/src/metrics.rs
-  visibility: internal
-- name: mz_stashed_peek_seconds_sum
-  help: Time spent reading a peek result and stashing it in the peek result stash (aka. persist blob).
-  labels:
-  - worker_id
-  source: src/compute/src/metrics.rs
-  visibility: internal
 - name: mz_statement_logging_actual_bytes
   help: The total amount of SQL text that was logged by statement logging.
   source: src/adapter/src/metrics.rs

doc/user/data/metrics.yml

@@ -832,8 +832,6 @@ def get_default_system_parameters(
     "compute_index_peek_permit_fraction",
     "compute_peek_response_stash_read_batch_size_bytes",
     "compute_peek_response_stash_read_memory_budget_bytes",
-    "compute_peek_stash_num_batches",
-    "compute_peek_stash_batch_size",
     "storage_statistics_retention_duration",
     "enable_paused_cluster_readhold_downgrade",
     "kafka_retry_backoff",

misc/python/materialize/mzcompose/__init__.py

@@ -3437,8 +3437,6 @@ def __init__(
             "mz_metrics_lgalloc_refresh_interval",
             "mz_metrics_rusage_refresh_interval",
             "mz_metrics_usage_refresh_interval",
-            "compute_peek_stash_num_batches",
-            "compute_peek_stash_batch_size",
             "compute_peek_response_stash_batch_max_runs",
             "compute_peek_response_stash_read_batch_size_bytes",
             "compute_peek_response_stash_read_memory_budget_bytes",

misc/python/materialize/parallel_workload/action.py

@@ -1111,7 +1111,9 @@ impl crate::coord::Coordinator {
                     // - ProtoBatch is lost in flight
                     // - ProtoBatch is lost because when combining PeekResponse
                     // from workers a cancellation or error "overrides" other
-                    // results, meaning we drop them
+                    // results, meaning we drop them, which includes the
+                    // `max_result_size` rejection the controller raises once it
+                    // has summed the workers' stashed rows
                     // - This task here is not run to completion before it can
                     // delete all batches
                     //

src/adapter/src/coord/peek.rs

@@ -218,6 +218,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.
@@ -291,7 +303,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.
@@ -322,6 +342,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/protocol/response.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,25 @@ 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: an error merged over a stashed response drops the batches that response
+            // names, and nothing else reaches them: only the reader deletes a stashed answer's
+            // parts, and it never sees this response. Counting stashed rows here makes that
+            // reachable for a peek over the ceiling, where before the check could not fire on a
+            // stashed answer at all. It is the leak DB-50 already reports for a cancelled peek,
+            // which drops the same batches at the same merge, and reclaiming either needs the
+            // batches to survive the merge rather than a bound that avoids producing them.
+            //
             // 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 +646,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 +655,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 +683,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.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-client/src/service/tests.rs

@@ -626,23 +626,6 @@ pub const PEEK_RESPONSE_STASH_READ_MEMORY_BUDGET_BYTES: Config<usize> = Config::
     ParameterScope::Environment,
 );
 
-/// The number of batches to pump from the peek result iterator when stashing peek responses.
-pub const PEEK_STASH_NUM_BATCHES: Config<usize> = Config::new(
-    "compute_peek_stash_num_batches",
-    100,
-    "The number of batches to pump from the peek result iterator (in one iteration through the worker loop) when stashing peek responses.",
-    ParameterScope::Environment,
-);
-
-/// The size of each batch, as number of rows, pumped from the peek result
-/// iterator when stashing peek responses.
-pub const PEEK_STASH_BATCH_SIZE: Config<usize> = Config::new(
-    "compute_peek_stash_batch_size",
-    100000,
-    "The size, as number of rows, of each batch pumped from the peek result iterator (in one iteration through the worker loop) when stashing peek responses.",
-    ParameterScope::Environment,
-);
-
 /// Whether compute should stop peeks that iterate over too many rows.
 pub const ENABLE_PEEK_ROW_ITERATION_LIMIT: Config<bool> = Config::new(
     "enable_compute_peek_row_iteration_limit",
@@ -652,19 +635,26 @@ pub const ENABLE_PEEK_ROW_ITERATION_LIMIT: Config<bool> = Config::new(
 );
 
 /// The maximum number of rows a peek may iterate over on each worker.
+///
+/// The count spans a peek's whole walk of its arrangement, rows written to the peek stash
+/// included, because a peek walks its arrangement once and the count travels with that walk.
 pub const PEEK_ROW_ITERATION_LIMIT: Config<usize> = Config::new(
     "compute_peek_row_iteration_limit",
     1000,
-    "The maximum number of rows a peek may iterate over on each worker when enable_compute_peek_row_iteration_limit is enabled. Does not apply once a peek's results move to the peek stash.",
+    "The maximum number of rows a peek may iterate over on each worker when enable_compute_peek_row_iteration_limit is enabled. The count spans the peek's whole walk, rows written to the peek stash included.",
     ParameterScope::Environment,
 );
 
-/// Whether a fast-path index peek may move its walk off the timely worker.
+/// Whether a fast-path index peek may move its walk off the timely worker for latency.
 ///
-/// Off, a peek walks to completion on the worker that owns it, delaying every other message that
-/// worker serves. On, a peek that outruns [`INDEX_PEEK_INLINE_BUDGET`] finishes away from it.
+/// Off, a peek walks on the worker that owns it until it answers, delaying every other message
+/// that worker serves. On, a peek that outruns [`INDEX_PEEK_INLINE_BUDGET`] finishes away from it.
 ///
-/// The kill switch for the whole mechanism, off by default until the path earns trust.
+/// This gates latency offload only. A peek whose rows outgrow an inline answer is offloaded either
+/// way, because the driver that writes to the peek stash is the offloaded one. Off means an ordinary
+/// peek runs where it used to, not that none leaves the worker.
+///
+/// Off by default until the path earns trust.
 pub const ENABLE_INDEX_PEEK_OFFLOAD: Config<bool> = Config::new(
     "enable_compute_index_peek_offload",
     false,
@@ -713,8 +703,12 @@ pub const INDEX_PEEK_ACTIVATION_BUDGET: Config<usize> = Config::new(
 ///
 /// At a plausible 100ns to 1us per position this bounds cancellation latency to single-digit
 /// milliseconds. Larger than [`INDEX_PEEK_INLINE_BUDGET`] because an offloaded scan is off the
-/// worker's critical path, so its slices answer to cancellation latency rather than to the
-/// worker's availability.
+/// worker's critical path, so its slices answer to cancellation latency, not to the worker's
+/// availability.
+///
+/// An upper bound, not a period: a walk bound for the peek stash suspends once its accumulation
+/// crosses `peek_response_stash_threshold_bytes`, by far the smaller trigger at that threshold's
+/// default, and unspent fuel is not carried over.
 pub const INDEX_PEEK_YIELD_GRANULARITY: Config<usize> = Config::new(
     "compute_index_peek_yield_granularity",
     10000,
@@ -827,8 +821,6 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet {
         .add(&PEEK_RESPONSE_STASH_BATCH_MAX_RUNS)
         .add(&PEEK_RESPONSE_STASH_READ_BATCH_SIZE_BYTES)
         .add(&PEEK_RESPONSE_STASH_READ_MEMORY_BUDGET_BYTES)
-        .add(&PEEK_STASH_NUM_BATCHES)
-        .add(&PEEK_STASH_BATCH_SIZE)
         .add(&ENABLE_PEEK_ROW_ITERATION_LIMIT)
         .add(&PEEK_ROW_ITERATION_LIMIT)
         .add(&ENABLE_INDEX_PEEK_OFFLOAD)

src/compute-types/src/dyncfgs.rs

@@ -33,6 +33,7 @@ mz-expr = { path = "../expr" }
 mz-metrics = { path = "../metrics" }
 mz-ore = { path = "../ore", features = ["async", "process", "tracing", "columnar", "differential-dataflow", "pager", "region"] }
 mz-proto = { path = "../proto" }
+mz-persist = { path = "../persist" }
 mz-persist-client = { path = "../persist-client" }
 mz-persist-types = { path = "../persist-types" }
 mz-repr = { path = "../repr" }
@@ -47,6 +48,7 @@ prost.workspace = true
 scopeguard.workspace = true
 serde.workspace = true
 smallvec = { workspace = true, features = ["serde", "union"] }
+thiserror.workspace = true
 timely.workspace = true
 tokio.workspace = true
 tracing.workspace = true