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 enough updates have been staged to fill a
91//! `Chunk`, 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 — the number of sub-chains grows with the number of distinct times,
143//! so we instead materialize the affected updates, advance their times, and sort and consolidate
144//! them in one O(U log U) pass.
145
146use std::cmp::Ordering;
147use std::collections::{BinaryHeap, VecDeque};
148use std::fmt;
149use std::rc::Rc;
150use std::sync::{Mutex, OnceLock};
151
152use columnar::{Columnar, Index, Len, Ref};
153use mz_ore::cast::CastLossy;
154use mz_ore::soft_assert_or_log;
155use mz_persist_client::metrics::{SinkMetrics, SinkWorkerMetrics, UpdateDelta};
156use mz_repr::{Diff, Timestamp};
157use mz_timely_util::column_pager::{self, PagedColumn};
158use mz_timely_util::columnar::Column;
159use mz_timely_util::temporal::{Bucket, BucketChain};
160use timely::PartialOrder;
161use timely::dataflow::channels::ContainerBytes;
162use timely::progress::Antichain;
163
164use crate::sink::correction::{ChannelLogging, SizeMetrics};
165
166/// Convenient alias for use in data trait bounds.
167///
168/// `D` is constrained to be `Columnar`, so that updates can be stored in a single columnar
169/// region per chunk, and the variable-length payload (e.g. `Row` bytes) lives in the same
170/// allocation as the rest of the chunk. The `Ref`-level `Eq + Ord` bounds let the merge/heap
171/// code compare updates directly through the columnar borrow, avoiding `into_owned` clones
172/// on the hot path.
173pub trait Data:
174    differential_dataflow::Data
175    + Columnar<Container: Send + Sync + Clone + for<'a> columnar::Borrow<Ref<'a>: Eq + Ord>>
176    + Send
177    + Sync
178{
179}
180impl<D> Data for D where
181    D: differential_dataflow::Data
182        + Columnar<Container: Send + Sync + Clone + for<'a> columnar::Borrow<Ref<'a>: Eq + Ord>>
183        + Send
184        + Sync
185{
186}
187
188/// A data structure used to store corrections in the MV sink implementation.
189///
190/// In contrast to `CorrectionV1`, this implementation stores updates in columnation regions,
191/// allowing their memory to be transparently spilled to disk.
192#[derive(Debug)]
193pub struct CorrectionV2<D: Data> {
194    /// Bucketed storage for updates at times at or beyond `boundary`.
195    ///
196    /// Buckets cover exponentially growing time ranges, so reads only touch the buckets below
197    /// their `upper`, and far-future updates (e.g. retractions produced by temporal filters) are
198    /// rarely touched.
199    chain: BucketChain<ChainBucket<D>>,
200    /// Chains at times below `boundary` that were not yet emitted.
201    ///
202    /// Filled by inserts at times below the boundary (mostly persist feedback) and by the
203    /// remainders of `emitted` when a read uses a smaller `upper` than the previous one. Merged
204    /// into `emitted` by the next read.
205    pending_low: Vec<Chain<D>>,
206    /// Updates that were emitted by `updates_before` but not yet cancelled by persist feedback.
207    ///
208    /// Sorted and consolidated, with all times advanced to the `since`.
209    emitted: Chain<D>,
210    /// A staging area for updates, to speed up small inserts.
211    stage: Stage<D>,
212    /// The lower bound of times stored in `chain`. Only ever advances.
213    ///
214    /// Times below the boundary have been peeled off the bucket chain and can only be stored in
215    /// `pending_low` or `emitted`.
216    boundary: Antichain<Timestamp>,
217    /// The frontier by which all contained times are advanced.
218    since: Antichain<Timestamp>,
219
220    /// Total count of updates in the correction buffer.
221    ///
222    /// Tracked to compute deltas in `update_metrics`.
223    prev_update_count: usize,
224    /// Total heap size used by the correction buffer.
225    ///
226    /// Tracked to compute deltas in `update_metrics`.
227    prev_size: SizeMetrics,
228    /// Global persist sink metrics.
229    metrics: SinkMetrics,
230    /// Per-worker persist sink metrics.
231    worker_metrics: SinkWorkerMetrics,
232    /// Introspection logging.
233    logging: Option<ChannelLogging>,
234}
235
236/// Fuel for restoring the bucket chain invariant after peeling.
237///
238/// Bounds the restoration work per buffer operation. The bucket chain remains functional when
239/// restoration is incomplete -- peeling and finding work on ill-formed chains, at the cost of
240/// more in-line splitting -- so leftover restoration is simply picked up by the next operation.
241///
242/// `restore` spends one unit of fuel per bucket split, and a single `peel` leaves at most
243/// `BucketTimestamp::DOMAIN` (64) buckets to re-split, so this budget completes restoration in one
244/// call for any realistic buffer. It is deliberately generous: the "incomplete restoration is
245/// picked up next op" path is a correctness safety net for pathological bucket counts, not a hot
246/// path we expect to exercise. Lower it if restoration ever needs to interleave with other work.
247const RESTORE_FUEL: i64 = 1_000_000;
248
249impl<D: Data> CorrectionV2<D> {
250    /// Construct a new [`CorrectionV2`] instance.
251    pub fn new(
252        metrics: SinkMetrics,
253        worker_metrics: SinkWorkerMetrics,
254        logging: Option<ChannelLogging>,
255        chain_proportionality: f64,
256        chunk_size: usize,
257    ) -> Self {
258        let update_size = std::mem::size_of::<(D, Timestamp, Diff)>();
259        let chunk_capacity = std::cmp::max(chunk_size / update_size, 1);
260
261        Self {
262            chain: BucketChain::new(ChainBucket::new(chain_proportionality, logging.clone())),
263            pending_low: Vec::new(),
264            emitted: Chain::new(),
265            stage: Stage::new(logging.clone(), chunk_capacity),
266            boundary: Antichain::from_elem(Timestamp::MIN),
267            since: Antichain::from_elem(Timestamp::MIN),
268            prev_update_count: 0,
269            prev_size: Default::default(),
270            metrics,
271            worker_metrics,
272            logging,
273        }
274    }
275
276    /// Insert a batch of updates.
277    pub fn insert(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) {
278        let Some(since_ts) = self.since.as_option() else {
279            // If the since is the empty frontier, discard all updates.
280            updates.clear();
281            return;
282        };
283
284        for (_, time, _) in &mut *updates {
285            *time = std::cmp::max(*time, *since_ts);
286        }
287
288        self.insert_inner(updates);
289    }
290
291    /// Insert a batch of updates, after negating their diffs.
292    pub fn insert_negated(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) {
293        let Some(since_ts) = self.since.as_option() else {
294            // If the since is the empty frontier, discard all updates.
295            updates.clear();
296            return;
297        };
298
299        for (_, time, diff) in &mut *updates {
300            *time = std::cmp::max(*time, *since_ts);
301            *diff = -*diff;
302        }
303
304        self.insert_inner(updates);
305    }
306
307    /// Insert a batch of updates into the stage, flushing it when full.
308    ///
309    /// All times are expected to be >= the `since`.
310    fn insert_inner(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) {
311        debug_assert!(updates.iter().all(|(_, t, _)| self.since.less_equal(t)));
312
313        if let Some(mut ready) = self.stage.insert(updates) {
314            self.route(&mut ready);
315        }
316
317        self.update_metrics();
318    }
319
320    /// Route a batch of sorted, consolidated updates to `pending_low` or their chain buckets.
321    fn route(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) {
322        // Updates at times below the boundary become a pending low chain.
323        let idx = updates.partition_point(|(_, t, _)| !self.boundary.less_equal(t));
324        if idx > 0 {
325            let mut builder = ChainBuilder::default();
326            builder.extend(updates.drain(..idx));
327            let chain = builder.finish();
328            if !chain.is_empty() {
329                self.log_chain_created(&chain);
330                self.pending_low.push(chain);
331            }
332        }
333
334        // Updates at times at or beyond the boundary go into their chain buckets. Walk ranges of
335        // times that fall into the same bucket, to push batches of updates at once.
336        let mut drain = updates.drain(..).peekable();
337        while let Some(update) = drain.next() {
338            let time = update.1;
339            let range = self
340                .chain
341                .range_of(&time)
342                .expect("bucket chain covers all times at or beyond the boundary");
343            let mut builder = ChainBuilder::default();
344            builder.extend(std::iter::once(update));
345            while let Some(update) = drain.next_if(|(_, t, _)| range.contains(t)) {
346                builder.extend(std::iter::once(update));
347            }
348            let bucket = self
349                .chain
350                .find_mut(&range.start)
351                .expect("bucket chain covers all times at or beyond the boundary");
352            bucket.push_chain(builder.finish());
353        }
354    }
355
356    /// Return consolidated updates before the given `upper`.
357    pub fn updates_before<'a>(
358        &'a mut self,
359        upper: &Antichain<Timestamp>,
360    ) -> impl Iterator<Item = (D, Timestamp, Diff)> + Send + 'a {
361        self.consolidate_before(upper);
362        self.consolidated_updates_before(upper)
363    }
364
365    /// Return the updates before the given `upper`, as consolidated by a preceding
366    /// [`CorrectionV2::consolidate_before`] call.
367    ///
368    /// The caller must have invoked `consolidate_before` with the same `upper` and must not have
369    /// mutated the buffer since. Otherwise the returned updates are neither consolidated nor
370    /// necessarily complete.
371    pub fn consolidated_updates_before<'a>(
372        &'a self,
373        upper: &Antichain<Timestamp>,
374    ) -> impl Iterator<Item = (D, Timestamp, Diff)> + Send + use<'a, D> {
375        // All contained times are advanced to at least the `since`, so a read at an `upper` that
376        // is not beyond the `since` is always empty. This mirrors the short-circuit in
377        // `consolidate_before`, which leaves `emitted` untouched in that case.
378        if !PartialOrder::less_than(&self.since, upper) {
379            return None.into_iter().flatten();
380        }
381
382        // After `consolidate_before`, `emitted` holds exactly the updates before `upper`: every
383        // path that populates it splits at `upper` (pushing the remainder to `pending_low`), and
384        // the guard above guarantees `upper > since`, so advancing stale times to the `since`
385        // cannot lift them to or beyond `upper`. We can therefore yield all of `emitted`. Guard
386        // the invariant: a violation would write updates beyond the batch upper to persist.
387        soft_assert_or_log!(
388            self.emitted
389                .last()
390                .is_none_or(|(_, t, _)| !upper.less_equal(&t)),
391            "emitted contains times at or beyond the upper",
392        );
393        Some(self.emitted.iter()).into_iter().flatten()
394    }
395
396    /// Consolidate all updates before the given `upper` into the `emitted` chain.
397    ///
398    /// Once this method returns, `emitted` contains all updates at times before `upper`,
399    /// consolidated.
400    ///
401    /// Does nothing if `upper` is not beyond the `since`: all contained times are advanced to at
402    /// least the `since`, so such a read is empty anyway, and skipping avoids an eager peel,
403    /// merge, and `boundary` advancement. Normal reads and `consolidate_at_since` always pass an
404    /// `upper` beyond the `since`.
405    pub fn consolidate_before(&mut self, upper: &Antichain<Timestamp>) {
406        if !PartialOrder::less_than(&self.since, upper) {
407            return;
408        }
409
410        if let Some(mut ready) = self.stage.flush() {
411            self.route(&mut ready);
412        }
413
414        let Some(&since_ts) = self.since.as_option() else {
415            // If the since is the empty frontier, discard all updates.
416            let peeled = self.chain.peel(Antichain::new().borrow());
417            for bucket in peeled {
418                for chain in bucket.into_chains() {
419                    self.log_chain_dropped(&chain);
420                }
421            }
422            for chain in std::mem::take(&mut self.pending_low) {
423                self.log_chain_dropped(&chain);
424            }
425            let emitted = std::mem::replace(&mut self.emitted, Chain::new());
426            if !emitted.is_empty() {
427                self.log_chain_dropped(&emitted);
428            }
429            self.update_metrics();
430            return;
431        };
432
433        // Peel the buckets below the upper off the bucket chain. Bucket splits during the peel
434        // only touch chunks around the upper; chunks wholly on either side are reused.
435        let peeled = self.chain.peel(upper.borrow());
436        if PartialOrder::less_than(&self.boundary, upper) {
437            self.boundary = upper.clone();
438        }
439
440        // Collect candidate chains: peeled bucket contents, pending low chains, and the previous
441        // emitted chain. All contain only times below the boundary.
442        let emitted = std::mem::replace(&mut self.emitted, Chain::new());
443        let mut candidates: Vec<Chain<D>> = Vec::new();
444        for bucket in peeled {
445            candidates.extend(bucket.into_chains());
446        }
447        candidates.append(&mut self.pending_low);
448        if !emitted.is_empty() {
449            candidates.push(emitted);
450        }
451
452        if candidates.is_empty() {
453            self.restore_chain();
454            self.update_metrics();
455            return;
456        }
457
458        candidates.iter().for_each(|c| self.log_chain_dropped(c));
459
460        // Split the candidates at the upper. Parts at or beyond the upper (possible when `upper`
461        // regresses below a previous one) stay pending.
462        let mut lowers = Vec::new();
463        for chain in candidates {
464            match upper.as_option() {
465                Some(&upper_ts) => {
466                    let (lower, remainder) = chain.split_at_time(upper_ts);
467                    if !lower.is_empty() {
468                        lowers.push(lower);
469                    }
470                    if !remainder.is_empty() {
471                        self.log_chain_created(&remainder);
472                        self.pending_low.push(remainder);
473                    }
474                }
475                // The empty upper is greater than all times.
476                None => lowers.push(chain),
477            }
478        }
479
480        // Merge the lower parts into the new emitted chain, advancing times below the since.
481        // Advancing times in a (time, data)-sorted chain can break its sort order, so chains
482        // containing stale times cannot be merged as they are. Stale times are expected in steady
483        // state: the previous emitted chain was written before the since advanced past it.
484        //
485        // Count the distinct stale times, up to a small cap. For few distinct stale times -- the
486        // steady state -- split cursors into runs that remain sorted under advancement and merge
487        // those. For many distinct stale times -- e.g. a since jump across many buffered
488        // timestamps when a sink restarts with an old as-of -- the number of runs and the cost of
489        // cloning cursor state per run grow with the number of distinct times, so materialize,
490        // advance, and consolidate in one O(U log U) pass instead.
491        const MAX_STALE_RUNS: usize = 32;
492        let mut stale_times = 0;
493        for chain in &lowers {
494            stale_times += chain.distinct_times_before(since_ts, MAX_STALE_RUNS - stale_times);
495            if stale_times >= MAX_STALE_RUNS {
496                break;
497            }
498        }
499
500        let merged = if stale_times == 0 {
501            let cursors: Vec<_> = lowers.into_iter().filter_map(Chain::into_cursor).collect();
502            merge_cursors(cursors)
503        } else if stale_times < MAX_STALE_RUNS {
504            let mut runs = Vec::new();
505            for chain in lowers {
506                if let Some(cursor) = chain.into_cursor() {
507                    runs.append(&mut cursor.advance_by(since_ts));
508                }
509            }
510            merge_cursors(runs)
511        } else {
512            let mut updates: Vec<_> = lowers.iter().flat_map(|c| c.iter()).collect();
513            for (_, time, _) in &mut updates {
514                *time = std::cmp::max(*time, since_ts);
515            }
516            consolidate(&mut updates);
517            let mut builder = ChainBuilder::default();
518            builder.extend(updates);
519            let chain = builder.finish();
520
521            // Advancement can move updates to or beyond the upper; such updates stay pending.
522            match upper.as_option() {
523                Some(&upper_ts) => {
524                    let (lower, remainder) = chain.split_at_time(upper_ts);
525                    if !remainder.is_empty() {
526                        self.log_chain_created(&remainder);
527                        self.pending_low.push(remainder);
528                    }
529                    lower
530                }
531                None => chain,
532            }
533        };
534
535        if !merged.is_empty() {
536            self.log_chain_created(&merged);
537        }
538        self.emitted = merged;
539
540        self.restore_chain();
541        self.update_metrics();
542    }
543
544    /// Perform a bounded amount of work towards restoring the bucket chain invariant.
545    ///
546    /// Restoration is allowed to remain incomplete: the bucket chain supports peeling and finding
547    /// on ill-formed chains, so any leftover work is picked up by subsequent operations. The fuel
548    /// bound keeps individual buffer operations from stalling the operator that owns the buffer.
549    fn restore_chain(&mut self) {
550        let mut fuel = RESTORE_FUEL;
551        self.chain.restore(&mut fuel);
552    }
553
554    /// Advance the since frontier.
555    ///
556    /// Time advancement of updates in the bucket chain is lazy: it happens when the updates are
557    /// consolidated by a read.
558    ///
559    /// # Panics
560    ///
561    /// Panics if the given `since` is less than the current since frontier.
562    pub fn advance_since(&mut self, since: Antichain<Timestamp>) {
563        assert!(PartialOrder::less_equal(&self.since, &since));
564        self.stage.advance_times(&since);
565        self.since = since;
566    }
567
568    /// Consolidate all updates at the current `since`.
569    pub fn consolidate_at_since(&mut self) {
570        let upper_ts = self.since.as_option().and_then(|t| t.try_step_forward());
571        if let Some(upper_ts) = upper_ts {
572            let upper = Antichain::from_elem(upper_ts);
573            self.consolidate_before(&upper);
574        }
575    }
576
577    fn log_chain_created(&self, chain: &Chain<D>) {
578        if let Some(logging) = &self.logging {
579            logging.chain_created(chain.update_count);
580        }
581    }
582
583    fn log_chain_dropped(&self, chain: &Chain<D>) {
584        if let Some(logging) = &self.logging {
585            logging.chain_dropped(chain.update_count);
586        }
587    }
588
589    /// Update persist sink metrics.
590    fn update_metrics(&mut self) {
591        let mut new_size = self.stage.get_size();
592        let mut new_length = self.stage.data.len();
593        for chain in &self.pending_low {
594            new_size += chain.get_size();
595            new_length += chain.update_count;
596        }
597        new_size += self.emitted.get_size();
598        new_length += self.emitted.update_count;
599        for bucket in self.chain.buckets() {
600            for chain in &bucket.chains {
601                new_size += chain.get_size();
602                new_length += chain.update_count;
603            }
604        }
605
606        self.update_metrics_inner(new_size, new_length);
607    }
608
609    /// Update persist sink metrics to the given new size and length.
610    fn update_metrics_inner(&mut self, new_size: SizeMetrics, new_length: usize) {
611        let old_size = self.prev_size;
612        let old_length = self.prev_update_count;
613        let len_delta = UpdateDelta::new(new_length, old_length);
614        let cap_delta = UpdateDelta::new(new_size.capacity, old_size.capacity);
615        self.metrics
616            .report_correction_update_deltas(len_delta, cap_delta);
617        self.worker_metrics
618            .report_correction_update_totals(new_length, new_size.capacity);
619
620        if let Some(logging) = &self.logging {
621            let i = |x: usize| isize::try_from(x).expect("must fit");
622            logging.report_size_diff(i(new_size.size) - i(old_size.size));
623            logging.report_capacity_diff(i(new_size.capacity) - i(old_size.capacity));
624            logging.report_allocations_diff(i(new_size.allocations) - i(old_size.allocations));
625        }
626
627        self.prev_size = new_size;
628        self.prev_update_count = new_length;
629    }
630}
631
632/// Merge the given cursors into one chain.
633fn merge_cursors<D: Data>(cursors: Vec<Cursor<D>>) -> Chain<D> {
634    match cursors.len() {
635        0 => Chain::new(),
636        1 => {
637            let [cur] = cursors.try_into().unwrap();
638            cur.into_chain()
639        }
640        2 => {
641            let [a, b] = cursors.try_into().unwrap();
642            merge_2(a, b)
643        }
644        _ => merge_many(cursors),
645    }
646}
647
648/// Merge the given two cursors using a 2-way merge.
649///
650/// This function is a specialization of `merge_many` that avoids the overhead of a binary heap.
651fn merge_2<D: Data>(cursor1: Cursor<D>, cursor2: Cursor<D>) -> Chain<D> {
652    let mut rest1 = Some(cursor1);
653    let mut rest2 = Some(cursor2);
654    let mut merged = ChainBuilder::default();
655
656    loop {
657        match (rest1, rest2) {
658            (Some(c1), Some(c2)) => {
659                let (d1, t1, r1) = c1.get();
660                let (d2, t2, r2) = c2.get();
661
662                match (t1, d1).cmp(&(t2, d2)) {
663                    Ordering::Less => {
664                        merged.push_ref((d1, t1, r1));
665                        rest1 = c1.step();
666                        rest2 = Some(c2);
667                    }
668                    Ordering::Greater => {
669                        merged.push_ref((d2, t2, r2));
670                        rest1 = Some(c1);
671                        rest2 = c2.step();
672                    }
673                    Ordering::Equal => {
674                        let r = r1 + r2;
675                        if r != Diff::ZERO {
676                            merged.push_ref((d1, t1, r));
677                        }
678                        rest1 = c1.step();
679                        rest2 = c2.step();
680                    }
681                }
682            }
683            (Some(c), None) | (None, Some(c)) => {
684                merged.push_cursor(c);
685                break;
686            }
687            (None, None) => break,
688        }
689    }
690
691    merged.finish()
692}
693
694/// Merge the given cursors using a k-way merge with a binary heap.
695fn merge_many<D: Data>(cursors: Vec<Cursor<D>>) -> Chain<D> {
696    let mut heap = MergeHeap::from_iter(cursors);
697    let mut merged = ChainBuilder::default();
698    while let Some(cursor1) = heap.pop() {
699        let (data, time, mut diff) = cursor1.get();
700
701        while let Some((cursor2, r)) = heap.pop_equal(data, time) {
702            diff += r;
703            if let Some(cursor2) = cursor2.step() {
704                heap.push(cursor2);
705            }
706        }
707
708        if diff != Diff::ZERO {
709            merged.push_ref((data, time, diff));
710        }
711        if let Some(cursor1) = cursor1.step() {
712            heap.push(cursor1);
713        }
714    }
715
716    merged.finish()
717}
718
719impl<D: Data> Drop for CorrectionV2<D> {
720    fn drop(&mut self) {
721        for bucket in self.chain.buckets() {
722            bucket.chains.iter().for_each(|c| self.log_chain_dropped(c));
723        }
724        self.pending_low
725            .iter()
726            .for_each(|c| self.log_chain_dropped(c));
727        if !self.emitted.is_empty() {
728            self.log_chain_dropped(&self.emitted);
729        }
730        self.update_metrics_inner(Default::default(), 0);
731    }
732}
733
734/// A bucket of `Chain`s, for use in a [`BucketChain`].
735///
736/// All chains are individually sorted by (time, data) and consolidated, but updates can appear in
737/// multiple chains, so consumers must merge the chains to obtain consolidated updates.
738struct ChainBucket<D: Data> {
739    /// The contained chains.
740    ///
741    /// Maintained with the chain invariant on pushes; splits can leave it violated until the next
742    /// push restores it.
743    chains: Vec<Chain<D>>,
744    /// The size factor of subsequent chains required by the chain invariant.
745    chain_proportionality: f64,
746    /// Introspection logging.
747    logging: Option<ChannelLogging>,
748}
749
750impl<D: Data> fmt::Debug for ChainBucket<D> {
751    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
752        f.debug_struct("ChainBucket")
753            .field("chains", &self.chains)
754            .finish_non_exhaustive()
755    }
756}
757
758impl<D: Data> ChainBucket<D> {
759    /// Construct a new, empty `ChainBucket`.
760    fn new(chain_proportionality: f64, logging: Option<ChannelLogging>) -> Self {
761        Self {
762            chains: Vec::new(),
763            chain_proportionality,
764            logging,
765        }
766    }
767
768    /// Push a chain onto the bucket, restoring the chain invariant.
769    fn push_chain(&mut self, chain: Chain<D>) {
770        if chain.is_empty() {
771            return;
772        }
773        if let Some(logging) = &self.logging {
774            logging.chain_created(chain.update_count);
775        }
776        self.chains.push(chain);
777
778        // Restore the chain invariant.
779        let prop = self.chain_proportionality;
780        let merge_needed = |chains: &[Chain<_>]| match chains {
781            [.., prev, last] => {
782                let last_len = f64::cast_lossy(last.update_count);
783                let prev_len = f64::cast_lossy(prev.update_count);
784                last_len * prop > prev_len
785            }
786            _ => false,
787        };
788
789        while merge_needed(&self.chains) {
790            let a = self.chains.pop().unwrap();
791            let b = self.chains.pop().unwrap();
792            if let Some(logging) = &self.logging {
793                logging.chain_dropped(a.update_count);
794                logging.chain_dropped(b.update_count);
795            }
796
797            let cursors = [a, b].into_iter().filter_map(Chain::into_cursor).collect();
798            let merged = merge_cursors(cursors);
799            if !merged.is_empty() {
800                if let Some(logging) = &self.logging {
801                    logging.chain_created(merged.update_count);
802                }
803                self.chains.push(merged);
804            }
805        }
806    }
807
808    /// Convert the bucket into its contained chains.
809    fn into_chains(self) -> Vec<Chain<D>> {
810        self.chains
811    }
812}
813
814impl<D: Data> Bucket for ChainBucket<D> {
815    type Timestamp = Timestamp;
816
817    fn split(self, timestamp: &Self::Timestamp, fuel: &mut i64) -> (Self, Self) {
818        let mut lower = Self::new(self.chain_proportionality, self.logging.clone());
819        let mut upper = Self::new(self.chain_proportionality, self.logging.clone());
820
821        for chain in self.chains {
822            // Whole chunks are reused; at most one chunk straddling the timestamp is copied per
823            // chain. Account fuel at chunk granularity.
824            *fuel = fuel.saturating_sub(i64::try_from(chain.chunks.len()).expect("must fit"));
825
826            if let Some(logging) = &self.logging {
827                logging.chain_dropped(chain.update_count);
828            }
829            let (lo, hi) = chain.split_at_time(*timestamp);
830            for (part, target) in [(lo, &mut lower), (hi, &mut upper)] {
831                if !part.is_empty() {
832                    if let Some(logging) = &self.logging {
833                        logging.chain_created(part.update_count);
834                    }
835                    target.chains.push(part);
836                }
837            }
838        }
839
840        (lower, upper)
841    }
842}
843
844/// A chain of [`Chunk`]s containing updates.
845///
846/// All updates in a chain are sorted by (time, data) and consolidated.
847///
848/// Note that, in contrast to [`Chunk`]s, chains can be empty. Though we generally try to avoid
849/// keeping around empty chains.
850#[derive(Debug)]
851struct Chain<D: Data> {
852    /// The contained chunks.
853    chunks: Vec<Chunk<D>>,
854    /// The number of updates contained in all chunks.
855    update_count: usize,
856}
857
858impl<D: Data> Chain<D> {
859    /// Construct an empty chain.
860    fn new() -> Self {
861        Self {
862            chunks: Default::default(),
863            update_count: 0,
864        }
865    }
866
867    /// Return whether the chain is empty.
868    fn is_empty(&self) -> bool {
869        self.chunks.is_empty()
870    }
871
872    /// Push a chunk onto the chain.
873    ///
874    /// All updates in the chunk must sort after all updates already in the chain, in
875    /// (time, data)-order, to ensure the chain remains sorted.
876    fn push_chunk(&mut self, chunk: Chunk<D>) {
877        mz_ore::soft_assert_no_log!(self.can_accept_chunk(&chunk));
878
879        self.update_count += chunk.len();
880        self.chunks.push(chunk);
881    }
882
883    /// Return whether the chain can accept the given chunk at its end while preserving
884    /// (time, data)-order.
885    ///
886    /// NOTE: The cached boundary times settle every case but a tie. On a tie the boundary updates
887    /// themselves are compared, which materializes both chunks and keeps them resident for the
888    /// rest of their lifetime. The only caller is the soft assertion in [`Chain::push_chunk`], and
889    /// soft assertions are live in any build started with `MZ_SOFT_ASSERTIONS` set, so this cost is
890    /// not confined to debug builds. Ties are reached whenever a run of updates at a single
891    /// timestamp spans a chunk boundary, which [`ChunkBuilder`] produces for any such run larger
892    /// than its byte limit.
893    fn can_accept_chunk(&self, chunk: &Chunk<D>) -> bool {
894        match self.chunks.last() {
895            None => true,
896            Some(last) => match last.last_time().cmp(&chunk.first_time()) {
897                Ordering::Less => true,
898                Ordering::Greater => false,
899                Ordering::Equal => {
900                    let (dc, _, _) = last.last();
901                    let (d, _, _) = chunk.first();
902                    dc < d
903                }
904            },
905        }
906    }
907
908    /// Return the last update in the chain, if any.
909    fn last(&self) -> Option<Ref<'_, (D, Timestamp, Diff)>> {
910        self.chunks.last().map(|c| c.last())
911    }
912
913    /// Convert the chain into a cursor over the contained updates.
914    fn into_cursor(self) -> Option<Cursor<D>> {
915        let chunks = self.chunks.into_iter().map(Rc::new).collect();
916        Cursor::new(chunks)
917    }
918
919    /// Return an iterator over the contained updates.
920    fn iter(&self) -> impl Iterator<Item = (D, Timestamp, Diff)> + '_ {
921        self.chunks.iter().flat_map(|c| {
922            (0..c.len()).map(move |i| {
923                let (d, t, r) = c.index(i);
924                (D::into_owned(d), t, r)
925            })
926        })
927    }
928
929    /// Count the distinct times of updates at times before `time`, up to the given cap.
930    ///
931    /// The scan uses one binary search per distinct time, so its cost is bounded by
932    /// O(cap log chunks).
933    fn distinct_times_before(&self, time: Timestamp, cap: usize) -> usize {
934        let mut count = 0;
935        let mut chunk_idx = 0;
936        let mut offset = 0;
937        while count < cap && chunk_idx < self.chunks.len() {
938            let chunk = &self.chunks[chunk_idx];
939            let current = chunk.index(offset).1;
940            if current >= time {
941                break;
942            }
943            count += 1;
944            // Skip to the first update at a time greater than `current`.
945            match chunk.find_time_greater_than(current) {
946                Some(idx) => offset = idx,
947                None => {
948                    // All later updates at `current` are in subsequent chunks.
949                    chunk_idx += 1;
950                    offset = 0;
951                    while chunk_idx < self.chunks.len() {
952                        match self.chunks[chunk_idx].find_time_greater_than(current) {
953                            Some(idx) => {
954                                offset = idx;
955                                break;
956                            }
957                            None => chunk_idx += 1,
958                        }
959                    }
960                }
961            }
962        }
963        count
964    }
965
966    /// Split the chain at the given time.
967    ///
968    /// Returns two chains, the first containing all updates at times < `time`, the second
969    /// containing all updates at times >= `time`. Chunks fully on either side of `time` are
970    /// reused; only a chunk straddling `time` is copied.
971    fn split_at_time(mut self, time: Timestamp) -> (Self, Self) {
972        let mut lower = Self::new();
973        let mut upper = Self::new();
974
975        let Some(skip_ts) = time.step_back() else {
976            // Nothing sorts before `time`.
977            return (lower, self);
978        };
979
980        for chunk in self.chunks.drain(..) {
981            // Route whole chunks by cached boundary times, so a chunk that lands entirely on one
982            // side is moved without paging it in. Only a straddling chunk is materialized here.
983            // With soft assertions on, `push_chunk` can still page in a chunk whose boundary time
984            // ties the chain's last one, see `Chain::can_accept_chunk`.
985            if chunk.last_time() < time {
986                lower.push_chunk(chunk);
987            } else if chunk.first_time() >= time {
988                upper.push_chunk(chunk);
989            } else {
990                // The chunk straddles `time`; copy its two halves.
991                let idx = chunk
992                    .find_time_greater_than(skip_ts)
993                    .expect("straddles time");
994                let mut builder = ChainBuilder::default();
995                for i in 0..idx {
996                    builder.push_ref(chunk.index(i));
997                }
998                for part in builder.finish().chunks {
999                    lower.push_chunk(part);
1000                }
1001                let mut builder = ChainBuilder::default();
1002                for i in idx..chunk.len() {
1003                    builder.push_ref(chunk.index(i));
1004                }
1005                for part in builder.finish().chunks {
1006                    upper.push_chunk(part);
1007                }
1008            }
1009        }
1010
1011        (lower, upper)
1012    }
1013
1014    /// Return the size of the chain, for use in metrics.
1015    fn get_size(&self) -> SizeMetrics {
1016        let mut metrics = SizeMetrics::default();
1017        for chunk in &self.chunks {
1018            metrics += chunk.get_size();
1019        }
1020        metrics
1021    }
1022}
1023
1024/// A builder that constructs a [`Chain`] from a stream of updates.
1025///
1026/// Wraps a [`ChunkBuilder`] and drains its minted chunks into a [`Chain`]. Pushed updates must
1027/// arrive in (time, data) sorted order.
1028struct ChainBuilder<D: Data> {
1029    builder: ChunkBuilder<D>,
1030    chain: Chain<D>,
1031}
1032
1033impl<D: Data> Default for ChainBuilder<D> {
1034    fn default() -> Self {
1035        Self {
1036            builder: Default::default(),
1037            chain: Chain::new(),
1038        }
1039    }
1040}
1041
1042impl<D: Data> ChainBuilder<D> {
1043    /// Push a reference-form update into the builder.
1044    fn push_ref(&mut self, update: Ref<'_, (D, Timestamp, Diff)>) {
1045        self.builder.push(update);
1046        self.drain();
1047    }
1048
1049    /// Push an owned-form update into the builder.
1050    fn push_owned(&mut self, update: &(D, Timestamp, Diff)) {
1051        self.builder.push(update);
1052        self.drain();
1053    }
1054
1055    /// Push the updates produced by a cursor into the builder.
1056    fn push_cursor(&mut self, cursor: Cursor<D>) {
1057        let mut rest = Some(cursor);
1058        while let Some(cursor) = rest.take() {
1059            let update = cursor.get();
1060            self.push_ref(update);
1061            rest = cursor.step();
1062        }
1063    }
1064
1065    /// Move any minted chunks from the builder into the chain.
1066    fn drain(&mut self) {
1067        while let Some(chunk) = self.builder.pop() {
1068            self.chain.push_chunk(chunk);
1069        }
1070    }
1071
1072    /// Finish building, returning the assembled [`Chain`].
1073    fn finish(self) -> Chain<D> {
1074        let Self { builder, mut chain } = self;
1075        for chunk in builder.finish() {
1076            if chunk.len() > 0 {
1077                chain.push_chunk(chunk);
1078            }
1079        }
1080        chain
1081    }
1082}
1083
1084impl<D: Data> Extend<(D, Timestamp, Diff)> for ChainBuilder<D> {
1085    fn extend<I: IntoIterator<Item = (D, Timestamp, Diff)>>(&mut self, iter: I) {
1086        for update in iter {
1087            self.push_owned(&update);
1088        }
1089    }
1090}
1091
1092/// A cursor over updates in a chain.
1093///
1094/// A cursor provides two guarantees:
1095///  * Produced updates are ordered and consolidated.
1096///  * A cursor always yields at least one update.
1097///
1098/// The second guarantee is enforced through the type system: Every method that steps a cursor
1099/// forward consumes `self` and returns an `Option<Cursor>` that's `None` if the operation stepped
1100/// over the last update.
1101///
1102/// A cursor holds on to `Rc<Chunk>`s, allowing multiple cursors to produce updates from the same
1103/// chunks concurrently. As soon as a cursor is done producing updates from a [`Chunk`] it drops
1104/// its reference. Once the last cursor is done with a [`Chunk`] its memory can be reclaimed.
1105#[derive(Clone, Debug)]
1106struct Cursor<D: Data> {
1107    /// The chunks from which updates can still be produced.
1108    chunks: VecDeque<Rc<Chunk<D>>>,
1109    /// The current offset into `chunks.front()`.
1110    chunk_offset: usize,
1111    /// An optional limit for the number of updates the cursor will produce.
1112    limit: Option<usize>,
1113    /// An optional overwrite for the timestamp of produced updates.
1114    overwrite_ts: Option<Timestamp>,
1115}
1116
1117impl<D: Data> Cursor<D> {
1118    /// Construct a cursor over a list of chunks.
1119    ///
1120    /// Returns `None` if `chunks` is empty.
1121    fn new(chunks: VecDeque<Rc<Chunk<D>>>) -> Option<Self> {
1122        if chunks.is_empty() {
1123            return None;
1124        }
1125
1126        Some(Self {
1127            chunks,
1128            chunk_offset: 0,
1129            limit: None,
1130            overwrite_ts: None,
1131        })
1132    }
1133
1134    /// Set a limit for the number of updates this cursor will produce.
1135    ///
1136    /// # Panics
1137    ///
1138    /// Panics if there is already a limit lower than the new one.
1139    fn set_limit(mut self, limit: usize) -> Option<Self> {
1140        assert!(self.limit.is_none_or(|l| l >= limit));
1141
1142        if limit == 0 {
1143            return None;
1144        }
1145
1146        // Release chunks made unreachable by the limit.
1147        let mut count = 0;
1148        let mut idx = 0;
1149        let mut offset = self.chunk_offset;
1150        while idx < self.chunks.len() && count < limit {
1151            let chunk = &self.chunks[idx];
1152            count += chunk.len() - offset;
1153            idx += 1;
1154            offset = 0;
1155        }
1156        self.chunks.truncate(idx);
1157
1158        if count > limit {
1159            self.limit = Some(limit);
1160        }
1161
1162        Some(self)
1163    }
1164
1165    /// Get a reference to the current update.
1166    fn get(&self) -> Ref<'_, (D, Timestamp, Diff)> {
1167        let chunk = self.get_chunk();
1168        let (d, t, r) = chunk.index(self.chunk_offset);
1169        let t = self.overwrite_ts.unwrap_or(t);
1170        (d, t, r)
1171    }
1172
1173    /// Get a reference to the current chunk.
1174    fn get_chunk(&self) -> &Chunk<D> {
1175        &self.chunks[0]
1176    }
1177
1178    /// Step to the next update.
1179    ///
1180    /// Returns the stepped cursor, or `None` if the step was over the last update.
1181    fn step(mut self) -> Option<Self> {
1182        if self.chunk_offset == self.get_chunk().len() - 1 {
1183            return self.skip_chunk().map(|(c, _)| c);
1184        }
1185
1186        self.chunk_offset += 1;
1187
1188        if let Some(limit) = &mut self.limit {
1189            *limit -= 1;
1190            if *limit == 0 {
1191                return None;
1192            }
1193        }
1194
1195        Some(self)
1196    }
1197
1198    /// Skip the remainder of the current chunk.
1199    ///
1200    /// Returns the forwarded cursor and the number of updates skipped, or `None` if no chunks are
1201    /// left after the skip.
1202    fn skip_chunk(mut self) -> Option<(Self, usize)> {
1203        let chunk = self.chunks.pop_front().expect("cursor invariant");
1204
1205        if self.chunks.is_empty() {
1206            return None;
1207        }
1208
1209        let skipped = chunk.len() - self.chunk_offset;
1210        self.chunk_offset = 0;
1211
1212        if let Some(limit) = &mut self.limit {
1213            if skipped >= *limit {
1214                return None;
1215            }
1216            *limit -= skipped;
1217        }
1218
1219        Some((self, skipped))
1220    }
1221
1222    /// Skip all updates with times <= the given time.
1223    ///
1224    /// Returns the forwarded cursor and the number of updates skipped, or `None` if no updates are
1225    /// left after the skip.
1226    fn skip_time(mut self, time: Timestamp) -> Option<(Self, usize)> {
1227        if self.overwrite_ts.is_some_and(|ts| ts <= time) {
1228            return None;
1229        } else if self.get().1 > time {
1230            return Some((self, 0));
1231        }
1232
1233        let mut skipped = 0;
1234
1235        let new_offset = loop {
1236            let chunk = self.get_chunk();
1237            if let Some(index) = chunk.find_time_greater_than(time) {
1238                break index;
1239            }
1240
1241            let (cursor, count) = self.skip_chunk()?;
1242            self = cursor;
1243            skipped += count;
1244        };
1245
1246        skipped += new_offset - self.chunk_offset;
1247        self.chunk_offset = new_offset;
1248
1249        Some((self, skipped))
1250    }
1251
1252    /// Advance all updates in this cursor by the given `since_ts`.
1253    ///
1254    /// Returns a list of cursors, each of which yields ordered and consolidated updates that have
1255    /// been advanced by `since_ts`.
1256    fn advance_by(mut self, since_ts: Timestamp) -> Vec<Self> {
1257        // If the cursor has an `overwrite_ts`, all its updates are at the same time already. We
1258        // only need to advance the `overwrite_ts` by the `since_ts`.
1259        if let Some(ts) = self.overwrite_ts {
1260            if ts < since_ts {
1261                self.overwrite_ts = Some(since_ts);
1262            }
1263            return vec![self];
1264        }
1265
1266        // Otherwise we need to split the cursor so that each new cursor only yields runs of
1267        // updates that are correctly (time, data)-ordered when advanced by `since_ts`. We achieve
1268        // this by splitting the cursor at each time <= `since_ts`.
1269        let mut splits = Vec::new();
1270        let mut remaining = Some(self);
1271
1272        while let Some(cursor) = remaining.take() {
1273            let (_, time, _) = cursor.get();
1274            if time >= since_ts {
1275                splits.push(cursor);
1276                break;
1277            }
1278
1279            let mut current = cursor.clone();
1280            if let Some((cursor, skipped)) = cursor.skip_time(time) {
1281                remaining = Some(cursor);
1282                current = current.set_limit(skipped).expect("skipped at least 1");
1283            }
1284            current.overwrite_ts = Some(since_ts);
1285            splits.push(current);
1286        }
1287
1288        splits
1289    }
1290
1291    /// Drain the cursor into a [`Chain`].
1292    ///
1293    /// This reuses the underlying chunks if possible, and writes new ones otherwise.
1294    fn into_chain(self) -> Chain<D> {
1295        match self.try_unwrap() {
1296            Ok(chain) => chain,
1297            Err((_, cursor)) => {
1298                let mut builder = ChainBuilder::default();
1299                builder.push_cursor(cursor);
1300                builder.finish()
1301            }
1302        }
1303    }
1304
1305    /// Attempt to unwrap the cursor into a [`Chain`].
1306    ///
1307    /// This operation efficiently reuses chunks by directly inserting them into the output chain
1308    /// where possible.
1309    ///
1310    /// An unwrap is only successful if the cursor's `limit` and `overwrite_ts` are both `None` and
1311    /// the cursor has unique references to its chunks. If the unwrap fails, this method returns an
1312    /// `Err` containing the cursor in an unchanged state, allowing the caller to convert it into a
1313    /// chain by copying chunks rather than reusing them.
1314    fn try_unwrap(self) -> Result<Chain<D>, (&'static str, Self)> {
1315        if self.limit.is_some() {
1316            return Err(("cursor with limit", self));
1317        }
1318        if self.overwrite_ts.is_some() {
1319            return Err(("cursor with overwrite_ts", self));
1320        }
1321        if self.chunks.iter().any(|c| Rc::strong_count(c) != 1) {
1322            return Err(("cursor on shared chunks", self));
1323        }
1324
1325        let mut builder = ChainBuilder::default();
1326        let mut remaining = Some(self);
1327
1328        // We might be partway through the first chunk, in which case we can't reuse it but need to
1329        // allocate a new one to contain only the updates the cursor can still yield.
1330        while let Some(cursor) = remaining.take() {
1331            if cursor.chunk_offset == 0 {
1332                remaining = Some(cursor);
1333                break;
1334            }
1335            let update = cursor.get();
1336            builder.push_ref(update);
1337            remaining = cursor.step();
1338        }
1339
1340        let mut chain = builder.finish();
1341        if let Some(cursor) = remaining {
1342            for chunk in cursor.chunks {
1343                let chunk = Rc::into_inner(chunk).expect("checked above");
1344                chain.push_chunk(chunk);
1345            }
1346        }
1347
1348        Ok(chain)
1349    }
1350}
1351
1352/// A non-empty chunk of updates, backed by a columnar region.
1353///
1354/// All updates in a chunk are sorted by (time, data) and consolidated.
1355///
1356/// Chunks are immutable once created. They are produced by [`ChunkBuilder`], which mints a
1357/// new chunk whenever its in-progress columnar container reaches a fixed serialized byte
1358/// boundary (~2 MiB, matching the ship granularity used elsewhere in the codebase), so each
1359/// chunk corresponds to a single, predictably sized allocation.
1360struct Chunk<D: Data> {
1361    /// The paged-out form, taken on first materialization.
1362    ///
1363    /// A `Mutex` (not `RefCell`) keeps the chunk `Sync`: cursors hold chunks behind a shared
1364    /// `Rc`, and the iterator returned by [`CorrectionV2::updates_before`] borrows them across
1365    /// the persist writer's `await`, so `&Chunk` must be `Send`. The lock is taken once, at
1366    /// materialization, and is otherwise uncontended (the sink runs single-threaded per worker).
1367    paged: Mutex<Option<PagedColumn<(D, Timestamp, Diff)>>>,
1368    /// The materialized form, populated lazily by [`Chunk::column`] on first access.
1369    ///
1370    /// An `OnceLock` (not `OnceCell`) for the same `Sync` reason. Once set the slot is never
1371    /// cleared, so its address is stable and [`Chunk::index`] can hand out `Ref<'_>` borrows tied
1372    /// to `&self`. The allocation is freed when the chunk drops, which bounds resident memory to
1373    /// the chunks under an active merge front.
1374    resident: OnceLock<Column<(D, Timestamp, Diff)>>,
1375    /// Number of updates, cached so `len` and chain bookkeeping never page the chunk in.
1376    len: usize,
1377    /// Time of the first update, cached so boundary checks (`split_at_time`, `can_accept`) route
1378    /// a resting chunk without materializing it.
1379    first_time: Timestamp,
1380    /// Time of the last update, cached likewise.
1381    last_time: Timestamp,
1382}
1383
1384impl<D: Data> fmt::Debug for Chunk<D> {
1385    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1386        write!(f, "Chunk(<{}>)", self.len())
1387    }
1388}
1389
1390impl<D: Data> Chunk<D> {
1391    /// Page the given non-empty column out into a chunk.
1392    ///
1393    /// Reads the cached metadata (length, boundary times) while the column is still resident, then
1394    /// hands it to the global column pager. The policy decides whether it actually spills; either
1395    /// way the chunk is born paged and materializes lazily on first read.
1396    ///
1397    /// # Panics
1398    ///
1399    /// Panics if the column is empty. Chunks are non-empty by construction; [`ChunkBuilder`] only
1400    /// ever builds a chunk from a populated column.
1401    fn from_column(mut data: Column<(D, Timestamp, Diff)>) -> Self {
1402        let (len, first_time, last_time) = {
1403            let borrowed = data.borrow();
1404            let len = borrowed.len();
1405            assert!(len > 0, "chunks are non-empty");
1406            (len, borrowed.get(0).1, borrowed.get(len - 1).1)
1407        };
1408
1409        let paged = column_pager::global_pager().page(&mut data);
1410        Self {
1411            paged: Mutex::new(Some(paged)),
1412            resident: OnceLock::new(),
1413            len,
1414            first_time,
1415            last_time,
1416        }
1417    }
1418
1419    /// Materialize the chunk's column, paging it in on first access.
1420    ///
1421    /// The returned reference is valid for as long as `&self`: the `OnceLock` slot is never
1422    /// cleared once populated, so its contents have a stable address.
1423    fn column(&self) -> &Column<(D, Timestamp, Diff)> {
1424        self.resident.get_or_init(|| {
1425            let paged = self
1426                .paged
1427                .lock()
1428                .expect("pager mutex poisoned")
1429                .take()
1430                .expect("paged form present until materialized");
1431            column_pager::global_pager().take(paged)
1432        })
1433    }
1434
1435    /// Return the number of updates in the chunk.
1436    fn len(&self) -> usize {
1437        self.len
1438    }
1439
1440    /// Return the update at the given index, paging the chunk in if necessary.
1441    ///
1442    /// # Panics
1443    ///
1444    /// Panics if the given index is not populated.
1445    fn index(&self, idx: usize) -> Ref<'_, (D, Timestamp, Diff)> {
1446        self.column().borrow().get(idx)
1447    }
1448
1449    /// Return the first update in the chunk, paging the chunk in if necessary.
1450    fn first(&self) -> Ref<'_, (D, Timestamp, Diff)> {
1451        self.index(0)
1452    }
1453
1454    /// Return the last update in the chunk, paging the chunk in if necessary.
1455    fn last(&self) -> Ref<'_, (D, Timestamp, Diff)> {
1456        self.index(self.len - 1)
1457    }
1458
1459    /// Return the time of the first update, without materializing the chunk.
1460    fn first_time(&self) -> Timestamp {
1461        self.first_time
1462    }
1463
1464    /// Return the time of the last update, without materializing the chunk.
1465    fn last_time(&self) -> Timestamp {
1466        self.last_time
1467    }
1468
1469    /// Return the index of the first update at a time greater than `time`, or `None` if no such
1470    /// update exists.
1471    ///
1472    /// The early-out uses the cached last time, so a chunk whose updates are all at or before
1473    /// `time` is skipped without paging it in.
1474    fn find_time_greater_than(&self, time: Timestamp) -> Option<usize> {
1475        if self.last_time <= time {
1476            return None;
1477        }
1478
1479        let mut lower = 0;
1480        let mut upper = self.len;
1481        while lower < upper {
1482            let idx = (lower + upper) / 2;
1483            if self.index(idx).1 > time {
1484                upper = idx;
1485            } else {
1486                lower = idx + 1;
1487            }
1488        }
1489
1490        Some(lower)
1491    }
1492
1493    /// Return the size of the chunk, for use in metrics.
1494    ///
1495    /// Reports resident bytes only: a chunk still spilled (on swap or in a pager file) is not part
1496    /// of RSS and contributes nothing, matching the accounting in
1497    /// [`mz_timely_util::columnar::merge_batcher`].
1498    fn get_size(&self) -> SizeMetrics {
1499        let resident = |col: &Column<(D, Timestamp, Diff)>| {
1500            let bytes = col.length_in_bytes();
1501            SizeMetrics {
1502                size: bytes,
1503                capacity: bytes,
1504                allocations: 1,
1505            }
1506        };
1507
1508        if let Some(col) = self.resident.get() {
1509            return resident(col);
1510        }
1511        // Not yet materialized: a policy that kept the column resident still occupies RSS, so
1512        // account for it; a genuinely spilled column does not.
1513        match &*self.paged.lock().expect("pager mutex poisoned") {
1514            Some(PagedColumn::Resident(col, _)) => resident(col),
1515            _ => SizeMetrics::default(),
1516        }
1517    }
1518}
1519
1520/// Builder that produces a stream of fixed-size [`Chunk`]s.
1521///
1522/// Wraps [`mz_timely_util::columnar::builder::ColumnBuilder`], which mints a new
1523/// [`Column::Align`] chunk whenever its in-progress columnar container reaches a fixed
1524/// serialized byte boundary (~2 MiB, matching the ship granularity used elsewhere in the
1525/// codebase). Each minted chunk is therefore a single, predictably-sized aligned allocation.
1526struct ChunkBuilder<D: Data> {
1527    inner: mz_timely_util::columnar::builder::ColumnBuilder<(D, Timestamp, Diff)>,
1528}
1529
1530impl<D: Data> Default for ChunkBuilder<D> {
1531    fn default() -> Self {
1532        Self {
1533            inner: Default::default(),
1534        }
1535    }
1536}
1537
1538impl<D: Data> ChunkBuilder<D> {
1539    /// Push an update into the builder.
1540    ///
1541    /// Accepts whatever the inner [`ColumnBuilder`]'s [`PushInto`] impl accepts — both the
1542    /// `Ref<'_, (D, T, R)>` refs produced by cursors and `&(D, T, R)` references to owned
1543    /// tuples drained from the staging buffer.
1544    ///
1545    /// [`ColumnBuilder`]: mz_timely_util::columnar::builder::ColumnBuilder
1546    /// [`PushInto`]: timely::container::PushInto
1547    #[inline]
1548    fn push<T>(&mut self, item: T)
1549    where
1550        mz_timely_util::columnar::builder::ColumnBuilder<(D, Timestamp, Diff)>:
1551            timely::container::PushInto<T>,
1552    {
1553        timely::container::PushInto::push_into(&mut self.inner, item);
1554    }
1555
1556    /// Pop a finished chunk, if one is available.
1557    fn pop(&mut self) -> Option<Chunk<D>> {
1558        use timely::container::ContainerBuilder;
1559        // `ColumnBuilder::extract` stashes the popped chunk in its `finished` slot so the
1560        // caller can read it through `&mut`; move it out with `mem::take` so we own it
1561        // (leaves `Column::Typed(Default::default())` behind, which the next `extract`
1562        // overwrites).
1563        self.inner
1564            .extract()
1565            .map(|c| Chunk::from_column(std::mem::take(c)))
1566    }
1567
1568    /// Finalize the builder: flush any in-progress updates as a typed chunk and drain pending.
1569    fn finish(mut self) -> impl Iterator<Item = Chunk<D>> {
1570        use timely::container::ContainerBuilder;
1571        // `ColumnBuilder::finish` flushes the in-progress container into the pending queue
1572        // (as `Column::Typed`) and returns the first pending entry. Subsequent calls drain
1573        // the rest until `None`. Translate that into an owning iterator.
1574        //
1575        // `finish` can hand back an empty column (e.g. when the last shipped chunk landed exactly
1576        // on the boundary). Skip those: `Chunk::from_column` requires a non-empty column, and an
1577        // empty chunk would needlessly engage the pager.
1578        std::iter::from_fn(move || {
1579            loop {
1580                let col = std::mem::take(self.inner.finish()?);
1581                if !col.is_empty() {
1582                    return Some(Chunk::from_column(col));
1583                }
1584            }
1585        })
1586    }
1587}
1588
1589/// A buffer for staging updates before they are inserted into the sorted chains.
1590#[derive(Debug)]
1591struct Stage<D> {
1592    /// The contained updates.
1593    ///
1594    /// This vector has a fixed capacity equal to the [`Chunk`] capacity.
1595    data: Vec<(D, Timestamp, Diff)>,
1596    /// Introspection logging.
1597    ///
1598    /// We want to report the number of records in the stage. To do so, we pretend that the stage
1599    /// is a chain, and every time the number of updates inside changes, the chain gets dropped and
1600    /// re-created.
1601    logging: Option<ChannelLogging>,
1602}
1603
1604impl<D: Data> Stage<D> {
1605    fn new(logging: Option<ChannelLogging>, chunk_capacity: usize) -> Self {
1606        // For logging, we pretend the stage consists of a single chain.
1607        if let Some(logging) = &logging {
1608            logging.chain_created(0);
1609        }
1610
1611        Self {
1612            data: Vec::with_capacity(chunk_capacity),
1613            logging,
1614        }
1615    }
1616
1617    /// Insert a batch of updates, possibly producing a batch of sorted, consolidated updates
1618    /// ready to be stored.
1619    fn insert(
1620        &mut self,
1621        updates: &mut Vec<(D, Timestamp, Diff)>,
1622    ) -> Option<Vec<(D, Timestamp, Diff)>> {
1623        if updates.is_empty() {
1624            return None;
1625        }
1626
1627        let prev_length = self.ilen();
1628
1629        // Determine how many chunks we can fill with the available updates.
1630        let update_count = self.data.len() + updates.len();
1631        let chunk_capacity = self.data.capacity();
1632        let chunk_count = update_count / chunk_capacity;
1633
1634        let mut new_updates = updates.drain(..);
1635
1636        // If we have enough shipable updates, collect them and consolidate.
1637        let maybe_ready = if chunk_count > 0 {
1638            let ship_count = chunk_count * chunk_capacity;
1639            let mut buffer = Vec::with_capacity(ship_count);
1640
1641            buffer.append(&mut self.data);
1642            while buffer.len() < ship_count {
1643                let update = new_updates.next().unwrap();
1644                buffer.push(update);
1645            }
1646
1647            consolidate(&mut buffer);
1648
1649            Some(buffer)
1650        } else {
1651            None
1652        };
1653
1654        // Stage the remaining updates.
1655        Extend::extend(&mut self.data, new_updates);
1656
1657        self.log_length_diff(self.ilen() - prev_length);
1658
1659        maybe_ready
1660    }
1661
1662    /// Flush all currently staged updates, returning them sorted and consolidated.
1663    fn flush(&mut self) -> Option<Vec<(D, Timestamp, Diff)>> {
1664        self.log_length_diff(-self.ilen());
1665
1666        consolidate(&mut self.data);
1667
1668        if self.data.is_empty() {
1669            return None;
1670        }
1671
1672        let capacity = self.data.capacity();
1673        let data = std::mem::replace(&mut self.data, Vec::with_capacity(capacity));
1674        Some(data)
1675    }
1676
1677    /// Advance the times of staged updates by the given `since`.
1678    fn advance_times(&mut self, since: &Antichain<Timestamp>) {
1679        let Some(since_ts) = since.as_option() else {
1680            // If the since is the empty frontier, discard all updates.
1681            self.log_length_diff(-self.ilen());
1682            self.data.clear();
1683            return;
1684        };
1685
1686        for (_, time, _) in &mut self.data {
1687            *time = std::cmp::max(*time, *since_ts);
1688        }
1689    }
1690
1691    /// Return the size of the stage, for use in metrics.
1692    ///
1693    /// Note: We don't follow pointers here, so the returned `size` and `capacity` values are
1694    /// under-estimates. That's fine as the stage should always be small.
1695    fn get_size(&self) -> SizeMetrics {
1696        SizeMetrics {
1697            size: self.data.len() * std::mem::size_of::<(D, Timestamp, Diff)>(),
1698            capacity: self.data.capacity() * std::mem::size_of::<(D, Timestamp, Diff)>(),
1699            allocations: 1,
1700        }
1701    }
1702
1703    /// Return the number of updates in the stage, as an `isize`.
1704    fn ilen(&self) -> isize {
1705        self.data.len().try_into().expect("must fit")
1706    }
1707
1708    fn log_length_diff(&self, diff: isize) {
1709        let Some(logging) = &self.logging else { return };
1710        if diff > 0 {
1711            let count = usize::try_from(diff).expect("must fit");
1712            logging.chain_created(count);
1713            logging.chain_dropped(0);
1714        } else if diff < 0 {
1715            let count = usize::try_from(-diff).expect("must fit");
1716            logging.chain_created(0);
1717            logging.chain_dropped(count);
1718        }
1719    }
1720}
1721
1722impl<D> Drop for Stage<D> {
1723    fn drop(&mut self) {
1724        if let Some(logging) = &self.logging {
1725            logging.chain_dropped(self.data.len());
1726        }
1727    }
1728}
1729
1730/// Sort and consolidate the given list of updates.
1731///
1732/// This function is the same as [`differential_dataflow::consolidation::consolidate_updates`],
1733/// except that it sorts updates by (time, data) instead of (data, time).
1734fn consolidate<D: Data>(updates: &mut Vec<(D, Timestamp, Diff)>) {
1735    if updates.len() <= 1 {
1736        return;
1737    }
1738
1739    let diff = |update: &(_, _, Diff)| update.2;
1740
1741    updates.sort_unstable_by(|(d1, t1, _), (d2, t2, _)| (t1, d1).cmp(&(t2, d2)));
1742
1743    let mut offset = 0;
1744    let mut accum = diff(&updates[0]);
1745
1746    for idx in 1..updates.len() {
1747        let this = &updates[idx];
1748        let prev = &updates[idx - 1];
1749        if this.0 == prev.0 && this.1 == prev.1 {
1750            accum += diff(&updates[idx]);
1751        } else {
1752            if accum != Diff::ZERO {
1753                updates.swap(offset, idx - 1);
1754                updates[offset].2 = accum;
1755                offset += 1;
1756            }
1757            accum = diff(&updates[idx]);
1758        }
1759    }
1760
1761    if accum != Diff::ZERO {
1762        let len = updates.len();
1763        updates.swap(offset, len - 1);
1764        updates[offset].2 = accum;
1765        offset += 1;
1766    }
1767
1768    updates.truncate(offset);
1769}
1770
1771/// Compare two columnar refs that have unrelated input lifetimes.
1772///
1773/// `<D::Container as Borrow>::Ref<'a>` is an associated-type projection through a trait, so
1774/// the compiler treats it as invariant in `'a` and won't auto-shorten the inputs by variance.
1775/// We instead explicitly reborrow both to a fresh, local lifetime `'x` via
1776/// [`Columnar::reborrow`] before letting the inner `==` pick up the `for<'a> Ref<'a>: Eq`
1777/// bound on [`Data`].
1778#[inline]
1779fn refs_eq<D: Data>(a: Ref<'_, D>, b: Ref<'_, D>) -> bool {
1780    #[inline]
1781    fn eq<'x, D: Data>(a: Ref<'x, D>, b: Ref<'x, D>) -> bool {
1782        a == b
1783    }
1784    eq::<D>(D::reborrow(a), D::reborrow(b))
1785}
1786
1787/// A binary heap specialized for merging [`Cursor`]s.
1788struct MergeHeap<D: Data>(BinaryHeap<MergeCursor<D>>);
1789
1790impl<D: Data> FromIterator<Cursor<D>> for MergeHeap<D> {
1791    fn from_iter<I: IntoIterator<Item = Cursor<D>>>(cursors: I) -> Self {
1792        let inner = cursors.into_iter().map(MergeCursor).collect();
1793        Self(inner)
1794    }
1795}
1796
1797impl<D: Data> MergeHeap<D> {
1798    /// Pop the next cursor (the one yielding the least update) from the heap.
1799    fn pop(&mut self) -> Option<Cursor<D>> {
1800        self.0.pop().map(|MergeCursor(c)| c)
1801    }
1802
1803    /// Pop the next cursor from the heap, provided the data and time of its current update are
1804    /// equal to the given values.
1805    ///
1806    /// Returns both the cursor and the diff corresponding to `data` and `time`.
1807    fn pop_equal(&mut self, data: Ref<'_, D>, time: Timestamp) -> Option<(Cursor<D>, Diff)> {
1808        let r = {
1809            let MergeCursor(cursor) = self.0.peek()?;
1810            let (d, t, r) = cursor.get();
1811            if t != time || !refs_eq::<D>(d, data) {
1812                return None;
1813            }
1814            r
1815        };
1816        let cursor = self.pop().expect("checked above");
1817        Some((cursor, r))
1818    }
1819
1820    /// Push a cursor onto the heap.
1821    fn push(&mut self, cursor: Cursor<D>) {
1822        self.0.push(MergeCursor(cursor));
1823    }
1824}
1825
1826/// A wrapper for [`Cursor`]s on a [`MergeHeap`].
1827///
1828/// Implements the cursor ordering required for merging cursors.
1829struct MergeCursor<D: Data>(Cursor<D>);
1830
1831impl<D: Data> PartialEq for MergeCursor<D> {
1832    fn eq(&self, other: &Self) -> bool {
1833        self.cmp(other).is_eq()
1834    }
1835}
1836
1837impl<D: Data> Eq for MergeCursor<D> {}
1838
1839impl<D: Data> PartialOrd for MergeCursor<D> {
1840    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1841        Some(self.cmp(other))
1842    }
1843}
1844
1845impl<D: Data> Ord for MergeCursor<D> {
1846    fn cmp(&self, other: &Self) -> Ordering {
1847        let (d1, t1, _) = self.0.get();
1848        let (d2, t2, _) = other.0.get();
1849        (t1, d1).cmp(&(t2, d2)).reverse()
1850    }
1851}
1852
1853#[cfg(test)]
1854mod tests {
1855    use mz_ore::metrics::MetricsRegistry;
1856    use mz_persist_client::cfg::PersistConfig;
1857    use mz_persist_client::metrics::Metrics;
1858    use mz_repr::{Diff, Timestamp};
1859
1860    use super::*;
1861    use crate::sink::correction::CorrectionV1;
1862
1863    #[mz_ore::test]
1864    fn chain_builder_update_count_matches_items() {
1865        let mut builder = ChainBuilder::<i64>::default();
1866        for i in 0..10_u64 {
1867            let d = i64::try_from(i).expect("fits");
1868            builder.push_owned(&(d, Timestamp::new(i), Diff::ONE));
1869        }
1870        let chain = builder.finish();
1871        assert_eq!(chain.update_count, chain.iter().count());
1872    }
1873
1874    /// Push enough updates to cross at least one `mint()` boundary, forcing the
1875    /// `Align` encode -> `from_bytes` decode roundtrip (the spilling path this data
1876    /// structure exists to support), and assert `iter()` roundtrips values, order,
1877    /// and diffs across the spill boundary.
1878    #[mz_ore::test]
1879    #[cfg_attr(miri, ignore)] // too slow: crossing the ~2 MiB mint boundary needs ~200k updates
1880    fn chain_builder_roundtrips_across_mint_boundary() {
1881        // A single `mint()` fires near the ~2 MiB (`SHIP_WORDS`) serialized boundary. With
1882        // three 8-byte columns per update that's tens of thousands of updates; pushing 200k
1883        // comfortably forces multiple mints.
1884        let count = 200_000_u64;
1885
1886        let mut builder = ChainBuilder::<i64>::default();
1887        for i in 0..count {
1888            let d = i64::try_from(i).expect("fits");
1889            builder.push_owned(&(d, Timestamp::new(i), Diff::ONE));
1890        }
1891        let chain = builder.finish();
1892
1893        // Crossing the mint boundary must have produced more than one chunk; otherwise the spill
1894        // path (each minted chunk is paged out and read back through the pager) wouldn't be
1895        // exercised. The chunk payload itself is now behind the pager (see [`Chunk`]), so we
1896        // assert on chunk count rather than inspecting the column variant directly.
1897        assert!(
1898            chain.chunks.len() > 1,
1899            "expected multiple minted chunks, got {} chunk(s): {:?}",
1900            chain.chunks.len(),
1901            chain.chunks,
1902        );
1903
1904        // `iter()` must roundtrip every update, in order, with correct diffs.
1905        assert_eq!(chain.update_count, usize::try_from(count).expect("fits"));
1906        let mut expected = 0_u64;
1907        for (d, t, r) in chain.iter() {
1908            assert_eq!(d, i64::try_from(expected).expect("fits"));
1909            assert_eq!(t, Timestamp::new(expected));
1910            assert_eq!(r, Diff::ONE);
1911            expected += 1;
1912        }
1913        assert_eq!(expected, count);
1914    }
1915
1916    fn sink_metrics() -> SinkMetrics {
1917        let registry = MetricsRegistry::new();
1918        let metrics = Metrics::new(&PersistConfig::new_for_tests(), &registry);
1919        metrics.sink.clone()
1920    }
1921
1922    /// Run the same stepwise-drain workload through `CorrectionV1` and `CorrectionV2` and assert
1923    /// that they emit the same updates at every step.
1924    ///
1925    /// Models the `write_batches` operator catching up through many distinct timestamps: the
1926    /// desired input runs ahead, batches are written one timestamp at a time, and written updates
1927    /// come back negated through the persist feedback.
1928    #[mz_ore::test]
1929    // Columnation regions are not Stacked Borrows compliant: later pushes invalidate the
1930    // provenance of previously stored items under Miri.
1931    #[cfg_attr(miri, ignore)]
1932    fn equivalence_with_v1() {
1933        let sink_metrics = sink_metrics();
1934
1935        let mut v1 =
1936            CorrectionV1::<String>::new(sink_metrics.clone(), sink_metrics.for_worker(0), 1);
1937        let mut v2 = CorrectionV2::<String>::new(
1938            sink_metrics.clone(),
1939            sink_metrics.for_worker(0),
1940            None,
1941            3.0,
1942            8 * 1024,
1943        );
1944
1945        let num_ts = 50;
1946        let keys = 4;
1947
1948        // Upsert-style input: every timestamp updates each key, retracting the previous value.
1949        let batch = |t: u64| -> Vec<(String, Timestamp, Diff)> {
1950            (0..keys)
1951                .flat_map(|k| {
1952                    let addition = (format!("{k}-{t}"), Timestamp::from(t), Diff::ONE);
1953                    let retraction = t
1954                        .checked_sub(1)
1955                        .map(|p| (format!("{k}-{p}"), Timestamp::from(t), -Diff::ONE));
1956                    std::iter::once(addition).chain(retraction)
1957                })
1958                .collect()
1959        };
1960
1961        // Pre-fill both with all batches, like a catch-up where the input runs ahead.
1962        for t in 0..num_ts {
1963            v1.insert(&mut batch(t));
1964            v2.insert(&mut batch(t));
1965        }
1966
1967        // Drain stepwise, with persist feedback, comparing emissions.
1968        for t in 0..num_ts {
1969            let upper = Antichain::from_elem(Timestamp::from(t + 1));
1970
1971            let mut out1: Vec<_> = v1.updates_before(&upper).collect();
1972            let mut out2: Vec<_> = v2.updates_before(&upper).collect();
1973            out1.sort();
1974            out2.sort();
1975            assert_eq!(out1, out2, "diverged at t={t}");
1976
1977            v1.insert_negated(&mut out1.clone());
1978            v2.insert_negated(&mut out2);
1979            v1.advance_since(upper.clone());
1980            v2.advance_since(upper);
1981        }
1982
1983        // Compare the final state at the since.
1984        let upper = Antichain::from_elem(Timestamp::from(num_ts + 1));
1985        v1.consolidate_at_since();
1986        v2.consolidate_at_since();
1987        let mut out1: Vec<_> = v1.updates_before(&upper).collect();
1988        let mut out2: Vec<_> = v2.updates_before(&upper).collect();
1989        out1.sort();
1990        out2.sort();
1991        assert_eq!(out1, out2);
1992    }
1993
1994    /// A since jump across many distinct buffered timestamps must collapse them onto the since.
1995    #[mz_ore::test]
1996    // Columnation regions are not Stacked Borrows compliant: later pushes invalidate the
1997    // provenance of previously stored items under Miri.
1998    #[cfg_attr(miri, ignore)]
1999    fn since_jump() {
2000        let sink_metrics = sink_metrics();
2001        let mut v2 = CorrectionV2::<String>::new(
2002            sink_metrics.clone(),
2003            sink_metrics.for_worker(0),
2004            None,
2005            3.0,
2006            8 * 1024,
2007        );
2008
2009        let num_ts = 100;
2010        for t in 0..num_ts {
2011            v2.insert(&mut vec![
2012                (format!("a-{t}"), Timestamp::from(t), Diff::ONE),
2013                (format!("a-{t}"), Timestamp::from(t), -Diff::ONE),
2014                (format!("b-{t}"), Timestamp::from(t), Diff::ONE),
2015            ]);
2016        }
2017
2018        v2.advance_since(Antichain::from_elem(Timestamp::from(num_ts)));
2019        v2.consolidate_at_since();
2020
2021        let upper = Antichain::from_elem(Timestamp::from(num_ts + 1));
2022        let out: Vec<_> = v2.updates_before(&upper).collect();
2023        assert_eq!(out.len(), usize::try_from(num_ts).unwrap());
2024        assert!(
2025            out.iter()
2026                .all(|(_, t, r)| *t == Timestamp::from(num_ts) && *r == Diff::ONE)
2027        );
2028    }
2029
2030    /// Reads must not observe updates at or beyond their `upper`, even when the `upper` is not
2031    /// beyond the `since`.
2032    #[mz_ore::test]
2033    // Columnation regions are not Stacked Borrows compliant: later pushes invalidate the
2034    // provenance of previously stored items under Miri.
2035    #[cfg_attr(miri, ignore)]
2036    fn upper_not_beyond_since() {
2037        let sink_metrics = sink_metrics();
2038        let mut v2 = CorrectionV2::<String>::new(
2039            sink_metrics.clone(),
2040            sink_metrics.for_worker(0),
2041            None,
2042            3.0,
2043            8 * 1024,
2044        );
2045
2046        v2.insert(&mut vec![(
2047            "a".to_owned(),
2048            Timestamp::from(5_u64),
2049            Diff::ONE,
2050        )]);
2051        v2.advance_since(Antichain::from_elem(Timestamp::from(10_u64)));
2052
2053        // The update logically lives at time 10 now, so a read before 7 must be empty.
2054        let upper = Antichain::from_elem(Timestamp::from(7_u64));
2055        assert_eq!(v2.updates_before(&upper).count(), 0);
2056
2057        // A read before 11 must emit it, advanced to the since.
2058        let upper = Antichain::from_elem(Timestamp::from(11_u64));
2059        let out: Vec<_> = v2.updates_before(&upper).collect();
2060        assert_eq!(
2061            out,
2062            vec![("a".to_owned(), Timestamp::from(10_u64), Diff::ONE)]
2063        );
2064    }
2065
2066    /// A [`PagingPolicy`] that always spills to the swap backend, uncompressed.
2067    ///
2068    /// The default global pager keeps every chunk resident; installing this drives the actual
2069    /// spill path so the tests exercise [`Chunk::column`]'s page-in through [`mz_ore::pager`].
2070    ///
2071    /// [`PagingPolicy`]: column_pager::PagingPolicy
2072    struct ForceSwap;
2073
2074    impl column_pager::PagingPolicy for ForceSwap {
2075        fn decide(&self, _hint: column_pager::PageHint) -> column_pager::PageDecision {
2076            column_pager::PageDecision::Page {
2077                backend: mz_ore::pager::Backend::Swap,
2078                codec: None,
2079            }
2080        }
2081        fn record(&self, _event: column_pager::PageEvent) {}
2082    }
2083
2084    /// Install a global pager that spills every chunk to swap for the duration of `f`, then
2085    /// restore the default (disabled) pager. The global pager is process-wide; concurrent tests
2086    /// only ever observe a correct round-trip regardless of backend, so racing on it is benign.
2087    fn with_swap_pager<R>(f: impl FnOnce() -> R) -> R {
2088        use std::sync::Arc;
2089        column_pager::set_global_pager(column_pager::ColumnPager::new(Arc::new(ForceSwap)));
2090        let result = f();
2091        column_pager::set_global_pager(column_pager::ColumnPager::disabled());
2092        result
2093    }
2094
2095    /// Build a chain crossing the mint boundary while every chunk is spilled to swap, then assert
2096    /// `iter()` (the read path behind `updates_before`) pages each chunk back in and roundtrips
2097    /// values, order, and diffs.
2098    #[mz_ore::test]
2099    #[cfg_attr(miri, ignore)] // madvise on the swap backend is unsupported under miri
2100    fn iter_roundtrips_through_swap_backend() {
2101        let count = 200_000_u64;
2102        with_swap_pager(|| {
2103            let mut builder = ChainBuilder::<i64>::default();
2104            for i in 0..count {
2105                let d = i64::try_from(i).expect("fits");
2106                builder.push_owned(&(d, Timestamp::new(i), Diff::ONE));
2107            }
2108            let chain = builder.finish();
2109            assert!(chain.chunks.len() > 1, "expected multiple minted chunks");
2110            assert_eq!(chain.update_count, usize::try_from(count).expect("fits"));
2111
2112            let mut expected = 0_u64;
2113            for (d, t, r) in chain.iter() {
2114                assert_eq!(d, i64::try_from(expected).expect("fits"));
2115                assert_eq!(t, Timestamp::new(expected));
2116                assert_eq!(r, Diff::ONE);
2117                expected += 1;
2118            }
2119            assert_eq!(expected, count);
2120        });
2121    }
2122
2123    /// Drive a [`Cursor`] over a spilled, multi-chunk chain to completion (the access pattern
2124    /// merges use). Each step pages the front chunk back in via [`Chunk::column`]; assert the
2125    /// cursor yields every update in order.
2126    #[mz_ore::test]
2127    #[cfg_attr(miri, ignore)] // madvise on the swap backend is unsupported under miri
2128    fn cursor_steps_through_swap_backend() {
2129        let count = 200_000_u64;
2130        with_swap_pager(|| {
2131            let mut builder = ChainBuilder::<i64>::default();
2132            for i in 0..count {
2133                let d = i64::try_from(i).expect("fits");
2134                builder.push_owned(&(d, Timestamp::new(i), Diff::ONE));
2135            }
2136            let chain = builder.finish();
2137            assert!(chain.chunks.len() > 1, "expected multiple minted chunks");
2138
2139            let mut rest = chain.into_cursor();
2140            let mut expected = 0_u64;
2141            while let Some(cursor) = rest.take() {
2142                let (d, t, r) = cursor.get();
2143                assert_eq!(i64::into_owned(d), i64::try_from(expected).expect("fits"));
2144                assert_eq!(t, Timestamp::new(expected));
2145                assert_eq!(r, Diff::ONE);
2146                expected += 1;
2147                rest = cursor.step();
2148            }
2149            assert_eq!(expected, count);
2150        });
2151    }
2152}