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