Skip to main content

mz_compute/sink/
correction_v2.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! An implementation of the `Correction` data structure used by the MV sink's `write_batches`
11//! operator to stash updates before they are written.
12//!
13//! The `Correction` data structure provides methods to:
14//!  * insert new updates
15//!  * advance the compaction frontier (called `since`)
16//!  * obtain an iterator over consolidated updates before some `upper`
17//!  * force consolidation of updates before some `upper`
18//!
19//! The goal is to provide good performance for each of these operations, even in the presence of
20//! future updates. MVs downstream of temporal filters might have to deal with large amounts of
21//! retractions for future times and we want those to be handled efficiently as well.
22//!
23//! Note that `Correction` does not provide a method to directly remove updates. Instead updates
24//! are removed by inserting their retractions so that they consolidate away to nothing.
25//!
26//! ## Storage of Updates
27//!
28//! Stored updates are of the form `(data, time, diff)`, where `time` and `diff` are fixed to
29//! [`mz_repr::Timestamp`] and [`mz_repr::Diff`], respectively.
30//!
31//! [`CorrectionV2`] holds onto a list of `Chain`s containing `Chunk`s of stashed updates. Each
32//! `Chunk` is a columnation region containing a fixed maximum number of updates. All updates in
33//! a chunk, and all updates in a chain, are ordered by (time, data) and consolidated.
34//!
35//! Chains live in three places:
36//!
37//!  * A [`BucketChain`] partitions times at or beyond the `boundary` (the largest read `upper`
38//!    seen so far) into buckets of exponentially growing time ranges, each holding a list of
39//!    chains. Reads only touch the buckets below their `upper`, so the bulk of the buffered
40//!    updates — in particular far-future retractions produced by temporal filters — is left
41//!    alone.
42//!  * `pending_low` holds chains at times below the `boundary`, mostly insertions arriving
43//!    through the persist feedback.
44//!  * `emitted` is a single chain holding the updates returned by the last read. Updates must
45//!    stay in the buffer until their feedback retractions arrive, and keeping them separate from
46//!    the bucket chain means reads never have to re-merge future updates.
47//!
48//! ```text
49//!       chain[0]   |   chain[1]   |   chain[2]
50//!                  |              |
51//!     chunk[0]     | chunk[0]     | chunk[0]
52//!       (a, 1, +1) |   (a, 1, +1) |   (d, 3, +1)
53//!       (b, 1, +1) |   (b, 2, -1) |   (d, 4, -1)
54//!     chunk[1]     | chunk[1]     |
55//!       (c, 1, +1) |   (c, 2, -2) |
56//!       (a, 2, -1) |   (c, 4, -1) |
57//!     chunk[2]     |              |
58//!       (b, 2, +1) |              |
59//!       (c, 2, +1) |              |
60//!     chunk[3]     |              |
61//!       (b, 3, -1) |              |
62//!       (c, 3, +1) |              |
63//! ```
64//!
65//! The "chain invariant" states that each chain in a bucket has at least `chain_proportionality` times as
66//! many updates as the next one. This means that chain sizes will often be powers of
67//! `chain_proportionality`, but they don't have to be. For example, for a proportionality of 2,
68//! the chain sizes `[11, 5, 2, 1]` would satisfy the chain invariant.
69//!
70//! Note that the invariant is maintained on update counts, not chunk counts. Chunks are
71//! byte-bounded (see `ChunkBuilder`), so chunk count is not proportional to update count and
72//! would be a poor proxy: any chain below the chunk byte boundary is a single chunk regardless
73//! of how many updates it holds, which would let the geometric invariant collapse and break the
74//! O(log N) amortization of inserts.
75//!
76//! Choosing the `chain_proportionality` value allows tuning the trade-off between memory and CPU
77//! resources required to maintain corrections. A higher proportionality forces more frequent chain
78//! merges, and therefore consolidation, reducing memory usage but increasing CPU usage.
79//!
80//! ## Inserting Updates
81//!
82//! A batch of updates is routed by time: updates below the `boundary` become a `pending_low`
83//! chain, the rest is appended as new chains to their respective buckets. Appending to a bucket
84//! merges chains until the chain invariant is restored.
85//!
86//! Inserting an update into the correction buffer can be expensive: It involves allocating a new
87//! chunk, copying the update in, and then likely merging with an existing chain to restore the
88//! chain invariant. If updates trickle in in small batches, this can cause a considerable
89//! overhead. To amortize this overhead, new updates aren't immediately inserted into the sorted
90//! chains but instead stored in a `Stage` buffer. Once the staged updates reach the configured
91//! byte size, they are sorted and routed.
92//!
93//! The insert operation has an amortized complexity of O(log N), with N being the current number
94//! of updates stored.
95//!
96//! ## Retrieving Consolidated Updates
97//!
98//! Retrieving consolidated updates before a given `upper` works by peeling all buckets below the
99//! `upper` off the bucket chain, splitting their chains, the pending low chains, and the previous
100//! `emitted` chain at the `upper`, merging the parts below the `upper` into the new `emitted`
101//! chain, and returning an iterator over that chain.
102//!
103//! Because each chain contains updates ordered by time first, splitting a chain at the `upper`
104//! reuses whole chunks and copies at most one chunk straddling the split point. Updates at times
105//! at or beyond the `upper` are never touched, no matter how many the buffer holds. The
106//! complexity of a read is O(U log K), with U being the number of updates before `upper` and K
107//! the number of chains containing them.
108//!
109//! ## Merging Chains
110//!
111//! Merging multiple chains into a single chain is done using a k-way merge. As the input chains
112//! are sorted by (time, data) and consolidated, the same properties hold for the output chain. The
113//! complexity of a merge of K chains containing N updates is O(N log K).
114//!
115//! There is a twist though: Merging also has to respect the `since` frontier, which determines how
116//! far the times of updates should be advanced. Advancing times in a sorted chain of updates
117//! can make them become unsorted, so we cannot just merge the chains from top to bottom.
118//!
119//! For example, consider these two chains, assuming `since = [2]`:
120//!   chain 1: [(c, 1, +1), (b, 2, -1), (a, 3, -1)]
121//!   chain 2: [(b, 1, +1), (a, 2, +1), (c, 2, -1)]
122//! After time advancement, the chains look like this:
123//!   chain 1: [(c, 2, +1), (b, 2, -1), (a, 3, -1)]
124//!   chain 2: [(b, 2, +1), (a, 2, +1), (c, 2, -1)]
125//! Merging them naively yields [(b, 2, +1), (a, 2, +1), (b, 2, -1), (a, 3, -1)], a chain that's
126//! neither sorted nor consolidated.
127//!
128//! Times below the `since` can only exist in chains read by `consolidate_before`, and only if
129//! the `since` advanced past buffered times since the previous read. For few distinct stale
130//! times — the steady state, where the previously emitted chain was written just before the
131//! since advanced past it — we merge sub-chains, one for each distinct time that's before or at
132//! the `since`. Each of these sub-chains retains the (time, data) ordering after the time
133//! advancement to `since`, so merging those yields the expected result.
134//!
135//! For the above example, the chains we would merge are:
136//!   chain 1.a: [(c, 2, +1)]
137//!   chain 1.b: [(b, 2, -1), (a, 3, -1)]
138//!   chain 2.a: [(b, 2, +1)],
139//!   chain 2.b: [(a, 2, +1), (c, 2, -1)]
140//!
141//! For many distinct stale times — e.g. a since jump across many buffered timestamps when a sink
142//! restarts with an old as-of, or when the first batch of a hydrating sink lands — the number of
143//! sub-chains grows with the number of distinct times, so we instead materialize the affected
144//! updates, advance their times, and sort and consolidate them in bounded runs that we then merge.
145
146use std::cmp::{Ordering, Reverse};
147use std::collections::{BinaryHeap, VecDeque};
148use std::fmt;
149use std::rc::Rc;
150use std::sync::atomic::{self, AtomicUsize};
151use std::sync::{Arc, Mutex, OnceLock};
152
153use columnar::bytes::indexed;
154use columnar::{Columnar, Index, Len, Ref};
155use itertools::Itertools;
156use mz_ore::cast::CastLossy;
157use mz_ore::pool::ChunkHandle;
158use mz_ore::soft_assert_or_log;
159use mz_persist_client::metrics::{SinkMetrics, SinkWorkerMetrics, UpdateDelta};
160use mz_repr::{Diff, Row, Timestamp};
161use mz_timely_util::columnar::body::ColumnBody;
162use mz_timely_util::columnar::chunk;
163use mz_timely_util::temporal::{Bucket, BucketChain};
164use timely::PartialOrder;
165use timely::container::PushInto;
166use timely::progress::Antichain;
167
168use crate::sink::correction::{ChannelLogging, SizeMetrics};
169
170/// Convenient alias for use in data trait bounds.
171///
172/// `D` is constrained to be `Columnar`, so that updates can be stored in a single columnar
173/// region per chunk, and the variable-length payload (e.g. `Row` bytes) lives in the same
174/// allocation as the rest of the chunk. [`DataContainer`] carries the bounds on that container.
175pub trait Data:
176    differential_dataflow::Data + Columnar<Container: DataContainer> + DataBytes + Send + Sync
177{
178}
179impl<D> Data for D where
180    D: differential_dataflow::Data + Columnar<Container: DataContainer> + DataBytes + Send + Sync
181{
182}
183
184/// The bytes a datum occupies, so the staging area can be sized in bytes.
185///
186/// Counts `size_of::<Self>()` plus what the datum owns on the heap. An estimate suffices: the
187/// figure paces when staged updates ship and feeds the buffer's size metrics, nothing else.
188pub trait DataBytes {
189    /// The bytes this datum occupies, heap included.
190    fn data_bytes(&self) -> usize;
191}
192
193impl DataBytes for Row {
194    fn data_bytes(&self) -> usize {
195        self.byte_len()
196    }
197}
198
199/// The bounds [`Data`] places on its columnar container.
200///
201/// The `Ref`-level `Eq + Ord` bounds let the merge/heap code compare updates directly through
202/// the columnar borrow, avoiding `into_owned` clones on the hot path. The `Borrowed`-level
203/// `Send` bound lets a hoisted `Chunk::view` travel with the iterators that
204/// [`CorrectionV2::updates_before`] hands across the persist writer's `await`.
205pub trait DataContainer:
206    Send + Sync + Clone + for<'a> columnar::Borrow<Ref<'a>: Eq + Ord, Borrowed<'a>: Send>
207{
208}
209impl<C> DataContainer for C where
210    C: Send + Sync + Clone + for<'a> columnar::Borrow<Ref<'a>: Eq + Ord, Borrowed<'a>: Send>
211{
212}
213
214/// A borrowed view over a [`Chunk`]'s column.
215///
216/// Obtained from [`Chunk::view`] and indexed with `get`.
217type ChunkView<'a, D> =
218    <<(D, Timestamp, Diff) as Columnar>::Container as columnar::Borrow>::Borrowed<'a>;
219
220/// A data structure used to store corrections in the MV sink implementation.
221///
222/// In contrast to `CorrectionV1`, this implementation stores updates in columnation regions,
223/// allowing their memory to be transparently spilled to disk.
224#[derive(Debug)]
225pub struct CorrectionV2<D: Data> {
226    /// Bucketed storage for updates at times at or beyond `boundary`.
227    ///
228    /// Buckets cover exponentially growing time ranges, so reads only touch the buckets below
229    /// their `upper`, and far-future updates (e.g. retractions produced by temporal filters) are
230    /// rarely touched.
231    chain: BucketChain<ChainBucket<D>>,
232    /// Chains at times below `boundary` that were not yet emitted.
233    ///
234    /// Filled by inserts at times below the boundary (mostly persist feedback) and by the
235    /// remainders of `emitted` when a read uses a smaller `upper` than the previous one. Merged
236    /// into `emitted` by the next read.
237    pending_low: Vec<Chain<D>>,
238    /// Updates that were emitted by `updates_before` but not yet cancelled by persist feedback.
239    ///
240    /// Sorted and consolidated, with all times advanced to the `since`.
241    emitted: Chain<D>,
242    /// A staging area for updates, to speed up small inserts.
243    stage: Stage<D>,
244    /// The lower bound of times stored in `chain`. Only ever advances.
245    ///
246    /// Times below the boundary have been peeled off the bucket chain and can only be stored in
247    /// `pending_low` or `emitted`.
248    boundary: Antichain<Timestamp>,
249    /// The frontier by which all contained times are advanced.
250    since: Antichain<Timestamp>,
251
252    /// Total count of updates last reported to metrics.
253    ///
254    /// Tracked to compute deltas in `update_metrics`.
255    prev_update_count: usize,
256    /// Total size last reported to metrics.
257    ///
258    /// Tracked to compute deltas in `update_metrics`.
259    prev_size: SizeMetrics,
260    /// Global persist sink metrics.
261    metrics: SinkMetrics,
262    /// Per-worker persist sink metrics.
263    worker_metrics: SinkWorkerMetrics,
264    /// Running totals and introspection logging.
265    accounting: Accounting,
266}
267
268/// Fuel for restoring the bucket chain invariant after peeling.
269///
270/// Bounds the restoration work per buffer operation. The bucket chain remains functional when
271/// restoration is incomplete -- peeling and finding work on ill-formed chains, at the cost of
272/// more in-line splitting -- so leftover restoration is simply picked up by the next operation.
273///
274/// `restore` spends one unit of fuel per bucket split, and a single `peel` leaves at most
275/// `BucketTimestamp::DOMAIN` (64) buckets to re-split, so this budget completes restoration in one
276/// call for any realistic buffer. It is deliberately generous: the "incomplete restoration is
277/// picked up next op" path is a correctness safety net for pathological bucket counts, not a hot
278/// path we expect to exercise. Lower it if restoration ever needs to interleave with other work.
279const RESTORE_FUEL: i64 = 1_000_000;
280
281impl<D: Data> CorrectionV2<D> {
282    /// Construct a new [`CorrectionV2`] instance.
283    pub fn new(
284        metrics: SinkMetrics,
285        worker_metrics: SinkWorkerMetrics,
286        logging: Option<ChannelLogging>,
287        chain_proportionality: f64,
288        chunk_size: usize,
289    ) -> Self {
290        let accounting = Accounting::new(logging);
291
292        Self {
293            chain: BucketChain::new(ChainBucket::new(chain_proportionality, accounting.clone())),
294            pending_low: Vec::new(),
295            emitted: Chain::new(),
296            stage: Stage::new(accounting.clone(), chunk_size),
297            boundary: Antichain::from_elem(Timestamp::MIN),
298            since: Antichain::from_elem(Timestamp::MIN),
299            prev_update_count: 0,
300            prev_size: Default::default(),
301            metrics,
302            worker_metrics,
303            accounting,
304        }
305    }
306
307    /// Insert a batch of updates.
308    pub fn insert(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) {
309        let Some(since_ts) = self.since.as_option() else {
310            // If the since is the empty frontier, discard all updates.
311            updates.clear();
312            return;
313        };
314
315        for (_, time, _) in &mut *updates {
316            *time = std::cmp::max(*time, *since_ts);
317        }
318
319        self.insert_inner(updates);
320    }
321
322    /// Insert a batch of updates, after negating their diffs.
323    pub fn insert_negated(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) {
324        let Some(since_ts) = self.since.as_option() else {
325            // If the since is the empty frontier, discard all updates.
326            updates.clear();
327            return;
328        };
329
330        for (_, time, diff) in &mut *updates {
331            *time = std::cmp::max(*time, *since_ts);
332            *diff = -*diff;
333        }
334
335        self.insert_inner(updates);
336    }
337
338    /// Insert a batch of updates into the stage, flushing it when full.
339    ///
340    /// All times are expected to be >= the `since`.
341    fn insert_inner(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) {
342        debug_assert!(updates.iter().all(|(_, t, _)| self.since.less_equal(t)));
343
344        if let Some(mut ready) = self.stage.insert(updates) {
345            self.route(&mut ready);
346        }
347
348        self.update_metrics();
349    }
350
351    /// Route a batch of sorted, consolidated updates to `pending_low` or their chain buckets.
352    fn route(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) {
353        // Updates at times below the boundary become a pending low chain.
354        let idx = updates.partition_point(|(_, t, _)| !self.boundary.less_equal(t));
355        if idx > 0 {
356            let mut builder = ChainBuilder::default();
357            builder.extend(updates.drain(..idx));
358            let chain = builder.finish();
359            if !chain.is_empty() {
360                self.account_chain_created(&chain);
361                self.pending_low.push(chain);
362            }
363        }
364
365        // Updates at times at or beyond the boundary go into their chain buckets. Walk ranges of
366        // times that fall into the same bucket, to push batches of updates at once.
367        let mut drain = updates.drain(..).peekable();
368        while let Some(update) = drain.next() {
369            let time = update.1;
370            let range = self
371                .chain
372                .range_of(&time)
373                .expect("bucket chain covers all times at or beyond the boundary");
374            let mut builder = ChainBuilder::default();
375            builder.extend(std::iter::once(update));
376            while let Some(update) = drain.next_if(|(_, t, _)| range.contains(t)) {
377                builder.extend(std::iter::once(update));
378            }
379            let bucket = self
380                .chain
381                .find_mut(&range.start)
382                .expect("bucket chain covers all times at or beyond the boundary");
383            bucket.push_chain(builder.finish());
384        }
385    }
386
387    /// Return consolidated updates before the given `upper`.
388    pub fn updates_before<'a>(
389        &'a mut self,
390        upper: &Antichain<Timestamp>,
391    ) -> impl Iterator<Item = (D, Timestamp, Diff)> + Send + 'a {
392        self.consolidate_before(upper);
393        self.consolidated_updates_before(upper)
394    }
395
396    /// Return the updates before the given `upper`, as consolidated by a preceding
397    /// [`CorrectionV2::consolidate_before`] call.
398    ///
399    /// The caller must have invoked `consolidate_before` with the same `upper` and must not have
400    /// mutated the buffer since. Otherwise the returned updates are neither consolidated nor
401    /// necessarily complete.
402    pub fn consolidated_updates_before<'a>(
403        &'a self,
404        upper: &Antichain<Timestamp>,
405    ) -> impl Iterator<Item = (D, Timestamp, Diff)> + Send + use<'a, D> {
406        // All contained times are advanced to at least the `since`, so a read at an `upper` that
407        // is not beyond the `since` is always empty. This mirrors the short-circuit in
408        // `consolidate_before`, which leaves `emitted` untouched in that case.
409        if !PartialOrder::less_than(&self.since, upper) {
410            return None.into_iter().flatten();
411        }
412
413        // After `consolidate_before`, `emitted` holds exactly the updates before `upper`: every
414        // path that populates it splits at `upper` (pushing the remainder to `pending_low`), and
415        // the guard above guarantees `upper > since`, so advancing stale times to the `since`
416        // cannot lift them to or beyond `upper`. We can therefore yield all of `emitted`. Guard
417        // the invariant: a violation would write updates beyond the batch upper to persist.
418        soft_assert_or_log!(
419            self.emitted
420                .chunks
421                .last()
422                .is_none_or(|c| !upper.less_equal(&c.last_time())),
423            "emitted contains times at or beyond the upper",
424        );
425        Some(self.emitted.iter()).into_iter().flatten()
426    }
427
428    /// Consolidate all updates before the given `upper` into the `emitted` chain.
429    ///
430    /// Once this method returns, `emitted` contains all updates at times before `upper`,
431    /// consolidated.
432    ///
433    /// Does nothing if `upper` is not beyond the `since`: all contained times are advanced to at
434    /// least the `since`, so such a read is empty anyway, and skipping avoids an eager peel,
435    /// merge, and `boundary` advancement. Normal reads and `consolidate_at_since` always pass an
436    /// `upper` beyond the `since`.
437    pub fn consolidate_before(&mut self, upper: &Antichain<Timestamp>) {
438        if !PartialOrder::less_than(&self.since, upper) {
439            return;
440        }
441
442        if let Some(mut ready) = self.stage.flush() {
443            self.route(&mut ready);
444        }
445
446        let Some(&since_ts) = self.since.as_option() else {
447            // If the since is the empty frontier, discard all updates.
448            let peeled = self.chain.peel(Antichain::new().borrow());
449            for bucket in peeled {
450                for chain in bucket.into_chains() {
451                    self.account_chain_dropped(&chain);
452                }
453            }
454            for chain in std::mem::take(&mut self.pending_low) {
455                self.account_chain_dropped(&chain);
456            }
457            let emitted = std::mem::replace(&mut self.emitted, Chain::new());
458            if !emitted.is_empty() {
459                self.account_chain_dropped(&emitted);
460            }
461            self.update_metrics();
462            return;
463        };
464
465        // Peel the buckets below the upper off the bucket chain. Bucket splits during the peel
466        // only touch chunks around the upper; chunks wholly on either side are reused.
467        let peeled = self.chain.peel(upper.borrow());
468        if PartialOrder::less_than(&self.boundary, upper) {
469            self.boundary = upper.clone();
470        }
471
472        // Collect candidate chains: peeled bucket contents, pending low chains, and the previous
473        // emitted chain. All contain only times below the boundary.
474        let emitted = std::mem::replace(&mut self.emitted, Chain::new());
475        let mut candidates: Vec<Chain<D>> = Vec::new();
476        for bucket in peeled {
477            candidates.extend(bucket.into_chains());
478        }
479        candidates.append(&mut self.pending_low);
480        if !emitted.is_empty() {
481            candidates.push(emitted);
482        }
483
484        if candidates.is_empty() {
485            self.restore_chain();
486            self.update_metrics();
487            return;
488        }
489
490        candidates
491            .iter()
492            .for_each(|c| self.account_chain_dropped(c));
493
494        // Split the candidates at the upper. Parts at or beyond the upper (possible when `upper`
495        // regresses below a previous one) stay pending.
496        let mut lowers = Vec::new();
497        for chain in candidates {
498            match upper.as_option() {
499                Some(&upper_ts) => {
500                    let (lower, remainder) = chain.split_at_time(upper_ts);
501                    if !lower.is_empty() {
502                        lowers.push(lower);
503                    }
504                    if !remainder.is_empty() {
505                        self.account_chain_created(&remainder);
506                        self.pending_low.push(remainder);
507                    }
508                }
509                // The empty upper is greater than all times.
510                None => lowers.push(chain),
511            }
512        }
513
514        // Merge the lower parts into the new emitted chain, advancing times below the since.
515        // Advancing times in a (time, data)-sorted chain can break its sort order, so chains
516        // containing stale times cannot be merged as they are. Stale times are expected in steady
517        // state: the previous emitted chain was written before the since advanced past it.
518        //
519        // Count the distinct stale times, up to a small cap. For few distinct stale times -- the
520        // steady state -- split cursors into runs that remain sorted under advancement and merge
521        // those. For many distinct stale times -- e.g. a since jump across many buffered
522        // timestamps when a sink restarts with an old as-of, or when a hydrating sink's first batch
523        // lands -- the number of runs and the cost of cloning cursor state per run grow with the
524        // number of distinct times, so materialize, advance, and sort in bounded runs instead.
525        const MAX_STALE_RUNS: usize = 32;
526        let mut stale_times = 0;
527        for chain in &lowers {
528            stale_times += chain.distinct_times_before(since_ts, MAX_STALE_RUNS - stale_times);
529            if stale_times >= MAX_STALE_RUNS {
530                break;
531            }
532        }
533
534        // The merged chain becomes `emitted`, which the caller reads next and the next
535        // consolidation merges with the feedback retractions, so the chunks the merge writes are
536        // the youngest generation whatever the depth of their inputs. A lone input chain is
537        // reused as is and keeps its chunks' depths, since re-spilling them would cost a copy.
538        let merged = if stale_times == 0 {
539            let cursors: Vec<_> = lowers.into_iter().filter_map(Chain::into_cursor).collect();
540            merge_cursors(cursors, 0)
541        } else if stale_times < MAX_STALE_RUNS {
542            let mut runs = Vec::new();
543            for chain in lowers {
544                if let Some(cursor) = chain.into_cursor() {
545                    runs.append(&mut cursor.advance_by(since_ts));
546                }
547            }
548            merge_cursors(runs, 0)
549        } else {
550            let chain = sort_advanced(lowers, since_ts, STALE_SORT_RUN);
551
552            // Advancement can move updates to or beyond the upper; such updates stay pending.
553            match upper.as_option() {
554                Some(&upper_ts) => {
555                    let (lower, remainder) = chain.split_at_time(upper_ts);
556                    if !remainder.is_empty() {
557                        self.account_chain_created(&remainder);
558                        self.pending_low.push(remainder);
559                    }
560                    lower
561                }
562                None => chain,
563            }
564        };
565
566        if !merged.is_empty() {
567            self.account_chain_created(&merged);
568        }
569        self.emitted = merged;
570
571        self.restore_chain();
572        self.update_metrics();
573    }
574
575    /// Perform a bounded amount of work towards restoring the bucket chain invariant.
576    ///
577    /// Restoration is allowed to remain incomplete: the bucket chain supports peeling and finding
578    /// on ill-formed chains, so any leftover work is picked up by subsequent operations. The fuel
579    /// bound keeps individual buffer operations from stalling the operator that owns the buffer.
580    fn restore_chain(&mut self) {
581        let mut fuel = RESTORE_FUEL;
582        self.chain.restore(&mut fuel);
583    }
584
585    /// Advance the since frontier.
586    ///
587    /// Time advancement of updates in the bucket chain is lazy: it happens when the updates are
588    /// consolidated by a read.
589    ///
590    /// # Panics
591    ///
592    /// Panics if the given `since` is less than the current since frontier.
593    pub fn advance_since(&mut self, since: Antichain<Timestamp>) {
594        assert!(PartialOrder::less_equal(&self.since, &since));
595        self.stage.advance_times(&since);
596        self.since = since;
597    }
598
599    /// Consolidate all updates at the current `since`.
600    pub fn consolidate_at_since(&mut self) {
601        let upper_ts = self.since.as_option().and_then(|t| t.try_step_forward());
602        if let Some(upper_ts) = upper_ts {
603            let upper = Antichain::from_elem(upper_ts);
604            self.consolidate_before(&upper);
605        }
606    }
607
608    fn account_chain_created(&self, chain: &Chain<D>) {
609        self.accounting.chain_created(chain);
610    }
611
612    fn account_chain_dropped(&self, chain: &Chain<D>) {
613        self.accounting.chain_dropped(chain);
614    }
615
616    /// Update persist sink metrics.
617    ///
618    /// Reads the running totals maintained by [`Accounting`], so its cost is independent of how
619    /// much the buffer holds. Nothing here walks chains or pages a chunk in.
620    fn update_metrics(&mut self) {
621        let (new_length, new_size) = self.accounting.totals();
622        self.update_metrics_inner(new_size, new_length);
623    }
624
625    /// Update persist sink metrics to the given new size and length.
626    fn update_metrics_inner(&mut self, new_size: SizeMetrics, new_length: usize) {
627        let old_size = self.prev_size;
628        let old_length = self.prev_update_count;
629        let len_delta = UpdateDelta::new(new_length, old_length);
630        let cap_delta = UpdateDelta::new(new_size.capacity, old_size.capacity);
631        self.metrics
632            .report_correction_update_deltas(len_delta, cap_delta);
633        self.worker_metrics
634            .report_correction_update_totals(new_length, new_size.capacity);
635
636        self.accounting.report_size_metrics(new_size, old_size);
637
638        self.prev_size = new_size;
639        self.prev_update_count = new_length;
640    }
641}
642
643/// The number of updates [`sort_advanced`] sorts on the heap at a time.
644///
645/// The heap a since jump costs per worker is one run of owned updates, the chunks the merges are
646/// building, and the input chunk being read. For rows of TPC-H lineitem's width, about 190 bytes
647/// owned, a run takes about 12 MiB, and a benchmark of 8 Mi such updates peaked at about 14 MiB.
648/// Shorter runs barely lower that peak, because the merges' chunks remain, while each halving adds
649/// a merge level. Consolidation costs per update, so the bound is a count rather than bytes.
650const STALE_SORT_RUN: usize = 1 << 16;
651
652/// Advance the times of the updates in `chains` to at least `since_ts` and return them as one
653/// sorted and consolidated chain.
654///
655/// Sorts runs of at most `run_len` updates on the heap and mints each as a chain before merging
656/// them, so the heap holds one run of owned updates rather than all of them. Each input chunk is
657/// dropped once read.
658fn sort_advanced<D: Data>(chains: Vec<Chain<D>>, since_ts: Timestamp, run_len: usize) -> Chain<D> {
659    let mut runs = Vec::new();
660    let mut run = Vec::with_capacity(run_len);
661    for chain in chains {
662        for chunk in chain.chunks {
663            chunk.with_view(|view| {
664                for i in 0..chunk.len() {
665                    let (data, time, diff) = view.get(i);
666                    run.push((D::into_owned(data), std::cmp::max(time, since_ts), diff));
667                    if run.len() >= run_len {
668                        consolidate(&mut run);
669                        // A run that consolidation left at most half full keeps accumulating, so
670                        // updates that cancel within the heap cost no chunks, and each
671                        // consolidation either halves the run or mints it.
672                        if run.len() * 2 > run_len {
673                            push_run(&mut runs, mint_run(&mut run));
674                        }
675                    }
676                }
677            });
678        }
679    }
680    consolidate(&mut run);
681    push_run(&mut runs, mint_run(&mut run));
682
683    let cursors = runs.into_iter().filter_map(Chain::into_cursor).collect();
684    merge_cursors(cursors, 0)
685}
686
687/// Mint the sorted and consolidated `run` as a chain, emptying it.
688fn mint_run<D: Data>(run: &mut Vec<(D, Timestamp, Diff)>) -> Chain<D> {
689    let mut builder = ChainBuilder::default();
690    builder.extend(run.drain(..));
691    builder.finish()
692}
693
694/// Push `run` onto the stack of `runs`, merging it into its predecessors while it holds at least
695/// half as many updates as they do.
696///
697/// The stack then holds runs of geometrically decreasing size, a number logarithmic in the
698/// updates, so the final merge stays narrow. Updates that cancel do so at the first merge that
699/// brings them together rather than in the final one. A merged run is one generation deeper than
700/// its inputs, so the pool can compress the runs that wait longest for the final merge.
701fn push_run<D: Data>(runs: &mut Vec<Chain<D>>, mut run: Chain<D>) {
702    while let Some(prev) = runs.last() {
703        if run.update_count * 2 < prev.update_count {
704            break;
705        }
706        let prev = runs.pop().expect("checked above");
707        let depth = prev.depth.max(run.depth).saturating_add(1);
708        let cursors = [prev, run]
709            .into_iter()
710            .filter_map(Chain::into_cursor)
711            .collect();
712        run = merge_cursors(cursors, depth);
713    }
714    if !run.is_empty() {
715        runs.push(run);
716    }
717}
718
719/// Merge the given cursors into one chain, writing new chunks at `depth`.
720fn merge_cursors<D: Data>(cursors: Vec<Cursor<D>>, depth: u8) -> Chain<D> {
721    match cursors.len() {
722        0 => Chain::new(),
723        1 => {
724            let [cur] = cursors.try_into().unwrap();
725            cur.into_chain(depth)
726        }
727        2 => {
728            let [a, b] = cursors.try_into().unwrap();
729            merge_2(a, b, depth)
730        }
731        _ => merge_many(cursors, depth),
732    }
733}
734
735/// Merge the given two cursors using a 2-way merge.
736///
737/// This function is a specialization of `merge_many` that avoids the overhead of a binary heap.
738fn merge_2<D: Data>(cursor1: Cursor<D>, cursor2: Cursor<D>, depth: u8) -> Chain<D> {
739    let mut rest1 = Some(cursor1);
740    let mut rest2 = Some(cursor2);
741    let mut merged = ChainBuilder::at_depth(depth);
742
743    // One borrow per chunk pair, not per update: `Chunk::view` re-decodes the column header on
744    // every call. The inner loop runs until either cursor crosses into its next chunk, at which
745    // point the outer loop re-borrows both.
746    while rest1.is_some() && rest2.is_some() {
747        let chunk1 = rest1.as_ref().expect("checked above").chunk_handle();
748        let chunk2 = rest2.as_ref().expect("checked above").chunk_handle();
749        let view1 = chunk1.view();
750        let view2 = chunk2.view();
751
752        loop {
753            let (Some(c1), Some(c2)) = (rest1.as_ref(), rest2.as_ref()) else {
754                break;
755            };
756            if !c1.reads_from(&chunk1) || !c2.reads_from(&chunk2) {
757                break;
758            }
759
760            let (d1, t1, r1) = c1.get_with(&view1);
761            let (d2, t2, r2) = c2.get_with(&view2);
762
763            match refs_cmp::<D>((t1, d1), (t2, d2)) {
764                Ordering::Less => {
765                    merged.push_ref((d1, t1, r1));
766                    rest1 = rest1.take().expect("checked above").step();
767                }
768                Ordering::Greater => {
769                    merged.push_ref((d2, t2, r2));
770                    rest2 = rest2.take().expect("checked above").step();
771                }
772                Ordering::Equal => {
773                    let r = r1 + r2;
774                    if r != Diff::ZERO {
775                        merged.push_ref((d1, t1, r));
776                    }
777                    rest1 = rest1.take().expect("checked above").step();
778                    rest2 = rest2.take().expect("checked above").step();
779                }
780            }
781        }
782    }
783
784    match (rest1, rest2) {
785        (Some(c), None) | (None, Some(c)) => merged.push_cursor(c),
786        (Some(_), Some(_)) => unreachable!("loop runs while both cursors are live"),
787        (None, None) => (),
788    }
789
790    merged.finish()
791}
792
793/// Merge the given cursors using a k-way merge with a binary heap.
794fn merge_many<D: Data>(cursors: Vec<Cursor<D>>, depth: u8) -> Chain<D> {
795    let mut cursors: Vec<Option<Cursor<D>>> = cursors.into_iter().map(Some).collect();
796    let mut merged = ChainBuilder::at_depth(depth);
797
798    // One borrow per chunk, not per update, as in `merge_2`. Each round borrows the current chunk
799    // of every live cursor, and merges until a cursor crosses into its next chunk, whose updates
800    // the round's heap cannot hold. The next round re-borrows.
801    while cursors.iter().any(Option::is_some) {
802        let chunks: Vec<Option<Rc<Chunk<D>>>> = cursors
803            .iter()
804            .map(|c| c.as_ref().map(Cursor::chunk_handle))
805            .collect();
806        let views: Vec<Option<ChunkView<'_, D>>> = chunks
807            .iter()
808            .map(|c| c.as_deref().map(Chunk::view))
809            .collect();
810
811        // Keyed by `(time, data)`, the merge order, with the cursor index as a tie breaker.
812        let mut heap = BinaryHeap::new();
813        for (i, (cursor, view)) in cursors.iter().zip_eq(&views).enumerate() {
814            if let (Some(cursor), Some(view)) = (cursor, view) {
815                let (d, t, _) = cursor.get_with(view);
816                heap.push(Reverse((t, d, i)));
817            }
818        }
819
820        while let Some(Reverse((time, data, first))) = heap.pop() {
821            let mut diff = Diff::ZERO;
822            let mut crossed = false;
823            let mut next = Some(first);
824            while let Some(i) = next {
825                let view = views[i].as_ref().expect("live cursors have a view");
826                let cursor = cursors[i].take().expect("heap entries are live cursors");
827                diff += cursor.get_with(view).2;
828                cursors[i] = cursor.step();
829                if let Some(cursor) = &cursors[i] {
830                    let chunk = chunks[i].as_ref().expect("live cursors have a chunk");
831                    if cursor.reads_from(chunk) {
832                        let (d, t, _) = cursor.get_with(view);
833                        heap.push(Reverse((t, d, i)));
834                    } else {
835                        crossed = true;
836                    }
837                }
838
839                // Every cursor at the same update contributes to it before it is pushed, so a
840                // round only ends between distinct updates.
841                next = match heap.peek() {
842                    Some(Reverse((t, d, _))) if *t == time && *d == data => {
843                        heap.pop().map(|Reverse((_, _, j))| j)
844                    }
845                    _ => None,
846                };
847            }
848
849            if diff != Diff::ZERO {
850                merged.push_ref((data, time, diff));
851            }
852            if crossed {
853                break;
854            }
855        }
856    }
857
858    merged.finish()
859}
860
861impl<D: Data> Drop for CorrectionV2<D> {
862    fn drop(&mut self) {
863        for bucket in self.chain.buckets() {
864            bucket
865                .chains
866                .iter()
867                .for_each(|c| self.account_chain_dropped(c));
868        }
869        self.pending_low
870            .iter()
871            .for_each(|c| self.account_chain_dropped(c));
872        if !self.emitted.is_empty() {
873            self.account_chain_dropped(&self.emitted);
874        }
875        self.update_metrics_inner(Default::default(), 0);
876    }
877}
878
879/// Running totals over the buffer's contents, plus the introspection logging that reports them.
880///
881/// Shared by the buffer and by every structure that holds its chains, so the totals are
882/// maintained where chains are minted and retired rather than by a walk over the buffer.
883///
884/// Correctness rests on a discipline the code already keeps for the introspection logging: a
885/// chain the buffer holds has been announced exactly once through [`Accounting::chain_created`],
886/// and is retired exactly once, through [`Accounting::chain_dropped`], before it is dropped or
887/// consumed. Chains are immutable once announced, so the totals a chain contributes at
888/// retirement are the ones it contributed at announcement. `metrics_totals_match_walk` checks
889/// the result against a walk.
890#[derive(Clone, Debug)]
891struct Accounting {
892    /// The totals, shared with every clone.
893    totals: Arc<Totals>,
894    /// Introspection logging, absent when the sink is not logging.
895    logging: Option<ChannelLogging>,
896}
897
898/// The running totals behind an [`Accounting`].
899///
900/// Atomics because the buffer must stay `Send`, not because the counters are contended: the sink
901/// touches its buffer from one thread at a time, which is why `Relaxed` suffices.
902#[derive(Debug, Default)]
903struct Totals {
904    /// Number of updates the buffer holds.
905    records: AtomicUsize,
906    /// Serialized size of what the buffer holds, in bytes.
907    size: AtomicUsize,
908    /// Number of allocations holding it: one per chunk, plus the staging vector while it has one.
909    allocations: AtomicUsize,
910}
911
912impl Accounting {
913    /// Construct accounting that reports to the given logging, if any.
914    fn new(logging: Option<ChannelLogging>) -> Self {
915        Self {
916            totals: Default::default(),
917            logging,
918        }
919    }
920
921    /// Return the current totals as `(records, size metrics)`.
922    ///
923    /// The reported capacity equals the size: chunk bodies are exactly sized when minted, and the
924    /// staging vector's slack is bounded by one chunk's worth of updates and not tracked.
925    fn totals(&self) -> (usize, SizeMetrics) {
926        let size = self.totals.size.load(atomic::Ordering::Relaxed);
927        let metrics = SizeMetrics {
928            size,
929            capacity: size,
930            allocations: self.totals.allocations.load(atomic::Ordering::Relaxed),
931        };
932        (self.totals.records.load(atomic::Ordering::Relaxed), metrics)
933    }
934
935    /// Account for a chain the buffer now holds.
936    fn chain_created<D: Data>(&self, chain: &Chain<D>) {
937        let relaxed = atomic::Ordering::Relaxed;
938        self.totals.records.fetch_add(chain.update_count, relaxed);
939        self.totals.size.fetch_add(chain.size, relaxed);
940        self.totals
941            .allocations
942            .fetch_add(chain.chunks.len(), relaxed);
943        if let Some(logging) = &self.logging {
944            logging.chain_created(chain.update_count);
945        }
946    }
947
948    /// Account for a chain the buffer no longer holds.
949    fn chain_dropped<D: Data>(&self, chain: &Chain<D>) {
950        sub_checked(&self.totals.records, chain.update_count);
951        sub_checked(&self.totals.size, chain.size);
952        sub_checked(&self.totals.allocations, chain.chunks.len());
953        if let Some(logging) = &self.logging {
954            logging.chain_dropped(chain.update_count);
955        }
956    }
957
958    /// Announce the staging area as an empty chain, which is how its population is reported.
959    fn stage_created(&self) {
960        if let Some(logging) = &self.logging {
961            logging.chain_created(0);
962        }
963    }
964
965    /// Retire the staging area, which holds `records` updates, `bytes` bytes, and `allocations`
966    /// allocations.
967    ///
968    /// Balances [`Accounting::stage_created`]: the chain the stage is reported as is dropped at
969    /// its current length.
970    fn stage_dropped(&self, records: usize, bytes: usize, allocations: usize) {
971        sub_checked(&self.totals.records, records);
972        sub_checked(&self.totals.size, bytes);
973        sub_checked(&self.totals.allocations, allocations);
974        if let Some(logging) = &self.logging {
975            logging.chain_dropped(records);
976        }
977    }
978
979    /// Account for the staging area changing by `records` updates, `bytes` bytes, and
980    /// `allocations` allocations.
981    fn stage_diff(&self, records: isize, bytes: isize, allocations: isize) {
982        add_signed(&self.totals.records, records);
983        add_signed(&self.totals.size, bytes);
984        add_signed(&self.totals.allocations, allocations);
985
986        // The stage is reported as a chain that is dropped and re-created at its new length.
987        let Some(logging) = &self.logging else { return };
988        if records > 0 {
989            logging.chain_created(usize::try_from(records).expect("positive"));
990            logging.chain_dropped(0);
991        } else if records < 0 {
992            logging.chain_created(0);
993            logging.chain_dropped(usize::try_from(-records).expect("positive"));
994        }
995    }
996
997    /// Report the change from `old` to `new` size metrics to introspection.
998    fn report_size_metrics(&self, new: SizeMetrics, old: SizeMetrics) {
999        let Some(logging) = &self.logging else { return };
1000        let i = |x: usize| isize::try_from(x).expect("must fit");
1001        logging.report_size_diff(i(new.size) - i(old.size));
1002        logging.report_capacity_diff(i(new.capacity) - i(old.capacity));
1003        logging.report_allocations_diff(i(new.allocations) - i(old.allocations));
1004    }
1005}
1006
1007/// Add a signed `diff` to `counter`.
1008///
1009/// # Panics
1010///
1011/// Panics if the result would be negative, which means a retirement was never matched by an
1012/// announcement.
1013fn add_signed(counter: &AtomicUsize, diff: isize) {
1014    let relaxed = atomic::Ordering::Relaxed;
1015    if diff >= 0 {
1016        counter.fetch_add(usize::try_from(diff).expect("non-negative"), relaxed);
1017    } else {
1018        sub_checked(counter, usize::try_from(-diff).expect("positive"));
1019    }
1020}
1021
1022/// Subtract `sub` from `counter`.
1023///
1024/// # Panics
1025///
1026/// Panics if the result would be negative, which means a retirement was never matched by an
1027/// announcement.
1028fn sub_checked(counter: &AtomicUsize, sub: usize) {
1029    let prev = counter.fetch_sub(sub, atomic::Ordering::Relaxed);
1030    assert!(prev >= sub, "total retired below zero");
1031}
1032
1033/// A bucket of `Chain`s, for use in a [`BucketChain`].
1034///
1035/// All chains are individually sorted by (time, data) and consolidated, but updates can appear in
1036/// multiple chains, so consumers must merge the chains to obtain consolidated updates.
1037struct ChainBucket<D: Data> {
1038    /// The contained chains.
1039    ///
1040    /// Maintained with the chain invariant on pushes; splits can leave it violated until the next
1041    /// push restores it.
1042    chains: Vec<Chain<D>>,
1043    /// The size factor of subsequent chains required by the chain invariant.
1044    chain_proportionality: f64,
1045    /// Running totals and introspection logging.
1046    accounting: Accounting,
1047}
1048
1049impl<D: Data> fmt::Debug for ChainBucket<D> {
1050    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1051        f.debug_struct("ChainBucket")
1052            .field("chains", &self.chains)
1053            .finish_non_exhaustive()
1054    }
1055}
1056
1057impl<D: Data> ChainBucket<D> {
1058    /// Construct a new, empty `ChainBucket`.
1059    fn new(chain_proportionality: f64, accounting: Accounting) -> Self {
1060        Self {
1061            chains: Vec::new(),
1062            chain_proportionality,
1063            accounting,
1064        }
1065    }
1066
1067    /// Push a chain onto the bucket, restoring the chain invariant.
1068    fn push_chain(&mut self, chain: Chain<D>) {
1069        if chain.is_empty() {
1070            return;
1071        }
1072        self.accounting.chain_created(&chain);
1073        self.chains.push(chain);
1074
1075        // Restore the chain invariant.
1076        let prop = self.chain_proportionality;
1077        let merge_needed = |chains: &[Chain<_>]| match chains {
1078            [.., prev, last] => {
1079                let last_len = f64::cast_lossy(last.update_count);
1080                let prev_len = f64::cast_lossy(prev.update_count);
1081                last_len * prop > prev_len
1082            }
1083            _ => false,
1084        };
1085
1086        while merge_needed(&self.chains) {
1087            let a = self.chains.pop().unwrap();
1088            let b = self.chains.pop().unwrap();
1089            self.accounting.chain_dropped(&a);
1090            self.accounting.chain_dropped(&b);
1091
1092            // A merged chain waits about `chain_proportionality` times longer for its next merge
1093            // than its inputs did, so it is one generation deeper.
1094            let depth = a.depth.max(b.depth).saturating_add(1);
1095            let cursors = [a, b].into_iter().filter_map(Chain::into_cursor).collect();
1096            let merged = merge_cursors(cursors, depth);
1097            if !merged.is_empty() {
1098                self.accounting.chain_created(&merged);
1099                self.chains.push(merged);
1100            }
1101        }
1102    }
1103
1104    /// Convert the bucket into its contained chains.
1105    fn into_chains(self) -> Vec<Chain<D>> {
1106        self.chains
1107    }
1108}
1109
1110impl<D: Data> Bucket for ChainBucket<D> {
1111    type Timestamp = Timestamp;
1112
1113    fn split(self, timestamp: &Self::Timestamp, fuel: &mut i64) -> (Self, Self) {
1114        let mut lower = Self::new(self.chain_proportionality, self.accounting.clone());
1115        let mut upper = Self::new(self.chain_proportionality, self.accounting.clone());
1116
1117        for chain in self.chains {
1118            // Whole chunks are reused; at most one chunk straddling the timestamp is copied per
1119            // chain. Account fuel at chunk granularity.
1120            *fuel = fuel.saturating_sub(i64::try_from(chain.chunks.len()).expect("must fit"));
1121
1122            self.accounting.chain_dropped(&chain);
1123            let (lo, hi) = chain.split_at_time(*timestamp);
1124            for (part, target) in [(lo, &mut lower), (hi, &mut upper)] {
1125                if !part.is_empty() {
1126                    target.accounting.chain_created(&part);
1127                    target.chains.push(part);
1128                }
1129            }
1130        }
1131
1132        (lower, upper)
1133    }
1134}
1135
1136/// A chain of [`Chunk`]s containing updates.
1137///
1138/// All updates in a chain are sorted by (time, data) and consolidated.
1139///
1140/// Note that, in contrast to [`Chunk`]s, chains can be empty. Though we generally try to avoid
1141/// keeping around empty chains.
1142#[derive(Debug)]
1143struct Chain<D: Data> {
1144    /// The contained chunks.
1145    chunks: Vec<Chunk<D>>,
1146    /// The number of updates contained in all chunks.
1147    update_count: usize,
1148    /// The serialized size of all chunks, in bytes.
1149    ///
1150    /// Maintained as chunks are pushed, so metrics never walk the chunks.
1151    size: usize,
1152    /// The deepest generation among the chunks, see [`Chunk::depth`].
1153    depth: u8,
1154}
1155
1156impl<D: Data> Chain<D> {
1157    /// Construct an empty chain.
1158    fn new() -> Self {
1159        Self {
1160            chunks: Default::default(),
1161            update_count: 0,
1162            size: 0,
1163            depth: 0,
1164        }
1165    }
1166
1167    /// Return whether the chain is empty.
1168    fn is_empty(&self) -> bool {
1169        self.chunks.is_empty()
1170    }
1171
1172    /// Push a chunk onto the chain.
1173    ///
1174    /// All updates in the chunk must sort after all updates already in the chain, in
1175    /// (time, data)-order, to ensure the chain remains sorted.
1176    fn push_chunk(&mut self, chunk: Chunk<D>) {
1177        mz_ore::soft_assert_no_log!(self.can_accept_chunk(&chunk));
1178
1179        self.update_count += chunk.len();
1180        self.size += chunk.size();
1181        self.depth = self.depth.max(chunk.depth);
1182        self.chunks.push(chunk);
1183    }
1184
1185    /// Return whether the chain can accept the given chunk at its end while preserving
1186    /// (time, data)-order.
1187    ///
1188    /// NOTE: The cached boundary times settle every case but a tie. On a tie the boundary updates
1189    /// themselves are compared, which reads both chunks. The reads are scoped, so neither chunk
1190    /// stays materialized, but a spilled chunk costs a copy out of the pool. The only caller is the
1191    /// soft assertion in [`Chain::push_chunk`], and soft assertions are live in any build started
1192    /// with `MZ_SOFT_ASSERTIONS` set, so this cost is not confined to debug builds. Ties are
1193    /// reached whenever a run of updates at a single timestamp spans a chunk boundary, which
1194    /// [`ChunkBuilder`] produces for any such run larger than its byte limit.
1195    fn can_accept_chunk(&self, chunk: &Chunk<D>) -> bool {
1196        match self.chunks.last() {
1197            None => true,
1198            Some(last) => match last.last_time().cmp(&chunk.first_time()) {
1199                Ordering::Less => true,
1200                Ordering::Greater => false,
1201                Ordering::Equal => last.with_view(|last_view| {
1202                    chunk.with_view(|view| {
1203                        let time = chunk.first_time();
1204                        let (dc, _, _) = last_view.get(last.len() - 1);
1205                        let (d, _, _) = view.get(0);
1206                        refs_cmp::<D>((time, dc), (time, d)).is_lt()
1207                    })
1208                }),
1209            },
1210        }
1211    }
1212
1213    /// Convert the chain into a cursor over the contained updates.
1214    fn into_cursor(self) -> Option<Cursor<D>> {
1215        let chunks = self.chunks.into_iter().map(Rc::new).collect();
1216        Cursor::new(chunks)
1217    }
1218
1219    /// Return an iterator over the contained updates.
1220    ///
1221    /// Reads each chunk scoped, so a spilled chunk stays spillable after the read. This is the
1222    /// read behind [`CorrectionV2::updates_before`], after which `emitted` rests until the next
1223    /// consolidation merges it: a batch as large as the snapshot after hydration would otherwise
1224    /// sit on the heap for that whole time. Each chunk's updates are made owned in one pass, so
1225    /// the iterator holds at most one chunk's worth of owned updates.
1226    fn iter(&self) -> impl Iterator<Item = (D, Timestamp, Diff)> + '_ {
1227        self.chunks.iter().flat_map(|c| {
1228            c.with_view(|view| {
1229                (0..c.len())
1230                    .map(|i| {
1231                        let (d, t, r) = view.get(i);
1232                        (D::into_owned(d), t, r)
1233                    })
1234                    .collect::<Vec<_>>()
1235            })
1236        })
1237    }
1238
1239    /// Count the distinct times of updates at times before `time`, up to the given cap.
1240    ///
1241    /// The scan uses one binary search per distinct time, so its cost is bounded by
1242    /// O(cap log chunks).
1243    fn distinct_times_before(&self, time: Timestamp, cap: usize) -> usize {
1244        let mut count = 0;
1245        let mut chunk_idx = 0;
1246        let mut offset = 0;
1247        while count < cap && chunk_idx < self.chunks.len() {
1248            let chunk = &self.chunks[chunk_idx];
1249            let current = chunk.index(offset).1;
1250            if current >= time {
1251                break;
1252            }
1253            count += 1;
1254            // Skip to the first update at a time greater than `current`.
1255            match chunk.find_time_greater_than(current) {
1256                Some(idx) => offset = idx,
1257                None => {
1258                    // All later updates at `current` are in subsequent chunks.
1259                    chunk_idx += 1;
1260                    offset = 0;
1261                    while chunk_idx < self.chunks.len() {
1262                        match self.chunks[chunk_idx].find_time_greater_than(current) {
1263                            Some(idx) => {
1264                                offset = idx;
1265                                break;
1266                            }
1267                            None => chunk_idx += 1,
1268                        }
1269                    }
1270                }
1271            }
1272        }
1273        count
1274    }
1275
1276    /// Split the chain at the given time.
1277    ///
1278    /// Returns two chains, the first containing all updates at times < `time`, the second
1279    /// containing all updates at times >= `time`. Chunks fully on either side of `time` are
1280    /// reused; only a chunk straddling `time` is copied.
1281    fn split_at_time(mut self, time: Timestamp) -> (Self, Self) {
1282        let mut lower = Self::new();
1283        let mut upper = Self::new();
1284
1285        let Some(skip_ts) = time.step_back() else {
1286            // Nothing sorts before `time`.
1287            return (lower, self);
1288        };
1289
1290        for chunk in self.chunks.drain(..) {
1291            // Route whole chunks by cached boundary times, so a chunk that lands entirely on one
1292            // side is moved without paging it in. Only a straddling chunk is materialized here.
1293            // With soft assertions on, `push_chunk` can still page in a chunk whose boundary time
1294            // ties the chain's last one, see `Chain::can_accept_chunk`.
1295            if chunk.last_time() < time {
1296                lower.push_chunk(chunk);
1297            } else if chunk.first_time() >= time {
1298                upper.push_chunk(chunk);
1299            } else {
1300                // The chunk straddles `time`; copy its two halves.
1301                let idx = chunk
1302                    .find_time_greater_than(skip_ts)
1303                    .expect("straddles time");
1304                let view = chunk.view();
1305                let mut builder = ChainBuilder::at_depth(chunk.depth);
1306                for i in 0..idx {
1307                    builder.push_ref(view.get(i));
1308                }
1309                for part in builder.finish().chunks {
1310                    lower.push_chunk(part);
1311                }
1312                let mut builder = ChainBuilder::at_depth(chunk.depth);
1313                for i in idx..chunk.len() {
1314                    builder.push_ref(view.get(i));
1315                }
1316                for part in builder.finish().chunks {
1317                    upper.push_chunk(part);
1318                }
1319            }
1320        }
1321
1322        (lower, upper)
1323    }
1324}
1325
1326/// A builder that constructs a [`Chain`] from a stream of updates.
1327///
1328/// Wraps a [`ChunkBuilder`] and drains its minted chunks into a [`Chain`]. Pushed updates must
1329/// arrive in (time, data) sorted order.
1330struct ChainBuilder<D: Data> {
1331    builder: ChunkBuilder<D>,
1332    chain: Chain<D>,
1333}
1334
1335impl<D: Data> Default for ChainBuilder<D> {
1336    fn default() -> Self {
1337        Self::at_depth(0)
1338    }
1339}
1340
1341impl<D: Data> ChainBuilder<D> {
1342    /// A builder whose chunks are minted at the given generational depth.
1343    fn at_depth(depth: u8) -> Self {
1344        Self {
1345            builder: ChunkBuilder::at_depth(depth),
1346            chain: Chain::new(),
1347        }
1348    }
1349
1350    /// Push a reference-form update into the builder.
1351    fn push_ref(&mut self, update: Ref<'_, (D, Timestamp, Diff)>) {
1352        if let Some(chunk) = self.builder.push(update) {
1353            self.chain.push_chunk(chunk);
1354        }
1355    }
1356
1357    /// Push an owned-form update into the builder.
1358    fn push_owned(&mut self, update: &(D, Timestamp, Diff)) {
1359        if let Some(chunk) = self.builder.push(update) {
1360            self.chain.push_chunk(chunk);
1361        }
1362    }
1363
1364    /// Push the updates produced by a cursor into the builder.
1365    fn push_cursor(&mut self, cursor: Cursor<D>) {
1366        let mut rest = Some(cursor);
1367        // One borrow per chunk: see `Chunk::view` for why this must not move into the inner loop.
1368        while let Some(cursor) = rest.take() {
1369            let chunk = cursor.chunk_handle();
1370            let view = chunk.view();
1371            rest = Some(cursor);
1372
1373            while let Some(cursor) = rest.as_ref() {
1374                if !cursor.reads_from(&chunk) {
1375                    break;
1376                }
1377                self.push_ref(cursor.get_with(&view));
1378                rest = rest.take().expect("checked above").step();
1379            }
1380        }
1381    }
1382
1383    /// Finish building, returning the assembled [`Chain`].
1384    fn finish(self) -> Chain<D> {
1385        let Self { builder, mut chain } = self;
1386        if let Some(chunk) = builder.finish() {
1387            chain.push_chunk(chunk);
1388        }
1389        chain
1390    }
1391}
1392
1393impl<D: Data> Extend<(D, Timestamp, Diff)> for ChainBuilder<D> {
1394    fn extend<I: IntoIterator<Item = (D, Timestamp, Diff)>>(&mut self, iter: I) {
1395        for update in iter {
1396            self.push_owned(&update);
1397        }
1398    }
1399}
1400
1401/// A cursor over updates in a chain.
1402///
1403/// A cursor provides two guarantees:
1404///  * Produced updates are ordered and consolidated.
1405///  * A cursor always yields at least one update.
1406///
1407/// The second guarantee is enforced through the type system: Every method that steps a cursor
1408/// forward consumes `self` and returns an `Option<Cursor>` that's `None` if the operation stepped
1409/// over the last update.
1410///
1411/// A cursor holds on to `Rc<Chunk>`s, allowing multiple cursors to produce updates from the same
1412/// chunks concurrently. As soon as a cursor is done producing updates from a [`Chunk`] it drops
1413/// its reference. Once the last cursor is done with a [`Chunk`] its memory can be reclaimed.
1414#[derive(Clone, Debug)]
1415struct Cursor<D: Data> {
1416    /// The chunks from which updates can still be produced.
1417    chunks: VecDeque<Rc<Chunk<D>>>,
1418    /// The current offset into `chunks.front()`.
1419    chunk_offset: usize,
1420    /// An optional limit for the number of updates the cursor will produce.
1421    limit: Option<usize>,
1422    /// An optional overwrite for the timestamp of produced updates.
1423    overwrite_ts: Option<Timestamp>,
1424}
1425
1426impl<D: Data> Cursor<D> {
1427    /// Construct a cursor over a list of chunks.
1428    ///
1429    /// Returns `None` if `chunks` is empty.
1430    fn new(chunks: VecDeque<Rc<Chunk<D>>>) -> Option<Self> {
1431        if chunks.is_empty() {
1432            return None;
1433        }
1434
1435        Some(Self {
1436            chunks,
1437            chunk_offset: 0,
1438            limit: None,
1439            overwrite_ts: None,
1440        })
1441    }
1442
1443    /// Set a limit for the number of updates this cursor will produce.
1444    ///
1445    /// # Panics
1446    ///
1447    /// Panics if there is already a limit lower than the new one.
1448    fn set_limit(mut self, limit: usize) -> Option<Self> {
1449        assert!(self.limit.is_none_or(|l| l >= limit));
1450
1451        if limit == 0 {
1452            return None;
1453        }
1454
1455        // Release chunks made unreachable by the limit.
1456        let mut count = 0;
1457        let mut idx = 0;
1458        let mut offset = self.chunk_offset;
1459        while idx < self.chunks.len() && count < limit {
1460            let chunk = &self.chunks[idx];
1461            count += chunk.len() - offset;
1462            idx += 1;
1463            offset = 0;
1464        }
1465        self.chunks.truncate(idx);
1466
1467        if count > limit {
1468            self.limit = Some(limit);
1469        }
1470
1471        Some(self)
1472    }
1473
1474    /// Get a reference to the current update.
1475    ///
1476    /// Single-access only. A loop over a cursor must hoist [`Cursor::chunk_handle`]'s
1477    /// [`Chunk::view`] and read through [`Cursor::get_with`] instead.
1478    fn get(&self) -> Ref<'_, (D, Timestamp, Diff)> {
1479        let chunk = self.get_chunk();
1480        let (d, t, r) = chunk.index(self.chunk_offset);
1481        let t = self.overwrite_ts.unwrap_or(t);
1482        (d, t, r)
1483    }
1484
1485    /// Get a reference to the current update, reading through an already-borrowed view.
1486    ///
1487    /// # Panics
1488    ///
1489    /// Panics if `view` is not a view of the cursor's current chunk. Guard loops with
1490    /// [`Cursor::reads_from`], which is how a caller learns the cursor has crossed into the next
1491    /// chunk and the view must be refreshed.
1492    fn get_with<'a>(&self, view: &ChunkView<'a, D>) -> Ref<'a, (D, Timestamp, Diff)> {
1493        debug_assert_eq!(view.len(), self.get_chunk().len(), "view of another chunk");
1494        let (d, t, r) = view.get(self.chunk_offset);
1495        let t = self.overwrite_ts.unwrap_or(t);
1496        (d, t, r)
1497    }
1498
1499    /// A shared handle on the chunk the cursor currently reads from.
1500    ///
1501    /// Held by callers that hoist a [`Chunk::view`], so the view's borrow outlives the cursor
1502    /// steps taken against it.
1503    fn chunk_handle(&self) -> Rc<Chunk<D>> {
1504        Rc::clone(&self.chunks[0])
1505    }
1506
1507    /// Whether the cursor still reads from `chunk`.
1508    fn reads_from(&self, chunk: &Rc<Chunk<D>>) -> bool {
1509        Rc::ptr_eq(&self.chunks[0], chunk)
1510    }
1511
1512    /// Get a reference to the current chunk.
1513    fn get_chunk(&self) -> &Chunk<D> {
1514        &self.chunks[0]
1515    }
1516
1517    /// Step to the next update.
1518    ///
1519    /// Returns the stepped cursor, or `None` if the step was over the last update.
1520    fn step(mut self) -> Option<Self> {
1521        if self.chunk_offset == self.get_chunk().len() - 1 {
1522            return self.skip_chunk().map(|(c, _)| c);
1523        }
1524
1525        self.chunk_offset += 1;
1526
1527        if let Some(limit) = &mut self.limit {
1528            *limit -= 1;
1529            if *limit == 0 {
1530                return None;
1531            }
1532        }
1533
1534        Some(self)
1535    }
1536
1537    /// Skip the remainder of the current chunk.
1538    ///
1539    /// Returns the forwarded cursor and the number of updates skipped, or `None` if no chunks are
1540    /// left after the skip.
1541    fn skip_chunk(mut self) -> Option<(Self, usize)> {
1542        let chunk = self.chunks.pop_front().expect("cursor invariant");
1543
1544        if self.chunks.is_empty() {
1545            return None;
1546        }
1547
1548        let skipped = chunk.len() - self.chunk_offset;
1549        self.chunk_offset = 0;
1550
1551        if let Some(limit) = &mut self.limit {
1552            if skipped >= *limit {
1553                return None;
1554            }
1555            *limit -= skipped;
1556        }
1557
1558        Some((self, skipped))
1559    }
1560
1561    /// Skip all updates with times <= the given time.
1562    ///
1563    /// Returns the forwarded cursor and the number of updates skipped, or `None` if no updates are
1564    /// left after the skip.
1565    fn skip_time(mut self, time: Timestamp) -> Option<(Self, usize)> {
1566        if self.overwrite_ts.is_some_and(|ts| ts <= time) {
1567            return None;
1568        } else if self.get().1 > time {
1569            return Some((self, 0));
1570        }
1571
1572        let mut skipped = 0;
1573
1574        let new_offset = loop {
1575            let chunk = self.get_chunk();
1576            if let Some(index) = chunk.find_time_greater_than(time) {
1577                break index;
1578            }
1579
1580            let (cursor, count) = self.skip_chunk()?;
1581            self = cursor;
1582            skipped += count;
1583        };
1584
1585        skipped += new_offset - self.chunk_offset;
1586        self.chunk_offset = new_offset;
1587
1588        Some((self, skipped))
1589    }
1590
1591    /// Advance all updates in this cursor by the given `since_ts`.
1592    ///
1593    /// Returns a list of cursors, each of which yields ordered and consolidated updates that have
1594    /// been advanced by `since_ts`.
1595    fn advance_by(mut self, since_ts: Timestamp) -> Vec<Self> {
1596        // If the cursor has an `overwrite_ts`, all its updates are at the same time already. We
1597        // only need to advance the `overwrite_ts` by the `since_ts`.
1598        if let Some(ts) = self.overwrite_ts {
1599            if ts < since_ts {
1600                self.overwrite_ts = Some(since_ts);
1601            }
1602            return vec![self];
1603        }
1604
1605        // Otherwise we need to split the cursor so that each new cursor only yields runs of
1606        // updates that are correctly (time, data)-ordered when advanced by `since_ts`. We achieve
1607        // this by splitting the cursor at each time <= `since_ts`.
1608        let mut splits = Vec::new();
1609        let mut remaining = Some(self);
1610
1611        while let Some(cursor) = remaining.take() {
1612            let (_, time, _) = cursor.get();
1613            if time >= since_ts {
1614                splits.push(cursor);
1615                break;
1616            }
1617
1618            let mut current = cursor.clone();
1619            if let Some((cursor, skipped)) = cursor.skip_time(time) {
1620                remaining = Some(cursor);
1621                current = current.set_limit(skipped).expect("skipped at least 1");
1622            }
1623            current.overwrite_ts = Some(since_ts);
1624            splits.push(current);
1625        }
1626
1627        splits
1628    }
1629
1630    /// Drain the cursor into a [`Chain`].
1631    ///
1632    /// This reuses the underlying chunks if possible, and writes new ones at `depth` otherwise.
1633    /// Reused chunks keep their own depth.
1634    fn into_chain(self, depth: u8) -> Chain<D> {
1635        match self.try_unwrap(depth) {
1636            Ok(chain) => chain,
1637            Err((_, cursor)) => {
1638                let mut builder = ChainBuilder::at_depth(depth);
1639                builder.push_cursor(cursor);
1640                builder.finish()
1641            }
1642        }
1643    }
1644
1645    /// Attempt to unwrap the cursor into a [`Chain`].
1646    ///
1647    /// This operation efficiently reuses chunks by directly inserting them into the output chain
1648    /// where possible.
1649    ///
1650    /// An unwrap is only successful if the cursor's `limit` and `overwrite_ts` are both `None` and
1651    /// the cursor has unique references to its chunks. If the unwrap fails, this method returns an
1652    /// `Err` containing the cursor in an unchanged state, allowing the caller to convert it into a
1653    /// chain by copying chunks rather than reusing them.
1654    fn try_unwrap(self, depth: u8) -> Result<Chain<D>, (&'static str, Self)> {
1655        if self.limit.is_some() {
1656            return Err(("cursor with limit", self));
1657        }
1658        if self.overwrite_ts.is_some() {
1659            return Err(("cursor with overwrite_ts", self));
1660        }
1661        if self.chunks.iter().any(|c| Rc::strong_count(c) != 1) {
1662            return Err(("cursor on shared chunks", self));
1663        }
1664
1665        let mut builder = ChainBuilder::at_depth(depth);
1666        let mut remaining = Some(self);
1667
1668        // We might be partway through the first chunk, in which case we can't reuse it but need to
1669        // allocate a new one to contain only the updates the cursor can still yield.
1670        while let Some(cursor) = remaining.take() {
1671            if cursor.chunk_offset == 0 {
1672                remaining = Some(cursor);
1673                break;
1674            }
1675            let update = cursor.get();
1676            builder.push_ref(update);
1677            remaining = cursor.step();
1678        }
1679
1680        let mut chain = builder.finish();
1681        if let Some(cursor) = remaining {
1682            for chunk in cursor.chunks {
1683                let chunk = Rc::into_inner(chunk).expect("checked above");
1684                chain.push_chunk(chunk);
1685            }
1686        }
1687
1688        Ok(chain)
1689    }
1690}
1691
1692/// A non-empty chunk of updates, backed by a columnar region.
1693///
1694/// All updates in a chunk are sorted by (time, data) and consolidated.
1695///
1696/// Chunks are immutable once created. They are produced by [`ChunkBuilder`], which mints a
1697/// new chunk whenever its in-progress columnar container reaches a fixed serialized byte
1698/// boundary (~2 MiB, matching the ship granularity used elsewhere in the codebase), so each
1699/// chunk corresponds to a single, predictably sized allocation.
1700struct Chunk<D: Data> {
1701    /// The body in the buffer pool, which may spill it.
1702    ///
1703    /// Empty when the body stays on the heap, which is what
1704    /// [`mz_timely_util::columnar::chunk::try_spill_ref`] decides, and emptied again by
1705    /// [`Chunk::body`]: a materialized chunk holds its body on the heap for the rest of its
1706    /// life, so keeping the pool's copy alive would only double the footprint and hold budget
1707    /// the pool could give to a chunk that still needs it. [`Chunk::with_view`] reads a copy and
1708    /// leaves the body here.
1709    ///
1710    /// A `Mutex` (not `RefCell`) keeps the chunk `Sync`: cursors hold chunks behind a shared
1711    /// `Rc`, and the iterator returned by [`CorrectionV2::updates_before`] borrows them across
1712    /// the persist writer's `await`, so `&Chunk` must be `Send`. The lock is taken at
1713    /// materialization and for scoped reads, and is otherwise uncontended (the sink runs
1714    /// single-threaded per worker).
1715    pooled: Mutex<Option<ChunkHandle>>,
1716    /// The materialized form, populated lazily by [`Chunk::body`] on first access.
1717    ///
1718    /// An `OnceLock` for the same `Sync` reason as `pooled`. Once set the slot is never
1719    /// cleared, so its address is stable and [`Chunk::index`] can hand out `Ref<'_>` borrows
1720    /// tied to `&self`. The allocation is freed when the chunk drops, which bounds resident
1721    /// memory to the chunks under an active merge front.
1722    resident: OnceLock<ColumnBody<(D, Timestamp, Diff)>>,
1723    /// Number of updates, cached so `len` and chain bookkeeping never page the chunk in.
1724    len: usize,
1725    /// Serialized size of the body in words, cached so size accounting never pages it in.
1726    body_words: usize,
1727    /// Time of the first update, cached so boundary checks (`split_at_time`, `can_accept`) route
1728    /// a resting chunk without materializing it.
1729    first_time: Timestamp,
1730    /// Time of the last update, cached likewise.
1731    last_time: Timestamp,
1732    /// The generational depth: 0 for chunks built from staged updates or written by a read,
1733    /// one more than the deepest input for chunks written by a bucket's chain merge. The pool
1734    /// treats deeper chunks as colder, and compresses them past the floor set by
1735    /// [`chunk::set_compress_min_depth`].
1736    depth: u8,
1737}
1738
1739impl<D: Data> fmt::Debug for Chunk<D> {
1740    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1741        write!(f, "Chunk(<{}>)", self.len())
1742    }
1743}
1744
1745impl<D: Data> Chunk<D> {
1746    /// Mint a chunk from the given non-empty body, emptying it.
1747    ///
1748    /// Reads the metadata a resting chunk must answer (length, size, boundary times) while the
1749    /// body is still in hand, then offers it to the buffer pool. A body the pool takes is
1750    /// encoded straight into its slot, which leaves the typed allocation behind for the caller
1751    /// to refill, and materializes lazily on first read. A body the pool declines is
1752    /// serialized out of the typed allocation instead, because a resident chunk has to own it.
1753    ///
1754    /// # Panics
1755    ///
1756    /// Panics if the body is empty. Chunks are non-empty by construction; [`ChunkBuilder`]
1757    /// only ever mints from a populated body.
1758    fn mint(body: &mut ColumnBody<(D, Timestamp, Diff)>, depth: u8) -> Self {
1759        let (len, first_time, last_time) = {
1760            let borrowed = body.borrow();
1761            let len = borrowed.len();
1762            assert!(len > 0, "chunks are non-empty");
1763            (len, borrowed.get(0).1, borrowed.get(len - 1).1)
1764        };
1765        let body_words = body.length_in_bytes() / std::mem::size_of::<u64>();
1766
1767        // TODO: chunks resting in far-future buckets go untouched for longer than their merge
1768        // depth suggests, and could spill deeper still.
1769        let (pooled, resident) = match chunk::try_spill_ref(body, depth) {
1770            Some(handle) => {
1771                body.clear();
1772                (Mutex::new(Some(handle)), OnceLock::new())
1773            }
1774            None => {
1775                // Serialize into an exactly sized vector, which leaves the typed allocation
1776                // in place for the caller and keeps its spare capacity out of a body that is
1777                // held for the rest of the chunk's life.
1778                let mut words = Vec::with_capacity(body_words);
1779                indexed::encode(&mut words, &body.borrow());
1780                body.clear();
1781                let cell = OnceLock::new();
1782                if cell.set(ColumnBody::Words(words)).is_err() {
1783                    unreachable!("cell is fresh");
1784                }
1785                (Mutex::new(None), cell)
1786            }
1787        };
1788        Self {
1789            pooled,
1790            body_words,
1791            resident,
1792            len,
1793            first_time,
1794            last_time,
1795            depth,
1796        }
1797    }
1798
1799    /// Materialize the chunk's body, taking it out of the pool on first access.
1800    ///
1801    /// The returned reference is valid for as long as `&self`: the `OnceLock` slot is never
1802    /// cleared once populated, so its contents have a stable address. Taking the body frees the
1803    /// pool's copy, so the chunk holds exactly one.
1804    ///
1805    /// For reads by merges and splits, whose output replaces the chunk. A read after which the
1806    /// chunk survives uses [`Chunk::with_view`], which keeps the body spillable.
1807    fn body(&self) -> &ColumnBody<(D, Timestamp, Diff)> {
1808        self.resident.get_or_init(|| {
1809            let handle = self
1810                .pooled
1811                .lock()
1812                .expect("pool handle mutex poisoned")
1813                .take()
1814                .expect("a chunk the pool declined is materialized at construction");
1815            let mut words = Vec::new();
1816            handle.take(&mut words);
1817            ColumnBody::Words(words)
1818        })
1819    }
1820
1821    /// Call `f` with a view of the chunk's body, leaving the body where it is.
1822    ///
1823    /// A materialized chunk is read in place. Otherwise the body is copied out of the pool for
1824    /// the call and the copy dropped after it, so the pool keeps its slot and may still spill
1825    /// the body.
1826    fn with_view<R>(&self, f: impl FnOnce(ChunkView<'_, D>) -> R) -> R {
1827        if let Some(body) = self.resident.get() {
1828            return f(body.borrow());
1829        }
1830        let mut words = Vec::new();
1831        {
1832            let pooled = self.pooled.lock().expect("pool handle mutex poisoned");
1833            match pooled.as_ref() {
1834                Some(handle) => handle.read_into(&mut words),
1835                // Materialized between the check above and the lock.
1836                None => {
1837                    drop(pooled);
1838                    let body = self.resident.get().expect("a taken body is materialized");
1839                    return f(body.borrow());
1840                }
1841            }
1842        }
1843        let body = ColumnBody::<(D, Timestamp, Diff)>::Words(words);
1844        f(body.borrow())
1845    }
1846
1847    /// Return the number of updates in the chunk.
1848    fn len(&self) -> usize {
1849        self.len
1850    }
1851
1852    /// Borrow the chunk's body, paging it in if necessary.
1853    ///
1854    /// Any caller that touches more than one update must hoist this out of its loop and index the
1855    /// returned view. `ColumnBody::borrow` on a serialized body rebuilds the struct-of-arrays
1856    /// view from the serialized header on every call, so borrowing per element pays that decode
1857    /// per element.
1858    fn view(&self) -> ChunkView<'_, D> {
1859        self.body().borrow()
1860    }
1861
1862    /// Return the update at the given index, paging the chunk in if necessary.
1863    ///
1864    /// Single-access only. Indexing a hoisted [`Chunk::view`] is the loop form.
1865    ///
1866    /// # Panics
1867    ///
1868    /// Panics if the given index is not populated.
1869    fn index(&self, idx: usize) -> Ref<'_, (D, Timestamp, Diff)> {
1870        self.view().get(idx)
1871    }
1872
1873    /// Return the time of the first update, without materializing the chunk.
1874    fn first_time(&self) -> Timestamp {
1875        self.first_time
1876    }
1877
1878    /// Return the time of the last update, without materializing the chunk.
1879    fn last_time(&self) -> Timestamp {
1880        self.last_time
1881    }
1882
1883    /// Return the index of the first update at a time greater than `time`, or `None` if no such
1884    /// update exists.
1885    ///
1886    /// The early-out uses the cached last time, so a chunk whose updates are all at or before
1887    /// `time` is skipped without paging it in.
1888    fn find_time_greater_than(&self, time: Timestamp) -> Option<usize> {
1889        if self.last_time <= time {
1890            return None;
1891        }
1892
1893        let view = self.view();
1894        let mut lower = 0;
1895        let mut upper = self.len;
1896        while lower < upper {
1897            let idx = (lower + upper) / 2;
1898            if view.get(idx).1 > time {
1899                upper = idx;
1900            } else {
1901                lower = idx + 1;
1902            }
1903        }
1904
1905        Some(lower)
1906    }
1907
1908    /// Return the serialized size of the chunk's body in bytes, for use in metrics.
1909    ///
1910    /// The chunk holds exactly one copy of its body, in the pool or on the heap, so this is its
1911    /// size either way. A scoped read's copy lives only for the read. It is not a resident-byte
1912    /// number: what the pool has evicted is the pool's business, and it publishes
1913    /// `resident_bytes` itself (`mz_ore::pool::PoolStats`).
1914    fn size(&self) -> usize {
1915        self.body_words * std::mem::size_of::<u64>()
1916    }
1917}
1918
1919/// Builder that mints fixed-size [`Chunk`]s from a stream of updates.
1920///
1921/// Updates accumulate in one body, which is minted into a chunk as soon as it reports itself at
1922/// capacity, so every chunk is a single, predictably sized body. Minting a spilled body leaves
1923/// the typed allocation in place, so a builder that mints many chunks grows one allocation once.
1924struct ChunkBuilder<D: Data> {
1925    /// The updates pushed since the last mint.
1926    current: ColumnBody<(D, Timestamp, Diff)>,
1927    /// The generational depth minted chunks carry.
1928    depth: u8,
1929}
1930
1931impl<D: Data> ChunkBuilder<D> {
1932    /// A builder whose chunks are minted at the given generational depth.
1933    fn at_depth(depth: u8) -> Self {
1934        Self {
1935            current: Default::default(),
1936            depth,
1937        }
1938    }
1939
1940    /// Push an update, returning a chunk if the push completed one.
1941    ///
1942    /// Accepts whatever [`ColumnBody`]'s [`PushInto`] impl accepts, both the
1943    /// `Ref<'_, (D, T, R)>` refs produced by cursors and `&(D, T, R)` references to owned
1944    /// tuples drained from the staging buffer.
1945    ///
1946    /// [`PushInto`]: timely::container::PushInto
1947    #[inline]
1948    fn push<T>(&mut self, item: T) -> Option<Chunk<D>>
1949    where
1950        ColumnBody<(D, Timestamp, Diff)>: PushInto<T>,
1951    {
1952        PushInto::push_into(&mut self.current, item);
1953        // The ship test walks the body's slice lengths, so it costs per push, not per byte.
1954        self.current
1955            .at_capacity()
1956            .then(|| Chunk::mint(&mut self.current, self.depth))
1957    }
1958
1959    /// Mint the updates pushed since the last mint, if any.
1960    fn finish(mut self) -> Option<Chunk<D>> {
1961        (!self.current.is_empty()).then(|| Chunk::mint(&mut self.current, self.depth))
1962    }
1963}
1964
1965/// A buffer for staging updates before they are inserted into the sorted chains.
1966#[derive(Debug)]
1967struct Stage<D> {
1968    /// The contained updates.
1969    ///
1970    /// Grows on demand rather than being allocated at the ship size, so a sink that never
1971    /// stages that much never holds it. One of these exists per sink and worker, which is why
1972    /// the eager allocation is worth avoiding even though a staging area is small.
1973    data: Vec<(D, Timestamp, Diff)>,
1974    /// Bytes held by `data`: the tuples plus what they own on the heap.
1975    bytes: usize,
1976    /// How many bytes to accumulate before shipping a batch, from
1977    /// `compute_correction_v2_chunk_size`.
1978    ///
1979    /// Shipping less often costs staging memory and saves inserts: it is the number of chains
1980    /// minted per update, and every chain minted is a chain some later read has to merge.
1981    /// Counted in bytes with the heap included, so the staging memory is bounded by this
1982    /// setting regardless of row width, and a wide row fills the stage in fewer updates.
1983    chunk_size: usize,
1984    /// Running totals and introspection logging.
1985    ///
1986    /// We want to report the number of records in the stage. To do so, we pretend that the stage
1987    /// is a chain, and every time the number of updates inside changes, the chain gets dropped and
1988    /// re-created.
1989    accounting: Accounting,
1990}
1991
1992impl<D: Data> Stage<D> {
1993    fn new(accounting: Accounting, chunk_size: usize) -> Self {
1994        // For logging, we pretend the stage consists of a single chain.
1995        accounting.stage_created();
1996
1997        Self {
1998            data: Vec::new(),
1999            bytes: 0,
2000            chunk_size,
2001            accounting,
2002        }
2003    }
2004
2005    /// Bytes an update occupies: its datum, heap included, plus time and diff.
2006    fn update_bytes(update: &(D, Timestamp, Diff)) -> usize {
2007        update.0.data_bytes() + std::mem::size_of::<(Timestamp, Diff)>()
2008    }
2009
2010    /// Insert a batch of updates, possibly producing a batch of sorted, consolidated updates
2011    /// ready to be stored.
2012    ///
2013    /// Everything staged ships together, the new batch included, once the staged bytes reach
2014    /// `chunk_size`.
2015    fn insert(
2016        &mut self,
2017        updates: &mut Vec<(D, Timestamp, Diff)>,
2018    ) -> Option<Vec<(D, Timestamp, Diff)>> {
2019        if updates.is_empty() {
2020            return None;
2021        }
2022
2023        let before = self.snapshot();
2024
2025        self.bytes += updates.iter().map(Self::update_bytes).sum::<usize>();
2026        self.data.append(updates);
2027
2028        let maybe_ready = if self.bytes >= self.chunk_size {
2029            let mut ready = std::mem::take(&mut self.data);
2030            self.bytes = 0;
2031            consolidate(&mut ready);
2032            Some(ready)
2033        } else {
2034            None
2035        };
2036
2037        self.account_since(before);
2038
2039        maybe_ready
2040    }
2041
2042    /// Flush all currently staged updates, returning them sorted and consolidated.
2043    fn flush(&mut self) -> Option<Vec<(D, Timestamp, Diff)>> {
2044        let before = self.snapshot();
2045
2046        consolidate(&mut self.data);
2047        self.bytes = 0;
2048        let data = (!self.data.is_empty()).then(|| std::mem::take(&mut self.data));
2049
2050        self.account_since(before);
2051        data
2052    }
2053
2054    /// Advance the times of staged updates by the given `since`.
2055    fn advance_times(&mut self, since: &Antichain<Timestamp>) {
2056        let Some(since_ts) = since.as_option() else {
2057            // If the since is the empty frontier, discard all updates.
2058            let before = self.snapshot();
2059            self.data.clear();
2060            self.bytes = 0;
2061            self.account_since(before);
2062            return;
2063        };
2064
2065        for (_, time, _) in &mut self.data {
2066            *time = std::cmp::max(*time, *since_ts);
2067        }
2068    }
2069
2070    /// The staged length, bytes, and allocation count, taken before a mutation for
2071    /// [`Stage::account_since`].
2072    fn snapshot(&self) -> (isize, isize, isize) {
2073        let i = |x: usize| isize::try_from(x).expect("must fit");
2074        let allocated = isize::from(self.data.capacity() > 0);
2075        (i(self.data.len()), i(self.bytes), allocated)
2076    }
2077
2078    /// Account for the change to the stage since `before`, a [`Stage::snapshot`].
2079    fn account_since(&self, before: (isize, isize, isize)) {
2080        let (len, bytes, allocations) = self.snapshot();
2081        self.accounting
2082            .stage_diff(len - before.0, bytes - before.1, allocations - before.2);
2083    }
2084}
2085
2086impl<D> Drop for Stage<D> {
2087    fn drop(&mut self) {
2088        let allocated = usize::from(self.data.capacity() > 0);
2089        self.accounting
2090            .stage_dropped(self.data.len(), self.bytes, allocated);
2091    }
2092}
2093
2094/// Sort and consolidate the given list of updates.
2095///
2096/// This function is the same as [`differential_dataflow::consolidation::consolidate_updates`],
2097/// except that it sorts updates by (time, data) instead of (data, time).
2098fn consolidate<D: Data>(updates: &mut Vec<(D, Timestamp, Diff)>) {
2099    if updates.len() <= 1 {
2100        return;
2101    }
2102
2103    let diff = |update: &(_, _, Diff)| update.2;
2104
2105    updates.sort_unstable_by(|(d1, t1, _), (d2, t2, _)| (t1, d1).cmp(&(t2, d2)));
2106
2107    let mut offset = 0;
2108    let mut accum = diff(&updates[0]);
2109
2110    for idx in 1..updates.len() {
2111        let this = &updates[idx];
2112        let prev = &updates[idx - 1];
2113        if this.0 == prev.0 && this.1 == prev.1 {
2114            accum += diff(&updates[idx]);
2115        } else {
2116            if accum != Diff::ZERO {
2117                updates.swap(offset, idx - 1);
2118                updates[offset].2 = accum;
2119                offset += 1;
2120            }
2121            accum = diff(&updates[idx]);
2122        }
2123    }
2124
2125    if accum != Diff::ZERO {
2126        let len = updates.len();
2127        updates.swap(offset, len - 1);
2128        updates[offset].2 = accum;
2129        offset += 1;
2130    }
2131
2132    updates.truncate(offset);
2133}
2134
2135/// Compare two `(time, data)` pairs of columnar refs that have unrelated input lifetimes.
2136///
2137/// `<D::Container as Borrow>::Ref<'a>` is an associated-type projection through a trait, so
2138/// the compiler treats it as invariant in `'a` and won't auto-shorten the inputs by variance.
2139/// We instead explicitly reborrow both to a fresh, local lifetime `'x` via
2140/// [`Columnar::reborrow`] before letting the inner `cmp` pick up the `for<'a> Ref<'a>: Ord`
2141/// bound on [`Data`].
2142#[inline]
2143fn refs_cmp<D: Data>(a: (Timestamp, Ref<'_, D>), b: (Timestamp, Ref<'_, D>)) -> Ordering {
2144    #[inline]
2145    fn cmp<'x, D: Data>(a: (Timestamp, Ref<'x, D>), b: (Timestamp, Ref<'x, D>)) -> Ordering {
2146        a.cmp(&b)
2147    }
2148    cmp::<D>((a.0, D::reborrow(a.1)), (b.0, D::reborrow(b.1)))
2149}
2150
2151#[cfg(test)]
2152mod tests {
2153    use mz_ore::metrics::MetricsRegistry;
2154    use mz_persist_client::cfg::PersistConfig;
2155    use mz_persist_client::metrics::Metrics;
2156    use mz_repr::{Diff, Timestamp};
2157
2158    use super::*;
2159    use crate::sink::correction::CorrectionV1;
2160
2161    #[mz_ore::test]
2162    fn chain_builder_update_count_matches_items() {
2163        let mut builder = ChainBuilder::<i64>::default();
2164        for i in 0..10_u64 {
2165            let d = i64::try_from(i).expect("fits");
2166            builder.push_owned(&(d, Timestamp::new(i), Diff::ONE));
2167        }
2168        let chain = builder.finish();
2169        assert_eq!(chain.update_count, chain.iter().count());
2170    }
2171
2172    /// Push enough updates to cross at least one `mint()` boundary, forcing the
2173    /// `Align` encode -> `from_bytes` decode roundtrip (the spilling path this data
2174    /// structure exists to support), and assert `iter()` roundtrips values, order,
2175    /// and diffs across the spill boundary.
2176    #[mz_ore::test]
2177    #[cfg_attr(miri, ignore)] // too slow: crossing the ~2 MiB mint boundary needs ~200k updates
2178    fn chain_builder_roundtrips_across_mint_boundary() {
2179        // A single `mint()` fires near the ~2 MiB (`SHIP_WORDS`) serialized boundary. With
2180        // three 8-byte columns per update that's tens of thousands of updates; pushing 200k
2181        // comfortably forces multiple mints.
2182        let count = 200_000_u64;
2183
2184        let mut builder = ChainBuilder::<i64>::default();
2185        for i in 0..count {
2186            let d = i64::try_from(i).expect("fits");
2187            builder.push_owned(&(d, Timestamp::new(i), Diff::ONE));
2188        }
2189        let chain = builder.finish();
2190
2191        // Crossing the mint boundary must have produced more than one chunk; otherwise the spill
2192        // path (each minted chunk is paged out and read back through the pager) wouldn't be
2193        // exercised. The chunk payload itself is now behind the pager (see [`Chunk`]), so we
2194        // assert on chunk count rather than inspecting the column variant directly.
2195        assert!(
2196            chain.chunks.len() > 1,
2197            "expected multiple minted chunks, got {} chunk(s): {:?}",
2198            chain.chunks.len(),
2199            chain.chunks,
2200        );
2201
2202        // `iter()` must roundtrip every update, in order, with correct diffs.
2203        assert_eq!(chain.update_count, usize::try_from(count).expect("fits"));
2204        let mut expected = 0_u64;
2205        for (d, t, r) in chain.iter() {
2206            assert_eq!(d, i64::try_from(expected).expect("fits"));
2207            assert_eq!(t, Timestamp::new(expected));
2208            assert_eq!(r, Diff::ONE);
2209            expected += 1;
2210        }
2211        assert_eq!(expected, count);
2212    }
2213
2214    impl DataBytes for String {
2215        fn data_bytes(&self) -> usize {
2216            std::mem::size_of::<Self>() + self.len()
2217        }
2218    }
2219
2220    impl DataBytes for i64 {
2221        fn data_bytes(&self) -> usize {
2222            std::mem::size_of::<Self>()
2223        }
2224    }
2225
2226    /// A k-way merge over chains of several chunks each consolidates across all of them, with
2227    /// its rounds ending wherever any input crosses a chunk boundary.
2228    #[mz_ore::test]
2229    #[cfg_attr(miri, ignore)] // slow under Miri
2230    fn merge_many_consolidates_across_chunks() {
2231        // Large enough that the sparsest chain still spans several chunks.
2232        let count = 600_000_i64;
2233        // Chain `k` holds the updates of `inputs[k]`, sorted by (time, data).
2234        let inputs: [Vec<(i64, Timestamp, Diff)>; 3] = [
2235            (0..count)
2236                .map(|d| (d, Timestamp::new(0), Diff::ONE))
2237                .collect(),
2238            (0..count)
2239                .filter(|d| d % 2 == 0)
2240                .map(|d| (d, Timestamp::new(0), Diff::ONE))
2241                .collect(),
2242            (0..count)
2243                .filter(|d| d % 3 == 0)
2244                .map(|d| (d, Timestamp::new(0), -Diff::ONE))
2245                .collect(),
2246        ];
2247
2248        let mut expected = std::collections::BTreeMap::new();
2249        for (d, _, r) in inputs.iter().flatten() {
2250            *expected.entry(*d).or_insert(Diff::ZERO) += *r;
2251        }
2252        expected.retain(|_, r| *r != Diff::ZERO);
2253
2254        let cursors = inputs
2255            .iter()
2256            .map(|updates| {
2257                let mut builder = ChainBuilder::<i64>::default();
2258                for update in updates {
2259                    builder.push_owned(update);
2260                }
2261                let chain = builder.finish();
2262                assert!(chain.chunks.len() > 1, "expected multiple minted chunks");
2263                chain.into_cursor().expect("non-empty")
2264            })
2265            .collect();
2266        let merged = merge_many(cursors, 0);
2267
2268        let actual: Vec<_> = merged.iter().map(|(d, _, r)| (d, r)).collect();
2269        let expected: Vec<_> = expected.into_iter().collect();
2270        assert_eq!(actual, expected);
2271    }
2272
2273    fn sink_metrics() -> SinkMetrics {
2274        let registry = MetricsRegistry::new();
2275        let metrics = Metrics::new(&PersistConfig::new_for_tests(), &registry);
2276        metrics.sink.clone()
2277    }
2278
2279    /// Insert single updates of the given payload width until the stage ships, returning how
2280    /// many updates that took.
2281    fn updates_until_ship(width: usize, chunk_size: usize) -> usize {
2282        let metrics = sink_metrics();
2283        let mut v2 = CorrectionV2::<String>::new(
2284            metrics.clone(),
2285            metrics.for_worker(0),
2286            None,
2287            3.0,
2288            chunk_size,
2289        );
2290        for i in 0..chunk_size {
2291            let datum = format!("{i:0width$}");
2292            v2.insert(&mut vec![(datum, Timestamp::from(0), Diff::ONE)]);
2293            if v2.stage.data.is_empty() {
2294                assert_eq!(v2.stage.bytes, 0, "a shipped stage holds no bytes");
2295                return i + 1;
2296            }
2297        }
2298        panic!("the stage never shipped");
2299    }
2300
2301    /// The stage ships by bytes, heap included, so wider rows ship after fewer updates.
2302    #[mz_ore::test]
2303    #[cfg_attr(miri, ignore)] // slow under Miri
2304    fn stage_ships_by_bytes() {
2305        let chunk_size = 16 * 1024;
2306        let narrow = updates_until_ship(8, chunk_size);
2307        let wide = updates_until_ship(1024, chunk_size);
2308
2309        let per_update = |width: usize| {
2310            std::mem::size_of::<String>() + width + std::mem::size_of::<(Timestamp, Diff)>()
2311        };
2312        assert_eq!(narrow, chunk_size.div_ceil(per_update(8)));
2313        assert_eq!(wide, chunk_size.div_ceil(per_update(1024)));
2314        assert!(wide < narrow);
2315    }
2316
2317    /// Run the same stepwise-drain workload through `CorrectionV1` and `CorrectionV2` and assert
2318    /// that they emit the same updates at every step.
2319    ///
2320    /// Models the `write_batches` operator catching up through many distinct timestamps: the
2321    /// desired input runs ahead, batches are written one timestamp at a time, and written updates
2322    /// come back negated through the persist feedback.
2323    #[mz_ore::test]
2324    // Columnation regions are not Stacked Borrows compliant: later pushes invalidate the
2325    // provenance of previously stored items under Miri.
2326    #[cfg_attr(miri, ignore)]
2327    fn equivalence_with_v1() {
2328        let sink_metrics = sink_metrics();
2329
2330        let mut v1 =
2331            CorrectionV1::<String>::new(sink_metrics.clone(), sink_metrics.for_worker(0), 1);
2332        let mut v2 = CorrectionV2::<String>::new(
2333            sink_metrics.clone(),
2334            sink_metrics.for_worker(0),
2335            None,
2336            3.0,
2337            8 * 1024,
2338        );
2339
2340        let num_ts = 50;
2341        let keys = 4;
2342
2343        // Upsert-style input: every timestamp updates each key, retracting the previous value.
2344        let batch = |t: u64| -> Vec<(String, Timestamp, Diff)> {
2345            (0..keys)
2346                .flat_map(|k| {
2347                    let addition = (format!("{k}-{t}"), Timestamp::from(t), Diff::ONE);
2348                    let retraction = t
2349                        .checked_sub(1)
2350                        .map(|p| (format!("{k}-{p}"), Timestamp::from(t), -Diff::ONE));
2351                    std::iter::once(addition).chain(retraction)
2352                })
2353                .collect()
2354        };
2355
2356        // Pre-fill both with all batches, like a catch-up where the input runs ahead.
2357        for t in 0..num_ts {
2358            v1.insert(&mut batch(t));
2359            v2.insert(&mut batch(t));
2360        }
2361
2362        // Drain stepwise, with persist feedback, comparing emissions.
2363        for t in 0..num_ts {
2364            let upper = Antichain::from_elem(Timestamp::from(t + 1));
2365
2366            let mut out1: Vec<_> = v1.updates_before(&upper).collect();
2367            let mut out2: Vec<_> = v2.updates_before(&upper).collect();
2368            out1.sort();
2369            out2.sort();
2370            assert_eq!(out1, out2, "diverged at t={t}");
2371
2372            v1.insert_negated(&mut out1.clone());
2373            v2.insert_negated(&mut out2);
2374            v1.advance_since(upper.clone());
2375            v2.advance_since(upper);
2376        }
2377
2378        // Compare the final state at the since.
2379        let upper = Antichain::from_elem(Timestamp::from(num_ts + 1));
2380        v1.consolidate_at_since();
2381        v2.consolidate_at_since();
2382        let mut out1: Vec<_> = v1.updates_before(&upper).collect();
2383        let mut out2: Vec<_> = v2.updates_before(&upper).collect();
2384        out1.sort();
2385        out2.sort();
2386        assert_eq!(out1, out2);
2387    }
2388
2389    #[mz_ore::test]
2390    #[cfg_attr(miri, ignore)] // slow under Miri, and the sort and merge it checks have no unsafe code
2391    fn sort_advanced_consolidates_across_runs() {
2392        // Each time carries an insertion and, at a later time, its retraction, so with runs of
2393        // a few updates the pair lands in different runs and only the merge can cancel it.
2394        let since_ts = Timestamp::from(100);
2395        let mut updates = Vec::new();
2396        for i in 0..50_i64 {
2397            let t = u64::try_from(i).expect("fits");
2398            updates.push((i, Timestamp::from(t), Diff::ONE));
2399            updates.push((i, Timestamp::from(t + 50), -Diff::ONE));
2400            updates.push((i, Timestamp::from(t), Diff::ONE));
2401        }
2402        // Times beyond the since keep their own time.
2403        updates.push((7, Timestamp::from(200), Diff::ONE));
2404        let merged = sort_advanced(chains_of(&updates, 9), since_ts, 4);
2405
2406        let mut expected: Vec<_> = updates
2407            .into_iter()
2408            .map(|(d, t, r)| (d, std::cmp::max(t, since_ts), r))
2409            .collect();
2410        consolidate(&mut expected);
2411        assert_eq!(merged.iter().collect::<Vec<_>>(), expected);
2412        assert_eq!(
2413            merged.update_count, 51,
2414            "50 surviving insertions and the late update"
2415        );
2416    }
2417
2418    /// Builds one chain per `per_chain` updates, each sorted and consolidated on its own.
2419    fn chains_of(updates: &[(i64, Timestamp, Diff)], per_chain: usize) -> Vec<Chain<i64>> {
2420        updates
2421            .chunks(per_chain)
2422            .map(|part| {
2423                let mut part = part.to_vec();
2424                consolidate(&mut part);
2425                mint_run(&mut part)
2426            })
2427            .collect()
2428    }
2429
2430    #[mz_ore::test]
2431    #[cfg_attr(miri, ignore)] // slow under Miri, and the sort and merge it checks have no unsafe code
2432    fn sort_advanced_matches_one_sort() {
2433        // Insertions spread over many times, a share of them retracted at a later time, so runs
2434        // both survive consolidation and cancel across runs at several merge levels.
2435        let since_ts = Timestamp::from(500);
2436        let mut updates = Vec::new();
2437        for i in 0..2_000_i64 {
2438            let t = u64::try_from(i % 997).expect("fits");
2439            updates.push((i % 300, Timestamp::from(t), Diff::ONE));
2440            if i % 3 == 0 {
2441                updates.push((i % 300, Timestamp::from(t + 3), -Diff::ONE));
2442            }
2443        }
2444        let mut expected: Vec<_> = updates
2445            .iter()
2446            .map(|(d, t, r)| (*d, std::cmp::max(*t, since_ts), *r))
2447            .collect();
2448        consolidate(&mut expected);
2449
2450        for run_len in [1, 2, 7, 64, 100_000] {
2451            let merged = sort_advanced(chains_of(&updates, 37), since_ts, run_len);
2452            assert_eq!(
2453                merged.iter().collect::<Vec<_>>(),
2454                expected,
2455                "run_len {run_len}"
2456            );
2457        }
2458    }
2459
2460    #[mz_ore::test]
2461    #[cfg_attr(miri, ignore)] // the pool's mapped regions are unsupported under miri
2462    fn sort_advanced_matches_one_sort_through_pool() {
2463        // Enough updates that input chains and runs mint chunks large enough to spill, so the sort
2464        // reads spilled input and merges spilled runs.
2465        let since_ts = Timestamp::from(5_000);
2466        let mut updates = Vec::new();
2467        for i in 0..400_000_i64 {
2468            let t = u64::try_from(i % 9_973).expect("fits");
2469            updates.push((i % 50_000, Timestamp::from(t), Diff::ONE));
2470            if i % 3 == 0 {
2471                updates.push((i % 50_000, Timestamp::from(t + 7), -Diff::ONE));
2472            }
2473        }
2474        let mut expected: Vec<_> = updates
2475            .iter()
2476            .map(|(d, t, r)| (*d, std::cmp::max(*t, since_ts), *r))
2477            .collect();
2478        consolidate(&mut expected);
2479
2480        with_spill_pool(|| {
2481            let chains = chains_of(&updates, 100_000);
2482            for chain in &chains {
2483                assert_all_spilled(chain);
2484            }
2485            let merged = sort_advanced(chains, since_ts, 1 << 12);
2486            assert_eq!(merged.iter().collect::<Vec<_>>(), expected);
2487        });
2488    }
2489
2490    #[mz_ore::test]
2491    #[cfg_attr(miri, ignore)] // slow under Miri, and the sort and merge it checks have no unsafe code
2492    fn sort_advanced_cancels_to_nothing() {
2493        let since_ts = Timestamp::from(10);
2494        let updates: Vec<_> = (0..100_i64)
2495            .flat_map(|i| {
2496                let t = Timestamp::from(u64::try_from(i % 10).expect("fits"));
2497                [
2498                    (i, t, Diff::ONE),
2499                    (i, t.max(Timestamp::from(5)), -Diff::ONE),
2500                ]
2501            })
2502            .collect();
2503        let merged = sort_advanced(chains_of(&updates, 13), since_ts, 8);
2504        assert!(merged.is_empty(), "every insertion meets its retraction");
2505    }
2506
2507    /// A since jump across many distinct buffered timestamps must collapse them onto the since.
2508    #[mz_ore::test]
2509    // Columnation regions are not Stacked Borrows compliant: later pushes invalidate the
2510    // provenance of previously stored items under Miri.
2511    #[cfg_attr(miri, ignore)]
2512    fn since_jump() {
2513        let sink_metrics = sink_metrics();
2514        let mut v2 = CorrectionV2::<String>::new(
2515            sink_metrics.clone(),
2516            sink_metrics.for_worker(0),
2517            None,
2518            3.0,
2519            8 * 1024,
2520        );
2521
2522        let num_ts = 100;
2523        for t in 0..num_ts {
2524            v2.insert(&mut vec![
2525                (format!("a-{t}"), Timestamp::from(t), Diff::ONE),
2526                (format!("a-{t}"), Timestamp::from(t), -Diff::ONE),
2527                (format!("b-{t}"), Timestamp::from(t), Diff::ONE),
2528            ]);
2529        }
2530
2531        v2.advance_since(Antichain::from_elem(Timestamp::from(num_ts)));
2532        v2.consolidate_at_since();
2533
2534        let upper = Antichain::from_elem(Timestamp::from(num_ts + 1));
2535        let out: Vec<_> = v2.updates_before(&upper).collect();
2536        assert_eq!(out.len(), usize::try_from(num_ts).unwrap());
2537        assert!(
2538            out.iter()
2539                .all(|(_, t, r)| *t == Timestamp::from(num_ts) && *r == Diff::ONE)
2540        );
2541    }
2542
2543    /// The maintained record, size, and allocation totals must equal a walk over the chunks.
2544    ///
2545    /// Guards the delta accounting in [`Chain::push_chunk`] and [`Stage::account_since`]: a
2546    /// site that grows a chain or the stage without adjusting the totals shows up here.
2547    #[mz_ore::test]
2548    #[cfg_attr(miri, ignore)] // slow under Miri, and the bookkeeping it checks has no unsafe code
2549    fn metrics_totals_match_walk() {
2550        let sink_metrics = sink_metrics();
2551        let mut v2 = CorrectionV2::<String>::new(
2552            sink_metrics.clone(),
2553            sink_metrics.for_worker(0),
2554            None,
2555            3.0,
2556            8 * 1024,
2557        );
2558
2559        let num_ts = 200_u64;
2560        for t in 0..num_ts {
2561            v2.insert(&mut vec![
2562                (format!("a-{t}"), Timestamp::from(t), Diff::ONE),
2563                (format!("b-{t}"), Timestamp::from(t), Diff::ONE),
2564            ]);
2565        }
2566
2567        // Drain part of the buffer, so chains sit in `emitted` and `pending_low` as well as in
2568        // the bucket chain.
2569        let upper = Antichain::from_elem(Timestamp::from(num_ts / 2));
2570        let drained = v2.updates_before(&upper).count();
2571        assert_eq!(drained, usize::try_from(num_ts).expect("fits"));
2572
2573        // Advancing the since and consolidating there exercises the peel, split, and merge paths,
2574        // each of which retires and mints chains.
2575        v2.advance_since(Antichain::from_elem(Timestamp::from(num_ts / 2)));
2576        v2.consolidate_at_since();
2577        v2.insert(&mut vec![(
2578            "late".to_owned(),
2579            Timestamp::from(num_ts),
2580            Diff::ONE,
2581        )]);
2582
2583        let mut chains: Vec<&Chain<String>> = vec![&v2.emitted];
2584        chains.extend(&v2.pending_low);
2585        for bucket in v2.chain.buckets() {
2586            chains.extend(&bucket.chains);
2587        }
2588
2589        let mut size = v2.stage.bytes;
2590        let mut records = v2.stage.data.len();
2591        let mut allocations = usize::from(v2.stage.data.capacity() > 0);
2592        for chain in chains {
2593            size += chain.chunks.iter().map(Chunk::size).sum::<usize>();
2594            records += chain.chunks.iter().map(Chunk::len).sum::<usize>();
2595            allocations += chain.chunks.len();
2596        }
2597
2598        assert!(size > 0, "workload must leave chunks behind");
2599        assert_eq!(v2.prev_size.size, size);
2600        assert_eq!(v2.prev_size.capacity, size);
2601        assert_eq!(v2.prev_size.allocations, allocations);
2602        assert_eq!(v2.prev_update_count, records);
2603    }
2604
2605    #[mz_ore::test]
2606    #[cfg_attr(miri, ignore)] // slow under Miri
2607    fn bucket_merges_deepen_and_reads_reset_depth() {
2608        let sink_metrics = sink_metrics();
2609        let mut v2 = CorrectionV2::<String>::new(
2610            sink_metrics.clone(),
2611            sink_metrics.for_worker(0),
2612            None,
2613            3.0,
2614            8 * 1024,
2615        );
2616
2617        // Many stage flushes at one time land in one bucket, whose chain merges deepen.
2618        let mut i = 0;
2619        let deepest = loop {
2620            v2.insert(&mut vec![(
2621                format!("{i:08}"),
2622                Timestamp::from(0),
2623                Diff::ONE,
2624            )]);
2625            i += 1;
2626            let chains: Vec<&Chain<String>> = v2.chain.buckets().flat_map(|b| &b.chains).collect();
2627            for chain in &chains {
2628                let max = chain.chunks.iter().map(|c| c.depth).max().unwrap_or(0);
2629                assert_eq!(chain.depth, max, "a chain's depth is its deepest chunk's");
2630            }
2631            let deepest = chains.iter().map(|c| c.depth).max().unwrap_or(0);
2632            if deepest >= 2 && chains.len() >= 2 {
2633                break deepest;
2634            }
2635            assert!(i < 1_000_000, "bucket merges never deepened");
2636        };
2637        assert!(deepest >= 2);
2638
2639        let upper = Antichain::from_elem(Timestamp::from(1));
2640        let read = v2.updates_before(&upper).count();
2641        assert_eq!(read, i);
2642        assert!(
2643            v2.emitted.chunks.iter().all(|c| c.depth == 0),
2644            "a read writes the youngest generation",
2645        );
2646    }
2647
2648    /// Every chain announced to introspection is retired by the time the buffer drops, the
2649    /// stage's included, whether or not the stage still holds updates.
2650    #[mz_ore::test]
2651    #[cfg_attr(miri, ignore)] // slow under Miri
2652    fn logged_chains_balance_on_drop() {
2653        use crate::sink::correction::LoggingEvent;
2654
2655        for staged in [false, true] {
2656            let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2657            let sink_metrics = sink_metrics();
2658            let mut v2 = CorrectionV2::<String>::new(
2659                sink_metrics.clone(),
2660                sink_metrics.for_worker(0),
2661                Some(ChannelLogging::new(tx)),
2662                3.0,
2663                // One byte ships every insert, one mebibyte ships none of them.
2664                if staged { 1 << 20 } else { 1 },
2665            );
2666            for t in 0..100_u64 {
2667                v2.insert(&mut vec![(format!("{t}"), Timestamp::from(t), Diff::ONE)]);
2668            }
2669            assert_eq!(v2.stage.data.is_empty(), !staged);
2670            drop(v2);
2671
2672            let (mut batches, mut records) = (0_isize, 0_isize);
2673            while let Ok(event) = rx.try_recv() {
2674                match event {
2675                    LoggingEvent::ChainCreated(n) => {
2676                        batches += 1;
2677                        records += isize::try_from(n).expect("fits");
2678                    }
2679                    LoggingEvent::ChainDropped(n) => {
2680                        batches -= 1;
2681                        records -= isize::try_from(n).expect("fits");
2682                    }
2683                    _ => {}
2684                }
2685            }
2686            assert_eq!((batches, records), (0, 0), "staged={staged}");
2687        }
2688    }
2689
2690    /// Reads must not observe updates at or beyond their `upper`, even when the `upper` is not
2691    /// beyond the `since`.
2692    #[mz_ore::test]
2693    // Columnation regions are not Stacked Borrows compliant: later pushes invalidate the
2694    // provenance of previously stored items under Miri.
2695    #[cfg_attr(miri, ignore)]
2696    fn upper_not_beyond_since() {
2697        let sink_metrics = sink_metrics();
2698        let mut v2 = CorrectionV2::<String>::new(
2699            sink_metrics.clone(),
2700            sink_metrics.for_worker(0),
2701            None,
2702            3.0,
2703            8 * 1024,
2704        );
2705
2706        v2.insert(&mut vec![(
2707            "a".to_owned(),
2708            Timestamp::from(5_u64),
2709            Diff::ONE,
2710        )]);
2711        v2.advance_since(Antichain::from_elem(Timestamp::from(10_u64)));
2712
2713        // The update logically lives at time 10 now, so a read before 7 must be empty.
2714        let upper = Antichain::from_elem(Timestamp::from(7_u64));
2715        assert_eq!(v2.updates_before(&upper).count(), 0);
2716
2717        // A read before 11 must emit it, advanced to the since.
2718        let upper = Antichain::from_elem(Timestamp::from(11_u64));
2719        let out: Vec<_> = v2.updates_before(&upper).collect();
2720        assert_eq!(
2721            out,
2722            vec![("a".to_owned(), Timestamp::from(10_u64), Diff::ONE)]
2723        );
2724    }
2725
2726    /// One pool shared by every test in the module. A pool reserves a large slab of address
2727    /// space, so one per test exhausts the VM map under parallel test threads.
2728    fn test_pool() -> mz_ore::pool::Pool {
2729        static POOL: std::sync::OnceLock<mz_ore::pool::Pool> = std::sync::OnceLock::new();
2730        POOL.get_or_init(|| mz_ore::pool::Pool::new().expect("pool reservation"))
2731            .clone()
2732    }
2733
2734    /// Route this thread's chunk spills through the test pool for the duration of `f`.
2735    ///
2736    /// Without an override no pool is installed and every chunk stays resident, so the tests
2737    /// would not exercise [`Chunk::body`]'s read-back at all. The override is thread-scoped,
2738    /// so concurrently running tests do not race on it.
2739    fn with_spill_pool<R>(f: impl FnOnce() -> R) -> R {
2740        chunk::set_spill_override(Some(test_pool()));
2741        let result = f();
2742        chunk::set_spill_override(None);
2743        result
2744    }
2745
2746    /// Assert every chunk in the chain went to the pool, so a read must fetch it back.
2747    fn assert_all_spilled<D: Data>(chain: &Chain<D>) {
2748        assert!(
2749            chain
2750                .chunks
2751                .iter()
2752                .all(|c| c.pooled.lock().expect("not poisoned").is_some()),
2753            "expected every chunk to spill, got {:?}",
2754            chain.chunks,
2755        );
2756    }
2757
2758    /// Build a chain crossing the mint boundary while every chunk spills, then assert `iter()`
2759    /// (the read path behind `updates_before`) reads each chunk back, roundtrips values, order,
2760    /// and diffs, and leaves every body in the pool.
2761    #[mz_ore::test]
2762    #[cfg_attr(miri, ignore)] // the pool's mapped regions are unsupported under miri
2763    fn iter_roundtrips_through_pool() {
2764        let count = 200_000_u64;
2765        with_spill_pool(|| {
2766            let mut builder = ChainBuilder::<i64>::default();
2767            for i in 0..count {
2768                let d = i64::try_from(i).expect("fits");
2769                builder.push_owned(&(d, Timestamp::new(i), Diff::ONE));
2770            }
2771            let chain = builder.finish();
2772            assert!(chain.chunks.len() > 1, "expected multiple minted chunks");
2773            assert_eq!(chain.update_count, usize::try_from(count).expect("fits"));
2774            assert_all_spilled(&chain);
2775
2776            let mut expected = 0_u64;
2777            for (d, t, r) in chain.iter() {
2778                assert_eq!(d, i64::try_from(expected).expect("fits"));
2779                assert_eq!(t, Timestamp::new(expected));
2780                assert_eq!(r, Diff::ONE);
2781                expected += 1;
2782            }
2783            assert_eq!(expected, count);
2784            // The read is scoped: every body is still in the pool and none was materialized.
2785            assert_all_spilled(&chain);
2786            assert!(chain.chunks.iter().all(|c| c.resident.get().is_none()));
2787        });
2788    }
2789
2790    /// Drive a [`Cursor`] over a spilled, multi-chunk chain to completion (the access pattern
2791    /// merges use). Each step reads the front chunk back via [`Chunk::body`]; assert the
2792    /// cursor yields every update in order.
2793    #[mz_ore::test]
2794    #[cfg_attr(miri, ignore)] // the pool's mapped regions are unsupported under miri
2795    fn cursor_steps_through_pool() {
2796        let count = 200_000_u64;
2797        with_spill_pool(|| {
2798            let mut builder = ChainBuilder::<i64>::default();
2799            for i in 0..count {
2800                let d = i64::try_from(i).expect("fits");
2801                builder.push_owned(&(d, Timestamp::new(i), Diff::ONE));
2802            }
2803            let chain = builder.finish();
2804            assert!(chain.chunks.len() > 1, "expected multiple minted chunks");
2805            assert_all_spilled(&chain);
2806
2807            let mut rest = chain.into_cursor();
2808            let mut expected = 0_u64;
2809            while let Some(cursor) = rest.take() {
2810                let (d, t, r) = cursor.get();
2811                assert_eq!(i64::into_owned(d), i64::try_from(expected).expect("fits"));
2812                assert_eq!(t, Timestamp::new(expected));
2813                assert_eq!(r, Diff::ONE);
2814                expected += 1;
2815                rest = cursor.step();
2816            }
2817            assert_eq!(expected, count);
2818        });
2819    }
2820}