compute: coalesce MV sink batches across workers

[antiguru]

Continues #35411 (draft by @bkirwi), rebased onto main.

With many workers each MV sink worker writes its own small batch per interval, so the sink emits many small parts and a large consensus state diff. This coalesces the parts the workers in one process write for a batch interval into a single, larger shared batch.

SharedBatches hands out a SharedBatchBuilder per batch id; handles for the same id feed one process-global builder task, and the last handle to finish receives the batch. The sync (v2) sink is wired to it, best-effort and default off (enable_compute_sync_mv_sink_shared_batches). On an 8-worker replica this cut persist blob PUTs ~67% and consensus state-diff bytes ~54%, at unchanged consensus command count.

Best-effort means a worker only merges its part into the shared batch when its write overlaps the others still building in the same interval, so parts per append land around 1.4 rather than the ideal 1.0. The barrier that closes that gap to exactly 1.0 is split into the follow-up #37617.

The v1 sink is untouched apart from carrying a shared batch id on BatchDescription. @bkirwi's draft also rewired v1 onto the shared builder, but that path is the default sink and the rewiring is not flag-gated, so it is left out here.

Original authorship by @bkirwi is preserved in the first commit.

[def-]

QA LLM Review

1. HIGH -- Cleanup sweep's Weak::upgrade can orphan a finished shared batch

src/storage-operators/src/persist.rs:98

The map-cleanup sweep in SharedBatches::builder probes liveness with weak.upgrade(), which takes a temporary strong reference. SharedBatchBuilder::finish decides who owns the batch with Arc::into_inner, which returns None whenever the strong count is above one. If the sweep's probe overlaps the last real handle's finish, every worker returns None, nobody emits the batch, and the append operator advances the shard upper over an interval whose updates were never written.

Details

The window: builder runs the sweep whenever state_map.len() > last_retained_len * 2, which is most calls once last_retained_len settles low, and each sweep upgrades every live entry in the map. A worker entering builder for sink X therefore probes sink Y's in-flight entry while another worker is finishing it. Both operations are unsynchronized with respect to each other: finish never takes self.data's mutex.

What follows from a hit:

• Arc::into_inner at src/storage-operators/src/persist.rs:218 sees count 2, returns None, and drops the last real handle.
• The sweep's temporary then drops the BatchState, closing tx. The builder task completes normally and produces Some(batch), but its JoinHandle lives in the dropped BatchState, so the return value is discarded. The parts stay in blob until leaked-blob cleanup (with an un-consumed Batch ... dangling blob keys warning).
• Every worker's write_shared_batch returned None, so apply_command sends WriteResponse { batch: None } and no worker emits into batches. In append::maybe_append_batches, batch_description is still Some(desc) and batches is empty, so compare_and_append_batch(&mut [], lower, upper, true) seals [lower, upper) with no data.
• The corrections are not lost (consolidated_updates_before borrows rather than drains), but the subsequent advance_since(upper) rounds them to upper, so they are re-written at upper. The MV's contents at every timestamp in [lower, upper) are permanently missing those updates. That is observable by AS OF reads, by SUBSCRIBE, and by downstream indexes/MVs that already computed at those times; the self-correcting feedback loop cannot repair a sealed timestamp.

Fix — probe the count instead of materializing a reference:

-            state_map.retain(|_, weak| weak.upgrade().is_some());
+            // Do not `upgrade()` here: the temporary strong reference it creates can make a
+            // concurrent `SharedBatchBuilder::finish` observe a strong count above one and
+            // return `None`, orphaning a batch that nobody then appends.
+            state_map.retain(|_, weak| weak.strong_count() > 0);

Weak::strong_count is racy in the harmless direction only: a stale 0 cannot become live again, and a stale non-zero just keeps a dead entry for one more sweep.

2. MEDIUM -- add_part recomputes accumulated part sizes on every call, so a flush window costs O(n²)

src/persist-client/src/batch.rs:673

in_progress_goodbytes walks all of finished_parts and calls Part::goodbytes() on each, and add_part calls it twice per part while finished_parts grows monotonically until the blob target is reached. With the default 128 MiB persist_blob_target_size and the sink's 1 KiB-update parts, a narrow MV accumulates thousands of parts per flush window, so the shared builder task makes on the order of n² Part::goodbytes() calls per window — millions of them — before a single part is flushed.

Details

Part::goodbytes() is not free: it constructs an ArrayOrd for the key and the val on each call (src/persist-types/src/part.rs:43), which allocates a Vec per struct level and clones an ArrayRef per column. This is confined to the shared-batch path — add_part has no other caller, and in the add path finished_parts is always empty — but that is the path this PR introduces, and the cost concentrates exactly where batches are largest: hydration and backfill of a large MV.

Keeping a running byte count alongside the vector removes it:

• add a finished_parts_goodbytes: usize field, bump it in finish_part_builder and in add_part's push, and zero it in flush_parts next to finished_parts.clear();
• in_progress_goodbytes then becomes self.part_builder.goodbytes() + self.finished_parts_goodbytes.

[antiguru]

Both confirmed and fixed in the first commit. Branch is now at ae3b03bc7c (pr-35411) and a5cf162608 (pr-35411-barrier).

1. Cleanup sweep orphaning a finished batch. Confirmed, including the consequence. The sweep holds the registry mutex, finish never takes it, so the two are unsynchronized, and Arc::into_inner returns None on any strong count above one. Corrections survive the empty append because consolidated_updates_before borrows rather than drains, but the following advance_since(upper) rounds them to upper, so the interval's timestamps are permanently missing those updates rather than eventually repaired.

Applied the suggested fix, with the reasoning recorded at the call site:

// Probe the count rather than `upgrade()`: the temporary strong reference an upgrade
// creates can make a concurrent `SharedBatchBuilder::finish` see a strong count above
// one and hand the batch to nobody, which seals the interval with no data. A stale
// count is harmless in both directions, since zero cannot become live again and a
// stale non-zero only keeps a dead entry until the next sweep.
inner.states.retain(|_, weak| weak.strong_count() > 0);

Two other Weak::upgrade sites were checked and left alone. upgrade_or_init materializes a reference for a worker that is a genuine participant, so a finish losing the race to it still delivers the batch through the joining handle. note_skip (in the barrier follow-up) has the same shape but only runs in barrier mode, where finish_barrier resolves through the barrier rather than Arc::into_inner.

2. Quadratic in_progress_goodbytes. Confirmed: Part::goodbytes builds an ArrayOrd for the key and the val on every call, and add_part calls the sum twice per part. Added a finished_parts_goodbytes running counter, maintained in a new push_finished_part helper (used by both finish_part_builder and add_part) and zeroed next to finished_parts.clear() in flush_parts, so in_progress_goodbytes is part_builder.goodbytes() + finished_parts_goodbytes.

Two things also changed since the reviewed revision, both from the same pass:

• The v1 sink rewiring is dropped. It was not flag-gated, and enable_compute_sync_mv_sink defaults off, so it would have changed the default sink's behavior. v1 now only carries PersistApi.persist_batches and BatchDescription.shared_id, both read solely by v2.
• Dropping it exposed two defects that also affected the v2 shared path. PersistClient::batch_builder hardcoded "peek_stash" as the shard-metrics name, which is first-writer-wins and would have mislabeled a user shard for its lifetime; it now takes a shard_name. And the sink assembled Schemas { id: None, .. }, so parts were written with schema_id: None and could not be decoded through the shard's schema registry; the id now comes from WriteHandle::schema_id().

Reviewed and applied by Claude Code.