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

[antiguru]

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. This is the layer that makes the one-path claim in the design document true.

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, 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, which is the trade the design says the effort is making. 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. The dyncfg description said otherwise and no longer 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. A Linear issue tracks the persist-side fix that would remove the need for the catch.

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.

What an upload can delete is not always everything it wrote. Persist merges runs once a builder holds more than peek_response_stash_batch_max_runs of them, and a merge writes a fresh output and drops the inputs it read without entering them into shard state. Parts flush at persist_blob_target_size, so an upload small enough never to merge deletes all of what it wrote and a large one deletes the output of its last merge. A reader deleting a response it has finished with reaches exactly the same parts, so this is a property of the builder rather than of abandonment, and it is unchanged by this PR.

Nothing bounds the total size of a stashed answer, which is also true today. max_result_size bounds the prefix a single scan retains between batches and never the sum across them. The only ceiling is the reader's own budget.

Known gaps, none of which this PR closes:

• No sqllogictest or testdrive exercises a real stashed fast-path peek through the adapter reader. inline_rows was always empty from a worker on the old path and is now non-empty up to the threshold, and the adapter merges it across workers, so that merge has no coverage in CI.
• Nothing bounds how many abandoned uploads hold a part at once. Each can hold up to persist_blob_target_size, and the abandon task outlives the walk's permit.
• The caught persist panic is verified by construction rather than by a test, because the panic handler that turns it into an abort is not installed under cargo test.

🤖 Opened by Claude Code on behalf of @antiguru

Replaces #38482, which GitHub closed when the design document moved from the bottom of the stack to the top and its branch was force-pushed past these commits. Same content, new base.

[def-]

QA LLM Review

1. MEDIUM -- Every peek over the stash threshold now needs a promotion permit, and holds it across the blob write

src/compute/src/compute_state/peek_offload.rs:187

Routing the stash through the promoted driver puts every streamable peek whose result exceeds compute_peek_response_stash_threshold_bytes (10 KB) behind compute_index_peek_permits, whose default is one per timely worker, and the permit is held until the persist batch has been finished. On a single-worker replica process those peeks now serialize on one permit across a blob round-trip each, where the previous stash upload ran as an ungated task. This lands with enable_index_peek_offload off, so the kill switch does not cover it.

Details

OffloadedPeek::promote acquires a permit unconditionally; the flag only chooses between ActivationBudget::Unbounded and Bounded. With the switch off, unbounded fuel still leaves the batch-ready suspension, which src/compute/src/compute_state.rs:2002 now turns into a promotion. So the sequence for an ordinary SELECT with no ORDER BY returning more than 10 KB is: walk ~10 KB inline, suspend, spawn, queue for a permit, and only then open the shard and write. _permit is a parameter of walk, so it is not released until stashed_answer has awaited StashUpload::finish, which is where the buffered part reaches blob storage.

PeekPermits::new(workers_per_process) (src/compute/src/server.rs:193) with INDEX_PEEK_PERMITS defaulting to 0 means one permit on the smallest replica sizes. The dyncfg's own justification is still purely about CPU -- "A peek that walks on its worker occupies one core, so peek CPU is capped today at one core per worker" -- but the work it now gates is mostly waiting on persist, so the bound both idles the core it reserves and caps concurrent uploads at one. Two second-order effects: the PR's claim that "with the switch off, ordinary peeks behave exactly as they do today" does not hold for the >10 KB ones, which are not rare; and with the switch on, latency-motivated promotions contend for the same permits against a stream of ordinary stashing peeks.

Worth deciding explicitly rather than inheriting: either release the permit once the walk stops producing rows and let finish run unpermitted, or size the pool for the stash workload rather than for core parity (and update the INDEX_PEEK_PERMITS rationale to say what it now bounds).

2. LOW -- PersistLocation is cloned on every visit to a pending index peek

src/compute/src/compute_state.rs:1254

The stash location is now cloned before seek_fulfillment runs its frontier check, so a peek that is still NotReady pays two Url clones (two heap allocations) on every sweep, not just once when it is promoted. The previous code only read .is_some() here. process_peeks is on the worker loop, which the surrounding comments treat as hot enough to justify an early return before the pending map is even walked.

