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(), ®istry);
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}