From 210aff9dbc772269112dbb7973bba588536ef345 Mon Sep 17 00:00:00 2001 From: Moritz Hoffmann Date: Wed, 26 Aug 2026 13:14:21 +0200 Subject: [PATCH 1/2] compute: remove the correction v1 buffer The `enable_compute_correction_v2` flag has served `true` in production for every build at or beyond 26.7.0-rc.1, so the v1 correction buffer is dead code in every deployment we support. This change deletes it along with the flag that selected between the two implementations. `CorrectionV2` becomes the only `Correction`, and its file moves to `sink/correction.rs`. The logging and metrics helpers that both implementations shared move to `sink/correction/logging.rs`, which keeps the buffer file from growing further. `Correction::new` now reads the chain proportionality and chunk size from the `ConfigSet`, matching the constructor the sink previously called on the enum wrapper; `Correction::with_params` takes the two values directly for tests and benchmarks. Removing v1 also removes its only consumer of `ConsolidatingVec`, so that type and the `consolidating_vec_growth_dampener` dyncfg that tuned it are gone as well. Both flag keys move to the LaunchDarkly consistency check's stale list, because the last published release still synchronizes them. The `equivalence_with_v1` unit test compared the two implementations step by step. It is replaced by `equivalence_with_reference`, which runs the same upsert-and-feedback workload against a naive in-test buffer that keeps every update in a flat vector, plus an assertion that each step emits something so the comparison cannot pass vacuously. The correction benchmark loses its version-dispatch wrapper and now measures the single implementation. The dyncfg keys `compute_correction_v2_chain_proportionality` and `compute_correction_v2_chunk_size` keep their names. Renaming them would drop the production overrides currently set against those keys. Co-Authored-By: Claude Opus 5 (1M context) --- misc/python/materialize/mzcompose/__init__.py | 2 - .../materialize/parallel_workload/action.py | 2 - src/compute-types/src/dyncfgs.rs | 19 - src/compute/benches/correction.rs | 165 +- src/compute/src/sink.rs | 4 - src/compute/src/sink/correction.rs | 2684 +++++++++++++---- src/compute/src/sink/correction/logging.rs | 283 ++ src/compute/src/sink/correction_v2.rs | 2152 ------------- src/compute/src/sink/materialized_view.rs | 3 +- src/compute/src/sink/materialized_view_v2.rs | 3 +- .../mzcompose.py | 4 +- 11 files changed, 2394 insertions(+), 2927 deletions(-) create mode 100644 src/compute/src/sink/correction/logging.rs delete mode 100644 src/compute/src/sink/correction_v2.rs diff --git a/misc/python/materialize/mzcompose/__init__.py b/misc/python/materialize/mzcompose/__init__.py index 9ce47a13dc0a5..f3175cde6b1ac 100644 --- a/misc/python/materialize/mzcompose/__init__.py +++ b/misc/python/materialize/mzcompose/__init__.py @@ -96,7 +96,6 @@ def get_minimal_system_parameters( "enable_cast_elimination": "true", "enable_coalesce_case_transform": "true", "enable_columnation_lgalloc": "false", - "enable_compute_correction_v2": "true", "enable_compute_logical_backpressure": "true", "enable_connection_validation_syntax": "true", "enable_create_table_from_source": "true", @@ -643,7 +642,6 @@ def get_default_system_parameters( "compute_dataflow_max_inflight_bytes_cc", "compute_flat_map_fuel", "compute_temporal_bucketing_summary", - "consolidating_vec_growth_dampener", "copy_to_s3_parquet_row_group_file_ratio", "copy_to_s3_arrow_builder_buffer_ratio", "copy_to_s3_multipart_part_size_bytes", diff --git a/misc/python/materialize/parallel_workload/action.py b/misc/python/materialize/parallel_workload/action.py index 3e2751d164899..abc8e256e9ef3 100644 --- a/misc/python/materialize/parallel_workload/action.py +++ b/misc/python/materialize/parallel_workload/action.py @@ -3187,7 +3187,6 @@ def __init__( "persist_blob_hedged_get_warm_interval", "enable_compute_half_join2", "enable_mz_join_core", - "enable_compute_correction_v2", "linear_join_yielding", "enable_lgalloc", "enable_lgalloc_eager_reclamation", @@ -3207,7 +3206,6 @@ def __init__( "compute_dataflow_max_inflight_bytes_cc", "subscribe_max_buffered_bytes", "compute_flat_map_fuel", - "consolidating_vec_growth_dampener", "compute_hydration_concurrency", "copy_to_s3_parquet_row_group_file_ratio", "copy_to_s3_arrow_builder_buffer_ratio", diff --git a/src/compute-types/src/dyncfgs.rs b/src/compute-types/src/dyncfgs.rs index 2f7a2c7586549..0a8c7150c6d07 100644 --- a/src/compute-types/src/dyncfgs.rs +++ b/src/compute-types/src/dyncfgs.rs @@ -262,14 +262,6 @@ pub const ENABLE_SYNC_MV_SINK: Config = Config::new( ParameterScope::Environment, ); -/// Whether rendering should use the new MV sink correction buffer implementation. -pub const ENABLE_CORRECTION_V2: Config = Config::new( - "enable_compute_correction_v2", - true, - "Whether compute should use the new MV sink correction buffer implementation.", - ParameterScope::Environment, -); - /// The size factor of subsequent chains in the correction V2 buffer. pub const CORRECTION_V2_CHAIN_PROPORTIONALITY: Config = Config::new( "compute_correction_v2_chain_proportionality", @@ -422,15 +414,6 @@ pub const DATAFLOW_MAX_INFLIGHT_BYTES_CC: Config> = Config::new( ParameterScope::Replica, ); -/// The term `n` in the growth rate `1 + 1/(n + 1)` for `ConsolidatingVec`. -/// The smallest value `0` corresponds to the greatest allowed growth, of doubling. -pub const CONSOLIDATING_VEC_GROWTH_DAMPENER: Config = Config::new( - "consolidating_vec_growth_dampener", - 1, - "Dampener in growth rate for consolidating vector size", - ParameterScope::Replica, -); - /// The number of dataflows that may hydrate concurrently. /// /// Enforced in `environmentd`, by the controller's per-replica hydration @@ -688,7 +671,6 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet { .add(&ENABLE_ERROR_DISTINCT) .add(&ENABLE_MZ_JOIN_CORE) .add(&ENABLE_SYNC_MV_SINK) - .add(&ENABLE_CORRECTION_V2) .add(&CORRECTION_V2_CHAIN_PROPORTIONALITY) .add(&CORRECTION_V2_CHUNK_SIZE) .add(&ENABLE_COMPUTE_TEMPORAL_BUCKETING) @@ -715,7 +697,6 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet { .add(&COMPUTE_REPLICA_EXPIRATION_OFFSET) .add(&COMPUTE_APPLY_COLUMN_DEMANDS) .add(&COMPUTE_FLAT_MAP_FUEL) - .add(&CONSOLIDATING_VEC_GROWTH_DAMPENER) .add(&ENABLE_COMPUTE_RENDER_FUELED_AS_SPECIFIC_COLLECTION) .add(&ENABLE_COMPUTE_LOGICAL_BACKPRESSURE) .add(&COMPUTE_LOGICAL_BACKPRESSURE_MAX_RETAINED_CAPABILITIES) diff --git a/src/compute/benches/correction.rs b/src/compute/benches/correction.rs index eccc74737fae7..199ac38071d0d 100644 --- a/src/compute/benches/correction.rs +++ b/src/compute/benches/correction.rs @@ -7,9 +7,9 @@ // the Business Source License, use of this software will be governed // by the Apache License, Version 2.0. -//! Micro-benchmark comparing `CorrectionV1` and `CorrectionV2` on hydration-style -//! workloads, i.e., workloads where the input catches up with the current time by -//! passing through many distinct timestamps. +//! Micro-benchmark for the MV sink `Correction` buffer on hydration-style workloads, +//! i.e., workloads where the input catches up with the current time by passing through +//! many distinct timestamps. //! //! The scenario this models: an MV sink restarts with an old as-of and the desired //! input replays through `T` distinct timestamps while persist writes (and thus @@ -30,8 +30,7 @@ use criterion::measurement::WallTime; use criterion::{ BatchSize, BenchmarkGroup, BenchmarkId, Criterion, criterion_group, criterion_main, }; -use mz_compute::sink::correction::CorrectionV1; -use mz_compute::sink::correction_v2::CorrectionV2; +use mz_compute::sink::correction::Correction; use mz_ore::metrics::MetricsRegistry; use mz_persist_client::cfg::PersistConfig; use mz_persist_client::metrics::{Metrics, SinkMetrics}; @@ -42,8 +41,6 @@ use timely::progress::Antichain; const CHAIN_PROPORTIONALITY: f64 = 3.0; /// Default value of `compute_correction_v2_chunk_size`. const CHUNK_SIZE: usize = 8 * 1024; -/// Default value of `consolidating_vec_growth_dampener`. -const GROWTH_DAMPENER: usize = 1; /// Number of updates inserted per distinct timestamp. const UPDATES_PER_TS: u64 = 16; @@ -51,69 +48,6 @@ const UPDATES_PER_TS: u64 = 16; /// Time offset of far-future retractions in the temporal-filter pattern. const TEMPORAL_OFFSET: u64 = 1 << 40; -#[derive(Clone, Copy)] -enum Version { - V1, - V2, -} - -impl Version { - fn name(self) -> &'static str { - match self { - Self::V1 => "v1", - Self::V2 => "v2", - } - } -} - -/// Local dispatch over the correction buffer implementations. -/// -/// The production `Correction` enum only contains v1 and v2, so the bench carries its own. -enum Corr { - V1(CorrectionV1), - V2(CorrectionV2), -} - -impl Corr { - fn insert(&mut self, updates: &mut Vec<(Row, Timestamp, Diff)>) { - match self { - Self::V1(c) => c.insert(updates), - Self::V2(c) => c.insert(updates), - } - } - - fn insert_negated(&mut self, updates: &mut Vec<(Row, Timestamp, Diff)>) { - match self { - Self::V1(c) => c.insert_negated(updates), - Self::V2(c) => c.insert_negated(updates), - } - } - - fn updates_before( - &mut self, - upper: &Antichain, - ) -> Box + '_> { - match self { - Self::V1(c) => Box::new(c.updates_before(upper)), - Self::V2(c) => Box::new(c.updates_before(upper)), - } - } - - fn advance_since(&mut self, since: Antichain) { - match self { - Self::V1(c) => c.advance_since(since), - Self::V2(c) => c.advance_since(since), - } - } - - fn consolidate_at_since(&mut self) { - match self { - Self::V1(c) => c.consolidate_at_since(), - Self::V2(c) => c.consolidate_at_since(), - } - } -} - /// The shape of the update stream fed into the correction buffer. #[derive(Clone, Copy)] enum Pattern { @@ -145,22 +79,14 @@ fn sink_metrics() -> SinkMetrics { metrics.sink.clone() } -fn make_correction(version: Version, metrics: &SinkMetrics) -> Corr { - let worker_metrics = metrics.for_worker(0); - match version { - Version::V1 => Corr::V1(CorrectionV1::new( - metrics.clone(), - worker_metrics, - GROWTH_DAMPENER, - )), - Version::V2 => Corr::V2(CorrectionV2::new( - metrics.clone(), - worker_metrics, - None, - CHAIN_PROPORTIONALITY, - CHUNK_SIZE, - )), - } +fn make_correction(metrics: &SinkMetrics) -> Correction { + Correction::with_params( + metrics.clone(), + metrics.for_worker(0), + None, + CHAIN_PROPORTIONALITY, + CHUNK_SIZE, + ) } fn row(key: u64, value: u64) -> Row { @@ -219,20 +145,28 @@ fn make_batches(num_ts: u64, pattern: Pattern) -> Vec], -) -> Corr { - let mut correction = make_correction(version, metrics); +) -> Correction { + let mut correction = make_correction(metrics); for batch in batches { correction.insert(&mut batch.clone()); } correction } -impl Corr { +/// Benchmark helpers modelling the `write_batches` operator's use of the buffer. +/// +/// A trait rather than an inherent `impl`, because `Correction` is defined in another crate. +trait WriteSteps { /// Emulate one `write_batches` write: read the updates before `upper` and feed back their /// negations, like the persist input does. + fn write_step(&mut self, upper: &Antichain) -> usize; + /// Read and count the updates before `upper`, without persist feedback. + fn read_count(&mut self, upper: &Antichain) -> usize; +} + +impl WriteSteps for Correction { fn write_step(&mut self, upper: &Antichain) -> usize { let mut written: Vec<_> = self.updates_before(upper).collect(); let count = written.len(); @@ -240,7 +174,6 @@ impl Corr { count } - /// Read and count the updates before `upper`, without persist feedback. fn read_count(&mut self, upper: &Antichain) -> usize { self.updates_before(upper).count() } @@ -253,7 +186,7 @@ impl Corr { /// updates come back through the persist input and are removed by `insert_negated`. We /// model that feedback here, otherwise the buffer re-emits all previous updates on every /// step and clone costs drown out the structure-management costs we want to measure. -fn drain_stepwise(mut correction: Corr, num_ts: u64) -> Corr { +fn drain_stepwise(mut correction: Correction, num_ts: u64) -> Correction { for t in 0..num_ts { let upper = Antichain::from_elem(Timestamp::from(t + 1)); let count = correction.write_step(&upper); @@ -265,7 +198,7 @@ fn drain_stepwise(mut correction: Corr, num_ts: u64) -> Corr { /// Advance the since across all buffered times at once, then consolidate and drain. /// This exercises `Cursor::advance_by` with many distinct times below the since. -fn advance_jump(mut correction: Corr, num_ts: u64) -> Corr { +fn advance_jump(mut correction: Correction, num_ts: u64) -> Correction { correction.advance_since(Antichain::from_elem(Timestamp::from(num_ts))); correction.consolidate_at_since(); let upper = Antichain::from_elem(Timestamp::from(num_ts + 1)); @@ -281,23 +214,21 @@ fn bench_scenario( num_ts: u64, routine: F, ) where - F: Fn(Corr, u64) -> Corr, + F: Fn(Correction, u64) -> Correction, { let batches = make_batches(num_ts, pattern); - for version in [Version::V1, Version::V2] { - group.bench_function(BenchmarkId::new(version.name(), num_ts), |b| { - b.iter_batched( - || filled_correction(version, metrics, &batches), - |correction| routine(correction, num_ts), - BatchSize::PerIteration, - ) - }); - } + group.bench_function(BenchmarkId::from_parameter(num_ts), |b| { + b.iter_batched( + || filled_correction(metrics, &batches), + |correction| routine(correction, num_ts), + BatchSize::PerIteration, + ) + }); } /// Configure a benchmark group for short wall-clock time. /// -/// The interesting effects here are order-of-magnitude differences between implementations, so we +/// The interesting effects here are order-of-magnitude differences across problem sizes, so we /// trade statistical rigor for execution speed. fn configure(group: &mut BenchmarkGroup) { group @@ -329,20 +260,18 @@ fn bench_correction(c: &mut Criterion) { configure(&mut group); for num_ts in num_ts_values { let batches = make_batches(num_ts, pattern); - for version in [Version::V1, Version::V2] { - group.bench_function(BenchmarkId::new(version.name(), num_ts), |b| { - b.iter_batched( - || (make_correction(version, &metrics), batches.clone()), - |(mut correction, mut batches)| { - for batch in &mut batches { - correction.insert(batch); - } - correction - }, - BatchSize::PerIteration, - ) - }); - } + group.bench_function(BenchmarkId::from_parameter(num_ts), |b| { + b.iter_batched( + || (make_correction(&metrics), batches.clone()), + |(mut correction, mut batches)| { + for batch in &mut batches { + correction.insert(batch); + } + correction + }, + BatchSize::PerIteration, + ) + }); } group.finish(); } diff --git a/src/compute/src/sink.rs b/src/compute/src/sink.rs index 002e19738b05e..2f8e856328c52 100644 --- a/src/compute/src/sink.rs +++ b/src/compute/src/sink.rs @@ -12,10 +12,6 @@ mod copy_to_s3_oneshot; pub mod correction; #[cfg(not(feature = "bench"))] mod correction; -#[cfg(feature = "bench")] -pub mod correction_v2; -#[cfg(not(feature = "bench"))] -mod correction_v2; mod materialized_view; mod materialized_view_v2; mod metric_sink; diff --git a/src/compute/src/sink/correction.rs b/src/compute/src/sink/correction.rs index e046f46da27ce..72a94b69649cc 100644 --- a/src/compute/src/sink/correction.rs +++ b/src/compute/src/sink/correction.rs @@ -7,212 +7,300 @@ // the Business Source License, use of this software will be governed // by the Apache License, Version 2.0. -//! The `Correction` data structure used by `persist_sink::write_batches` to stash updates before -//! they are written into batches. - -use std::collections::BTreeMap; +//! The `Correction` data structure used by the MV sink's `write_batches` operator to stash +//! updates before they are written. +//! +//! The `Correction` data structure provides methods to: +//! * insert new updates +//! * advance the compaction frontier (called `since`) +//! * obtain an iterator over consolidated updates before some `upper` +//! * force consolidation of updates before some `upper` +//! +//! The goal is to provide good performance for each of these operations, even in the presence of +//! future updates. MVs downstream of temporal filters might have to deal with large amounts of +//! retractions for future times and we want those to be handled efficiently as well. +//! +//! Note that `Correction` does not provide a method to directly remove updates. Instead updates +//! are removed by inserting their retractions so that they consolidate away to nothing. +//! +//! ## Storage of Updates +//! +//! Stored updates are of the form `(data, time, diff)`, where `time` and `diff` are fixed to +//! [`mz_repr::Timestamp`] and [`mz_repr::Diff`], respectively. +//! +//! [`Correction`] holds onto a list of `Chain`s containing `Chunk`s of stashed updates. Each +//! `Chunk` is a columnation region containing a fixed maximum number of updates. All updates in +//! a chunk, and all updates in a chain, are ordered by (time, data) and consolidated. +//! +//! Chains live in three places: +//! +//! * A [`BucketChain`] partitions times at or beyond the `boundary` (the largest read `upper` +//! seen so far) into buckets of exponentially growing time ranges, each holding a list of +//! chains. Reads only touch the buckets below their `upper`, so the bulk of the buffered +//! updates — in particular far-future retractions produced by temporal filters — is left +//! alone. +//! * `pending_low` holds chains at times below the `boundary`, mostly insertions arriving +//! through the persist feedback. +//! * `emitted` is a single chain holding the updates returned by the last read. Updates must +//! stay in the buffer until their feedback retractions arrive, and keeping them separate from +//! the bucket chain means reads never have to re-merge future updates. +//! +//! ```text +//! chain[0] | chain[1] | chain[2] +//! | | +//! chunk[0] | chunk[0] | chunk[0] +//! (a, 1, +1) | (a, 1, +1) | (d, 3, +1) +//! (b, 1, +1) | (b, 2, -1) | (d, 4, -1) +//! chunk[1] | chunk[1] | +//! (c, 1, +1) | (c, 2, -2) | +//! (a, 2, -1) | (c, 4, -1) | +//! chunk[2] | | +//! (b, 2, +1) | | +//! (c, 2, +1) | | +//! chunk[3] | | +//! (b, 3, -1) | | +//! (c, 3, +1) | | +//! ``` +//! +//! The "chain invariant" states that each chain in a bucket has at least `chain_proportionality` times as +//! many updates as the next one. This means that chain sizes will often be powers of +//! `chain_proportionality`, but they don't have to be. For example, for a proportionality of 2, +//! the chain sizes `[11, 5, 2, 1]` would satisfy the chain invariant. +//! +//! Note that the invariant is maintained on update counts, not chunk counts. Chunks are +//! byte-bounded (see `ChunkBuilder`), so chunk count is not proportional to update count and +//! would be a poor proxy: any chain below the chunk byte boundary is a single chunk regardless +//! of how many updates it holds, which would let the geometric invariant collapse and break the +//! O(log N) amortization of inserts. +//! +//! Choosing the `chain_proportionality` value allows tuning the trade-off between memory and CPU +//! resources required to maintain corrections. A higher proportionality forces more frequent chain +//! merges, and therefore consolidation, reducing memory usage but increasing CPU usage. +//! +//! ## Inserting Updates +//! +//! A batch of updates is routed by time: updates below the `boundary` become a `pending_low` +//! chain, the rest is appended as new chains to their respective buckets. Appending to a bucket +//! merges chains until the chain invariant is restored. +//! +//! Inserting an update into the correction buffer can be expensive: It involves allocating a new +//! chunk, copying the update in, and then likely merging with an existing chain to restore the +//! chain invariant. If updates trickle in in small batches, this can cause a considerable +//! overhead. To amortize this overhead, new updates aren't immediately inserted into the sorted +//! chains but instead stored in a `Stage` buffer. Once enough updates have been staged to fill a +//! `Chunk`, they are sorted and routed. +//! +//! The insert operation has an amortized complexity of O(log N), with N being the current number +//! of updates stored. +//! +//! ## Retrieving Consolidated Updates +//! +//! Retrieving consolidated updates before a given `upper` works by peeling all buckets below the +//! `upper` off the bucket chain, splitting their chains, the pending low chains, and the previous +//! `emitted` chain at the `upper`, merging the parts below the `upper` into the new `emitted` +//! chain, and returning an iterator over that chain. +//! +//! Because each chain contains updates ordered by time first, splitting a chain at the `upper` +//! reuses whole chunks and copies at most one chunk straddling the split point. Updates at times +//! at or beyond the `upper` are never touched, no matter how many the buffer holds. The +//! complexity of a read is O(U log K), with U being the number of updates before `upper` and K +//! the number of chains containing them. +//! +//! ## Merging Chains +//! +//! Merging multiple chains into a single chain is done using a k-way merge. As the input chains +//! are sorted by (time, data) and consolidated, the same properties hold for the output chain. The +//! complexity of a merge of K chains containing N updates is O(N log K). +//! +//! There is a twist though: Merging also has to respect the `since` frontier, which determines how +//! far the times of updates should be advanced. Advancing times in a sorted chain of updates +//! can make them become unsorted, so we cannot just merge the chains from top to bottom. +//! +//! For example, consider these two chains, assuming `since = [2]`: +//! chain 1: [(c, 1, +1), (b, 2, -1), (a, 3, -1)] +//! chain 2: [(b, 1, +1), (a, 2, +1), (c, 2, -1)] +//! After time advancement, the chains look like this: +//! chain 1: [(c, 2, +1), (b, 2, -1), (a, 3, -1)] +//! chain 2: [(b, 2, +1), (a, 2, +1), (c, 2, -1)] +//! Merging them naively yields [(b, 2, +1), (a, 2, +1), (b, 2, -1), (a, 3, -1)], a chain that's +//! neither sorted nor consolidated. +//! +//! Times below the `since` can only exist in chains read by `consolidate_before`, and only if +//! the `since` advanced past buffered times since the previous read. For few distinct stale +//! times — the steady state, where the previously emitted chain was written just before the +//! since advanced past it — we merge sub-chains, one for each distinct time that's before or at +//! the `since`. Each of these sub-chains retains the (time, data) ordering after the time +//! advancement to `since`, so merging those yields the expected result. +//! +//! For the above example, the chains we would merge are: +//! chain 1.a: [(c, 2, +1)] +//! chain 1.b: [(b, 2, -1), (a, 3, -1)] +//! chain 2.a: [(b, 2, +1)], +//! chain 2.b: [(a, 2, +1), (c, 2, -1)] +//! +//! For many distinct stale times — e.g. a since jump across many buffered timestamps when a sink +//! restarts with an old as-of — the number of sub-chains grows with the number of distinct times, +//! so we instead materialize the affected updates, advance their times, and sort and consolidate +//! them in one O(U log U) pass. + +pub mod logging; + +use std::cmp::Ordering; +use std::collections::{BinaryHeap, VecDeque}; use std::fmt; -use std::num::NonZeroIsize; -use std::ops::{AddAssign, Bound, RangeBounds, SubAssign}; - -use differential_dataflow::consolidation::{consolidate, consolidate_updates}; -use differential_dataflow::logging::{BatchEvent, DropEvent}; -use itertools::{Either, Itertools}; -use mz_compute_types::dyncfgs::{ - CONSOLIDATING_VEC_GROWTH_DAMPENER, CORRECTION_V2_CHAIN_PROPORTIONALITY, - CORRECTION_V2_CHUNK_SIZE, ENABLE_CORRECTION_V2, -}; +use std::rc::Rc; +use std::sync::{Mutex, OnceLock}; + +use columnar::{Columnar, Index, Len, Ref}; +use mz_compute_types::dyncfgs::{CORRECTION_V2_CHAIN_PROPORTIONALITY, CORRECTION_V2_CHUNK_SIZE}; use mz_dyncfg::ConfigSet; +use mz_ore::cast::CastLossy; +use mz_ore::soft_assert_or_log; use mz_persist_client::metrics::{SinkMetrics, SinkWorkerMetrics, UpdateDelta}; use mz_repr::{Diff, Timestamp}; +use mz_timely_util::column_pager::{self, PagedColumn}; +use mz_timely_util::columnar::Column; +use mz_timely_util::temporal::{Bucket, BucketChain}; use timely::PartialOrder; +use timely::dataflow::channels::ContainerBytes; use timely::progress::Antichain; -use tokio::sync::mpsc; -use crate::logging::compute::{ - ArrangementHeapAllocations, ArrangementHeapCapacity, ArrangementHeapSize, - ArrangementHeapSizeOperator, ArrangementHeapSizeOperatorDrop, ComputeEvent, - Logger as ComputeLogger, -}; -use crate::sink::correction_v2::{CorrectionV2, Data}; +use crate::sink::correction::logging::{ChannelLogging, SizeMetrics}; -/// A data structure suitable for storing updates in a self-correcting persist sink. +/// Convenient alias for use in data trait bounds. /// -/// Selects one of two correction buffer implementations. `V1` is the original simple -/// implementation that stores updates in non-spillable memory. `V2` improves on `V1` by supporting -/// spill-to-disk but is less battle-tested so for now we want to keep the option of reverting to -/// `V1` in a pinch. The plan is to remove `V1` eventually. -pub enum Correction { - /// Correction buffer based on a [`CorrectionV1`]. - V1(CorrectionV1), - /// Correction buffer based on a [`CorrectionV2`]. - V2(CorrectionV2), +/// `D` is constrained to be `Columnar`, so that updates can be stored in a single columnar +/// region per chunk, and the variable-length payload (e.g. `Row` bytes) lives in the same +/// allocation as the rest of the chunk. The `Ref`-level `Eq + Ord` bounds let the merge/heap +/// code compare updates directly through the columnar borrow, avoiding `into_owned` clones +/// on the hot path. +pub trait Data: + differential_dataflow::Data + + Columnar columnar::Borrow: Eq + Ord>> + + Send + + Sync +{ +} +impl Data for D where + D: differential_dataflow::Data + + Columnar columnar::Borrow: Eq + Ord>> + + Send + + Sync +{ } -impl Correction { - /// Construct a new `Correction` instance. - pub fn new( - metrics: SinkMetrics, - worker_metrics: SinkWorkerMetrics, - logging: Option, - config: &ConfigSet, - ) -> Self { - if ENABLE_CORRECTION_V2.get(config) { - let prop = CORRECTION_V2_CHAIN_PROPORTIONALITY.get(config); - let chunk_size = CORRECTION_V2_CHUNK_SIZE.get(config); - Self::V2(CorrectionV2::new( - metrics, - worker_metrics, - logging, - prop, - chunk_size, - )) - } else { - let growth_dampener = CONSOLIDATING_VEC_GROWTH_DAMPENER.get(config); - Self::V1(CorrectionV1::new(metrics, worker_metrics, growth_dampener)) - } - } - - /// Insert a batch of updates. - pub fn insert(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) { - match self { - Self::V1(c) => c.insert(updates), - Self::V2(c) => c.insert(updates), - } - } - - /// Insert a batch of updates, after negating their diffs. - pub fn insert_negated(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) { - match self { - Self::V1(c) => c.insert_negated(updates), - Self::V2(c) => c.insert_negated(updates), - } - } - - /// Consolidate and return updates before the given `upper`. - pub fn updates_before( - &mut self, - upper: &Antichain, - ) -> Box + Send + '_> { - match self { - Self::V1(c) => Box::new(c.updates_before(upper)), - Self::V2(c) => Box::new(c.updates_before(upper)), - } - } - - /// Consolidate the updates before the given `upper`. +/// A data structure used to store corrections in the MV sink implementation. +/// +/// Updates are stored in columnar chunks whose memory can be transparently spilled to disk. +#[derive(Debug)] +pub struct Correction { + /// Bucketed storage for updates at times at or beyond `boundary`. /// - /// This is the expensive half of [`Correction::updates_before`], split out so callers that - /// must not run unbounded CPU work inline can perform it elsewhere. - pub fn consolidate_before(&mut self, upper: &Antichain) { - match self { - Self::V1(c) => c.consolidate_before(upper), - Self::V2(c) => c.consolidate_before(upper), - } - } - - /// Return the updates before the given `upper`, as consolidated by a preceding - /// [`Correction::consolidate_before`] call. + /// Buckets cover exponentially growing time ranges, so reads only touch the buckets below + /// their `upper`, and far-future updates (e.g. retractions produced by temporal filters) are + /// rarely touched. + chain: BucketChain>, + /// Chains at times below `boundary` that were not yet emitted. /// - /// The caller must have invoked `consolidate_before` with the same `upper` and must not have - /// mutated the buffer since. Otherwise the returned updates are neither consolidated nor - /// necessarily complete. - pub fn consolidated_updates_before<'a>( - &'a self, - upper: &Antichain, - ) -> impl Iterator + Send + use<'a, D> { - match self { - Self::V1(c) => Either::Left(c.consolidated_updates_before(upper)), - Self::V2(c) => Either::Right(c.consolidated_updates_before(upper)), - } - } - - /// Advance the since frontier. + /// Filled by inserts at times below the boundary (mostly persist feedback) and by the + /// remainders of `emitted` when a read uses a smaller `upper` than the previous one. Merged + /// into `emitted` by the next read. + pending_low: Vec>, + /// Updates that were emitted by `updates_before` but not yet cancelled by persist feedback. /// - /// # Panics + /// Sorted and consolidated, with all times advanced to the `since`. + emitted: Chain, + /// A staging area for updates, to speed up small inserts. + stage: Stage, + /// The lower bound of times stored in `chain`. Only ever advances. /// - /// Panics if the given `since` is less than the current since frontier. - pub fn advance_since(&mut self, since: Antichain) { - match self { - Self::V1(c) => c.advance_since(since), - Self::V2(c) => c.advance_since(since), - } - } - - /// Consolidate all updates at the current `since`. - pub fn consolidate_at_since(&mut self) { - match self { - Self::V1(c) => c.consolidate_at_since(), - Self::V2(c) => c.consolidate_at_since(), - } - } -} - -/// A collection holding `persist_sink` updates. -/// -/// The `CorrectionV1` data structure is purpose-built for the `persist_sink::write_batches` -/// operator: -/// -/// * It stores updates by time, to enable efficient separation between updates that should -/// be written to a batch and updates whose time has not yet arrived. -/// * It eschews an interface for directly removing previously inserted updates. Instead, updates -/// are removed by inserting them again, with negated diffs. Stored updates are continuously -/// consolidated to give them opportunity to cancel each other out. -/// * It provides an interface for advancing all contained updates to a given frontier. -pub struct CorrectionV1 { - /// Stashed updates by time. - updates: BTreeMap>, - /// Frontier to which all update times are advanced. + /// Times below the boundary have been peeled off the bucket chain and can only be stored in + /// `pending_low` or `emitted`. + boundary: Antichain, + /// The frontier by which all contained times are advanced. since: Antichain, - /// Total length and capacity of vectors in `updates`. + /// Total count of updates in the correction buffer. /// - /// Tracked to maintain metrics. - total_size: LengthAndCapacity, + /// Tracked to compute deltas in `update_metrics`. + prev_update_count: usize, + /// Total heap size used by the correction buffer. + /// + /// Tracked to compute deltas in `update_metrics`. + prev_size: SizeMetrics, /// Global persist sink metrics. metrics: SinkMetrics, /// Per-worker persist sink metrics. worker_metrics: SinkWorkerMetrics, - /// Configuration for `ConsolidatingVec` driving the growth rate down from doubling. - growth_dampener: usize, + /// Introspection logging. + logging: Option, } -impl CorrectionV1 { - /// Construct a new `CorrectionV1` instance. +/// Fuel for restoring the bucket chain invariant after peeling. +/// +/// Bounds the restoration work per buffer operation. The bucket chain remains functional when +/// restoration is incomplete -- peeling and finding work on ill-formed chains, at the cost of +/// more in-line splitting -- so leftover restoration is simply picked up by the next operation. +/// +/// `restore` spends one unit of fuel per bucket split, and a single `peel` leaves at most +/// `BucketTimestamp::DOMAIN` (64) buckets to re-split, so this budget completes restoration in one +/// call for any realistic buffer. It is deliberately generous: the "incomplete restoration is +/// picked up next op" path is a correctness safety net for pathological bucket counts, not a hot +/// path we expect to exercise. Lower it if restoration ever needs to interleave with other work. +const RESTORE_FUEL: i64 = 1_000_000; + +impl Correction { + /// Construct a new [`Correction`] instance, tuned by the given configuration. pub fn new( metrics: SinkMetrics, worker_metrics: SinkWorkerMetrics, - growth_dampener: usize, + logging: Option, + config: &ConfigSet, ) -> Self { - Self { - updates: Default::default(), - since: Antichain::from_elem(Timestamp::MIN), - total_size: Default::default(), + let chain_proportionality = CORRECTION_V2_CHAIN_PROPORTIONALITY.get(config); + let chunk_size = CORRECTION_V2_CHUNK_SIZE.get(config); + Self::with_params( metrics, worker_metrics, - growth_dampener, - } + logging, + chain_proportionality, + chunk_size, + ) } - /// Update persist sink metrics to the given new length and capacity. - fn update_metrics(&mut self, new_size: LengthAndCapacity) { - let old_size = self.total_size; - let len_delta = UpdateDelta::new(new_size.length, old_size.length); - let cap_delta = UpdateDelta::new(new_size.capacity, old_size.capacity); - self.metrics - .report_correction_update_deltas(len_delta, cap_delta); - self.worker_metrics - .report_correction_update_totals(new_size.length, new_size.capacity); + /// Construct a new [`Correction`] instance with explicit tuning parameters. + /// + /// `chain_proportionality` is the size factor between subsequent chains, `chunk_size` the + /// byte size of a chunk. + pub fn with_params( + metrics: SinkMetrics, + worker_metrics: SinkWorkerMetrics, + logging: Option, + chain_proportionality: f64, + chunk_size: usize, + ) -> Self { + let update_size = std::mem::size_of::<(D, Timestamp, Diff)>(); + let chunk_capacity = std::cmp::max(chunk_size / update_size, 1); - self.total_size = new_size; + Self { + chain: BucketChain::new(ChainBucket::new(chain_proportionality, logging.clone())), + pending_low: Vec::new(), + emitted: Chain::new(), + stage: Stage::new(logging.clone(), chunk_capacity), + boundary: Antichain::from_elem(Timestamp::MIN), + since: Antichain::from_elem(Timestamp::MIN), + prev_update_count: 0, + prev_size: Default::default(), + metrics, + worker_metrics, + logging, + } } -} -impl CorrectionV1 { /// Insert a batch of updates. pub fn insert(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) { let Some(since_ts) = self.since.as_option() else { - // If the since frontier is empty, discard all updates. + // If the since is the empty frontier, discard all updates. updates.clear(); return; }; @@ -220,13 +308,14 @@ impl CorrectionV1 { for (_, time, _) in &mut *updates { *time = std::cmp::max(*time, *since_ts); } + self.insert_inner(updates); } /// Insert a batch of updates, after negating their diffs. pub fn insert_negated(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) { let Some(since_ts) = self.since.as_option() else { - // If the since frontier is empty, discard all updates. + // If the since is the empty frontier, discard all updates. updates.clear(); return; }; @@ -235,569 +324,1912 @@ impl CorrectionV1 { *time = std::cmp::max(*time, *since_ts); *diff = -*diff; } + self.insert_inner(updates); } - /// Insert a batch of updates. + /// Insert a batch of updates into the stage, flushing it when full. /// - /// The given `updates` must all have been advanced by `self.since`. + /// All times are expected to be >= the `since`. fn insert_inner(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) { - consolidate_updates(updates); - updates.sort_unstable_by_key(|(_, time, _)| *time); - - let mut new_size = self.total_size; - let mut updates = updates.drain(..).peekable(); - while let Some(&(_, time, _)) = updates.peek() { - mz_ore::soft_assert_no_log!( - self.since.less_equal(&time), - "update not advanced by `since`" - ); - - let data = updates - .peeking_take_while(|(_, t, _)| *t == time) - .map(|(d, _, r)| (d, r)); - - use std::collections::btree_map::Entry; - match self.updates.entry(time) { - Entry::Vacant(entry) => { - let mut vec: ConsolidatingVec<_> = data.collect(); - vec.growth_dampener = self.growth_dampener; - new_size += (vec.len(), vec.capacity()); - entry.insert(vec); - } - Entry::Occupied(mut entry) => { - let vec = entry.get_mut(); - new_size -= (vec.len(), vec.capacity()); - vec.extend(data); - new_size += (vec.len(), vec.capacity()); - } - } - } + debug_assert!(updates.iter().all(|(_, t, _)| self.since.less_equal(t))); - self.update_metrics(new_size); - } + if let Some(mut ready) = self.stage.insert(updates) { + self.route(&mut ready); + } - /// The range of stored times before the given `upper`. - fn range_before(upper: &Antichain) -> (Bound, Bound) { - let start = Bound::Included(Timestamp::MIN); - let end = match upper.as_option() { - Some(ts) => Bound::Excluded(*ts), - None => Bound::Unbounded, - }; - (start, end) + self.update_metrics(); } - /// Consolidate the updates before the given `upper`. - pub fn consolidate_before(&mut self, upper: &Antichain) { - let _ = self.consolidate(Self::range_before(upper)); - } + /// Route a batch of sorted, consolidated updates to `pending_low` or their chain buckets. + fn route(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) { + // Updates at times below the boundary become a pending low chain. + let idx = updates.partition_point(|(_, t, _)| !self.boundary.less_equal(t)); + if idx > 0 { + let mut builder = ChainBuilder::default(); + builder.extend(updates.drain(..idx)); + let chain = builder.finish(); + if !chain.is_empty() { + self.log_chain_created(&chain); + self.pending_low.push(chain); + } + } - /// Return the updates before the given `upper`, as consolidated by a preceding - /// [`CorrectionV1::consolidate_before`] call. - /// - /// The caller must have invoked `consolidate_before` with the same `upper` and must not have - /// mutated the buffer since. Otherwise the returned updates are not consolidated. - pub fn consolidated_updates_before<'a>( - &'a self, - upper: &Antichain, - ) -> impl Iterator + Send + use<'a, D> { - self.updates - .range(Self::range_before(upper)) - .flat_map(|(t, data)| data.iter().map(|(d, r)| (d.clone(), *t, *r))) + // Updates at times at or beyond the boundary go into their chain buckets. Walk ranges of + // times that fall into the same bucket, to push batches of updates at once. + let mut drain = updates.drain(..).peekable(); + while let Some(update) = drain.next() { + let time = update.1; + let range = self + .chain + .range_of(&time) + .expect("bucket chain covers all times at or beyond the boundary"); + let mut builder = ChainBuilder::default(); + builder.extend(std::iter::once(update)); + while let Some(update) = drain.next_if(|(_, t, _)| range.contains(t)) { + builder.extend(std::iter::once(update)); + } + let bucket = self + .chain + .find_mut(&range.start) + .expect("bucket chain covers all times at or beyond the boundary"); + bucket.push_chain(builder.finish()); + } } - /// Consolidate and return updates before the given `upper`. + /// Return consolidated updates before the given `upper`. pub fn updates_before<'a>( &'a mut self, upper: &Antichain, - ) -> impl Iterator + Send + use<'a, D> { + ) -> impl Iterator + Send + 'a { self.consolidate_before(upper); self.consolidated_updates_before(upper) } - /// Consolidate the updates at the times in the given range. + /// Return the updates before the given `upper`, as consolidated by a preceding + /// [`Correction::consolidate_before`] call. /// - /// Returns the number of updates remaining in the range afterwards. - fn consolidate(&mut self, range: R) -> usize - where - R: RangeBounds, - { - let mut new_size = self.total_size; - - let updates = self.updates.range_mut(range); - let count = updates.fold(0, |acc, (_, data)| { - new_size -= (data.len(), data.capacity()); - data.consolidate(); - new_size += (data.len(), data.capacity()); - acc + data.len() - }); + /// The caller must have invoked `consolidate_before` with the same `upper` and must not have + /// mutated the buffer since. Otherwise the returned updates are neither consolidated nor + /// necessarily complete. + pub fn consolidated_updates_before<'a>( + &'a self, + upper: &Antichain, + ) -> impl Iterator + Send + use<'a, D> { + // All contained times are advanced to at least the `since`, so a read at an `upper` that + // is not beyond the `since` is always empty. This mirrors the short-circuit in + // `consolidate_before`, which leaves `emitted` untouched in that case. + if !PartialOrder::less_than(&self.since, upper) { + return None.into_iter().flatten(); + } - self.update_metrics(new_size); - count + // After `consolidate_before`, `emitted` holds exactly the updates before `upper`: every + // path that populates it splits at `upper` (pushing the remainder to `pending_low`), and + // the guard above guarantees `upper > since`, so advancing stale times to the `since` + // cannot lift them to or beyond `upper`. We can therefore yield all of `emitted`. Guard + // the invariant: a violation would write updates beyond the batch upper to persist. + soft_assert_or_log!( + self.emitted + .last() + .is_none_or(|(_, t, _)| !upper.less_equal(&t)), + "emitted contains times at or beyond the upper", + ); + Some(self.emitted.iter()).into_iter().flatten() } - /// Advance the since frontier. + /// Consolidate all updates before the given `upper` into the `emitted` chain. /// - /// # Panics + /// Once this method returns, `emitted` contains all updates at times before `upper`, + /// consolidated. /// - /// Panics if the given `since` is less than the current since frontier. - pub fn advance_since(&mut self, since: Antichain) { - assert!(PartialOrder::less_equal(&self.since, &since)); + /// Does nothing if `upper` is not beyond the `since`: all contained times are advanced to at + /// least the `since`, so such a read is empty anyway, and skipping avoids an eager peel, + /// merge, and `boundary` advancement. Normal reads and `consolidate_at_since` always pass an + /// `upper` beyond the `since`. + pub fn consolidate_before(&mut self, upper: &Antichain) { + if !PartialOrder::less_than(&self.since, upper) { + return; + } - if since != self.since { - self.advance_by(&since); - self.since = since; + if let Some(mut ready) = self.stage.flush() { + self.route(&mut ready); } - } - /// Advance all contained updates by the given frontier. - /// - /// If the given frontier is empty, all remaining updates are discarded. - pub fn advance_by(&mut self, frontier: &Antichain) { - let Some(target_ts) = frontier.as_option() else { - self.updates.clear(); - self.update_metrics(Default::default()); + let Some(&since_ts) = self.since.as_option() else { + // If the since is the empty frontier, discard all updates. + let peeled = self.chain.peel(Antichain::new().borrow()); + for bucket in peeled { + for chain in bucket.into_chains() { + self.log_chain_dropped(&chain); + } + } + for chain in std::mem::take(&mut self.pending_low) { + self.log_chain_dropped(&chain); + } + let emitted = std::mem::replace(&mut self.emitted, Chain::new()); + if !emitted.is_empty() { + self.log_chain_dropped(&emitted); + } + self.update_metrics(); return; }; - let mut new_size = self.total_size; - while let Some((ts, data)) = self.updates.pop_first() { - if frontier.less_equal(&ts) { - // We have advanced all updates that can advance. - self.updates.insert(ts, data); - break; - } + // Peel the buckets below the upper off the bucket chain. Bucket splits during the peel + // only touch chunks around the upper; chunks wholly on either side are reused. + let peeled = self.chain.peel(upper.borrow()); + if PartialOrder::less_than(&self.boundary, upper) { + self.boundary = upper.clone(); + } - use std::collections::btree_map::Entry; - match self.updates.entry(*target_ts) { - Entry::Vacant(entry) => { - entry.insert(data); - } - Entry::Occupied(mut entry) => { - let vec = entry.get_mut(); - new_size -= (data.len(), data.capacity()); - new_size -= (vec.len(), vec.capacity()); - vec.extend(data); - new_size += (vec.len(), vec.capacity()); + // Collect candidate chains: peeled bucket contents, pending low chains, and the previous + // emitted chain. All contain only times below the boundary. + let emitted = std::mem::replace(&mut self.emitted, Chain::new()); + let mut candidates: Vec> = Vec::new(); + for bucket in peeled { + candidates.extend(bucket.into_chains()); + } + candidates.append(&mut self.pending_low); + if !emitted.is_empty() { + candidates.push(emitted); + } + + if candidates.is_empty() { + self.restore_chain(); + self.update_metrics(); + return; + } + + candidates.iter().for_each(|c| self.log_chain_dropped(c)); + + // Split the candidates at the upper. Parts at or beyond the upper (possible when `upper` + // regresses below a previous one) stay pending. + let mut lowers = Vec::new(); + for chain in candidates { + match upper.as_option() { + Some(&upper_ts) => { + let (lower, remainder) = chain.split_at_time(upper_ts); + if !lower.is_empty() { + lowers.push(lower); + } + if !remainder.is_empty() { + self.log_chain_created(&remainder); + self.pending_low.push(remainder); + } } + // The empty upper is greater than all times. + None => lowers.push(chain), } } - self.update_metrics(new_size); - } - - /// Consolidate all updates at the current `since`. - pub fn consolidate_at_since(&mut self) { - let Some(since_ts) = self.since.as_option() else { - return; - }; + // Merge the lower parts into the new emitted chain, advancing times below the since. + // Advancing times in a (time, data)-sorted chain can break its sort order, so chains + // containing stale times cannot be merged as they are. Stale times are expected in steady + // state: the previous emitted chain was written before the since advanced past it. + // + // Count the distinct stale times, up to a small cap. For few distinct stale times -- the + // steady state -- split cursors into runs that remain sorted under advancement and merge + // those. For many distinct stale times -- e.g. a since jump across many buffered + // timestamps when a sink restarts with an old as-of -- the number of runs and the cost of + // cloning cursor state per run grow with the number of distinct times, so materialize, + // advance, and consolidate in one O(U log U) pass instead. + const MAX_STALE_RUNS: usize = 32; + let mut stale_times = 0; + for chain in &lowers { + stale_times += chain.distinct_times_before(since_ts, MAX_STALE_RUNS - stale_times); + if stale_times >= MAX_STALE_RUNS { + break; + } + } - let start = Bound::Included(*since_ts); - let end = match since_ts.try_step_forward() { - Some(ts) => Bound::Excluded(ts), - None => Bound::Unbounded, + let merged = if stale_times == 0 { + let cursors: Vec<_> = lowers.into_iter().filter_map(Chain::into_cursor).collect(); + merge_cursors(cursors) + } else if stale_times < MAX_STALE_RUNS { + let mut runs = Vec::new(); + for chain in lowers { + if let Some(cursor) = chain.into_cursor() { + runs.append(&mut cursor.advance_by(since_ts)); + } + } + merge_cursors(runs) + } else { + let mut updates: Vec<_> = lowers.iter().flat_map(|c| c.iter()).collect(); + for (_, time, _) in &mut updates { + *time = std::cmp::max(*time, since_ts); + } + consolidate(&mut updates); + let mut builder = ChainBuilder::default(); + builder.extend(updates); + let chain = builder.finish(); + + // Advancement can move updates to or beyond the upper; such updates stay pending. + match upper.as_option() { + Some(&upper_ts) => { + let (lower, remainder) = chain.split_at_time(upper_ts); + if !remainder.is_empty() { + self.log_chain_created(&remainder); + self.pending_low.push(remainder); + } + lower + } + None => chain, + } }; - self.consolidate((start, end)); - } -} + if !merged.is_empty() { + self.log_chain_created(&merged); + } + self.emitted = merged; -impl Drop for CorrectionV1 { - fn drop(&mut self) { - self.update_metrics(Default::default()); + self.restore_chain(); + self.update_metrics(); } -} -/// Helper type for convenient tracking of length and capacity together. -#[derive(Clone, Copy, Debug, Default)] -pub(super) struct LengthAndCapacity { - pub length: usize, - pub capacity: usize, -} - -impl AddAssign for LengthAndCapacity { - fn add_assign(&mut self, size: Self) { - self.length += size.length; - self.capacity += size.capacity; + /// Perform a bounded amount of work towards restoring the bucket chain invariant. + /// + /// Restoration is allowed to remain incomplete: the bucket chain supports peeling and finding + /// on ill-formed chains, so any leftover work is picked up by subsequent operations. The fuel + /// bound keeps individual buffer operations from stalling the operator that owns the buffer. + fn restore_chain(&mut self) { + let mut fuel = RESTORE_FUEL; + self.chain.restore(&mut fuel); } -} -impl AddAssign<(usize, usize)> for LengthAndCapacity { - fn add_assign(&mut self, (len, cap): (usize, usize)) { - self.length += len; - self.capacity += cap; + /// Advance the since frontier. + /// + /// Time advancement of updates in the bucket chain is lazy: it happens when the updates are + /// consolidated by a read. + /// + /// # Panics + /// + /// Panics if the given `since` is less than the current since frontier. + pub fn advance_since(&mut self, since: Antichain) { + assert!(PartialOrder::less_equal(&self.since, &since)); + self.stage.advance_times(&since); + self.since = since; } -} -impl SubAssign<(usize, usize)> for LengthAndCapacity { - fn sub_assign(&mut self, (len, cap): (usize, usize)) { - self.length -= len; - self.capacity -= cap; + /// Consolidate all updates at the current `since`. + pub fn consolidate_at_since(&mut self) { + let upper_ts = self.since.as_option().and_then(|t| t.try_step_forward()); + if let Some(upper_ts) = upper_ts { + let upper = Antichain::from_elem(upper_ts); + self.consolidate_before(&upper); + } } -} - -/// A vector that consolidates its contents. -/// -/// The vector is filled with updates until it reaches capacity. At this point, the updates are -/// consolidated to free up space. This process repeats until the consolidation recovered less than -/// half of the vector's capacity, at which point the capacity is doubled. -#[derive(Debug)] -pub(crate) struct ConsolidatingVec { - data: Vec<(D, Diff)>, - /// A lower bound for how small we'll shrink the Vec's capacity. NB: The cap - /// might start smaller than this. - min_capacity: usize, - /// Dampener in the growth rate. 0 corresponds to doubling and in general `n` to `1+1/(n+1)`. - /// - /// If consolidation didn't free enough space, at least a linear amount, increase the capacity - /// Setting this to 0 results in doubling whenever the list is at least half full. - /// Larger numbers result in more conservative approaches that use more CPU, but less memory. - growth_dampener: usize, -} -impl ConsolidatingVec { - /// Return the length of the vector. - pub fn len(&self) -> usize { - self.data.len() + fn log_chain_created(&self, chain: &Chain) { + if let Some(logging) = &self.logging { + logging.chain_created(chain.update_count); + } } - /// Return the capacity of the vector. - pub fn capacity(&self) -> usize { - self.data.capacity() + fn log_chain_dropped(&self, chain: &Chain) { + if let Some(logging) = &self.logging { + logging.chain_dropped(chain.update_count); + } } - /// Pushes `item` into the vector. - /// - /// If the vector does not have sufficient capacity, we'll first consolidate and then increase - /// its capacity if the consolidated results still occupy a significant fraction of the vector. - /// - /// The worst-case cost of this function is O(n log n) in the number of items the vector stores, - /// but amortizes to O(log n). - pub fn push(&mut self, item: (D, Diff)) { - let capacity = self.data.capacity(); - if self.data.len() == capacity { - // The vector is full. First, consolidate to try to recover some space. - self.consolidate(); - - // We may need more capacity if our current capacity is within `1+1/(n+1)` of the length. - // This corresponds to `cap < len + len/(n+1)`, which is the logic we use. - let length = self.data.len(); - let dampener = self.growth_dampener; - if capacity < length + length / (dampener + 1) { - // We would like to increase the capacity by a factor of `1+1/(n+1)`, which involves - // determining the target capacity, and then reserving an amount that achieves this - // while working around the existing length. - let new_cap = capacity + capacity / (dampener + 1); - self.data.reserve_exact(new_cap - length); + /// Update persist sink metrics. + fn update_metrics(&mut self) { + let mut new_size = self.stage.get_size(); + let mut new_length = self.stage.data.len(); + for chain in &self.pending_low { + new_size += chain.get_size(); + new_length += chain.update_count; + } + new_size += self.emitted.get_size(); + new_length += self.emitted.update_count; + for bucket in self.chain.buckets() { + for chain in &bucket.chains { + new_size += chain.get_size(); + new_length += chain.update_count; } } - self.data.push(item); + self.update_metrics_inner(new_size, new_length); } - /// Consolidate the contents. - pub fn consolidate(&mut self) { - consolidate(&mut self.data); + /// Update persist sink metrics to the given new size and length. + fn update_metrics_inner(&mut self, new_size: SizeMetrics, new_length: usize) { + let old_size = self.prev_size; + let old_length = self.prev_update_count; + let len_delta = UpdateDelta::new(new_length, old_length); + let cap_delta = UpdateDelta::new(new_size.capacity, old_size.capacity); + self.metrics + .report_correction_update_deltas(len_delta, cap_delta); + self.worker_metrics + .report_correction_update_totals(new_length, new_size.capacity); - // We may have the opportunity to reclaim allocated memory. - // Given that `push` will at most double the capacity when the vector is more than half full, and - // we want to avoid entering into a resizing cycle, we choose to only shrink if the - // vector's length is less than one fourth of its capacity. - if self.data.len() < self.data.capacity() / 4 { - self.data.shrink_to(self.min_capacity); + if let Some(logging) = &self.logging { + let i = |x: usize| isize::try_from(x).expect("must fit"); + logging.report_size_diff(i(new_size.size) - i(old_size.size)); + logging.report_capacity_diff(i(new_size.capacity) - i(old_size.capacity)); + logging.report_allocations_diff(i(new_size.allocations) - i(old_size.allocations)); } - } - /// Return an iterator over the borrowed items. - pub fn iter(&self) -> impl Iterator { - self.data.iter() + self.prev_size = new_size; + self.prev_update_count = new_length; } } -impl IntoIterator for ConsolidatingVec { - type Item = (D, Diff); - type IntoIter = std::vec::IntoIter<(D, Diff)>; - - fn into_iter(self) -> Self::IntoIter { - self.data.into_iter() +/// Merge the given cursors into one chain. +fn merge_cursors(cursors: Vec>) -> Chain { + match cursors.len() { + 0 => Chain::new(), + 1 => { + let [cur] = cursors.try_into().unwrap(); + cur.into_chain() + } + 2 => { + let [a, b] = cursors.try_into().unwrap(); + merge_2(a, b) + } + _ => merge_many(cursors), } } -impl FromIterator<(D, Diff)> for ConsolidatingVec { - fn from_iter(iter: I) -> Self - where - I: IntoIterator, - { - Self { - data: Vec::from_iter(iter), - min_capacity: 0, - growth_dampener: 0, +/// Merge the given two cursors using a 2-way merge. +/// +/// This function is a specialization of `merge_many` that avoids the overhead of a binary heap. +fn merge_2(cursor1: Cursor, cursor2: Cursor) -> Chain { + let mut rest1 = Some(cursor1); + let mut rest2 = Some(cursor2); + let mut merged = ChainBuilder::default(); + + loop { + match (rest1, rest2) { + (Some(c1), Some(c2)) => { + let (d1, t1, r1) = c1.get(); + let (d2, t2, r2) = c2.get(); + + match (t1, d1).cmp(&(t2, d2)) { + Ordering::Less => { + merged.push_ref((d1, t1, r1)); + rest1 = c1.step(); + rest2 = Some(c2); + } + Ordering::Greater => { + merged.push_ref((d2, t2, r2)); + rest1 = Some(c1); + rest2 = c2.step(); + } + Ordering::Equal => { + let r = r1 + r2; + if r != Diff::ZERO { + merged.push_ref((d1, t1, r)); + } + rest1 = c1.step(); + rest2 = c2.step(); + } + } + } + (Some(c), None) | (None, Some(c)) => { + merged.push_cursor(c); + break; + } + (None, None) => break, } } + + merged.finish() } -impl Extend<(D, Diff)> for ConsolidatingVec { - fn extend(&mut self, iter: I) - where - I: IntoIterator, - { - for item in iter { - self.push(item); +/// Merge the given cursors using a k-way merge with a binary heap. +fn merge_many(cursors: Vec>) -> Chain { + let mut heap = MergeHeap::from_iter(cursors); + let mut merged = ChainBuilder::default(); + while let Some(cursor1) = heap.pop() { + let (data, time, mut diff) = cursor1.get(); + + while let Some((cursor2, r)) = heap.pop_equal(data, time) { + diff += r; + if let Some(cursor2) = cursor2.step() { + heap.push(cursor2); + } + } + + if diff != Diff::ZERO { + merged.push_ref((data, time, diff)); + } + if let Some(cursor1) = cursor1.step() { + heap.push(cursor1); } } -} -/// Helper type for convenient tracking of various size metrics together. -#[derive(Clone, Copy, Debug, Default)] -pub(super) struct SizeMetrics { - pub size: usize, - pub capacity: usize, - pub allocations: usize, + merged.finish() } -impl AddAssign for SizeMetrics { - fn add_assign(&mut self, other: Self) { - self.size += other.size; - self.capacity += other.capacity; - self.allocations += other.allocations; +impl Drop for Correction { + fn drop(&mut self) { + for bucket in self.chain.buckets() { + bucket.chains.iter().for_each(|c| self.log_chain_dropped(c)); + } + self.pending_low + .iter() + .for_each(|c| self.log_chain_dropped(c)); + if !self.emitted.is_empty() { + self.log_chain_dropped(&self.emitted); + } + self.update_metrics_inner(Default::default(), 0); } } -/// A logging event sent from the Tokio task back to the Timely thread. -#[derive(Debug)] -pub enum LoggingEvent { - /// A chain with the given number of updates was created. - ChainCreated(usize), - /// A chain with the given number of updates was dropped. - ChainDropped(usize), - /// The heap size of the correction buffer changed by the given amount. - SizeDiff(NonZeroIsize), - /// The heap capacity of the correction buffer changed by the given amount. - CapacityDiff(NonZeroIsize), - /// The number of allocations of the correction buffer changed by the given amount. - AllocationsDiff(NonZeroIsize), -} - -/// Channel-based logging for corrections on a Tokio task. `Send`-safe. +/// A bucket of `Chain`s, for use in a [`BucketChain`]. /// -/// Sends logging events to the Timely thread, where they are applied to the real `Logging` -/// instance. This allows corrections on the Tokio task to participate in introspection logging -/// without holding `Rc>`. -#[derive(Clone, Debug)] -pub struct ChannelLogging(mpsc::UnboundedSender); - -impl ChannelLogging { - /// Construct a new `ChannelLogging` sending events on the given channel. - pub fn new(tx: mpsc::UnboundedSender) -> Self { - Self(tx) - } - - /// Report the creation of a chain with the given number of updates. - pub fn chain_created(&self, updates: usize) { - let _ = self.0.send(LoggingEvent::ChainCreated(updates)); - } +/// All chains are individually sorted by (time, data) and consolidated, but updates can appear in +/// multiple chains, so consumers must merge the chains to obtain consolidated updates. +struct ChainBucket { + /// The contained chains. + /// + /// Maintained with the chain invariant on pushes; splits can leave it violated until the next + /// push restores it. + chains: Vec>, + /// The size factor of subsequent chains required by the chain invariant. + chain_proportionality: f64, + /// Introspection logging. + logging: Option, +} - /// Report the dropping of a chain with the given number of updates. - pub fn chain_dropped(&self, updates: usize) { - let _ = self.0.send(LoggingEvent::ChainDropped(updates)); +impl fmt::Debug for ChainBucket { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ChainBucket") + .field("chains", &self.chains) + .finish_non_exhaustive() } +} - /// Report a change in heap size by the given amount. - pub fn report_size_diff(&self, diff: isize) { - if let Some(diff) = NonZeroIsize::new(diff) { - let _ = self.0.send(LoggingEvent::SizeDiff(diff)); +impl ChainBucket { + /// Construct a new, empty `ChainBucket`. + fn new(chain_proportionality: f64, logging: Option) -> Self { + Self { + chains: Vec::new(), + chain_proportionality, + logging, } } - /// Report a change in heap capacity by the given amount. - pub fn report_capacity_diff(&self, diff: isize) { - if let Some(diff) = NonZeroIsize::new(diff) { - let _ = self.0.send(LoggingEvent::CapacityDiff(diff)); + /// Push a chain onto the bucket, restoring the chain invariant. + fn push_chain(&mut self, chain: Chain) { + if chain.is_empty() { + return; } - } + if let Some(logging) = &self.logging { + logging.chain_created(chain.update_count); + } + self.chains.push(chain); + + // Restore the chain invariant. + let prop = self.chain_proportionality; + let merge_needed = |chains: &[Chain<_>]| match chains { + [.., prev, last] => { + let last_len = f64::cast_lossy(last.update_count); + let prev_len = f64::cast_lossy(prev.update_count); + last_len * prop > prev_len + } + _ => false, + }; + + while merge_needed(&self.chains) { + let a = self.chains.pop().unwrap(); + let b = self.chains.pop().unwrap(); + if let Some(logging) = &self.logging { + logging.chain_dropped(a.update_count); + logging.chain_dropped(b.update_count); + } - /// Report a change in the number of allocations by the given amount. - pub fn report_allocations_diff(&self, diff: isize) { - if let Some(diff) = NonZeroIsize::new(diff) { - let _ = self.0.send(LoggingEvent::AllocationsDiff(diff)); + let cursors = [a, b].into_iter().filter_map(Chain::into_cursor).collect(); + let merged = merge_cursors(cursors); + if !merged.is_empty() { + if let Some(logging) = &self.logging { + logging.chain_created(merged.update_count); + } + self.chains.push(merged); + } } } -} -/// State for correction buffer logging on the Timely thread. -/// -/// Drains [`LoggingEvent`]s sent by [`ChannelLogging`] from the Tokio task and applies them -/// to the compute and differential loggers. Emits `ArrangementHeapSizeOperator` on construction -/// and `ArrangementHeapSizeOperatorDrop` on drop. -// TODO: Correction buffer logging currently reuses the arrangement batch and size logging. This -// isn't strictly correct as a correction buffer is not an arrangement. Consider refactoring this -// to be about "operator sizes" instead. -pub(super) struct CorrectionLogger { - compute_logger: ComputeLogger, - differential_logger: differential_dataflow::logging::Logger, - operator_id: usize, - rx: mpsc::UnboundedReceiver, - /// Net number of batches logged (BatchEvent - DropEvent). - net_batches: isize, - /// Net number of records logged across all batch/drop/merge events. - net_records: isize, - /// Cumulative heap size delta, for retraction on drop. - net_size: isize, - /// Cumulative heap capacity delta, for retraction on drop. - net_capacity: isize, - /// Cumulative heap allocations delta, for retraction on drop. - net_allocations: isize, -} - -impl fmt::Debug for CorrectionLogger { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("CorrectionLogger") - .field("operator_id", &self.operator_id) - .finish_non_exhaustive() + /// Convert the bucket into its contained chains. + fn into_chains(self) -> Vec> { + self.chains } } -impl CorrectionLogger { - pub fn new( - compute_logger: ComputeLogger, - differential_logger: differential_dataflow::logging::Logger, - operator_id: usize, - address: Vec, - rx: mpsc::UnboundedReceiver, - ) -> Self { - compute_logger.log(&ComputeEvent::ArrangementHeapSizeOperator( - ArrangementHeapSizeOperator { - operator_id, - address, - }, - )); +impl Bucket for ChainBucket { + type Timestamp = Timestamp; + + fn split(self, timestamp: &Self::Timestamp, fuel: &mut i64) -> (Self, Self) { + let mut lower = Self::new(self.chain_proportionality, self.logging.clone()); + let mut upper = Self::new(self.chain_proportionality, self.logging.clone()); + + for chain in self.chains { + // Whole chunks are reused; at most one chunk straddling the timestamp is copied per + // chain. Account fuel at chunk granularity. + *fuel = fuel.saturating_sub(i64::try_from(chain.chunks.len()).expect("must fit")); + + if let Some(logging) = &self.logging { + logging.chain_dropped(chain.update_count); + } + let (lo, hi) = chain.split_at_time(*timestamp); + for (part, target) in [(lo, &mut lower), (hi, &mut upper)] { + if !part.is_empty() { + if let Some(logging) = &self.logging { + logging.chain_created(part.update_count); + } + target.chains.push(part); + } + } + } + + (lower, upper) + } +} +/// A chain of [`Chunk`]s containing updates. +/// +/// All updates in a chain are sorted by (time, data) and consolidated. +/// +/// Note that, in contrast to [`Chunk`]s, chains can be empty. Though we generally try to avoid +/// keeping around empty chains. +#[derive(Debug)] +struct Chain { + /// The contained chunks. + chunks: Vec>, + /// The number of updates contained in all chunks. + update_count: usize, +} + +impl Chain { + /// Construct an empty chain. + fn new() -> Self { Self { - compute_logger, - differential_logger, - operator_id, - rx, - net_batches: 0, - net_records: 0, - net_size: 0, - net_capacity: 0, - net_allocations: 0, - } - } - - /// Drain logging events from the channel and apply them locally. - pub fn apply_events(&mut self) { - use LoggingEvent::*; - - while let Ok(event) = self.rx.try_recv() { - match event { - ChainCreated(length) => { - self.net_batches += 1; - self.net_records += isize::try_from(length).expect("must fit"); - self.differential_logger.log(BatchEvent { - operator: self.operator_id, - length, - }); + chunks: Default::default(), + update_count: 0, + } + } + + /// Return whether the chain is empty. + fn is_empty(&self) -> bool { + self.chunks.is_empty() + } + + /// Push a chunk onto the chain. + /// + /// All updates in the chunk must sort after all updates already in the chain, in + /// (time, data)-order, to ensure the chain remains sorted. + fn push_chunk(&mut self, chunk: Chunk) { + mz_ore::soft_assert_no_log!(self.can_accept_chunk(&chunk)); + + self.update_count += chunk.len(); + self.chunks.push(chunk); + } + + /// Return whether the chain can accept the given chunk at its end while preserving + /// (time, data)-order. + /// + /// NOTE: The cached boundary times settle every case but a tie. On a tie the boundary updates + /// themselves are compared, which materializes both chunks and keeps them resident for the + /// rest of their lifetime. The only caller is the soft assertion in [`Chain::push_chunk`], and + /// soft assertions are live in any build started with `MZ_SOFT_ASSERTIONS` set, so this cost is + /// not confined to debug builds. Ties are reached whenever a run of updates at a single + /// timestamp spans a chunk boundary, which [`ChunkBuilder`] produces for any such run larger + /// than its byte limit. + fn can_accept_chunk(&self, chunk: &Chunk) -> bool { + match self.chunks.last() { + None => true, + Some(last) => match last.last_time().cmp(&chunk.first_time()) { + Ordering::Less => true, + Ordering::Greater => false, + Ordering::Equal => { + let (dc, _, _) = last.last(); + let (d, _, _) = chunk.first(); + dc < d } - ChainDropped(length) => { - self.net_batches -= 1; - self.net_records -= isize::try_from(length).expect("must fit"); - self.differential_logger.log(DropEvent { - operator: self.operator_id, - length, - }); + }, + } + } + + /// Return the last update in the chain, if any. + fn last(&self) -> Option> { + self.chunks.last().map(|c| c.last()) + } + + /// Convert the chain into a cursor over the contained updates. + fn into_cursor(self) -> Option> { + let chunks = self.chunks.into_iter().map(Rc::new).collect(); + Cursor::new(chunks) + } + + /// Return an iterator over the contained updates. + fn iter(&self) -> impl Iterator + '_ { + self.chunks.iter().flat_map(|c| { + (0..c.len()).map(move |i| { + let (d, t, r) = c.index(i); + (D::into_owned(d), t, r) + }) + }) + } + + /// Count the distinct times of updates at times before `time`, up to the given cap. + /// + /// The scan uses one binary search per distinct time, so its cost is bounded by + /// O(cap log chunks). + fn distinct_times_before(&self, time: Timestamp, cap: usize) -> usize { + let mut count = 0; + let mut chunk_idx = 0; + let mut offset = 0; + while count < cap && chunk_idx < self.chunks.len() { + let chunk = &self.chunks[chunk_idx]; + let current = chunk.index(offset).1; + if current >= time { + break; + } + count += 1; + // Skip to the first update at a time greater than `current`. + match chunk.find_time_greater_than(current) { + Some(idx) => offset = idx, + None => { + // All later updates at `current` are in subsequent chunks. + chunk_idx += 1; + offset = 0; + while chunk_idx < self.chunks.len() { + match self.chunks[chunk_idx].find_time_greater_than(current) { + Some(idx) => { + offset = idx; + break; + } + None => chunk_idx += 1, + } + } + } + } + } + count + } + + /// Split the chain at the given time. + /// + /// Returns two chains, the first containing all updates at times < `time`, the second + /// containing all updates at times >= `time`. Chunks fully on either side of `time` are + /// reused; only a chunk straddling `time` is copied. + fn split_at_time(mut self, time: Timestamp) -> (Self, Self) { + let mut lower = Self::new(); + let mut upper = Self::new(); + + let Some(skip_ts) = time.step_back() else { + // Nothing sorts before `time`. + return (lower, self); + }; + + for chunk in self.chunks.drain(..) { + // Route whole chunks by cached boundary times, so a chunk that lands entirely on one + // side is moved without paging it in. Only a straddling chunk is materialized here. + // With soft assertions on, `push_chunk` can still page in a chunk whose boundary time + // ties the chain's last one, see `Chain::can_accept_chunk`. + if chunk.last_time() < time { + lower.push_chunk(chunk); + } else if chunk.first_time() >= time { + upper.push_chunk(chunk); + } else { + // The chunk straddles `time`; copy its two halves. + let idx = chunk + .find_time_greater_than(skip_ts) + .expect("straddles time"); + let mut builder = ChainBuilder::default(); + for i in 0..idx { + builder.push_ref(chunk.index(i)); } - SizeDiff(delta_size) => { - self.net_size += delta_size.get(); - self.compute_logger.log(&ComputeEvent::ArrangementHeapSize( - ArrangementHeapSize { - operator_id: self.operator_id, - delta_size: delta_size.get(), - }, - )); + for part in builder.finish().chunks { + lower.push_chunk(part); } - CapacityDiff(delta_capacity) => { - self.net_capacity += delta_capacity.get(); - self.compute_logger - .log(&ComputeEvent::ArrangementHeapCapacity( - ArrangementHeapCapacity { - operator_id: self.operator_id, - delta_capacity: delta_capacity.get(), - }, - )); + let mut builder = ChainBuilder::default(); + for i in idx..chunk.len() { + builder.push_ref(chunk.index(i)); } - AllocationsDiff(delta_allocations) => { - self.net_allocations += delta_allocations.get(); - self.compute_logger - .log(&ComputeEvent::ArrangementHeapAllocations( - ArrangementHeapAllocations { - operator_id: self.operator_id, - delta_allocations: delta_allocations.get(), - }, - )); + for part in builder.finish().chunks { + upper.push_chunk(part); } } } + + (lower, upper) + } + + /// Return the size of the chain, for use in metrics. + fn get_size(&self) -> SizeMetrics { + let mut metrics = SizeMetrics::default(); + for chunk in &self.chunks { + metrics += chunk.get_size(); + } + metrics } } -impl Drop for CorrectionLogger { - fn drop(&mut self) { - // Drain any events that arrived before the drop. Note that the Tokio task - // may still be running (abort is async), so some events may not have arrived - // yet. We retract any remaining batch/record counts below. - self.apply_events(); - - // Retract any outstanding batch and record counts that weren't balanced by - // ChainDropped events. This handles the case where the Tokio task is aborted - // and its Correction destructors haven't run yet (abort is async). - // - // Each DropEvent retracts one batch and `length` records, so we emit one per - // outstanding batch, with the first carrying all outstanding records. - for i in 0..self.net_batches { - let length = if i == 0 { - usize::try_from(self.net_records).unwrap_or(0) +/// A builder that constructs a [`Chain`] from a stream of updates. +/// +/// Wraps a [`ChunkBuilder`] and drains its minted chunks into a [`Chain`]. Pushed updates must +/// arrive in (time, data) sorted order. +struct ChainBuilder { + builder: ChunkBuilder, + chain: Chain, +} + +impl Default for ChainBuilder { + fn default() -> Self { + Self { + builder: Default::default(), + chain: Chain::new(), + } + } +} + +impl ChainBuilder { + /// Push a reference-form update into the builder. + fn push_ref(&mut self, update: Ref<'_, (D, Timestamp, Diff)>) { + self.builder.push(update); + self.drain(); + } + + /// Push an owned-form update into the builder. + fn push_owned(&mut self, update: &(D, Timestamp, Diff)) { + self.builder.push(update); + self.drain(); + } + + /// Push the updates produced by a cursor into the builder. + fn push_cursor(&mut self, cursor: Cursor) { + let mut rest = Some(cursor); + while let Some(cursor) = rest.take() { + let update = cursor.get(); + self.push_ref(update); + rest = cursor.step(); + } + } + + /// Move any minted chunks from the builder into the chain. + fn drain(&mut self) { + while let Some(chunk) = self.builder.pop() { + self.chain.push_chunk(chunk); + } + } + + /// Finish building, returning the assembled [`Chain`]. + fn finish(self) -> Chain { + let Self { builder, mut chain } = self; + for chunk in builder.finish() { + if chunk.len() > 0 { + chain.push_chunk(chunk); + } + } + chain + } +} + +impl Extend<(D, Timestamp, Diff)> for ChainBuilder { + fn extend>(&mut self, iter: I) { + for update in iter { + self.push_owned(&update); + } + } +} + +/// A cursor over updates in a chain. +/// +/// A cursor provides two guarantees: +/// * Produced updates are ordered and consolidated. +/// * A cursor always yields at least one update. +/// +/// The second guarantee is enforced through the type system: Every method that steps a cursor +/// forward consumes `self` and returns an `Option` that's `None` if the operation stepped +/// over the last update. +/// +/// A cursor holds on to `Rc`s, allowing multiple cursors to produce updates from the same +/// chunks concurrently. As soon as a cursor is done producing updates from a [`Chunk`] it drops +/// its reference. Once the last cursor is done with a [`Chunk`] its memory can be reclaimed. +#[derive(Clone, Debug)] +struct Cursor { + /// The chunks from which updates can still be produced. + chunks: VecDeque>>, + /// The current offset into `chunks.front()`. + chunk_offset: usize, + /// An optional limit for the number of updates the cursor will produce. + limit: Option, + /// An optional overwrite for the timestamp of produced updates. + overwrite_ts: Option, +} + +impl Cursor { + /// Construct a cursor over a list of chunks. + /// + /// Returns `None` if `chunks` is empty. + fn new(chunks: VecDeque>>) -> Option { + if chunks.is_empty() { + return None; + } + + Some(Self { + chunks, + chunk_offset: 0, + limit: None, + overwrite_ts: None, + }) + } + + /// Set a limit for the number of updates this cursor will produce. + /// + /// # Panics + /// + /// Panics if there is already a limit lower than the new one. + fn set_limit(mut self, limit: usize) -> Option { + assert!(self.limit.is_none_or(|l| l >= limit)); + + if limit == 0 { + return None; + } + + // Release chunks made unreachable by the limit. + let mut count = 0; + let mut idx = 0; + let mut offset = self.chunk_offset; + while idx < self.chunks.len() && count < limit { + let chunk = &self.chunks[idx]; + count += chunk.len() - offset; + idx += 1; + offset = 0; + } + self.chunks.truncate(idx); + + if count > limit { + self.limit = Some(limit); + } + + Some(self) + } + + /// Get a reference to the current update. + fn get(&self) -> Ref<'_, (D, Timestamp, Diff)> { + let chunk = self.get_chunk(); + let (d, t, r) = chunk.index(self.chunk_offset); + let t = self.overwrite_ts.unwrap_or(t); + (d, t, r) + } + + /// Get a reference to the current chunk. + fn get_chunk(&self) -> &Chunk { + &self.chunks[0] + } + + /// Step to the next update. + /// + /// Returns the stepped cursor, or `None` if the step was over the last update. + fn step(mut self) -> Option { + if self.chunk_offset == self.get_chunk().len() - 1 { + return self.skip_chunk().map(|(c, _)| c); + } + + self.chunk_offset += 1; + + if let Some(limit) = &mut self.limit { + *limit -= 1; + if *limit == 0 { + return None; + } + } + + Some(self) + } + + /// Skip the remainder of the current chunk. + /// + /// Returns the forwarded cursor and the number of updates skipped, or `None` if no chunks are + /// left after the skip. + fn skip_chunk(mut self) -> Option<(Self, usize)> { + let chunk = self.chunks.pop_front().expect("cursor invariant"); + + if self.chunks.is_empty() { + return None; + } + + let skipped = chunk.len() - self.chunk_offset; + self.chunk_offset = 0; + + if let Some(limit) = &mut self.limit { + if skipped >= *limit { + return None; + } + *limit -= skipped; + } + + Some((self, skipped)) + } + + /// Skip all updates with times <= the given time. + /// + /// Returns the forwarded cursor and the number of updates skipped, or `None` if no updates are + /// left after the skip. + fn skip_time(mut self, time: Timestamp) -> Option<(Self, usize)> { + if self.overwrite_ts.is_some_and(|ts| ts <= time) { + return None; + } else if self.get().1 > time { + return Some((self, 0)); + } + + let mut skipped = 0; + + let new_offset = loop { + let chunk = self.get_chunk(); + if let Some(index) = chunk.find_time_greater_than(time) { + break index; + } + + let (cursor, count) = self.skip_chunk()?; + self = cursor; + skipped += count; + }; + + skipped += new_offset - self.chunk_offset; + self.chunk_offset = new_offset; + + Some((self, skipped)) + } + + /// Advance all updates in this cursor by the given `since_ts`. + /// + /// Returns a list of cursors, each of which yields ordered and consolidated updates that have + /// been advanced by `since_ts`. + fn advance_by(mut self, since_ts: Timestamp) -> Vec { + // If the cursor has an `overwrite_ts`, all its updates are at the same time already. We + // only need to advance the `overwrite_ts` by the `since_ts`. + if let Some(ts) = self.overwrite_ts { + if ts < since_ts { + self.overwrite_ts = Some(since_ts); + } + return vec![self]; + } + + // Otherwise we need to split the cursor so that each new cursor only yields runs of + // updates that are correctly (time, data)-ordered when advanced by `since_ts`. We achieve + // this by splitting the cursor at each time <= `since_ts`. + let mut splits = Vec::new(); + let mut remaining = Some(self); + + while let Some(cursor) = remaining.take() { + let (_, time, _) = cursor.get(); + if time >= since_ts { + splits.push(cursor); + break; + } + + let mut current = cursor.clone(); + if let Some((cursor, skipped)) = cursor.skip_time(time) { + remaining = Some(cursor); + current = current.set_limit(skipped).expect("skipped at least 1"); + } + current.overwrite_ts = Some(since_ts); + splits.push(current); + } + + splits + } + + /// Drain the cursor into a [`Chain`]. + /// + /// This reuses the underlying chunks if possible, and writes new ones otherwise. + fn into_chain(self) -> Chain { + match self.try_unwrap() { + Ok(chain) => chain, + Err((_, cursor)) => { + let mut builder = ChainBuilder::default(); + builder.push_cursor(cursor); + builder.finish() + } + } + } + + /// Attempt to unwrap the cursor into a [`Chain`]. + /// + /// This operation efficiently reuses chunks by directly inserting them into the output chain + /// where possible. + /// + /// An unwrap is only successful if the cursor's `limit` and `overwrite_ts` are both `None` and + /// the cursor has unique references to its chunks. If the unwrap fails, this method returns an + /// `Err` containing the cursor in an unchanged state, allowing the caller to convert it into a + /// chain by copying chunks rather than reusing them. + fn try_unwrap(self) -> Result, (&'static str, Self)> { + if self.limit.is_some() { + return Err(("cursor with limit", self)); + } + if self.overwrite_ts.is_some() { + return Err(("cursor with overwrite_ts", self)); + } + if self.chunks.iter().any(|c| Rc::strong_count(c) != 1) { + return Err(("cursor on shared chunks", self)); + } + + let mut builder = ChainBuilder::default(); + let mut remaining = Some(self); + + // We might be partway through the first chunk, in which case we can't reuse it but need to + // allocate a new one to contain only the updates the cursor can still yield. + while let Some(cursor) = remaining.take() { + if cursor.chunk_offset == 0 { + remaining = Some(cursor); + break; + } + let update = cursor.get(); + builder.push_ref(update); + remaining = cursor.step(); + } + + let mut chain = builder.finish(); + if let Some(cursor) = remaining { + for chunk in cursor.chunks { + let chunk = Rc::into_inner(chunk).expect("checked above"); + chain.push_chunk(chunk); + } + } + + Ok(chain) + } +} + +/// A non-empty chunk of updates, backed by a columnar region. +/// +/// All updates in a chunk are sorted by (time, data) and consolidated. +/// +/// Chunks are immutable once created. They are produced by [`ChunkBuilder`], which mints a +/// new chunk whenever its in-progress columnar container reaches a fixed serialized byte +/// boundary (~2 MiB, matching the ship granularity used elsewhere in the codebase), so each +/// chunk corresponds to a single, predictably sized allocation. +struct Chunk { + /// The paged-out form, taken on first materialization. + /// + /// A `Mutex` (not `RefCell`) keeps the chunk `Sync`: cursors hold chunks behind a shared + /// `Rc`, and the iterator returned by [`Correction::updates_before`] borrows them across + /// the persist writer's `await`, so `&Chunk` must be `Send`. The lock is taken once, at + /// materialization, and is otherwise uncontended (the sink runs single-threaded per worker). + paged: Mutex>>, + /// The materialized form, populated lazily by [`Chunk::column`] on first access. + /// + /// An `OnceLock` (not `OnceCell`) for the same `Sync` reason. Once set the slot is never + /// cleared, so its address is stable and [`Chunk::index`] can hand out `Ref<'_>` borrows tied + /// to `&self`. The allocation is freed when the chunk drops, which bounds resident memory to + /// the chunks under an active merge front. + resident: OnceLock>, + /// Number of updates, cached so `len` and chain bookkeeping never page the chunk in. + len: usize, + /// Time of the first update, cached so boundary checks (`split_at_time`, `can_accept`) route + /// a resting chunk without materializing it. + first_time: Timestamp, + /// Time of the last update, cached likewise. + last_time: Timestamp, +} + +impl fmt::Debug for Chunk { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "Chunk(<{}>)", self.len()) + } +} + +impl Chunk { + /// Page the given non-empty column out into a chunk. + /// + /// Reads the cached metadata (length, boundary times) while the column is still resident, then + /// hands it to the global column pager. The policy decides whether it actually spills; either + /// way the chunk is born paged and materializes lazily on first read. + /// + /// # Panics + /// + /// Panics if the column is empty. Chunks are non-empty by construction; [`ChunkBuilder`] only + /// ever builds a chunk from a populated column. + fn from_column(mut data: Column<(D, Timestamp, Diff)>) -> Self { + let (len, first_time, last_time) = { + let borrowed = data.borrow(); + let len = borrowed.len(); + assert!(len > 0, "chunks are non-empty"); + (len, borrowed.get(0).1, borrowed.get(len - 1).1) + }; + + let paged = column_pager::global_pager().page(&mut data); + Self { + paged: Mutex::new(Some(paged)), + resident: OnceLock::new(), + len, + first_time, + last_time, + } + } + + /// Materialize the chunk's column, paging it in on first access. + /// + /// The returned reference is valid for as long as `&self`: the `OnceLock` slot is never + /// cleared once populated, so its contents have a stable address. + fn column(&self) -> &Column<(D, Timestamp, Diff)> { + self.resident.get_or_init(|| { + let paged = self + .paged + .lock() + .expect("pager mutex poisoned") + .take() + .expect("paged form present until materialized"); + column_pager::global_pager().take(paged) + }) + } + + /// Return the number of updates in the chunk. + fn len(&self) -> usize { + self.len + } + + /// Return the update at the given index, paging the chunk in if necessary. + /// + /// # Panics + /// + /// Panics if the given index is not populated. + fn index(&self, idx: usize) -> Ref<'_, (D, Timestamp, Diff)> { + self.column().borrow().get(idx) + } + + /// Return the first update in the chunk, paging the chunk in if necessary. + fn first(&self) -> Ref<'_, (D, Timestamp, Diff)> { + self.index(0) + } + + /// Return the last update in the chunk, paging the chunk in if necessary. + fn last(&self) -> Ref<'_, (D, Timestamp, Diff)> { + self.index(self.len - 1) + } + + /// Return the time of the first update, without materializing the chunk. + fn first_time(&self) -> Timestamp { + self.first_time + } + + /// Return the time of the last update, without materializing the chunk. + fn last_time(&self) -> Timestamp { + self.last_time + } + + /// Return the index of the first update at a time greater than `time`, or `None` if no such + /// update exists. + /// + /// The early-out uses the cached last time, so a chunk whose updates are all at or before + /// `time` is skipped without paging it in. + fn find_time_greater_than(&self, time: Timestamp) -> Option { + if self.last_time <= time { + return None; + } + + let mut lower = 0; + let mut upper = self.len; + while lower < upper { + let idx = (lower + upper) / 2; + if self.index(idx).1 > time { + upper = idx; } else { - 0 + lower = idx + 1; + } + } + + Some(lower) + } + + /// Return the size of the chunk, for use in metrics. + /// + /// Reports resident bytes only: a chunk still spilled (on swap or in a pager file) is not part + /// of RSS and contributes nothing, matching the accounting in + /// [`mz_timely_util::columnar::merge_batcher`]. + fn get_size(&self) -> SizeMetrics { + let resident = |col: &Column<(D, Timestamp, Diff)>| { + let bytes = col.length_in_bytes(); + SizeMetrics { + size: bytes, + capacity: bytes, + allocations: 1, + } + }; + + if let Some(col) = self.resident.get() { + return resident(col); + } + // Not yet materialized: a policy that kept the column resident still occupies RSS, so + // account for it; a genuinely spilled column does not. + match &*self.paged.lock().expect("pager mutex poisoned") { + Some(PagedColumn::Resident(col, _)) => resident(col), + _ => SizeMetrics::default(), + } + } +} + +/// Builder that produces a stream of fixed-size [`Chunk`]s. +/// +/// Wraps [`mz_timely_util::columnar::builder::ColumnBuilder`], which mints a new +/// [`Column::Align`] chunk whenever its in-progress columnar container reaches a fixed +/// serialized byte boundary (~2 MiB, matching the ship granularity used elsewhere in the +/// codebase). Each minted chunk is therefore a single, predictably-sized aligned allocation. +struct ChunkBuilder { + inner: mz_timely_util::columnar::builder::ColumnBuilder<(D, Timestamp, Diff)>, +} + +impl Default for ChunkBuilder { + fn default() -> Self { + Self { + inner: Default::default(), + } + } +} + +impl ChunkBuilder { + /// Push an update into the builder. + /// + /// Accepts whatever the inner [`ColumnBuilder`]'s [`PushInto`] impl accepts — both the + /// `Ref<'_, (D, T, R)>` refs produced by cursors and `&(D, T, R)` references to owned + /// tuples drained from the staging buffer. + /// + /// [`ColumnBuilder`]: mz_timely_util::columnar::builder::ColumnBuilder + /// [`PushInto`]: timely::container::PushInto + #[inline] + fn push(&mut self, item: T) + where + mz_timely_util::columnar::builder::ColumnBuilder<(D, Timestamp, Diff)>: + timely::container::PushInto, + { + timely::container::PushInto::push_into(&mut self.inner, item); + } + + /// Pop a finished chunk, if one is available. + fn pop(&mut self) -> Option> { + use timely::container::ContainerBuilder; + // `ColumnBuilder::extract` stashes the popped chunk in its `finished` slot so the + // caller can read it through `&mut`; move it out with `mem::take` so we own it + // (leaves `Column::Typed(Default::default())` behind, which the next `extract` + // overwrites). + self.inner + .extract() + .map(|c| Chunk::from_column(std::mem::take(c))) + } + + /// Finalize the builder: flush any in-progress updates as a typed chunk and drain pending. + fn finish(mut self) -> impl Iterator> { + use timely::container::ContainerBuilder; + // `ColumnBuilder::finish` flushes the in-progress container into the pending queue + // (as `Column::Typed`) and returns the first pending entry. Subsequent calls drain + // the rest until `None`. Translate that into an owning iterator. + // + // `finish` can hand back an empty column (e.g. when the last shipped chunk landed exactly + // on the boundary). Skip those: `Chunk::from_column` requires a non-empty column, and an + // empty chunk would needlessly engage the pager. + std::iter::from_fn(move || { + loop { + let col = std::mem::take(self.inner.finish()?); + if !col.is_empty() { + return Some(Chunk::from_column(col)); + } + } + }) + } +} + +/// A buffer for staging updates before they are inserted into the sorted chains. +#[derive(Debug)] +struct Stage { + /// The contained updates. + /// + /// This vector has a fixed capacity equal to the [`Chunk`] capacity. + data: Vec<(D, Timestamp, Diff)>, + /// Introspection logging. + /// + /// We want to report the number of records in the stage. To do so, we pretend that the stage + /// is a chain, and every time the number of updates inside changes, the chain gets dropped and + /// re-created. + logging: Option, +} + +impl Stage { + fn new(logging: Option, chunk_capacity: usize) -> Self { + // For logging, we pretend the stage consists of a single chain. + if let Some(logging) = &logging { + logging.chain_created(0); + } + + Self { + data: Vec::with_capacity(chunk_capacity), + logging, + } + } + + /// Insert a batch of updates, possibly producing a batch of sorted, consolidated updates + /// ready to be stored. + fn insert( + &mut self, + updates: &mut Vec<(D, Timestamp, Diff)>, + ) -> Option> { + if updates.is_empty() { + return None; + } + + let prev_length = self.ilen(); + + // Determine how many chunks we can fill with the available updates. + let update_count = self.data.len() + updates.len(); + let chunk_capacity = self.data.capacity(); + let chunk_count = update_count / chunk_capacity; + + let mut new_updates = updates.drain(..); + + // If we have enough shipable updates, collect them and consolidate. + let maybe_ready = if chunk_count > 0 { + let ship_count = chunk_count * chunk_capacity; + let mut buffer = Vec::with_capacity(ship_count); + + buffer.append(&mut self.data); + while buffer.len() < ship_count { + let update = new_updates.next().unwrap(); + buffer.push(update); + } + + consolidate(&mut buffer); + + Some(buffer) + } else { + None + }; + + // Stage the remaining updates. + Extend::extend(&mut self.data, new_updates); + + self.log_length_diff(self.ilen() - prev_length); + + maybe_ready + } + + /// Flush all currently staged updates, returning them sorted and consolidated. + fn flush(&mut self) -> Option> { + self.log_length_diff(-self.ilen()); + + consolidate(&mut self.data); + + if self.data.is_empty() { + return None; + } + + let capacity = self.data.capacity(); + let data = std::mem::replace(&mut self.data, Vec::with_capacity(capacity)); + Some(data) + } + + /// Advance the times of staged updates by the given `since`. + fn advance_times(&mut self, since: &Antichain) { + let Some(since_ts) = since.as_option() else { + // If the since is the empty frontier, discard all updates. + self.log_length_diff(-self.ilen()); + self.data.clear(); + return; + }; + + for (_, time, _) in &mut self.data { + *time = std::cmp::max(*time, *since_ts); + } + } + + /// Return the size of the stage, for use in metrics. + /// + /// Note: We don't follow pointers here, so the returned `size` and `capacity` values are + /// under-estimates. That's fine as the stage should always be small. + fn get_size(&self) -> SizeMetrics { + SizeMetrics { + size: self.data.len() * std::mem::size_of::<(D, Timestamp, Diff)>(), + capacity: self.data.capacity() * std::mem::size_of::<(D, Timestamp, Diff)>(), + allocations: 1, + } + } + + /// Return the number of updates in the stage, as an `isize`. + fn ilen(&self) -> isize { + self.data.len().try_into().expect("must fit") + } + + fn log_length_diff(&self, diff: isize) { + let Some(logging) = &self.logging else { return }; + if diff > 0 { + let count = usize::try_from(diff).expect("must fit"); + logging.chain_created(count); + logging.chain_dropped(0); + } else if diff < 0 { + let count = usize::try_from(-diff).expect("must fit"); + logging.chain_created(0); + logging.chain_dropped(count); + } + } +} + +impl Drop for Stage { + fn drop(&mut self) { + if let Some(logging) = &self.logging { + logging.chain_dropped(self.data.len()); + } + } +} + +/// Sort and consolidate the given list of updates. +/// +/// This function is the same as [`differential_dataflow::consolidation::consolidate_updates`], +/// except that it sorts updates by (time, data) instead of (data, time). +fn consolidate(updates: &mut Vec<(D, Timestamp, Diff)>) { + if updates.len() <= 1 { + return; + } + + let diff = |update: &(_, _, Diff)| update.2; + + updates.sort_unstable_by(|(d1, t1, _), (d2, t2, _)| (t1, d1).cmp(&(t2, d2))); + + let mut offset = 0; + let mut accum = diff(&updates[0]); + + for idx in 1..updates.len() { + let this = &updates[idx]; + let prev = &updates[idx - 1]; + if this.0 == prev.0 && this.1 == prev.1 { + accum += diff(&updates[idx]); + } else { + if accum != Diff::ZERO { + updates.swap(offset, idx - 1); + updates[offset].2 = accum; + offset += 1; + } + accum = diff(&updates[idx]); + } + } + + if accum != Diff::ZERO { + let len = updates.len(); + updates.swap(offset, len - 1); + updates[offset].2 = accum; + offset += 1; + } + + updates.truncate(offset); +} + +/// Compare two columnar refs that have unrelated input lifetimes. +/// +/// `::Ref<'a>` is an associated-type projection through a trait, so +/// the compiler treats it as invariant in `'a` and won't auto-shorten the inputs by variance. +/// We instead explicitly reborrow both to a fresh, local lifetime `'x` via +/// [`Columnar::reborrow`] before letting the inner `==` pick up the `for<'a> Ref<'a>: Eq` +/// bound on [`Data`]. +#[inline] +fn refs_eq(a: Ref<'_, D>, b: Ref<'_, D>) -> bool { + #[inline] + fn eq<'x, D: Data>(a: Ref<'x, D>, b: Ref<'x, D>) -> bool { + a == b + } + eq::(D::reborrow(a), D::reborrow(b)) +} + +/// A binary heap specialized for merging [`Cursor`]s. +struct MergeHeap(BinaryHeap>); + +impl FromIterator> for MergeHeap { + fn from_iter>>(cursors: I) -> Self { + let inner = cursors.into_iter().map(MergeCursor).collect(); + Self(inner) + } +} + +impl MergeHeap { + /// Pop the next cursor (the one yielding the least update) from the heap. + fn pop(&mut self) -> Option> { + self.0.pop().map(|MergeCursor(c)| c) + } + + /// Pop the next cursor from the heap, provided the data and time of its current update are + /// equal to the given values. + /// + /// Returns both the cursor and the diff corresponding to `data` and `time`. + fn pop_equal(&mut self, data: Ref<'_, D>, time: Timestamp) -> Option<(Cursor, Diff)> { + let r = { + let MergeCursor(cursor) = self.0.peek()?; + let (d, t, r) = cursor.get(); + if t != time || !refs_eq::(d, data) { + return None; + } + r + }; + let cursor = self.pop().expect("checked above"); + Some((cursor, r)) + } + + /// Push a cursor onto the heap. + fn push(&mut self, cursor: Cursor) { + self.0.push(MergeCursor(cursor)); + } +} + +/// A wrapper for [`Cursor`]s on a [`MergeHeap`]. +/// +/// Implements the cursor ordering required for merging cursors. +struct MergeCursor(Cursor); + +impl PartialEq for MergeCursor { + fn eq(&self, other: &Self) -> bool { + self.cmp(other).is_eq() + } +} + +impl Eq for MergeCursor {} + +impl PartialOrd for MergeCursor { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for MergeCursor { + fn cmp(&self, other: &Self) -> Ordering { + let (d1, t1, _) = self.0.get(); + let (d2, t2, _) = other.0.get(); + (t1, d1).cmp(&(t2, d2)).reverse() + } +} + +#[cfg(test)] +mod tests { + use mz_ore::metrics::MetricsRegistry; + use mz_persist_client::cfg::PersistConfig; + use mz_persist_client::metrics::Metrics; + use mz_repr::{Diff, Timestamp}; + + use super::*; + + /// A naive correction buffer, used as a reference for the real implementation. + /// + /// Keeps every update in a flat vector and does all the work (time advancement, + /// consolidation, filtering) on read. Obviously too slow for production, but simple enough + /// to be evidently correct. + struct ReferenceCorrection { + updates: Vec<(String, Timestamp, Diff)>, + since: Antichain, + } + + impl ReferenceCorrection { + fn new() -> Self { + Self { + updates: Vec::new(), + since: Antichain::from_elem(Timestamp::MIN), + } + } + + fn insert(&mut self, updates: &mut Vec<(String, Timestamp, Diff)>) { + let Some(since_ts) = self.since.as_option() else { + updates.clear(); + return; }; - self.differential_logger.log(DropEvent { - operator: self.operator_id, - length, - }); - } - - // Retract any outstanding heap size/capacity/allocations deltas. - if self.net_size != 0 { - self.compute_logger - .log(&ComputeEvent::ArrangementHeapSize(ArrangementHeapSize { - operator_id: self.operator_id, - delta_size: -self.net_size, - })); - } - if self.net_capacity != 0 { - self.compute_logger - .log(&ComputeEvent::ArrangementHeapCapacity( - ArrangementHeapCapacity { - operator_id: self.operator_id, - delta_capacity: -self.net_capacity, - }, - )); - } - if self.net_allocations != 0 { - self.compute_logger - .log(&ComputeEvent::ArrangementHeapAllocations( - ArrangementHeapAllocations { - operator_id: self.operator_id, - delta_allocations: -self.net_allocations, - }, - )); - } - - self.compute_logger - .log(&ComputeEvent::ArrangementHeapSizeOperatorDrop( - ArrangementHeapSizeOperatorDrop { - operator_id: self.operator_id, - }, - )); + for (d, t, r) in updates.drain(..) { + self.updates.push((d, std::cmp::max(t, *since_ts), r)); + } + } + + fn insert_negated(&mut self, updates: &mut Vec<(String, Timestamp, Diff)>) { + for (_, _, r) in &mut *updates { + *r = -*r; + } + self.insert(updates); + } + + fn advance_since(&mut self, since: Antichain) { + assert!(PartialOrder::less_equal(&self.since, &since)); + match since.as_option() { + Some(ts) => { + for (_, t, _) in &mut self.updates { + *t = std::cmp::max(*t, *ts); + } + } + None => self.updates.clear(), + } + self.since = since; + } + + fn updates_before(&self, upper: &Antichain) -> Vec<(String, Timestamp, Diff)> { + let mut out: Vec<_> = self + .updates + .iter() + .filter(|(_, t, _)| !upper.less_equal(t)) + .cloned() + .collect(); + differential_dataflow::consolidation::consolidate_updates(&mut out); + out + } + } + + #[mz_ore::test] + fn chain_builder_update_count_matches_items() { + let mut builder = ChainBuilder::::default(); + for i in 0..10_u64 { + let d = i64::try_from(i).expect("fits"); + builder.push_owned(&(d, Timestamp::new(i), Diff::ONE)); + } + let chain = builder.finish(); + assert_eq!(chain.update_count, chain.iter().count()); + } + + /// Push enough updates to cross at least one `mint()` boundary, forcing the + /// `Align` encode -> `from_bytes` decode roundtrip (the spilling path this data + /// structure exists to support), and assert `iter()` roundtrips values, order, + /// and diffs across the spill boundary. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow: crossing the ~2 MiB mint boundary needs ~200k updates + fn chain_builder_roundtrips_across_mint_boundary() { + // A single `mint()` fires near the ~2 MiB (`SHIP_WORDS`) serialized boundary. With + // three 8-byte columns per update that's tens of thousands of updates; pushing 200k + // comfortably forces multiple mints. + let count = 200_000_u64; + + let mut builder = ChainBuilder::::default(); + for i in 0..count { + let d = i64::try_from(i).expect("fits"); + builder.push_owned(&(d, Timestamp::new(i), Diff::ONE)); + } + let chain = builder.finish(); + + // Crossing the mint boundary must have produced more than one chunk; otherwise the spill + // path (each minted chunk is paged out and read back through the pager) wouldn't be + // exercised. The chunk payload itself is now behind the pager (see [`Chunk`]), so we + // assert on chunk count rather than inspecting the column variant directly. + assert!( + chain.chunks.len() > 1, + "expected multiple minted chunks, got {} chunk(s): {:?}", + chain.chunks.len(), + chain.chunks, + ); + + // `iter()` must roundtrip every update, in order, with correct diffs. + assert_eq!(chain.update_count, usize::try_from(count).expect("fits")); + let mut expected = 0_u64; + for (d, t, r) in chain.iter() { + assert_eq!(d, i64::try_from(expected).expect("fits")); + assert_eq!(t, Timestamp::new(expected)); + assert_eq!(r, Diff::ONE); + expected += 1; + } + assert_eq!(expected, count); + } + + fn sink_metrics() -> SinkMetrics { + let registry = MetricsRegistry::new(); + let metrics = Metrics::new(&PersistConfig::new_for_tests(), ®istry); + metrics.sink.clone() + } + + /// Run the same stepwise-drain workload through the reference implementation and + /// `Correction` and assert that they emit the same updates at every step. + /// + /// Models the `write_batches` operator catching up through many distinct timestamps: the + /// desired input runs ahead, batches are written one timestamp at a time, and written updates + /// come back negated through the persist feedback. + #[mz_ore::test] + // Columnation regions are not Stacked Borrows compliant: later pushes invalidate the + // provenance of previously stored items under Miri. + #[cfg_attr(miri, ignore)] + fn equivalence_with_reference() { + let sink_metrics = sink_metrics(); + + let mut reference = ReferenceCorrection::new(); + let mut correction = Correction::::with_params( + sink_metrics.clone(), + sink_metrics.for_worker(0), + None, + 3.0, + 8 * 1024, + ); + + let num_ts = 50; + let keys = 4; + + // Upsert-style input: every timestamp updates each key, retracting the previous value. + let batch = |t: u64| -> Vec<(String, Timestamp, Diff)> { + (0..keys) + .flat_map(|k| { + let addition = (format!("{k}-{t}"), Timestamp::from(t), Diff::ONE); + let retraction = t + .checked_sub(1) + .map(|p| (format!("{k}-{p}"), Timestamp::from(t), -Diff::ONE)); + std::iter::once(addition).chain(retraction) + }) + .collect() + }; + + // Pre-fill both with all batches, like a catch-up where the input runs ahead. + for t in 0..num_ts { + reference.insert(&mut batch(t)); + correction.insert(&mut batch(t)); + } + + // Drain stepwise, with persist feedback, comparing emissions. + for t in 0..num_ts { + let upper = Antichain::from_elem(Timestamp::from(t + 1)); + + let mut expected = reference.updates_before(&upper); + let mut actual: Vec<_> = correction.updates_before(&upper).collect(); + expected.sort(); + actual.sort(); + assert_eq!(expected, actual, "diverged at t={t}"); + // Guard against the comparison passing because both sides emit nothing. + assert!(!actual.is_empty(), "no updates emitted at t={t}"); + + reference.insert_negated(&mut expected.clone()); + correction.insert_negated(&mut actual); + reference.advance_since(upper.clone()); + correction.advance_since(upper); + } + + // Compare the final state at the since. + let upper = Antichain::from_elem(Timestamp::from(num_ts + 1)); + correction.consolidate_at_since(); + let mut expected = reference.updates_before(&upper); + let mut actual: Vec<_> = correction.updates_before(&upper).collect(); + expected.sort(); + actual.sort(); + assert_eq!(expected, actual); + } + + /// A since jump across many distinct buffered timestamps must collapse them onto the since. + #[mz_ore::test] + // Columnation regions are not Stacked Borrows compliant: later pushes invalidate the + // provenance of previously stored items under Miri. + #[cfg_attr(miri, ignore)] + fn since_jump() { + let sink_metrics = sink_metrics(); + let mut correction = Correction::::with_params( + sink_metrics.clone(), + sink_metrics.for_worker(0), + None, + 3.0, + 8 * 1024, + ); + + let num_ts = 100; + for t in 0..num_ts { + correction.insert(&mut vec![ + (format!("a-{t}"), Timestamp::from(t), Diff::ONE), + (format!("a-{t}"), Timestamp::from(t), -Diff::ONE), + (format!("b-{t}"), Timestamp::from(t), Diff::ONE), + ]); + } + + correction.advance_since(Antichain::from_elem(Timestamp::from(num_ts))); + correction.consolidate_at_since(); + + let upper = Antichain::from_elem(Timestamp::from(num_ts + 1)); + let out: Vec<_> = correction.updates_before(&upper).collect(); + assert_eq!(out.len(), usize::try_from(num_ts).unwrap()); + assert!( + out.iter() + .all(|(_, t, r)| *t == Timestamp::from(num_ts) && *r == Diff::ONE) + ); + } + + /// Reads must not observe updates at or beyond their `upper`, even when the `upper` is not + /// beyond the `since`. + #[mz_ore::test] + // Columnation regions are not Stacked Borrows compliant: later pushes invalidate the + // provenance of previously stored items under Miri. + #[cfg_attr(miri, ignore)] + fn upper_not_beyond_since() { + let sink_metrics = sink_metrics(); + let mut correction = Correction::::with_params( + sink_metrics.clone(), + sink_metrics.for_worker(0), + None, + 3.0, + 8 * 1024, + ); + + correction.insert(&mut vec![( + "a".to_owned(), + Timestamp::from(5_u64), + Diff::ONE, + )]); + correction.advance_since(Antichain::from_elem(Timestamp::from(10_u64))); + + // The update logically lives at time 10 now, so a read before 7 must be empty. + let upper = Antichain::from_elem(Timestamp::from(7_u64)); + assert_eq!(correction.updates_before(&upper).count(), 0); + + // A read before 11 must emit it, advanced to the since. + let upper = Antichain::from_elem(Timestamp::from(11_u64)); + let out: Vec<_> = correction.updates_before(&upper).collect(); + assert_eq!( + out, + vec![("a".to_owned(), Timestamp::from(10_u64), Diff::ONE)] + ); + } + + /// A [`PagingPolicy`] that always spills to the swap backend, uncompressed. + /// + /// The default global pager keeps every chunk resident; installing this drives the actual + /// spill path so the tests exercise [`Chunk::column`]'s page-in through [`mz_ore::pager`]. + /// + /// [`PagingPolicy`]: column_pager::PagingPolicy + struct ForceSwap; + + impl column_pager::PagingPolicy for ForceSwap { + fn decide(&self, _hint: column_pager::PageHint) -> column_pager::PageDecision { + column_pager::PageDecision::Page { + backend: mz_ore::pager::Backend::Swap, + codec: None, + } + } + fn record(&self, _event: column_pager::PageEvent) {} + } + + /// Install a global pager that spills every chunk to swap for the duration of `f`, then + /// restore the default (disabled) pager. The global pager is process-wide; concurrent tests + /// only ever observe a correct round-trip regardless of backend, so racing on it is benign. + fn with_swap_pager(f: impl FnOnce() -> R) -> R { + use std::sync::Arc; + column_pager::set_global_pager(column_pager::ColumnPager::new(Arc::new(ForceSwap))); + let result = f(); + column_pager::set_global_pager(column_pager::ColumnPager::disabled()); + result + } + + /// Build a chain crossing the mint boundary while every chunk is spilled to swap, then assert + /// `iter()` (the read path behind `updates_before`) pages each chunk back in and roundtrips + /// values, order, and diffs. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] // madvise on the swap backend is unsupported under miri + fn iter_roundtrips_through_swap_backend() { + let count = 200_000_u64; + with_swap_pager(|| { + let mut builder = ChainBuilder::::default(); + for i in 0..count { + let d = i64::try_from(i).expect("fits"); + builder.push_owned(&(d, Timestamp::new(i), Diff::ONE)); + } + let chain = builder.finish(); + assert!(chain.chunks.len() > 1, "expected multiple minted chunks"); + assert_eq!(chain.update_count, usize::try_from(count).expect("fits")); + + let mut expected = 0_u64; + for (d, t, r) in chain.iter() { + assert_eq!(d, i64::try_from(expected).expect("fits")); + assert_eq!(t, Timestamp::new(expected)); + assert_eq!(r, Diff::ONE); + expected += 1; + } + assert_eq!(expected, count); + }); + } + + /// Drive a [`Cursor`] over a spilled, multi-chunk chain to completion (the access pattern + /// merges use). Each step pages the front chunk back in via [`Chunk::column`]; assert the + /// cursor yields every update in order. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] // madvise on the swap backend is unsupported under miri + fn cursor_steps_through_swap_backend() { + let count = 200_000_u64; + with_swap_pager(|| { + let mut builder = ChainBuilder::::default(); + for i in 0..count { + let d = i64::try_from(i).expect("fits"); + builder.push_owned(&(d, Timestamp::new(i), Diff::ONE)); + } + let chain = builder.finish(); + assert!(chain.chunks.len() > 1, "expected multiple minted chunks"); + + let mut rest = chain.into_cursor(); + let mut expected = 0_u64; + while let Some(cursor) = rest.take() { + let (d, t, r) = cursor.get(); + assert_eq!(i64::into_owned(d), i64::try_from(expected).expect("fits")); + assert_eq!(t, Timestamp::new(expected)); + assert_eq!(r, Diff::ONE); + expected += 1; + rest = cursor.step(); + } + assert_eq!(expected, count); + }); } } diff --git a/src/compute/src/sink/correction/logging.rs b/src/compute/src/sink/correction/logging.rs new file mode 100644 index 0000000000000..4ae221e467601 --- /dev/null +++ b/src/compute/src/sink/correction/logging.rs @@ -0,0 +1,283 @@ +// Copyright Materialize, Inc. and contributors. All rights reserved. +// +// Use of this software is governed by the Business Source License +// included in the LICENSE file. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0. + +//! Introspection logging for the MV sink's correction buffer. +//! +//! The correction buffer lives on a Tokio task, while the introspection loggers are owned by the +//! Timely thread and are not `Send`. [`ChannelLogging`] bridges the two: the buffer reports size +//! and chain changes as [`LoggingEvent`]s over a channel, and [`CorrectionLogger`] drains them on +//! the Timely thread. + +use std::fmt; +use std::num::NonZeroIsize; +use std::ops::AddAssign; + +use differential_dataflow::logging::{BatchEvent, DropEvent}; +use tokio::sync::mpsc; + +use crate::logging::compute::{ + ArrangementHeapAllocations, ArrangementHeapCapacity, ArrangementHeapSize, + ArrangementHeapSizeOperator, ArrangementHeapSizeOperatorDrop, ComputeEvent, + Logger as ComputeLogger, +}; + +/// Helper type for convenient tracking of various size metrics together. +#[derive(Clone, Copy, Debug, Default)] +pub(super) struct SizeMetrics { + pub size: usize, + pub capacity: usize, + pub allocations: usize, +} + +impl AddAssign for SizeMetrics { + fn add_assign(&mut self, other: Self) { + self.size += other.size; + self.capacity += other.capacity; + self.allocations += other.allocations; + } +} + +/// A logging event sent from the Tokio task back to the Timely thread. +#[derive(Debug)] +pub enum LoggingEvent { + /// A chain with the given number of updates was created. + ChainCreated(usize), + /// A chain with the given number of updates was dropped. + ChainDropped(usize), + /// The heap size of the correction buffer changed by the given amount. + SizeDiff(NonZeroIsize), + /// The heap capacity of the correction buffer changed by the given amount. + CapacityDiff(NonZeroIsize), + /// The number of allocations of the correction buffer changed by the given amount. + AllocationsDiff(NonZeroIsize), +} + +/// Channel-based logging for corrections on a Tokio task. `Send`-safe. +/// +/// Sends logging events to the Timely thread, where they are applied to the real `Logging` +/// instance. This allows corrections on the Tokio task to participate in introspection logging +/// without holding `Rc>`. +#[derive(Clone, Debug)] +pub struct ChannelLogging(mpsc::UnboundedSender); + +impl ChannelLogging { + /// Construct a new `ChannelLogging` sending events on the given channel. + pub fn new(tx: mpsc::UnboundedSender) -> Self { + Self(tx) + } + + /// Report the creation of a chain with the given number of updates. + pub fn chain_created(&self, updates: usize) { + let _ = self.0.send(LoggingEvent::ChainCreated(updates)); + } + + /// Report the dropping of a chain with the given number of updates. + pub fn chain_dropped(&self, updates: usize) { + let _ = self.0.send(LoggingEvent::ChainDropped(updates)); + } + + /// Report a change in heap size by the given amount. + pub fn report_size_diff(&self, diff: isize) { + if let Some(diff) = NonZeroIsize::new(diff) { + let _ = self.0.send(LoggingEvent::SizeDiff(diff)); + } + } + + /// Report a change in heap capacity by the given amount. + pub fn report_capacity_diff(&self, diff: isize) { + if let Some(diff) = NonZeroIsize::new(diff) { + let _ = self.0.send(LoggingEvent::CapacityDiff(diff)); + } + } + + /// Report a change in the number of allocations by the given amount. + pub fn report_allocations_diff(&self, diff: isize) { + if let Some(diff) = NonZeroIsize::new(diff) { + let _ = self.0.send(LoggingEvent::AllocationsDiff(diff)); + } + } +} + +/// State for correction buffer logging on the Timely thread. +/// +/// Drains [`LoggingEvent`]s sent by [`ChannelLogging`] from the Tokio task and applies them +/// to the compute and differential loggers. Emits `ArrangementHeapSizeOperator` on construction +/// and `ArrangementHeapSizeOperatorDrop` on drop. +// TODO: Correction buffer logging currently reuses the arrangement batch and size logging. This +// isn't strictly correct as a correction buffer is not an arrangement. Consider refactoring this +// to be about "operator sizes" instead. +pub(crate) struct CorrectionLogger { + compute_logger: ComputeLogger, + differential_logger: differential_dataflow::logging::Logger, + operator_id: usize, + rx: mpsc::UnboundedReceiver, + /// Net number of batches logged (BatchEvent - DropEvent). + net_batches: isize, + /// Net number of records logged across all batch/drop/merge events. + net_records: isize, + /// Cumulative heap size delta, for retraction on drop. + net_size: isize, + /// Cumulative heap capacity delta, for retraction on drop. + net_capacity: isize, + /// Cumulative heap allocations delta, for retraction on drop. + net_allocations: isize, +} + +impl fmt::Debug for CorrectionLogger { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("CorrectionLogger") + .field("operator_id", &self.operator_id) + .finish_non_exhaustive() + } +} + +impl CorrectionLogger { + pub fn new( + compute_logger: ComputeLogger, + differential_logger: differential_dataflow::logging::Logger, + operator_id: usize, + address: Vec, + rx: mpsc::UnboundedReceiver, + ) -> Self { + compute_logger.log(&ComputeEvent::ArrangementHeapSizeOperator( + ArrangementHeapSizeOperator { + operator_id, + address, + }, + )); + + Self { + compute_logger, + differential_logger, + operator_id, + rx, + net_batches: 0, + net_records: 0, + net_size: 0, + net_capacity: 0, + net_allocations: 0, + } + } + + /// Drain logging events from the channel and apply them locally. + pub fn apply_events(&mut self) { + use LoggingEvent::*; + + while let Ok(event) = self.rx.try_recv() { + match event { + ChainCreated(length) => { + self.net_batches += 1; + self.net_records += isize::try_from(length).expect("must fit"); + self.differential_logger.log(BatchEvent { + operator: self.operator_id, + length, + }); + } + ChainDropped(length) => { + self.net_batches -= 1; + self.net_records -= isize::try_from(length).expect("must fit"); + self.differential_logger.log(DropEvent { + operator: self.operator_id, + length, + }); + } + SizeDiff(delta_size) => { + self.net_size += delta_size.get(); + self.compute_logger.log(&ComputeEvent::ArrangementHeapSize( + ArrangementHeapSize { + operator_id: self.operator_id, + delta_size: delta_size.get(), + }, + )); + } + CapacityDiff(delta_capacity) => { + self.net_capacity += delta_capacity.get(); + self.compute_logger + .log(&ComputeEvent::ArrangementHeapCapacity( + ArrangementHeapCapacity { + operator_id: self.operator_id, + delta_capacity: delta_capacity.get(), + }, + )); + } + AllocationsDiff(delta_allocations) => { + self.net_allocations += delta_allocations.get(); + self.compute_logger + .log(&ComputeEvent::ArrangementHeapAllocations( + ArrangementHeapAllocations { + operator_id: self.operator_id, + delta_allocations: delta_allocations.get(), + }, + )); + } + } + } + } +} + +impl Drop for CorrectionLogger { + fn drop(&mut self) { + // Drain any events that arrived before the drop. Note that the Tokio task + // may still be running (abort is async), so some events may not have arrived + // yet. We retract any remaining batch/record counts below. + self.apply_events(); + + // Retract any outstanding batch and record counts that weren't balanced by + // ChainDropped events. This handles the case where the Tokio task is aborted + // and its Correction destructors haven't run yet (abort is async). + // + // Each DropEvent retracts one batch and `length` records, so we emit one per + // outstanding batch, with the first carrying all outstanding records. + for i in 0..self.net_batches { + let length = if i == 0 { + usize::try_from(self.net_records).unwrap_or(0) + } else { + 0 + }; + self.differential_logger.log(DropEvent { + operator: self.operator_id, + length, + }); + } + + // Retract any outstanding heap size/capacity/allocations deltas. + if self.net_size != 0 { + self.compute_logger + .log(&ComputeEvent::ArrangementHeapSize(ArrangementHeapSize { + operator_id: self.operator_id, + delta_size: -self.net_size, + })); + } + if self.net_capacity != 0 { + self.compute_logger + .log(&ComputeEvent::ArrangementHeapCapacity( + ArrangementHeapCapacity { + operator_id: self.operator_id, + delta_capacity: -self.net_capacity, + }, + )); + } + if self.net_allocations != 0 { + self.compute_logger + .log(&ComputeEvent::ArrangementHeapAllocations( + ArrangementHeapAllocations { + operator_id: self.operator_id, + delta_allocations: -self.net_allocations, + }, + )); + } + + self.compute_logger + .log(&ComputeEvent::ArrangementHeapSizeOperatorDrop( + ArrangementHeapSizeOperatorDrop { + operator_id: self.operator_id, + }, + )); + } +} diff --git a/src/compute/src/sink/correction_v2.rs b/src/compute/src/sink/correction_v2.rs deleted file mode 100644 index b8913c38d5112..0000000000000 --- a/src/compute/src/sink/correction_v2.rs +++ /dev/null @@ -1,2152 +0,0 @@ -// Copyright Materialize, Inc. and contributors. All rights reserved. -// -// Use of this software is governed by the Business Source License -// included in the LICENSE file. -// -// As of the Change Date specified in that file, in accordance with -// the Business Source License, use of this software will be governed -// by the Apache License, Version 2.0. - -//! An implementation of the `Correction` data structure used by the MV sink's `write_batches` -//! operator to stash updates before they are written. -//! -//! The `Correction` data structure provides methods to: -//! * insert new updates -//! * advance the compaction frontier (called `since`) -//! * obtain an iterator over consolidated updates before some `upper` -//! * force consolidation of updates before some `upper` -//! -//! The goal is to provide good performance for each of these operations, even in the presence of -//! future updates. MVs downstream of temporal filters might have to deal with large amounts of -//! retractions for future times and we want those to be handled efficiently as well. -//! -//! Note that `Correction` does not provide a method to directly remove updates. Instead updates -//! are removed by inserting their retractions so that they consolidate away to nothing. -//! -//! ## Storage of Updates -//! -//! Stored updates are of the form `(data, time, diff)`, where `time` and `diff` are fixed to -//! [`mz_repr::Timestamp`] and [`mz_repr::Diff`], respectively. -//! -//! [`CorrectionV2`] holds onto a list of `Chain`s containing `Chunk`s of stashed updates. Each -//! `Chunk` is a columnation region containing a fixed maximum number of updates. All updates in -//! a chunk, and all updates in a chain, are ordered by (time, data) and consolidated. -//! -//! Chains live in three places: -//! -//! * A [`BucketChain`] partitions times at or beyond the `boundary` (the largest read `upper` -//! seen so far) into buckets of exponentially growing time ranges, each holding a list of -//! chains. Reads only touch the buckets below their `upper`, so the bulk of the buffered -//! updates — in particular far-future retractions produced by temporal filters — is left -//! alone. -//! * `pending_low` holds chains at times below the `boundary`, mostly insertions arriving -//! through the persist feedback. -//! * `emitted` is a single chain holding the updates returned by the last read. Updates must -//! stay in the buffer until their feedback retractions arrive, and keeping them separate from -//! the bucket chain means reads never have to re-merge future updates. -//! -//! ```text -//! chain[0] | chain[1] | chain[2] -//! | | -//! chunk[0] | chunk[0] | chunk[0] -//! (a, 1, +1) | (a, 1, +1) | (d, 3, +1) -//! (b, 1, +1) | (b, 2, -1) | (d, 4, -1) -//! chunk[1] | chunk[1] | -//! (c, 1, +1) | (c, 2, -2) | -//! (a, 2, -1) | (c, 4, -1) | -//! chunk[2] | | -//! (b, 2, +1) | | -//! (c, 2, +1) | | -//! chunk[3] | | -//! (b, 3, -1) | | -//! (c, 3, +1) | | -//! ``` -//! -//! The "chain invariant" states that each chain in a bucket has at least `chain_proportionality` times as -//! many updates as the next one. This means that chain sizes will often be powers of -//! `chain_proportionality`, but they don't have to be. For example, for a proportionality of 2, -//! the chain sizes `[11, 5, 2, 1]` would satisfy the chain invariant. -//! -//! Note that the invariant is maintained on update counts, not chunk counts. Chunks are -//! byte-bounded (see `ChunkBuilder`), so chunk count is not proportional to update count and -//! would be a poor proxy: any chain below the chunk byte boundary is a single chunk regardless -//! of how many updates it holds, which would let the geometric invariant collapse and break the -//! O(log N) amortization of inserts. -//! -//! Choosing the `chain_proportionality` value allows tuning the trade-off between memory and CPU -//! resources required to maintain corrections. A higher proportionality forces more frequent chain -//! merges, and therefore consolidation, reducing memory usage but increasing CPU usage. -//! -//! ## Inserting Updates -//! -//! A batch of updates is routed by time: updates below the `boundary` become a `pending_low` -//! chain, the rest is appended as new chains to their respective buckets. Appending to a bucket -//! merges chains until the chain invariant is restored. -//! -//! Inserting an update into the correction buffer can be expensive: It involves allocating a new -//! chunk, copying the update in, and then likely merging with an existing chain to restore the -//! chain invariant. If updates trickle in in small batches, this can cause a considerable -//! overhead. To amortize this overhead, new updates aren't immediately inserted into the sorted -//! chains but instead stored in a `Stage` buffer. Once enough updates have been staged to fill a -//! `Chunk`, they are sorted and routed. -//! -//! The insert operation has an amortized complexity of O(log N), with N being the current number -//! of updates stored. -//! -//! ## Retrieving Consolidated Updates -//! -//! Retrieving consolidated updates before a given `upper` works by peeling all buckets below the -//! `upper` off the bucket chain, splitting their chains, the pending low chains, and the previous -//! `emitted` chain at the `upper`, merging the parts below the `upper` into the new `emitted` -//! chain, and returning an iterator over that chain. -//! -//! Because each chain contains updates ordered by time first, splitting a chain at the `upper` -//! reuses whole chunks and copies at most one chunk straddling the split point. Updates at times -//! at or beyond the `upper` are never touched, no matter how many the buffer holds. The -//! complexity of a read is O(U log K), with U being the number of updates before `upper` and K -//! the number of chains containing them. -//! -//! ## Merging Chains -//! -//! Merging multiple chains into a single chain is done using a k-way merge. As the input chains -//! are sorted by (time, data) and consolidated, the same properties hold for the output chain. The -//! complexity of a merge of K chains containing N updates is O(N log K). -//! -//! There is a twist though: Merging also has to respect the `since` frontier, which determines how -//! far the times of updates should be advanced. Advancing times in a sorted chain of updates -//! can make them become unsorted, so we cannot just merge the chains from top to bottom. -//! -//! For example, consider these two chains, assuming `since = [2]`: -//! chain 1: [(c, 1, +1), (b, 2, -1), (a, 3, -1)] -//! chain 2: [(b, 1, +1), (a, 2, +1), (c, 2, -1)] -//! After time advancement, the chains look like this: -//! chain 1: [(c, 2, +1), (b, 2, -1), (a, 3, -1)] -//! chain 2: [(b, 2, +1), (a, 2, +1), (c, 2, -1)] -//! Merging them naively yields [(b, 2, +1), (a, 2, +1), (b, 2, -1), (a, 3, -1)], a chain that's -//! neither sorted nor consolidated. -//! -//! Times below the `since` can only exist in chains read by `consolidate_before`, and only if -//! the `since` advanced past buffered times since the previous read. For few distinct stale -//! times — the steady state, where the previously emitted chain was written just before the -//! since advanced past it — we merge sub-chains, one for each distinct time that's before or at -//! the `since`. Each of these sub-chains retains the (time, data) ordering after the time -//! advancement to `since`, so merging those yields the expected result. -//! -//! For the above example, the chains we would merge are: -//! chain 1.a: [(c, 2, +1)] -//! chain 1.b: [(b, 2, -1), (a, 3, -1)] -//! chain 2.a: [(b, 2, +1)], -//! chain 2.b: [(a, 2, +1), (c, 2, -1)] -//! -//! For many distinct stale times — e.g. a since jump across many buffered timestamps when a sink -//! restarts with an old as-of — the number of sub-chains grows with the number of distinct times, -//! so we instead materialize the affected updates, advance their times, and sort and consolidate -//! them in one O(U log U) pass. - -use std::cmp::Ordering; -use std::collections::{BinaryHeap, VecDeque}; -use std::fmt; -use std::rc::Rc; -use std::sync::{Mutex, OnceLock}; - -use columnar::{Columnar, Index, Len, Ref}; -use mz_ore::cast::CastLossy; -use mz_ore::soft_assert_or_log; -use mz_persist_client::metrics::{SinkMetrics, SinkWorkerMetrics, UpdateDelta}; -use mz_repr::{Diff, Timestamp}; -use mz_timely_util::column_pager::{self, PagedColumn}; -use mz_timely_util::columnar::Column; -use mz_timely_util::temporal::{Bucket, BucketChain}; -use timely::PartialOrder; -use timely::dataflow::channels::ContainerBytes; -use timely::progress::Antichain; - -use crate::sink::correction::{ChannelLogging, SizeMetrics}; - -/// Convenient alias for use in data trait bounds. -/// -/// `D` is constrained to be `Columnar`, so that updates can be stored in a single columnar -/// region per chunk, and the variable-length payload (e.g. `Row` bytes) lives in the same -/// allocation as the rest of the chunk. The `Ref`-level `Eq + Ord` bounds let the merge/heap -/// code compare updates directly through the columnar borrow, avoiding `into_owned` clones -/// on the hot path. -pub trait Data: - differential_dataflow::Data - + Columnar columnar::Borrow: Eq + Ord>> - + Send - + Sync -{ -} -impl Data for D where - D: differential_dataflow::Data - + Columnar columnar::Borrow: Eq + Ord>> - + Send - + Sync -{ -} - -/// A data structure used to store corrections in the MV sink implementation. -/// -/// In contrast to `CorrectionV1`, this implementation stores updates in columnation regions, -/// allowing their memory to be transparently spilled to disk. -#[derive(Debug)] -pub struct CorrectionV2 { - /// Bucketed storage for updates at times at or beyond `boundary`. - /// - /// Buckets cover exponentially growing time ranges, so reads only touch the buckets below - /// their `upper`, and far-future updates (e.g. retractions produced by temporal filters) are - /// rarely touched. - chain: BucketChain>, - /// Chains at times below `boundary` that were not yet emitted. - /// - /// Filled by inserts at times below the boundary (mostly persist feedback) and by the - /// remainders of `emitted` when a read uses a smaller `upper` than the previous one. Merged - /// into `emitted` by the next read. - pending_low: Vec>, - /// Updates that were emitted by `updates_before` but not yet cancelled by persist feedback. - /// - /// Sorted and consolidated, with all times advanced to the `since`. - emitted: Chain, - /// A staging area for updates, to speed up small inserts. - stage: Stage, - /// The lower bound of times stored in `chain`. Only ever advances. - /// - /// Times below the boundary have been peeled off the bucket chain and can only be stored in - /// `pending_low` or `emitted`. - boundary: Antichain, - /// The frontier by which all contained times are advanced. - since: Antichain, - - /// Total count of updates in the correction buffer. - /// - /// Tracked to compute deltas in `update_metrics`. - prev_update_count: usize, - /// Total heap size used by the correction buffer. - /// - /// Tracked to compute deltas in `update_metrics`. - prev_size: SizeMetrics, - /// Global persist sink metrics. - metrics: SinkMetrics, - /// Per-worker persist sink metrics. - worker_metrics: SinkWorkerMetrics, - /// Introspection logging. - logging: Option, -} - -/// Fuel for restoring the bucket chain invariant after peeling. -/// -/// Bounds the restoration work per buffer operation. The bucket chain remains functional when -/// restoration is incomplete -- peeling and finding work on ill-formed chains, at the cost of -/// more in-line splitting -- so leftover restoration is simply picked up by the next operation. -/// -/// `restore` spends one unit of fuel per bucket split, and a single `peel` leaves at most -/// `BucketTimestamp::DOMAIN` (64) buckets to re-split, so this budget completes restoration in one -/// call for any realistic buffer. It is deliberately generous: the "incomplete restoration is -/// picked up next op" path is a correctness safety net for pathological bucket counts, not a hot -/// path we expect to exercise. Lower it if restoration ever needs to interleave with other work. -const RESTORE_FUEL: i64 = 1_000_000; - -impl CorrectionV2 { - /// Construct a new [`CorrectionV2`] instance. - pub fn new( - metrics: SinkMetrics, - worker_metrics: SinkWorkerMetrics, - logging: Option, - chain_proportionality: f64, - chunk_size: usize, - ) -> Self { - let update_size = std::mem::size_of::<(D, Timestamp, Diff)>(); - let chunk_capacity = std::cmp::max(chunk_size / update_size, 1); - - Self { - chain: BucketChain::new(ChainBucket::new(chain_proportionality, logging.clone())), - pending_low: Vec::new(), - emitted: Chain::new(), - stage: Stage::new(logging.clone(), chunk_capacity), - boundary: Antichain::from_elem(Timestamp::MIN), - since: Antichain::from_elem(Timestamp::MIN), - prev_update_count: 0, - prev_size: Default::default(), - metrics, - worker_metrics, - logging, - } - } - - /// Insert a batch of updates. - pub fn insert(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) { - let Some(since_ts) = self.since.as_option() else { - // If the since is the empty frontier, discard all updates. - updates.clear(); - return; - }; - - for (_, time, _) in &mut *updates { - *time = std::cmp::max(*time, *since_ts); - } - - self.insert_inner(updates); - } - - /// Insert a batch of updates, after negating their diffs. - pub fn insert_negated(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) { - let Some(since_ts) = self.since.as_option() else { - // If the since is the empty frontier, discard all updates. - updates.clear(); - return; - }; - - for (_, time, diff) in &mut *updates { - *time = std::cmp::max(*time, *since_ts); - *diff = -*diff; - } - - self.insert_inner(updates); - } - - /// Insert a batch of updates into the stage, flushing it when full. - /// - /// All times are expected to be >= the `since`. - fn insert_inner(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) { - debug_assert!(updates.iter().all(|(_, t, _)| self.since.less_equal(t))); - - if let Some(mut ready) = self.stage.insert(updates) { - self.route(&mut ready); - } - - self.update_metrics(); - } - - /// Route a batch of sorted, consolidated updates to `pending_low` or their chain buckets. - fn route(&mut self, updates: &mut Vec<(D, Timestamp, Diff)>) { - // Updates at times below the boundary become a pending low chain. - let idx = updates.partition_point(|(_, t, _)| !self.boundary.less_equal(t)); - if idx > 0 { - let mut builder = ChainBuilder::default(); - builder.extend(updates.drain(..idx)); - let chain = builder.finish(); - if !chain.is_empty() { - self.log_chain_created(&chain); - self.pending_low.push(chain); - } - } - - // Updates at times at or beyond the boundary go into their chain buckets. Walk ranges of - // times that fall into the same bucket, to push batches of updates at once. - let mut drain = updates.drain(..).peekable(); - while let Some(update) = drain.next() { - let time = update.1; - let range = self - .chain - .range_of(&time) - .expect("bucket chain covers all times at or beyond the boundary"); - let mut builder = ChainBuilder::default(); - builder.extend(std::iter::once(update)); - while let Some(update) = drain.next_if(|(_, t, _)| range.contains(t)) { - builder.extend(std::iter::once(update)); - } - let bucket = self - .chain - .find_mut(&range.start) - .expect("bucket chain covers all times at or beyond the boundary"); - bucket.push_chain(builder.finish()); - } - } - - /// Return consolidated updates before the given `upper`. - pub fn updates_before<'a>( - &'a mut self, - upper: &Antichain, - ) -> impl Iterator + Send + 'a { - self.consolidate_before(upper); - self.consolidated_updates_before(upper) - } - - /// Return the updates before the given `upper`, as consolidated by a preceding - /// [`CorrectionV2::consolidate_before`] call. - /// - /// The caller must have invoked `consolidate_before` with the same `upper` and must not have - /// mutated the buffer since. Otherwise the returned updates are neither consolidated nor - /// necessarily complete. - pub fn consolidated_updates_before<'a>( - &'a self, - upper: &Antichain, - ) -> impl Iterator + Send + use<'a, D> { - // All contained times are advanced to at least the `since`, so a read at an `upper` that - // is not beyond the `since` is always empty. This mirrors the short-circuit in - // `consolidate_before`, which leaves `emitted` untouched in that case. - if !PartialOrder::less_than(&self.since, upper) { - return None.into_iter().flatten(); - } - - // After `consolidate_before`, `emitted` holds exactly the updates before `upper`: every - // path that populates it splits at `upper` (pushing the remainder to `pending_low`), and - // the guard above guarantees `upper > since`, so advancing stale times to the `since` - // cannot lift them to or beyond `upper`. We can therefore yield all of `emitted`. Guard - // the invariant: a violation would write updates beyond the batch upper to persist. - soft_assert_or_log!( - self.emitted - .last() - .is_none_or(|(_, t, _)| !upper.less_equal(&t)), - "emitted contains times at or beyond the upper", - ); - Some(self.emitted.iter()).into_iter().flatten() - } - - /// Consolidate all updates before the given `upper` into the `emitted` chain. - /// - /// Once this method returns, `emitted` contains all updates at times before `upper`, - /// consolidated. - /// - /// Does nothing if `upper` is not beyond the `since`: all contained times are advanced to at - /// least the `since`, so such a read is empty anyway, and skipping avoids an eager peel, - /// merge, and `boundary` advancement. Normal reads and `consolidate_at_since` always pass an - /// `upper` beyond the `since`. - pub fn consolidate_before(&mut self, upper: &Antichain) { - if !PartialOrder::less_than(&self.since, upper) { - return; - } - - if let Some(mut ready) = self.stage.flush() { - self.route(&mut ready); - } - - let Some(&since_ts) = self.since.as_option() else { - // If the since is the empty frontier, discard all updates. - let peeled = self.chain.peel(Antichain::new().borrow()); - for bucket in peeled { - for chain in bucket.into_chains() { - self.log_chain_dropped(&chain); - } - } - for chain in std::mem::take(&mut self.pending_low) { - self.log_chain_dropped(&chain); - } - let emitted = std::mem::replace(&mut self.emitted, Chain::new()); - if !emitted.is_empty() { - self.log_chain_dropped(&emitted); - } - self.update_metrics(); - return; - }; - - // Peel the buckets below the upper off the bucket chain. Bucket splits during the peel - // only touch chunks around the upper; chunks wholly on either side are reused. - let peeled = self.chain.peel(upper.borrow()); - if PartialOrder::less_than(&self.boundary, upper) { - self.boundary = upper.clone(); - } - - // Collect candidate chains: peeled bucket contents, pending low chains, and the previous - // emitted chain. All contain only times below the boundary. - let emitted = std::mem::replace(&mut self.emitted, Chain::new()); - let mut candidates: Vec> = Vec::new(); - for bucket in peeled { - candidates.extend(bucket.into_chains()); - } - candidates.append(&mut self.pending_low); - if !emitted.is_empty() { - candidates.push(emitted); - } - - if candidates.is_empty() { - self.restore_chain(); - self.update_metrics(); - return; - } - - candidates.iter().for_each(|c| self.log_chain_dropped(c)); - - // Split the candidates at the upper. Parts at or beyond the upper (possible when `upper` - // regresses below a previous one) stay pending. - let mut lowers = Vec::new(); - for chain in candidates { - match upper.as_option() { - Some(&upper_ts) => { - let (lower, remainder) = chain.split_at_time(upper_ts); - if !lower.is_empty() { - lowers.push(lower); - } - if !remainder.is_empty() { - self.log_chain_created(&remainder); - self.pending_low.push(remainder); - } - } - // The empty upper is greater than all times. - None => lowers.push(chain), - } - } - - // Merge the lower parts into the new emitted chain, advancing times below the since. - // Advancing times in a (time, data)-sorted chain can break its sort order, so chains - // containing stale times cannot be merged as they are. Stale times are expected in steady - // state: the previous emitted chain was written before the since advanced past it. - // - // Count the distinct stale times, up to a small cap. For few distinct stale times -- the - // steady state -- split cursors into runs that remain sorted under advancement and merge - // those. For many distinct stale times -- e.g. a since jump across many buffered - // timestamps when a sink restarts with an old as-of -- the number of runs and the cost of - // cloning cursor state per run grow with the number of distinct times, so materialize, - // advance, and consolidate in one O(U log U) pass instead. - const MAX_STALE_RUNS: usize = 32; - let mut stale_times = 0; - for chain in &lowers { - stale_times += chain.distinct_times_before(since_ts, MAX_STALE_RUNS - stale_times); - if stale_times >= MAX_STALE_RUNS { - break; - } - } - - let merged = if stale_times == 0 { - let cursors: Vec<_> = lowers.into_iter().filter_map(Chain::into_cursor).collect(); - merge_cursors(cursors) - } else if stale_times < MAX_STALE_RUNS { - let mut runs = Vec::new(); - for chain in lowers { - if let Some(cursor) = chain.into_cursor() { - runs.append(&mut cursor.advance_by(since_ts)); - } - } - merge_cursors(runs) - } else { - let mut updates: Vec<_> = lowers.iter().flat_map(|c| c.iter()).collect(); - for (_, time, _) in &mut updates { - *time = std::cmp::max(*time, since_ts); - } - consolidate(&mut updates); - let mut builder = ChainBuilder::default(); - builder.extend(updates); - let chain = builder.finish(); - - // Advancement can move updates to or beyond the upper; such updates stay pending. - match upper.as_option() { - Some(&upper_ts) => { - let (lower, remainder) = chain.split_at_time(upper_ts); - if !remainder.is_empty() { - self.log_chain_created(&remainder); - self.pending_low.push(remainder); - } - lower - } - None => chain, - } - }; - - if !merged.is_empty() { - self.log_chain_created(&merged); - } - self.emitted = merged; - - self.restore_chain(); - self.update_metrics(); - } - - /// Perform a bounded amount of work towards restoring the bucket chain invariant. - /// - /// Restoration is allowed to remain incomplete: the bucket chain supports peeling and finding - /// on ill-formed chains, so any leftover work is picked up by subsequent operations. The fuel - /// bound keeps individual buffer operations from stalling the operator that owns the buffer. - fn restore_chain(&mut self) { - let mut fuel = RESTORE_FUEL; - self.chain.restore(&mut fuel); - } - - /// Advance the since frontier. - /// - /// Time advancement of updates in the bucket chain is lazy: it happens when the updates are - /// consolidated by a read. - /// - /// # Panics - /// - /// Panics if the given `since` is less than the current since frontier. - pub fn advance_since(&mut self, since: Antichain) { - assert!(PartialOrder::less_equal(&self.since, &since)); - self.stage.advance_times(&since); - self.since = since; - } - - /// Consolidate all updates at the current `since`. - pub fn consolidate_at_since(&mut self) { - let upper_ts = self.since.as_option().and_then(|t| t.try_step_forward()); - if let Some(upper_ts) = upper_ts { - let upper = Antichain::from_elem(upper_ts); - self.consolidate_before(&upper); - } - } - - fn log_chain_created(&self, chain: &Chain) { - if let Some(logging) = &self.logging { - logging.chain_created(chain.update_count); - } - } - - fn log_chain_dropped(&self, chain: &Chain) { - if let Some(logging) = &self.logging { - logging.chain_dropped(chain.update_count); - } - } - - /// Update persist sink metrics. - fn update_metrics(&mut self) { - let mut new_size = self.stage.get_size(); - let mut new_length = self.stage.data.len(); - for chain in &self.pending_low { - new_size += chain.get_size(); - new_length += chain.update_count; - } - new_size += self.emitted.get_size(); - new_length += self.emitted.update_count; - for bucket in self.chain.buckets() { - for chain in &bucket.chains { - new_size += chain.get_size(); - new_length += chain.update_count; - } - } - - self.update_metrics_inner(new_size, new_length); - } - - /// Update persist sink metrics to the given new size and length. - fn update_metrics_inner(&mut self, new_size: SizeMetrics, new_length: usize) { - let old_size = self.prev_size; - let old_length = self.prev_update_count; - let len_delta = UpdateDelta::new(new_length, old_length); - let cap_delta = UpdateDelta::new(new_size.capacity, old_size.capacity); - self.metrics - .report_correction_update_deltas(len_delta, cap_delta); - self.worker_metrics - .report_correction_update_totals(new_length, new_size.capacity); - - if let Some(logging) = &self.logging { - let i = |x: usize| isize::try_from(x).expect("must fit"); - logging.report_size_diff(i(new_size.size) - i(old_size.size)); - logging.report_capacity_diff(i(new_size.capacity) - i(old_size.capacity)); - logging.report_allocations_diff(i(new_size.allocations) - i(old_size.allocations)); - } - - self.prev_size = new_size; - self.prev_update_count = new_length; - } -} - -/// Merge the given cursors into one chain. -fn merge_cursors(cursors: Vec>) -> Chain { - match cursors.len() { - 0 => Chain::new(), - 1 => { - let [cur] = cursors.try_into().unwrap(); - cur.into_chain() - } - 2 => { - let [a, b] = cursors.try_into().unwrap(); - merge_2(a, b) - } - _ => merge_many(cursors), - } -} - -/// Merge the given two cursors using a 2-way merge. -/// -/// This function is a specialization of `merge_many` that avoids the overhead of a binary heap. -fn merge_2(cursor1: Cursor, cursor2: Cursor) -> Chain { - let mut rest1 = Some(cursor1); - let mut rest2 = Some(cursor2); - let mut merged = ChainBuilder::default(); - - loop { - match (rest1, rest2) { - (Some(c1), Some(c2)) => { - let (d1, t1, r1) = c1.get(); - let (d2, t2, r2) = c2.get(); - - match (t1, d1).cmp(&(t2, d2)) { - Ordering::Less => { - merged.push_ref((d1, t1, r1)); - rest1 = c1.step(); - rest2 = Some(c2); - } - Ordering::Greater => { - merged.push_ref((d2, t2, r2)); - rest1 = Some(c1); - rest2 = c2.step(); - } - Ordering::Equal => { - let r = r1 + r2; - if r != Diff::ZERO { - merged.push_ref((d1, t1, r)); - } - rest1 = c1.step(); - rest2 = c2.step(); - } - } - } - (Some(c), None) | (None, Some(c)) => { - merged.push_cursor(c); - break; - } - (None, None) => break, - } - } - - merged.finish() -} - -/// Merge the given cursors using a k-way merge with a binary heap. -fn merge_many(cursors: Vec>) -> Chain { - let mut heap = MergeHeap::from_iter(cursors); - let mut merged = ChainBuilder::default(); - while let Some(cursor1) = heap.pop() { - let (data, time, mut diff) = cursor1.get(); - - while let Some((cursor2, r)) = heap.pop_equal(data, time) { - diff += r; - if let Some(cursor2) = cursor2.step() { - heap.push(cursor2); - } - } - - if diff != Diff::ZERO { - merged.push_ref((data, time, diff)); - } - if let Some(cursor1) = cursor1.step() { - heap.push(cursor1); - } - } - - merged.finish() -} - -impl Drop for CorrectionV2 { - fn drop(&mut self) { - for bucket in self.chain.buckets() { - bucket.chains.iter().for_each(|c| self.log_chain_dropped(c)); - } - self.pending_low - .iter() - .for_each(|c| self.log_chain_dropped(c)); - if !self.emitted.is_empty() { - self.log_chain_dropped(&self.emitted); - } - self.update_metrics_inner(Default::default(), 0); - } -} - -/// A bucket of `Chain`s, for use in a [`BucketChain`]. -/// -/// All chains are individually sorted by (time, data) and consolidated, but updates can appear in -/// multiple chains, so consumers must merge the chains to obtain consolidated updates. -struct ChainBucket { - /// The contained chains. - /// - /// Maintained with the chain invariant on pushes; splits can leave it violated until the next - /// push restores it. - chains: Vec>, - /// The size factor of subsequent chains required by the chain invariant. - chain_proportionality: f64, - /// Introspection logging. - logging: Option, -} - -impl fmt::Debug for ChainBucket { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("ChainBucket") - .field("chains", &self.chains) - .finish_non_exhaustive() - } -} - -impl ChainBucket { - /// Construct a new, empty `ChainBucket`. - fn new(chain_proportionality: f64, logging: Option) -> Self { - Self { - chains: Vec::new(), - chain_proportionality, - logging, - } - } - - /// Push a chain onto the bucket, restoring the chain invariant. - fn push_chain(&mut self, chain: Chain) { - if chain.is_empty() { - return; - } - if let Some(logging) = &self.logging { - logging.chain_created(chain.update_count); - } - self.chains.push(chain); - - // Restore the chain invariant. - let prop = self.chain_proportionality; - let merge_needed = |chains: &[Chain<_>]| match chains { - [.., prev, last] => { - let last_len = f64::cast_lossy(last.update_count); - let prev_len = f64::cast_lossy(prev.update_count); - last_len * prop > prev_len - } - _ => false, - }; - - while merge_needed(&self.chains) { - let a = self.chains.pop().unwrap(); - let b = self.chains.pop().unwrap(); - if let Some(logging) = &self.logging { - logging.chain_dropped(a.update_count); - logging.chain_dropped(b.update_count); - } - - let cursors = [a, b].into_iter().filter_map(Chain::into_cursor).collect(); - let merged = merge_cursors(cursors); - if !merged.is_empty() { - if let Some(logging) = &self.logging { - logging.chain_created(merged.update_count); - } - self.chains.push(merged); - } - } - } - - /// Convert the bucket into its contained chains. - fn into_chains(self) -> Vec> { - self.chains - } -} - -impl Bucket for ChainBucket { - type Timestamp = Timestamp; - - fn split(self, timestamp: &Self::Timestamp, fuel: &mut i64) -> (Self, Self) { - let mut lower = Self::new(self.chain_proportionality, self.logging.clone()); - let mut upper = Self::new(self.chain_proportionality, self.logging.clone()); - - for chain in self.chains { - // Whole chunks are reused; at most one chunk straddling the timestamp is copied per - // chain. Account fuel at chunk granularity. - *fuel = fuel.saturating_sub(i64::try_from(chain.chunks.len()).expect("must fit")); - - if let Some(logging) = &self.logging { - logging.chain_dropped(chain.update_count); - } - let (lo, hi) = chain.split_at_time(*timestamp); - for (part, target) in [(lo, &mut lower), (hi, &mut upper)] { - if !part.is_empty() { - if let Some(logging) = &self.logging { - logging.chain_created(part.update_count); - } - target.chains.push(part); - } - } - } - - (lower, upper) - } -} - -/// A chain of [`Chunk`]s containing updates. -/// -/// All updates in a chain are sorted by (time, data) and consolidated. -/// -/// Note that, in contrast to [`Chunk`]s, chains can be empty. Though we generally try to avoid -/// keeping around empty chains. -#[derive(Debug)] -struct Chain { - /// The contained chunks. - chunks: Vec>, - /// The number of updates contained in all chunks. - update_count: usize, -} - -impl Chain { - /// Construct an empty chain. - fn new() -> Self { - Self { - chunks: Default::default(), - update_count: 0, - } - } - - /// Return whether the chain is empty. - fn is_empty(&self) -> bool { - self.chunks.is_empty() - } - - /// Push a chunk onto the chain. - /// - /// All updates in the chunk must sort after all updates already in the chain, in - /// (time, data)-order, to ensure the chain remains sorted. - fn push_chunk(&mut self, chunk: Chunk) { - mz_ore::soft_assert_no_log!(self.can_accept_chunk(&chunk)); - - self.update_count += chunk.len(); - self.chunks.push(chunk); - } - - /// Return whether the chain can accept the given chunk at its end while preserving - /// (time, data)-order. - /// - /// NOTE: The cached boundary times settle every case but a tie. On a tie the boundary updates - /// themselves are compared, which materializes both chunks and keeps them resident for the - /// rest of their lifetime. The only caller is the soft assertion in [`Chain::push_chunk`], and - /// soft assertions are live in any build started with `MZ_SOFT_ASSERTIONS` set, so this cost is - /// not confined to debug builds. Ties are reached whenever a run of updates at a single - /// timestamp spans a chunk boundary, which [`ChunkBuilder`] produces for any such run larger - /// than its byte limit. - fn can_accept_chunk(&self, chunk: &Chunk) -> bool { - match self.chunks.last() { - None => true, - Some(last) => match last.last_time().cmp(&chunk.first_time()) { - Ordering::Less => true, - Ordering::Greater => false, - Ordering::Equal => { - let (dc, _, _) = last.last(); - let (d, _, _) = chunk.first(); - dc < d - } - }, - } - } - - /// Return the last update in the chain, if any. - fn last(&self) -> Option> { - self.chunks.last().map(|c| c.last()) - } - - /// Convert the chain into a cursor over the contained updates. - fn into_cursor(self) -> Option> { - let chunks = self.chunks.into_iter().map(Rc::new).collect(); - Cursor::new(chunks) - } - - /// Return an iterator over the contained updates. - fn iter(&self) -> impl Iterator + '_ { - self.chunks.iter().flat_map(|c| { - (0..c.len()).map(move |i| { - let (d, t, r) = c.index(i); - (D::into_owned(d), t, r) - }) - }) - } - - /// Count the distinct times of updates at times before `time`, up to the given cap. - /// - /// The scan uses one binary search per distinct time, so its cost is bounded by - /// O(cap log chunks). - fn distinct_times_before(&self, time: Timestamp, cap: usize) -> usize { - let mut count = 0; - let mut chunk_idx = 0; - let mut offset = 0; - while count < cap && chunk_idx < self.chunks.len() { - let chunk = &self.chunks[chunk_idx]; - let current = chunk.index(offset).1; - if current >= time { - break; - } - count += 1; - // Skip to the first update at a time greater than `current`. - match chunk.find_time_greater_than(current) { - Some(idx) => offset = idx, - None => { - // All later updates at `current` are in subsequent chunks. - chunk_idx += 1; - offset = 0; - while chunk_idx < self.chunks.len() { - match self.chunks[chunk_idx].find_time_greater_than(current) { - Some(idx) => { - offset = idx; - break; - } - None => chunk_idx += 1, - } - } - } - } - } - count - } - - /// Split the chain at the given time. - /// - /// Returns two chains, the first containing all updates at times < `time`, the second - /// containing all updates at times >= `time`. Chunks fully on either side of `time` are - /// reused; only a chunk straddling `time` is copied. - fn split_at_time(mut self, time: Timestamp) -> (Self, Self) { - let mut lower = Self::new(); - let mut upper = Self::new(); - - let Some(skip_ts) = time.step_back() else { - // Nothing sorts before `time`. - return (lower, self); - }; - - for chunk in self.chunks.drain(..) { - // Route whole chunks by cached boundary times, so a chunk that lands entirely on one - // side is moved without paging it in. Only a straddling chunk is materialized here. - // With soft assertions on, `push_chunk` can still page in a chunk whose boundary time - // ties the chain's last one, see `Chain::can_accept_chunk`. - if chunk.last_time() < time { - lower.push_chunk(chunk); - } else if chunk.first_time() >= time { - upper.push_chunk(chunk); - } else { - // The chunk straddles `time`; copy its two halves. - let idx = chunk - .find_time_greater_than(skip_ts) - .expect("straddles time"); - let mut builder = ChainBuilder::default(); - for i in 0..idx { - builder.push_ref(chunk.index(i)); - } - for part in builder.finish().chunks { - lower.push_chunk(part); - } - let mut builder = ChainBuilder::default(); - for i in idx..chunk.len() { - builder.push_ref(chunk.index(i)); - } - for part in builder.finish().chunks { - upper.push_chunk(part); - } - } - } - - (lower, upper) - } - - /// Return the size of the chain, for use in metrics. - fn get_size(&self) -> SizeMetrics { - let mut metrics = SizeMetrics::default(); - for chunk in &self.chunks { - metrics += chunk.get_size(); - } - metrics - } -} - -/// A builder that constructs a [`Chain`] from a stream of updates. -/// -/// Wraps a [`ChunkBuilder`] and drains its minted chunks into a [`Chain`]. Pushed updates must -/// arrive in (time, data) sorted order. -struct ChainBuilder { - builder: ChunkBuilder, - chain: Chain, -} - -impl Default for ChainBuilder { - fn default() -> Self { - Self { - builder: Default::default(), - chain: Chain::new(), - } - } -} - -impl ChainBuilder { - /// Push a reference-form update into the builder. - fn push_ref(&mut self, update: Ref<'_, (D, Timestamp, Diff)>) { - self.builder.push(update); - self.drain(); - } - - /// Push an owned-form update into the builder. - fn push_owned(&mut self, update: &(D, Timestamp, Diff)) { - self.builder.push(update); - self.drain(); - } - - /// Push the updates produced by a cursor into the builder. - fn push_cursor(&mut self, cursor: Cursor) { - let mut rest = Some(cursor); - while let Some(cursor) = rest.take() { - let update = cursor.get(); - self.push_ref(update); - rest = cursor.step(); - } - } - - /// Move any minted chunks from the builder into the chain. - fn drain(&mut self) { - while let Some(chunk) = self.builder.pop() { - self.chain.push_chunk(chunk); - } - } - - /// Finish building, returning the assembled [`Chain`]. - fn finish(self) -> Chain { - let Self { builder, mut chain } = self; - for chunk in builder.finish() { - if chunk.len() > 0 { - chain.push_chunk(chunk); - } - } - chain - } -} - -impl Extend<(D, Timestamp, Diff)> for ChainBuilder { - fn extend>(&mut self, iter: I) { - for update in iter { - self.push_owned(&update); - } - } -} - -/// A cursor over updates in a chain. -/// -/// A cursor provides two guarantees: -/// * Produced updates are ordered and consolidated. -/// * A cursor always yields at least one update. -/// -/// The second guarantee is enforced through the type system: Every method that steps a cursor -/// forward consumes `self` and returns an `Option` that's `None` if the operation stepped -/// over the last update. -/// -/// A cursor holds on to `Rc`s, allowing multiple cursors to produce updates from the same -/// chunks concurrently. As soon as a cursor is done producing updates from a [`Chunk`] it drops -/// its reference. Once the last cursor is done with a [`Chunk`] its memory can be reclaimed. -#[derive(Clone, Debug)] -struct Cursor { - /// The chunks from which updates can still be produced. - chunks: VecDeque>>, - /// The current offset into `chunks.front()`. - chunk_offset: usize, - /// An optional limit for the number of updates the cursor will produce. - limit: Option, - /// An optional overwrite for the timestamp of produced updates. - overwrite_ts: Option, -} - -impl Cursor { - /// Construct a cursor over a list of chunks. - /// - /// Returns `None` if `chunks` is empty. - fn new(chunks: VecDeque>>) -> Option { - if chunks.is_empty() { - return None; - } - - Some(Self { - chunks, - chunk_offset: 0, - limit: None, - overwrite_ts: None, - }) - } - - /// Set a limit for the number of updates this cursor will produce. - /// - /// # Panics - /// - /// Panics if there is already a limit lower than the new one. - fn set_limit(mut self, limit: usize) -> Option { - assert!(self.limit.is_none_or(|l| l >= limit)); - - if limit == 0 { - return None; - } - - // Release chunks made unreachable by the limit. - let mut count = 0; - let mut idx = 0; - let mut offset = self.chunk_offset; - while idx < self.chunks.len() && count < limit { - let chunk = &self.chunks[idx]; - count += chunk.len() - offset; - idx += 1; - offset = 0; - } - self.chunks.truncate(idx); - - if count > limit { - self.limit = Some(limit); - } - - Some(self) - } - - /// Get a reference to the current update. - fn get(&self) -> Ref<'_, (D, Timestamp, Diff)> { - let chunk = self.get_chunk(); - let (d, t, r) = chunk.index(self.chunk_offset); - let t = self.overwrite_ts.unwrap_or(t); - (d, t, r) - } - - /// Get a reference to the current chunk. - fn get_chunk(&self) -> &Chunk { - &self.chunks[0] - } - - /// Step to the next update. - /// - /// Returns the stepped cursor, or `None` if the step was over the last update. - fn step(mut self) -> Option { - if self.chunk_offset == self.get_chunk().len() - 1 { - return self.skip_chunk().map(|(c, _)| c); - } - - self.chunk_offset += 1; - - if let Some(limit) = &mut self.limit { - *limit -= 1; - if *limit == 0 { - return None; - } - } - - Some(self) - } - - /// Skip the remainder of the current chunk. - /// - /// Returns the forwarded cursor and the number of updates skipped, or `None` if no chunks are - /// left after the skip. - fn skip_chunk(mut self) -> Option<(Self, usize)> { - let chunk = self.chunks.pop_front().expect("cursor invariant"); - - if self.chunks.is_empty() { - return None; - } - - let skipped = chunk.len() - self.chunk_offset; - self.chunk_offset = 0; - - if let Some(limit) = &mut self.limit { - if skipped >= *limit { - return None; - } - *limit -= skipped; - } - - Some((self, skipped)) - } - - /// Skip all updates with times <= the given time. - /// - /// Returns the forwarded cursor and the number of updates skipped, or `None` if no updates are - /// left after the skip. - fn skip_time(mut self, time: Timestamp) -> Option<(Self, usize)> { - if self.overwrite_ts.is_some_and(|ts| ts <= time) { - return None; - } else if self.get().1 > time { - return Some((self, 0)); - } - - let mut skipped = 0; - - let new_offset = loop { - let chunk = self.get_chunk(); - if let Some(index) = chunk.find_time_greater_than(time) { - break index; - } - - let (cursor, count) = self.skip_chunk()?; - self = cursor; - skipped += count; - }; - - skipped += new_offset - self.chunk_offset; - self.chunk_offset = new_offset; - - Some((self, skipped)) - } - - /// Advance all updates in this cursor by the given `since_ts`. - /// - /// Returns a list of cursors, each of which yields ordered and consolidated updates that have - /// been advanced by `since_ts`. - fn advance_by(mut self, since_ts: Timestamp) -> Vec { - // If the cursor has an `overwrite_ts`, all its updates are at the same time already. We - // only need to advance the `overwrite_ts` by the `since_ts`. - if let Some(ts) = self.overwrite_ts { - if ts < since_ts { - self.overwrite_ts = Some(since_ts); - } - return vec![self]; - } - - // Otherwise we need to split the cursor so that each new cursor only yields runs of - // updates that are correctly (time, data)-ordered when advanced by `since_ts`. We achieve - // this by splitting the cursor at each time <= `since_ts`. - let mut splits = Vec::new(); - let mut remaining = Some(self); - - while let Some(cursor) = remaining.take() { - let (_, time, _) = cursor.get(); - if time >= since_ts { - splits.push(cursor); - break; - } - - let mut current = cursor.clone(); - if let Some((cursor, skipped)) = cursor.skip_time(time) { - remaining = Some(cursor); - current = current.set_limit(skipped).expect("skipped at least 1"); - } - current.overwrite_ts = Some(since_ts); - splits.push(current); - } - - splits - } - - /// Drain the cursor into a [`Chain`]. - /// - /// This reuses the underlying chunks if possible, and writes new ones otherwise. - fn into_chain(self) -> Chain { - match self.try_unwrap() { - Ok(chain) => chain, - Err((_, cursor)) => { - let mut builder = ChainBuilder::default(); - builder.push_cursor(cursor); - builder.finish() - } - } - } - - /// Attempt to unwrap the cursor into a [`Chain`]. - /// - /// This operation efficiently reuses chunks by directly inserting them into the output chain - /// where possible. - /// - /// An unwrap is only successful if the cursor's `limit` and `overwrite_ts` are both `None` and - /// the cursor has unique references to its chunks. If the unwrap fails, this method returns an - /// `Err` containing the cursor in an unchanged state, allowing the caller to convert it into a - /// chain by copying chunks rather than reusing them. - fn try_unwrap(self) -> Result, (&'static str, Self)> { - if self.limit.is_some() { - return Err(("cursor with limit", self)); - } - if self.overwrite_ts.is_some() { - return Err(("cursor with overwrite_ts", self)); - } - if self.chunks.iter().any(|c| Rc::strong_count(c) != 1) { - return Err(("cursor on shared chunks", self)); - } - - let mut builder = ChainBuilder::default(); - let mut remaining = Some(self); - - // We might be partway through the first chunk, in which case we can't reuse it but need to - // allocate a new one to contain only the updates the cursor can still yield. - while let Some(cursor) = remaining.take() { - if cursor.chunk_offset == 0 { - remaining = Some(cursor); - break; - } - let update = cursor.get(); - builder.push_ref(update); - remaining = cursor.step(); - } - - let mut chain = builder.finish(); - if let Some(cursor) = remaining { - for chunk in cursor.chunks { - let chunk = Rc::into_inner(chunk).expect("checked above"); - chain.push_chunk(chunk); - } - } - - Ok(chain) - } -} - -/// A non-empty chunk of updates, backed by a columnar region. -/// -/// All updates in a chunk are sorted by (time, data) and consolidated. -/// -/// Chunks are immutable once created. They are produced by [`ChunkBuilder`], which mints a -/// new chunk whenever its in-progress columnar container reaches a fixed serialized byte -/// boundary (~2 MiB, matching the ship granularity used elsewhere in the codebase), so each -/// chunk corresponds to a single, predictably sized allocation. -struct Chunk { - /// The paged-out form, taken on first materialization. - /// - /// A `Mutex` (not `RefCell`) keeps the chunk `Sync`: cursors hold chunks behind a shared - /// `Rc`, and the iterator returned by [`CorrectionV2::updates_before`] borrows them across - /// the persist writer's `await`, so `&Chunk` must be `Send`. The lock is taken once, at - /// materialization, and is otherwise uncontended (the sink runs single-threaded per worker). - paged: Mutex>>, - /// The materialized form, populated lazily by [`Chunk::column`] on first access. - /// - /// An `OnceLock` (not `OnceCell`) for the same `Sync` reason. Once set the slot is never - /// cleared, so its address is stable and [`Chunk::index`] can hand out `Ref<'_>` borrows tied - /// to `&self`. The allocation is freed when the chunk drops, which bounds resident memory to - /// the chunks under an active merge front. - resident: OnceLock>, - /// Number of updates, cached so `len` and chain bookkeeping never page the chunk in. - len: usize, - /// Time of the first update, cached so boundary checks (`split_at_time`, `can_accept`) route - /// a resting chunk without materializing it. - first_time: Timestamp, - /// Time of the last update, cached likewise. - last_time: Timestamp, -} - -impl fmt::Debug for Chunk { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "Chunk(<{}>)", self.len()) - } -} - -impl Chunk { - /// Page the given non-empty column out into a chunk. - /// - /// Reads the cached metadata (length, boundary times) while the column is still resident, then - /// hands it to the global column pager. The policy decides whether it actually spills; either - /// way the chunk is born paged and materializes lazily on first read. - /// - /// # Panics - /// - /// Panics if the column is empty. Chunks are non-empty by construction; [`ChunkBuilder`] only - /// ever builds a chunk from a populated column. - fn from_column(mut data: Column<(D, Timestamp, Diff)>) -> Self { - let (len, first_time, last_time) = { - let borrowed = data.borrow(); - let len = borrowed.len(); - assert!(len > 0, "chunks are non-empty"); - (len, borrowed.get(0).1, borrowed.get(len - 1).1) - }; - - let paged = column_pager::global_pager().page(&mut data); - Self { - paged: Mutex::new(Some(paged)), - resident: OnceLock::new(), - len, - first_time, - last_time, - } - } - - /// Materialize the chunk's column, paging it in on first access. - /// - /// The returned reference is valid for as long as `&self`: the `OnceLock` slot is never - /// cleared once populated, so its contents have a stable address. - fn column(&self) -> &Column<(D, Timestamp, Diff)> { - self.resident.get_or_init(|| { - let paged = self - .paged - .lock() - .expect("pager mutex poisoned") - .take() - .expect("paged form present until materialized"); - column_pager::global_pager().take(paged) - }) - } - - /// Return the number of updates in the chunk. - fn len(&self) -> usize { - self.len - } - - /// Return the update at the given index, paging the chunk in if necessary. - /// - /// # Panics - /// - /// Panics if the given index is not populated. - fn index(&self, idx: usize) -> Ref<'_, (D, Timestamp, Diff)> { - self.column().borrow().get(idx) - } - - /// Return the first update in the chunk, paging the chunk in if necessary. - fn first(&self) -> Ref<'_, (D, Timestamp, Diff)> { - self.index(0) - } - - /// Return the last update in the chunk, paging the chunk in if necessary. - fn last(&self) -> Ref<'_, (D, Timestamp, Diff)> { - self.index(self.len - 1) - } - - /// Return the time of the first update, without materializing the chunk. - fn first_time(&self) -> Timestamp { - self.first_time - } - - /// Return the time of the last update, without materializing the chunk. - fn last_time(&self) -> Timestamp { - self.last_time - } - - /// Return the index of the first update at a time greater than `time`, or `None` if no such - /// update exists. - /// - /// The early-out uses the cached last time, so a chunk whose updates are all at or before - /// `time` is skipped without paging it in. - fn find_time_greater_than(&self, time: Timestamp) -> Option { - if self.last_time <= time { - return None; - } - - let mut lower = 0; - let mut upper = self.len; - while lower < upper { - let idx = (lower + upper) / 2; - if self.index(idx).1 > time { - upper = idx; - } else { - lower = idx + 1; - } - } - - Some(lower) - } - - /// Return the size of the chunk, for use in metrics. - /// - /// Reports resident bytes only: a chunk still spilled (on swap or in a pager file) is not part - /// of RSS and contributes nothing, matching the accounting in - /// [`mz_timely_util::columnar::merge_batcher`]. - fn get_size(&self) -> SizeMetrics { - let resident = |col: &Column<(D, Timestamp, Diff)>| { - let bytes = col.length_in_bytes(); - SizeMetrics { - size: bytes, - capacity: bytes, - allocations: 1, - } - }; - - if let Some(col) = self.resident.get() { - return resident(col); - } - // Not yet materialized: a policy that kept the column resident still occupies RSS, so - // account for it; a genuinely spilled column does not. - match &*self.paged.lock().expect("pager mutex poisoned") { - Some(PagedColumn::Resident(col, _)) => resident(col), - _ => SizeMetrics::default(), - } - } -} - -/// Builder that produces a stream of fixed-size [`Chunk`]s. -/// -/// Wraps [`mz_timely_util::columnar::builder::ColumnBuilder`], which mints a new -/// [`Column::Align`] chunk whenever its in-progress columnar container reaches a fixed -/// serialized byte boundary (~2 MiB, matching the ship granularity used elsewhere in the -/// codebase). Each minted chunk is therefore a single, predictably-sized aligned allocation. -struct ChunkBuilder { - inner: mz_timely_util::columnar::builder::ColumnBuilder<(D, Timestamp, Diff)>, -} - -impl Default for ChunkBuilder { - fn default() -> Self { - Self { - inner: Default::default(), - } - } -} - -impl ChunkBuilder { - /// Push an update into the builder. - /// - /// Accepts whatever the inner [`ColumnBuilder`]'s [`PushInto`] impl accepts — both the - /// `Ref<'_, (D, T, R)>` refs produced by cursors and `&(D, T, R)` references to owned - /// tuples drained from the staging buffer. - /// - /// [`ColumnBuilder`]: mz_timely_util::columnar::builder::ColumnBuilder - /// [`PushInto`]: timely::container::PushInto - #[inline] - fn push(&mut self, item: T) - where - mz_timely_util::columnar::builder::ColumnBuilder<(D, Timestamp, Diff)>: - timely::container::PushInto, - { - timely::container::PushInto::push_into(&mut self.inner, item); - } - - /// Pop a finished chunk, if one is available. - fn pop(&mut self) -> Option> { - use timely::container::ContainerBuilder; - // `ColumnBuilder::extract` stashes the popped chunk in its `finished` slot so the - // caller can read it through `&mut`; move it out with `mem::take` so we own it - // (leaves `Column::Typed(Default::default())` behind, which the next `extract` - // overwrites). - self.inner - .extract() - .map(|c| Chunk::from_column(std::mem::take(c))) - } - - /// Finalize the builder: flush any in-progress updates as a typed chunk and drain pending. - fn finish(mut self) -> impl Iterator> { - use timely::container::ContainerBuilder; - // `ColumnBuilder::finish` flushes the in-progress container into the pending queue - // (as `Column::Typed`) and returns the first pending entry. Subsequent calls drain - // the rest until `None`. Translate that into an owning iterator. - // - // `finish` can hand back an empty column (e.g. when the last shipped chunk landed exactly - // on the boundary). Skip those: `Chunk::from_column` requires a non-empty column, and an - // empty chunk would needlessly engage the pager. - std::iter::from_fn(move || { - loop { - let col = std::mem::take(self.inner.finish()?); - if !col.is_empty() { - return Some(Chunk::from_column(col)); - } - } - }) - } -} - -/// A buffer for staging updates before they are inserted into the sorted chains. -#[derive(Debug)] -struct Stage { - /// The contained updates. - /// - /// This vector has a fixed capacity equal to the [`Chunk`] capacity. - data: Vec<(D, Timestamp, Diff)>, - /// Introspection logging. - /// - /// We want to report the number of records in the stage. To do so, we pretend that the stage - /// is a chain, and every time the number of updates inside changes, the chain gets dropped and - /// re-created. - logging: Option, -} - -impl Stage { - fn new(logging: Option, chunk_capacity: usize) -> Self { - // For logging, we pretend the stage consists of a single chain. - if let Some(logging) = &logging { - logging.chain_created(0); - } - - Self { - data: Vec::with_capacity(chunk_capacity), - logging, - } - } - - /// Insert a batch of updates, possibly producing a batch of sorted, consolidated updates - /// ready to be stored. - fn insert( - &mut self, - updates: &mut Vec<(D, Timestamp, Diff)>, - ) -> Option> { - if updates.is_empty() { - return None; - } - - let prev_length = self.ilen(); - - // Determine how many chunks we can fill with the available updates. - let update_count = self.data.len() + updates.len(); - let chunk_capacity = self.data.capacity(); - let chunk_count = update_count / chunk_capacity; - - let mut new_updates = updates.drain(..); - - // If we have enough shipable updates, collect them and consolidate. - let maybe_ready = if chunk_count > 0 { - let ship_count = chunk_count * chunk_capacity; - let mut buffer = Vec::with_capacity(ship_count); - - buffer.append(&mut self.data); - while buffer.len() < ship_count { - let update = new_updates.next().unwrap(); - buffer.push(update); - } - - consolidate(&mut buffer); - - Some(buffer) - } else { - None - }; - - // Stage the remaining updates. - Extend::extend(&mut self.data, new_updates); - - self.log_length_diff(self.ilen() - prev_length); - - maybe_ready - } - - /// Flush all currently staged updates, returning them sorted and consolidated. - fn flush(&mut self) -> Option> { - self.log_length_diff(-self.ilen()); - - consolidate(&mut self.data); - - if self.data.is_empty() { - return None; - } - - let capacity = self.data.capacity(); - let data = std::mem::replace(&mut self.data, Vec::with_capacity(capacity)); - Some(data) - } - - /// Advance the times of staged updates by the given `since`. - fn advance_times(&mut self, since: &Antichain) { - let Some(since_ts) = since.as_option() else { - // If the since is the empty frontier, discard all updates. - self.log_length_diff(-self.ilen()); - self.data.clear(); - return; - }; - - for (_, time, _) in &mut self.data { - *time = std::cmp::max(*time, *since_ts); - } - } - - /// Return the size of the stage, for use in metrics. - /// - /// Note: We don't follow pointers here, so the returned `size` and `capacity` values are - /// under-estimates. That's fine as the stage should always be small. - fn get_size(&self) -> SizeMetrics { - SizeMetrics { - size: self.data.len() * std::mem::size_of::<(D, Timestamp, Diff)>(), - capacity: self.data.capacity() * std::mem::size_of::<(D, Timestamp, Diff)>(), - allocations: 1, - } - } - - /// Return the number of updates in the stage, as an `isize`. - fn ilen(&self) -> isize { - self.data.len().try_into().expect("must fit") - } - - fn log_length_diff(&self, diff: isize) { - let Some(logging) = &self.logging else { return }; - if diff > 0 { - let count = usize::try_from(diff).expect("must fit"); - logging.chain_created(count); - logging.chain_dropped(0); - } else if diff < 0 { - let count = usize::try_from(-diff).expect("must fit"); - logging.chain_created(0); - logging.chain_dropped(count); - } - } -} - -impl Drop for Stage { - fn drop(&mut self) { - if let Some(logging) = &self.logging { - logging.chain_dropped(self.data.len()); - } - } -} - -/// Sort and consolidate the given list of updates. -/// -/// This function is the same as [`differential_dataflow::consolidation::consolidate_updates`], -/// except that it sorts updates by (time, data) instead of (data, time). -fn consolidate(updates: &mut Vec<(D, Timestamp, Diff)>) { - if updates.len() <= 1 { - return; - } - - let diff = |update: &(_, _, Diff)| update.2; - - updates.sort_unstable_by(|(d1, t1, _), (d2, t2, _)| (t1, d1).cmp(&(t2, d2))); - - let mut offset = 0; - let mut accum = diff(&updates[0]); - - for idx in 1..updates.len() { - let this = &updates[idx]; - let prev = &updates[idx - 1]; - if this.0 == prev.0 && this.1 == prev.1 { - accum += diff(&updates[idx]); - } else { - if accum != Diff::ZERO { - updates.swap(offset, idx - 1); - updates[offset].2 = accum; - offset += 1; - } - accum = diff(&updates[idx]); - } - } - - if accum != Diff::ZERO { - let len = updates.len(); - updates.swap(offset, len - 1); - updates[offset].2 = accum; - offset += 1; - } - - updates.truncate(offset); -} - -/// Compare two columnar refs that have unrelated input lifetimes. -/// -/// `::Ref<'a>` is an associated-type projection through a trait, so -/// the compiler treats it as invariant in `'a` and won't auto-shorten the inputs by variance. -/// We instead explicitly reborrow both to a fresh, local lifetime `'x` via -/// [`Columnar::reborrow`] before letting the inner `==` pick up the `for<'a> Ref<'a>: Eq` -/// bound on [`Data`]. -#[inline] -fn refs_eq(a: Ref<'_, D>, b: Ref<'_, D>) -> bool { - #[inline] - fn eq<'x, D: Data>(a: Ref<'x, D>, b: Ref<'x, D>) -> bool { - a == b - } - eq::(D::reborrow(a), D::reborrow(b)) -} - -/// A binary heap specialized for merging [`Cursor`]s. -struct MergeHeap(BinaryHeap>); - -impl FromIterator> for MergeHeap { - fn from_iter>>(cursors: I) -> Self { - let inner = cursors.into_iter().map(MergeCursor).collect(); - Self(inner) - } -} - -impl MergeHeap { - /// Pop the next cursor (the one yielding the least update) from the heap. - fn pop(&mut self) -> Option> { - self.0.pop().map(|MergeCursor(c)| c) - } - - /// Pop the next cursor from the heap, provided the data and time of its current update are - /// equal to the given values. - /// - /// Returns both the cursor and the diff corresponding to `data` and `time`. - fn pop_equal(&mut self, data: Ref<'_, D>, time: Timestamp) -> Option<(Cursor, Diff)> { - let r = { - let MergeCursor(cursor) = self.0.peek()?; - let (d, t, r) = cursor.get(); - if t != time || !refs_eq::(d, data) { - return None; - } - r - }; - let cursor = self.pop().expect("checked above"); - Some((cursor, r)) - } - - /// Push a cursor onto the heap. - fn push(&mut self, cursor: Cursor) { - self.0.push(MergeCursor(cursor)); - } -} - -/// A wrapper for [`Cursor`]s on a [`MergeHeap`]. -/// -/// Implements the cursor ordering required for merging cursors. -struct MergeCursor(Cursor); - -impl PartialEq for MergeCursor { - fn eq(&self, other: &Self) -> bool { - self.cmp(other).is_eq() - } -} - -impl Eq for MergeCursor {} - -impl PartialOrd for MergeCursor { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for MergeCursor { - fn cmp(&self, other: &Self) -> Ordering { - let (d1, t1, _) = self.0.get(); - let (d2, t2, _) = other.0.get(); - (t1, d1).cmp(&(t2, d2)).reverse() - } -} - -#[cfg(test)] -mod tests { - use mz_ore::metrics::MetricsRegistry; - use mz_persist_client::cfg::PersistConfig; - use mz_persist_client::metrics::Metrics; - use mz_repr::{Diff, Timestamp}; - - use super::*; - use crate::sink::correction::CorrectionV1; - - #[mz_ore::test] - fn chain_builder_update_count_matches_items() { - let mut builder = ChainBuilder::::default(); - for i in 0..10_u64 { - let d = i64::try_from(i).expect("fits"); - builder.push_owned(&(d, Timestamp::new(i), Diff::ONE)); - } - let chain = builder.finish(); - assert_eq!(chain.update_count, chain.iter().count()); - } - - /// Push enough updates to cross at least one `mint()` boundary, forcing the - /// `Align` encode -> `from_bytes` decode roundtrip (the spilling path this data - /// structure exists to support), and assert `iter()` roundtrips values, order, - /// and diffs across the spill boundary. - #[mz_ore::test] - #[cfg_attr(miri, ignore)] // too slow: crossing the ~2 MiB mint boundary needs ~200k updates - fn chain_builder_roundtrips_across_mint_boundary() { - // A single `mint()` fires near the ~2 MiB (`SHIP_WORDS`) serialized boundary. With - // three 8-byte columns per update that's tens of thousands of updates; pushing 200k - // comfortably forces multiple mints. - let count = 200_000_u64; - - let mut builder = ChainBuilder::::default(); - for i in 0..count { - let d = i64::try_from(i).expect("fits"); - builder.push_owned(&(d, Timestamp::new(i), Diff::ONE)); - } - let chain = builder.finish(); - - // Crossing the mint boundary must have produced more than one chunk; otherwise the spill - // path (each minted chunk is paged out and read back through the pager) wouldn't be - // exercised. The chunk payload itself is now behind the pager (see [`Chunk`]), so we - // assert on chunk count rather than inspecting the column variant directly. - assert!( - chain.chunks.len() > 1, - "expected multiple minted chunks, got {} chunk(s): {:?}", - chain.chunks.len(), - chain.chunks, - ); - - // `iter()` must roundtrip every update, in order, with correct diffs. - assert_eq!(chain.update_count, usize::try_from(count).expect("fits")); - let mut expected = 0_u64; - for (d, t, r) in chain.iter() { - assert_eq!(d, i64::try_from(expected).expect("fits")); - assert_eq!(t, Timestamp::new(expected)); - assert_eq!(r, Diff::ONE); - expected += 1; - } - assert_eq!(expected, count); - } - - fn sink_metrics() -> SinkMetrics { - let registry = MetricsRegistry::new(); - let metrics = Metrics::new(&PersistConfig::new_for_tests(), ®istry); - metrics.sink.clone() - } - - /// Run the same stepwise-drain workload through `CorrectionV1` and `CorrectionV2` and assert - /// that they emit the same updates at every step. - /// - /// Models the `write_batches` operator catching up through many distinct timestamps: the - /// desired input runs ahead, batches are written one timestamp at a time, and written updates - /// come back negated through the persist feedback. - #[mz_ore::test] - // Columnation regions are not Stacked Borrows compliant: later pushes invalidate the - // provenance of previously stored items under Miri. - #[cfg_attr(miri, ignore)] - fn equivalence_with_v1() { - let sink_metrics = sink_metrics(); - - let mut v1 = - CorrectionV1::::new(sink_metrics.clone(), sink_metrics.for_worker(0), 1); - let mut v2 = CorrectionV2::::new( - sink_metrics.clone(), - sink_metrics.for_worker(0), - None, - 3.0, - 8 * 1024, - ); - - let num_ts = 50; - let keys = 4; - - // Upsert-style input: every timestamp updates each key, retracting the previous value. - let batch = |t: u64| -> Vec<(String, Timestamp, Diff)> { - (0..keys) - .flat_map(|k| { - let addition = (format!("{k}-{t}"), Timestamp::from(t), Diff::ONE); - let retraction = t - .checked_sub(1) - .map(|p| (format!("{k}-{p}"), Timestamp::from(t), -Diff::ONE)); - std::iter::once(addition).chain(retraction) - }) - .collect() - }; - - // Pre-fill both with all batches, like a catch-up where the input runs ahead. - for t in 0..num_ts { - v1.insert(&mut batch(t)); - v2.insert(&mut batch(t)); - } - - // Drain stepwise, with persist feedback, comparing emissions. - for t in 0..num_ts { - let upper = Antichain::from_elem(Timestamp::from(t + 1)); - - let mut out1: Vec<_> = v1.updates_before(&upper).collect(); - let mut out2: Vec<_> = v2.updates_before(&upper).collect(); - out1.sort(); - out2.sort(); - assert_eq!(out1, out2, "diverged at t={t}"); - - v1.insert_negated(&mut out1.clone()); - v2.insert_negated(&mut out2); - v1.advance_since(upper.clone()); - v2.advance_since(upper); - } - - // Compare the final state at the since. - let upper = Antichain::from_elem(Timestamp::from(num_ts + 1)); - v1.consolidate_at_since(); - v2.consolidate_at_since(); - let mut out1: Vec<_> = v1.updates_before(&upper).collect(); - let mut out2: Vec<_> = v2.updates_before(&upper).collect(); - out1.sort(); - out2.sort(); - assert_eq!(out1, out2); - } - - /// A since jump across many distinct buffered timestamps must collapse them onto the since. - #[mz_ore::test] - // Columnation regions are not Stacked Borrows compliant: later pushes invalidate the - // provenance of previously stored items under Miri. - #[cfg_attr(miri, ignore)] - fn since_jump() { - let sink_metrics = sink_metrics(); - let mut v2 = CorrectionV2::::new( - sink_metrics.clone(), - sink_metrics.for_worker(0), - None, - 3.0, - 8 * 1024, - ); - - let num_ts = 100; - for t in 0..num_ts { - v2.insert(&mut vec![ - (format!("a-{t}"), Timestamp::from(t), Diff::ONE), - (format!("a-{t}"), Timestamp::from(t), -Diff::ONE), - (format!("b-{t}"), Timestamp::from(t), Diff::ONE), - ]); - } - - v2.advance_since(Antichain::from_elem(Timestamp::from(num_ts))); - v2.consolidate_at_since(); - - let upper = Antichain::from_elem(Timestamp::from(num_ts + 1)); - let out: Vec<_> = v2.updates_before(&upper).collect(); - assert_eq!(out.len(), usize::try_from(num_ts).unwrap()); - assert!( - out.iter() - .all(|(_, t, r)| *t == Timestamp::from(num_ts) && *r == Diff::ONE) - ); - } - - /// Reads must not observe updates at or beyond their `upper`, even when the `upper` is not - /// beyond the `since`. - #[mz_ore::test] - // Columnation regions are not Stacked Borrows compliant: later pushes invalidate the - // provenance of previously stored items under Miri. - #[cfg_attr(miri, ignore)] - fn upper_not_beyond_since() { - let sink_metrics = sink_metrics(); - let mut v2 = CorrectionV2::::new( - sink_metrics.clone(), - sink_metrics.for_worker(0), - None, - 3.0, - 8 * 1024, - ); - - v2.insert(&mut vec![( - "a".to_owned(), - Timestamp::from(5_u64), - Diff::ONE, - )]); - v2.advance_since(Antichain::from_elem(Timestamp::from(10_u64))); - - // The update logically lives at time 10 now, so a read before 7 must be empty. - let upper = Antichain::from_elem(Timestamp::from(7_u64)); - assert_eq!(v2.updates_before(&upper).count(), 0); - - // A read before 11 must emit it, advanced to the since. - let upper = Antichain::from_elem(Timestamp::from(11_u64)); - let out: Vec<_> = v2.updates_before(&upper).collect(); - assert_eq!( - out, - vec![("a".to_owned(), Timestamp::from(10_u64), Diff::ONE)] - ); - } - - /// A [`PagingPolicy`] that always spills to the swap backend, uncompressed. - /// - /// The default global pager keeps every chunk resident; installing this drives the actual - /// spill path so the tests exercise [`Chunk::column`]'s page-in through [`mz_ore::pager`]. - /// - /// [`PagingPolicy`]: column_pager::PagingPolicy - struct ForceSwap; - - impl column_pager::PagingPolicy for ForceSwap { - fn decide(&self, _hint: column_pager::PageHint) -> column_pager::PageDecision { - column_pager::PageDecision::Page { - backend: mz_ore::pager::Backend::Swap, - codec: None, - } - } - fn record(&self, _event: column_pager::PageEvent) {} - } - - /// Install a global pager that spills every chunk to swap for the duration of `f`, then - /// restore the default (disabled) pager. The global pager is process-wide; concurrent tests - /// only ever observe a correct round-trip regardless of backend, so racing on it is benign. - fn with_swap_pager(f: impl FnOnce() -> R) -> R { - use std::sync::Arc; - column_pager::set_global_pager(column_pager::ColumnPager::new(Arc::new(ForceSwap))); - let result = f(); - column_pager::set_global_pager(column_pager::ColumnPager::disabled()); - result - } - - /// Build a chain crossing the mint boundary while every chunk is spilled to swap, then assert - /// `iter()` (the read path behind `updates_before`) pages each chunk back in and roundtrips - /// values, order, and diffs. - #[mz_ore::test] - #[cfg_attr(miri, ignore)] // madvise on the swap backend is unsupported under miri - fn iter_roundtrips_through_swap_backend() { - let count = 200_000_u64; - with_swap_pager(|| { - let mut builder = ChainBuilder::::default(); - for i in 0..count { - let d = i64::try_from(i).expect("fits"); - builder.push_owned(&(d, Timestamp::new(i), Diff::ONE)); - } - let chain = builder.finish(); - assert!(chain.chunks.len() > 1, "expected multiple minted chunks"); - assert_eq!(chain.update_count, usize::try_from(count).expect("fits")); - - let mut expected = 0_u64; - for (d, t, r) in chain.iter() { - assert_eq!(d, i64::try_from(expected).expect("fits")); - assert_eq!(t, Timestamp::new(expected)); - assert_eq!(r, Diff::ONE); - expected += 1; - } - assert_eq!(expected, count); - }); - } - - /// Drive a [`Cursor`] over a spilled, multi-chunk chain to completion (the access pattern - /// merges use). Each step pages the front chunk back in via [`Chunk::column`]; assert the - /// cursor yields every update in order. - #[mz_ore::test] - #[cfg_attr(miri, ignore)] // madvise on the swap backend is unsupported under miri - fn cursor_steps_through_swap_backend() { - let count = 200_000_u64; - with_swap_pager(|| { - let mut builder = ChainBuilder::::default(); - for i in 0..count { - let d = i64::try_from(i).expect("fits"); - builder.push_owned(&(d, Timestamp::new(i), Diff::ONE)); - } - let chain = builder.finish(); - assert!(chain.chunks.len() > 1, "expected multiple minted chunks"); - - let mut rest = chain.into_cursor(); - let mut expected = 0_u64; - while let Some(cursor) = rest.take() { - let (d, t, r) = cursor.get(); - assert_eq!(i64::into_owned(d), i64::try_from(expected).expect("fits")); - assert_eq!(t, Timestamp::new(expected)); - assert_eq!(r, Diff::ONE); - expected += 1; - rest = cursor.step(); - } - assert_eq!(expected, count); - }); - } -} diff --git a/src/compute/src/sink/materialized_view.rs b/src/compute/src/sink/materialized_view.rs index e5d1a4064c1ba..47fe48680c688 100644 --- a/src/compute/src/sink/materialized_view.rs +++ b/src/compute/src/sink/materialized_view.rs @@ -155,7 +155,8 @@ use crate::compute_state::ComputeState; use crate::render::StartSignal; use crate::render::errors::DataflowErrorSer; use crate::render::sinks::SinkRender; -use crate::sink::correction::{ChannelLogging, Correction, CorrectionLogger}; +use crate::sink::correction::Correction; +use crate::sink::correction::logging::{ChannelLogging, CorrectionLogger}; use crate::sink::materialized_view_v2; use crate::sink::refresh::apply_refresh; diff --git a/src/compute/src/sink/materialized_view_v2.rs b/src/compute/src/sink/materialized_view_v2.rs index 5408d63142025..f688906c07f83 100644 --- a/src/compute/src/sink/materialized_view_v2.rs +++ b/src/compute/src/sink/materialized_view_v2.rs @@ -74,7 +74,8 @@ use tracing::trace; use crate::compute_state::ComputeState; use crate::render::StartSignal; use crate::render::errors::DataflowErrorSer; -use crate::sink::correction::{ChannelLogging, Correction, CorrectionLogger}; +use crate::sink::correction::Correction; +use crate::sink::correction::logging::{ChannelLogging, CorrectionLogger}; use crate::sink::materialized_view::{ BatchDescription, BatchesStream, DescsStream, DesiredStreams, OkErr, PersistApi, PersistStreams, SharedSinkFrontier, advance, operator_name, persist_source, diff --git a/test/launchdarkly-flag-consistency/mzcompose.py b/test/launchdarkly-flag-consistency/mzcompose.py index 51cfe8432f2d0..0c131410e5f79 100644 --- a/test/launchdarkly-flag-consistency/mzcompose.py +++ b/test/launchdarkly-flag-consistency/mzcompose.py @@ -459,9 +459,11 @@ clusterd_malloc_conf column_paged_batcher_spill_worker_count column_paged_batcher_use_pool + consolidating_vec_growth_dampener constraint_based_timestamp_selection enable_0dt_deployment enable_aws_msk_iam_auth + enable_compute_correction_v2 enable_consolidate_after_union_negate enable_continual_task_builtins enable_continual_task_transform @@ -532,7 +534,6 @@ # flag to a permanent category above. "compute_subscribe_snapshot_optimization", "enable_cast_elimination", - "enable_compute_correction_v2", "enable_compute_error_distinct", "enable_compute_temporal_bucketing", "enable_new_outer_join_lowering", @@ -557,7 +558,6 @@ enable_cluster_schedule_refresh enable_column_paged_batcher enable_column_paged_batcher_spill - enable_compute_correction_v2 enable_eager_delta_joins enable_index_options enable_join_prioritize_arranged From 937807eaf21f48e7467076ce1e1c6d53b15b9681 Mon Sep 17 00:00:00 2001 From: Moritz Hoffmann Date: Wed, 26 Aug 2026 13:19:06 +0200 Subject: [PATCH 2/2] compute: unlink private CorrectionLogger from module docs Rustdoc rejects an intra-doc link from a public module's documentation to a crate-private item. Co-Authored-By: Claude Opus 5 (1M context) --- src/compute/src/sink/correction/logging.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/compute/src/sink/correction/logging.rs b/src/compute/src/sink/correction/logging.rs index 4ae221e467601..647e7c2563d9a 100644 --- a/src/compute/src/sink/correction/logging.rs +++ b/src/compute/src/sink/correction/logging.rs @@ -11,7 +11,7 @@ //! //! The correction buffer lives on a Tokio task, while the introspection loggers are owned by the //! Timely thread and are not `Send`. [`ChannelLogging`] bridges the two: the buffer reports size -//! and chain changes as [`LoggingEvent`]s over a channel, and [`CorrectionLogger`] drains them on +//! and chain changes as [`LoggingEvent`]s over a channel, and `CorrectionLogger` drains them on //! the Timely thread. use std::fmt;