Details

The clone is only needed in the PeekStatus::Promote arm, which is where StashTarget::new consumes it. Keeping a bool here and cloning self.compute_state.peek_stash_persist_location inside that arm restores the old cost without changing behaviour.

[ggevay]

(Sorry, I don't have capacity to review today, because I need to urgently finish something else, and I'll be on PTO tomorrow. I can review next week, or you could ask Aljoscha, who originally authored the peek stash.)

[antiguru]

Review findings, posted by Claude Code on behalf of @antiguru. Line numbers are at 3cd9c996a3.

1. The size ceiling this layer added is per-worker, so a stashed answer still escapes it

This corrects the claim in d9da98bf6b's message and in the comment at peek_scan.rs:446-452, which says the check is measured against "everything the answer will contain [...] because a peek bound for the stash returns its rows to the client like any other". It is measured against everything this worker's answer will contain. Nothing sums the stashed bytes across workers:

• the controller's cross-worker check is PendingPeek::absorb (src/compute-client/src/service.rs:588-601), which accumulates response.inline_byte_len(), and PeekResponse::inline_byte_len() for Stashed counts only inline_rows (src/compute-client/src/protocol/response.rs:216). At the 10 KB default threshold that is at most 10 KB per worker, and the batches are invisible to it;
• the adapter's finish_incremental checks rows.byte_len() > max_result_size per chunk (src/expr/src/relation.rs:3711), and peek_stash_read_batch_size_bytes guarantees each chunk is small, so it never fires;
• the inline path is bounded across workers, because the coordinator checks the merged collection in RowSetFinishing::finish_inner.

The fixture in the tree demonstrates it. test/sqllogictest/max_result_size.slt:104-141 sets max_result_size to 1MB, and has to ALTER SYSTEM SET enable_compute_peek_response_stash TO 'false' to get the error, with a comment saying total in the message matters "because it indicates we failed when aggregating the result from multiple workers, as opposed to on any single worker". With the stash left on, each of eight workers accumulates roughly 140 KB, well under the 1 MiB per-worker ceiling, each stashes, the controller sums at most 80 KB of inline_rows, and the client receives about 1.1 MiB under a 1 MiB cap. The effective bound is workers × max_result_size.

What this layer did fix is real: a single worker's stashed answer was previously unbounded, because the prefix is handed away and reset on every batch. It does not restore the cross-worker bound. StashedPeekResponse::size_bytes() already exists, so feeding it into absorb's accumulator is the fix. Failing that, the comment should say it bounds one worker's share.

2. The rewrite ships on by default and the kill switch does not cover it

ENABLE_PEEK_RESPONSE_STASH defaults to true (src/compute-types/src/dyncfgs.rs:554-559). StashingPeek and PeekStatus::UsePeekStash are deleted, and ENABLE_INDEX_PEEK_OFFLOAD explicitly does not gate stash-bound promotion. Three things therefore reach every replica with no runtime revert:

• the new one-walk stash mechanism itself;
• max_result_size now failing stashed answers that previously succeeded, a user-visible error where there was none, and with a different message than the inline path produces (result exceeds max size of ..., without total);
• the row-iteration limit now spanning the stash walk. PEEK_ROW_ITERATION_LIMIT defaults to 1000 and is currently disabled, but if it is ever enabled every stashed answer dies, since a stashed answer needs only to exceed 10 KB to exist.

Turning enable_compute_peek_response_stash off is not a revert, it converts large results into max_result_size errors. The rest of this stack lands inert behind a switch, so this layer's posture is worth an explicit decision and a release note. Nothing in test/ pins the new error.

3. The panic catch this layer is built around is never exercised

peek_stash.rs:312-337. The catch exists because aborting a walk with a part write in flight leaves persist's Pending in Blocking, and the later finish hits panic!("block_until_ready cancelled?") in src/persist-client/src/internal/merge.rs:182, outside any catch, where the replica's handler aborts the process.

