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

doc/user/data/metrics.yml

@@ -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:
@@ -3709,25 +3713,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

misc/python/materialize/mzcompose/__init__.py

@@ -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/parallel_workload/action.py

@@ -3430,8 +3430,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",

src/adapter/src/coord/peek.rs

@@ -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/compute-client/src/protocol/response.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,14 @@ 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. `max_result_size` bounds
+    /// [`StashedPeekResponse::stashed_byte_len`] instead.
     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 +341,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,20 +589,24 @@ 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() {
-            // NOTE: Tests match on this exact message, so nothing else may produce it.
-            let error = PeekError::unstructured(format!(
-                "total result exceeds max size of {}",
-                ByteSize::b(max_result_size)
-            ));
+        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 only the reader, which never sees this response, deletes a stashed
+            // answer's parts. This is the leak DB-50 reports for a cancelled peek, which drops
+            // the same batches at the same merge.
+            let error = PeekError::ResultExceedsMaxSize {
+                max_result_size: max_result_size.cast_into(),
+            };
             let current = mem::replace(&mut self.response, PeekResponse::Canceled);
             self.response = merge_peek_responses(current, PeekResponse::Error(error));
         }
@@ -636,6 +640,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 +649,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 +677,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() {
@@ -60,10 +62,9 @@ fn peek_max_size_wins_over_row_iteration_limit_in_every_order() {
         [2, 0, 1],
         [2, 1, 0],
     ];
-    let expected = PeekResponse::Error(PeekError::unstructured(format!(
-        "total result exceeds max size of {}",
-        ByteSize::b(max_result_size)
-    )));
+    let expected = PeekResponse::Error(PeekError::ResultExceedsMaxSize {
+        max_result_size: max_result_size.cast_into(),
+    });
 
     for permutation in permutations {
         let mut pending = PendingPeek::new();
@@ -74,3 +75,43 @@ 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::ResultExceedsMaxSize {
+            max_result_size: max_result_size.cast_into(),
+        }),
+    );
+}

src/compute-types/src/dyncfgs.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,24 @@ 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.
+/// 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.
 pub const ENABLE_INDEX_PEEK_OFFLOAD: Config<bool> = Config::new(
     "enable_compute_index_peek_offload",
     false,
@@ -710,8 +698,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,
@@ -824,8 +816,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/Cargo.toml

@@ -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