The test harness sets persist_batch_builder_max_outstanding_parts = 8_192 (peek_stash/tests.rs:68-77) specifically to keep that stall out of the way, so both abort tests, an_aborted_walk_deletes_the_parts_its_upload_wrote and a_walk_cancelled_while_uploading_deletes_what_it_wrote, avoid the exact condition the catch is for. Delete the ore_catch_unwind and CI stays green while production gains a replica-abort path on peek cancellation.

The primitive is the right one: ore_catch_unwind scopes CATCHING_UNWIND_ASYNC, which the enhanced handler checks, and the panic is raised on the awaiting task inside the scope. It is the coverage that is missing. A test that leaves max_outstanding_parts at its default, aborts mid-write, and asserts the process survives and the warning fires would pin it.

Smaller, same area: the catch wraps only batch_builder.finish(upload). The batch.delete().await on the next line is outside it, so a persist panic there still aborts the replica, contrary to the module doc's "Cleanup holds itself to the same rule".

4. No test/-level coverage for the rewritten path

Everything is unit-level against in-memory persist. test/sqllogictest/max_result_size.slt is the obvious home for both the new ceiling error and a positive case, that a large answer still comes back intact with the stash on. Neither exists. For a change that rewrites the default-on large-result path, the unit tests do not cover the schema and encoding round trip through a real replica.

Minor

• compute_state.rs:1213 clones the whole PersistLocation, two SensitiveUrls and two allocations, on every visit to every pending index peek, including the point lookups this stack exists to speed up, and discards it unless the peek promotes. The base asked is_some(). Move the clone into the PeekStatus::Promote arm.
• peek_offload.rs:344-352: the Satisfied exit calls only observe_error_phase, so a stashed LIMIT peek contributes nothing to index_peek_row_iteration_seconds or _rows. Same on the push-error and open-error arms. Deliberate and pinned by a test, but the positions histogram systematically under-reports for exactly the peeks the stash exists for, and neither help text says so.
• peek_stash/tests.rs:73 says "See the finding recorded for this layer", which points outside the repo. CLAUDE.md wants a comment to stand on its own; name the Pending::Blocking panic instead.

Checked and clean

Blob lifetime. Every exit walked, no undocumented leak. Complete with an upload finishes it; Failed, a push error and an observed cancellation all discard; an open error and NO_STASH_LOCATION cannot have an upload; a task abort reaches Drop for StashUpload and abandon; an abort inside finish_batch's rx.await makes the detached task's send fail and DeliveredBatch::drop deletes; a panic between finish_batch and delivered.take() likewise. The two remaining leaks, a dropped PeekResponse::Stashed and a replica dying mid-upload or a shutting-down runtime, are the documented ones. finish_batch correctly detaches rather than aborts.

Satisfied against offset and limit. Sound, and not even a change. is_streamable requires an empty order_by and an identity projection, and stash eligibility is exactly is_streamable, so no ORDER BY can coexist with the stash and no top-N row can live past offset + limit. The base stopped at the same num_rows >= max_rows point including mid-batch. LIMIT 0 is handled, since demand() is checked before each row.

Old against new answer equivalence. Equal as a multiset with diffs preserved. num_rows_batches counts sum-of-diffs of what was added and excludes inline_rows, matching its documented contract, and StashedPeekResponse::num_rows adds both halves. The adapter sends inline_rows first and asserts is_streamable, so ordering is irrelevant by construction. Prefix thinning after a handoff can leave more than offset + limit rows, which the adapter's re-application absorbs, and never fewer.

Conventions. No as conversions, no std::collections::HashMap/HashSet, no side effects in debug_assert!, no em-dashes or structuring semicolons in added comments. doc/user/data/metrics.yml was regenerated for both the metric add and the mz_stashed_peek_seconds removal.

[antiguru]

Findings addressed, posted by Claude Code on behalf of @antiguru.

The result-size ceiling now bounds a stashed answer across workers. Each worker decides on its own whether its share is large enough to stash, so an oversized result arrives as several stashed responses each under the ceiling, and PendingPeek::absorb accumulated inline_byte_len(), which for Stashed counts only the inline rows. The batches were invisible to it. The response now carries the size of its stashed rows measured the way an inline answer measures its own, and the controller accumulates that. The encoded size the response already carried is what the batches cost in blob storage, a different number.

Corrections to my earlier report on this PR. I cited max_result_size.slt's enable_compute_peek_response_stash TO false line as proof the hole was reachable. It is not: that query is a persist fast-path peek, which never stashes, and removing the line leaves the file green either way. I checked before claiming the fix works.

The defect is real, and here is what actually demonstrates it. With an index on the table so the peek takes the fast-path index route, the stash threshold low enough that every worker's share reaches the stash, and compute_peek_response_stash_read_batch_size_bytes lowered so the adapter's per-chunk check on read-back cannot be what rejects the result: without this change the query returns roughly 1.1 MiB under a 1 MiB cap, and with it the peek fails as an inline one of that size would. That case is in max_result_size.slt now, and I verified it red before the fix and green after. The per-chunk check is why the gap was invisible: at the default read batch size it happened to catch this.

The panic catch is still untested, and now says so. The abort tests raise persist's outstanding-part bound to keep a stalled write out of the way, and that stall is the condition the catch exists for. Against in-memory persist a part write completes before an abort can land inside it, so the Blocking state is not reachable on demand. Stated with a TODO rather than papered over with a test that would not exercise it. The comment that pointed at a finding outside the repo names the hazard instead.

The stash location is off the sweep's hot path. It was cloned on every visit to every pending index peek, two SensitiveUrls, and discarded unless the peek promoted. A flag there, the clone in the arm that promotes.

Open for you. This layer ships default-on with no revert: enable_compute_peek_response_stash defaults true, so the one-walk stash reaches every replica whether or not the offload is armed, and turning it off converts large results into max_result_size errors rather than reverting. The design doc now says that under Kill switch. Whether that posture is right is a decision rather than a defect.

[def-]

QA LLM Review

1. MEDIUM -- A stashed row is charged 16 bytes more than the same row answered inline, so max_result_size rejects results that are under the cap

src/compute/src/compute_state/peek_stash.rs:156

StashUpload::push accounts a stashed row as Row::byte_len() + COUNT_BYTE_SIZE, while every inline row this PR sums it with is accounted as RowCollection::byte_len(). The two differ by a flat 16 bytes per row, and by up to 2x for narrow rows, so the aggregate ceiling the controller now applies to stashed answers fires well before the client would have received max_result_size bytes. Since enable_compute_peek_response_stash defaults on with a 10 KB threshold, this governs essentially every streamable fast-path index peek.

Details

Row::byte_len() (src/repr/src/row.rs:243) is size_of::<Row>(), statically asserted at 24 (src/repr/src/row.rs:301), plus the heap bytes only once the row spills past 23. RowCollection::byte_len() (src/expr/src/row/collection.rs:140) is the packed data length plus size_of::<usize>() plus size_of::<NonZeroUsize>(). So per entry the stash charges 24 + data + 8 where an inline answer charges data + 16: +16 bytes for any row wider than 23 bytes, and 32 against 21 for a five-byte row. PendingPeek::absorb (src/compute-client/src/service.rs:602) adds the two together and compares the sum to one ceiling.

The mismatch also sits inside a single response. StashedPeekResponse::answer_byte_len (src/compute-client/src/protocol/response.rs:347) adds stashed_byte_len to inline_rows.byte_len(), so the same row is charged differently depending on which side of the stash threshold the walk happened to cut, and a peek where some workers stash and others answer Rows mixes both units in one total. SELECT id FROM t over a narrow column measures roughly 1.5x what the client receives, so a max_result_size of 100 MB rejects at about 65 MB of result.

The scan's own per-worker check has always used this formula, so the ruler is not new; what is new is that it now reaches the aggregate check, which is the one that decides a multi-worker peek's fate, and that it is added there to bytes measured the other way. The doc on stashed_byte_len states the property that would make the bound mean what max_result_size means, "measured as an inline answer measures its own", and charging the row's packed data length rather than Row::byte_len() in push is what gives it. The regression note at the foot of test/sqllogictest/max_result_size.slt records an overcount of the same shape being fixed in the finishing path, so the direction is one this cap has been corrected for before.

[antiguru]

Confirmed and fixed. The mechanism is right on every step, and the fix goes wider than the one call site, because the site was not the only place using that ruler.

What the two rulers are. RowCollection::byte_len charges an entry the row's packed data plus 8 bytes of offset plus 8 bytes of count, so data + 16. The peek path charged Row::byte_len() + COUNT_BYTE_SIZE, which is the 24-byte Row struct plus heap bytes only once the row spills past 23, so 32 + data for a wide row and a flat 32 for a narrow one. The Row struct's own bytes are in no answer, which is why the finishing path never counted them: RowSetFinishing::finish_inner bounds max_result_size with rows.byte_len(), and the regression note at the foot of max_result_size.slt records this cap being corrected in that direction before.

One ruler now, in peek_scan. entry_byte_len(row) = row.data_len() + size_of::<usize>() + size_of::<NonZeroUsize>() replaces every use of the old formula: the scan's running total_size, the size it subtracts for rows thinning drops, and StashUpload::push. So the per-worker check, the aggregate check and stashed_byte_len all measure a row the way the answer that carries it does, and stashed_byte_len's documented "measured as an inline answer measures its own" is now true rather than aspirational.

Three test fixtures pinned thresholds as count * one_row_size taken from the first row of the fixture. That worked only because the old ruler charged every non-spilled row 32 bytes whatever it held. Datum::UInt8/Datum::UInt64 pack into the fewest bytes that hold the value, so ok_row(0) is a byte narrower than every other value the fixtures use, and the summed prefix of n rows now falls short of n times any single row's size. The thresholds are built from the widest row of the fixture where they have to hold for every batch of a walk, and summed over the rows they mean where they name one crossing.

The Unstructured audit that came with it. service.rs's aggregate ceiling reported an Unstructured error, which the adapter maps to AdapterError::Unstructured, an internal error. It now returns PeekError::ResultExceedsMaxSize, the variant the per-worker check already uses, which maps to AdapterError::ResultSize. Same for the persist fast-path peek in compute_state.rs, whose message was ResultExceedsMaxSize's Display verbatim.

That drops the word "total" from the peek message, which max_result_size.slt had been using to say the failure came from the aggregate rather than from one worker. The test does not lose its teeth: eight workers split 10000 rows of about 110 bytes, so no single share comes near a 1 MiB cap and only the total can reject the query. The comment says that now instead. SUBSCRIBE keeps "total result exceeds max size", since its check is a separate String path in stash.

The rest of the unstructured sites in the stack are genuine internal faults, each paired with a soft_panic_or_log! or a warn!: an offloaded walk with nowhere to write, a stash that would not open, a batch persist rejected, a task that died without an outcome, mismatched shard ids or relation descs across stashed responses. Two classes outside the stack would be worth a structured variant of their own, and are not one here: the negative-multiplicity errors, which have no structured form anywhere in the tree, and the "since has advanced" / "compaction frontier is beyond the read" pair, which are the same condition reported from the two peek paths with different text.

Posted by Claude Code on behalf of @antiguru.

[antiguru]

collect_finished_data's tail is one match now. The offloaded bool, the if !offloaded block and the two early returns are gone: ScanOutcome::Finished(result) counts the walk on the inline substrate and reports the error phase once for both answers, then a nested match on the Result adds the ok phase and builds the response or returns the error. Suspended is the arm that reports nothing, which is now visible from the arm itself rather than from a bool computed above the match.

One correction to the reasoning: observe_error_phase was already reached on the Ok path as well as the Err one, since the if !offloaded block ran ahead of the match. What is asymmetric is observe_ok_phase, and deliberately: an error answer can come from either walk, so the ok-walk numbers describe a finished phase only when the answer is rows. That is now stated at the arm that skips it. Behaviour is unchanged, which the 139 tests in mz-compute and mz-compute-client still pass.

This shape only works from this commit up. In peek/offload-driver the same code has three outcomes rather than two, because a suspension holding a full batch takes the peek to the stash by a different route, and both of those outcomes want the same accounting. The bool earns its keep there.

Posted by Claude Code on behalf of @antiguru.