Skip to main content

mz_persist_client/internal/
trace.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! An append-only collection of compactable update batches. The Spine below is
11//! a fork of Differential Dataflow's [Spine] with minimal modifications. The
12//! original Spine code is designed for incremental (via "fuel"ing) synchronous
13//! merge of in-memory batches. Persist doesn't want compaction to block
14//! incoming writes and, in fact, may in the future elect to push the work of
15//! compaction onto another machine entirely via RPC. As a result, we abuse the
16//! Spine code as follows:
17//!
18//! [Spine]: differential_dataflow::trace::implementations::spine_fueled::Spine
19//!
20//! - The normal Spine works in terms of [Batch] impls. A `Batch` is added to
21//!   the Spine. As progress is made, the Spine will merge two batches together
22//!   by: constructing a [Batch::Merger], giving it bits of fuel to
23//!   incrementally perform the merge (which spreads out the work, keeping
24//!   latencies even), and then once it's done fueling extracting the new single
25//!   output `Batch` and discarding the inputs.
26//! - Persist instead represents a batch of blob data with a [HollowBatch]
27//!   pointer which contains the normal `Batch` metadata plus the keys necessary
28//!   to retrieve the updates.
29//! - [SpineBatch] wraps `HollowBatch` and has a [FuelingMerge] companion
30//!   (analogous to `Batch::Merger`) that allows us to represent a merge as it
31//!   is fueling. Normally, this would represent real incremental compaction
32//!   progress, but in persist, it's simply a bookkeeping mechanism. Once fully
33//!   fueled, the `FuelingMerge` is turned into a fueled [SpineBatch],
34//!   which to the Spine is indistinguishable from a merged batch. At this
35//!   point, it is eligible for asynchronous compaction and a `FueledMergeReq`
36//!   is generated.
37//! - At any later point, this request may be answered via
38//!   [Trace::apply_merge_res_checked] or [Trace::apply_merge_res_unchecked].
39//!   This internally replaces the`SpineBatch`, which has no
40//!   effect on the structure of `Spine` but replaces the metadata
41//!   in persist's state to point at the new batch.
42//! - `SpineBatch` is explictly allowed to accumulate a list of `HollowBatch`s.
43//!   This decouples compaction from Spine progress and also allows us to reduce
44//!   write amplification by merging `N` batches at once where `N` can be
45//!   greater than 2.
46//!
47//! [Batch]: differential_dataflow::trace::Batch
48//! [Batch::Merger]: differential_dataflow::trace::Batch::Merger
49
50use std::cmp::Ordering;
51use std::collections::{BTreeMap, BTreeSet};
52use std::fmt::{Debug, Display};
53use std::mem;
54use std::ops::Range;
55use std::sync::Arc;
56
57use arrayvec::ArrayVec;
58use differential_dataflow::difference::Monoid;
59use differential_dataflow::lattice::Lattice;
60use differential_dataflow::trace::Description;
61use itertools::Itertools;
62use mz_ore::cast::CastFrom;
63use mz_persist::metrics::ColumnarMetrics;
64use mz_persist_types::Codec64;
65use serde::{Serialize, Serializer};
66use timely::PartialOrder;
67use timely::progress::frontier::AntichainRef;
68use timely::progress::{Antichain, Timestamp};
69use tracing::{error, warn};
70
71use crate::internal::paths::WriterKey;
72use crate::internal::state::{HollowBatch, RunId};
73
74use super::state::RunPart;
75
76#[derive(Debug, Clone, PartialEq)]
77pub struct FueledMergeReq<T> {
78    pub id: SpineId,
79    pub desc: Description<T>,
80    pub inputs: Vec<IdHollowBatch<T>>,
81}
82
83#[derive(Debug)]
84pub struct FueledMergeRes<T> {
85    pub output: HollowBatch<T>,
86    pub input: CompactionInput,
87    pub new_active_compaction: Option<ActiveCompaction>,
88}
89
90/// An append-only collection of compactable update batches.
91///
92/// In an effort to keep our fork of Spine as close as possible to the original,
93/// we push as many changes as possible into this wrapper.
94#[derive(Debug, Clone)]
95pub struct Trace<T> {
96    spine: Spine<T>,
97    pub(crate) roundtrip_structure: bool,
98}
99
100#[cfg(any(test, debug_assertions))]
101impl<T: PartialEq> PartialEq for Trace<T> {
102    fn eq(&self, other: &Self) -> bool {
103        // Deconstruct self and other so we get a compile failure if new fields
104        // are added.
105        let Trace {
106            spine: _,
107            roundtrip_structure: _,
108        } = self;
109        let Trace {
110            spine: _,
111            roundtrip_structure: _,
112        } = other;
113
114        // Intentionally use HollowBatches for this comparison so we ignore
115        // differences in spine layers.
116        self.batches().eq(other.batches())
117    }
118}
119
120impl<T: Timestamp + Lattice> Default for Trace<T> {
121    fn default() -> Self {
122        Self {
123            spine: Spine::new(),
124            roundtrip_structure: true,
125        }
126    }
127}
128
129#[derive(Clone, Debug, Serialize)]
130pub struct ThinSpineBatch<T> {
131    pub(crate) level: usize,
132    pub(crate) desc: Description<T>,
133    pub(crate) parts: Vec<SpineId>,
134    /// NB: this exists to validate legacy batch bounds during the migration;
135    /// it can be deleted once the roundtrip_structure flag is permanently rolled out.
136    pub(crate) descs: Vec<Description<T>>,
137}
138
139impl<T: PartialEq> PartialEq for ThinSpineBatch<T> {
140    fn eq(&self, other: &Self) -> bool {
141        // Ignore the temporary descs vector when comparing for equality.
142        (self.level, &self.desc, &self.parts).eq(&(other.level, &other.desc, &other.parts))
143    }
144}
145
146#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
147pub struct ThinMerge<T> {
148    pub(crate) since: Antichain<T>,
149    pub(crate) remaining_work: usize,
150    pub(crate) active_compaction: Option<ActiveCompaction>,
151}
152
153impl<T: Clone> ThinMerge<T> {
154    fn fueling(merge: &FuelingMerge<T>) -> Self {
155        ThinMerge {
156            since: merge.since.clone(),
157            remaining_work: merge.remaining_work,
158            active_compaction: None,
159        }
160    }
161
162    fn fueled(batch: &SpineBatch<T>) -> Self {
163        ThinMerge {
164            since: batch.desc.since().clone(),
165            remaining_work: 0,
166            active_compaction: batch.active_compaction.clone(),
167        }
168    }
169}
170
171/// This is a "flattened" representation of a Trace. Goals:
172/// - small updates to the trace should result in small differences in the `FlatTrace`;
173/// - two `FlatTrace`s should be efficient to diff;
174/// - converting to and from a `Trace` should be relatively straightforward.
175///
176/// These goals are all somewhat in tension, and the space of possible representations is pretty
177/// large. See individual fields for comments on some of the tradeoffs.
178#[derive(Clone, Debug)]
179pub struct FlatTrace<T> {
180    pub(crate) since: Antichain<T>,
181    /// Hollow batches without an associated ID. If this flattened trace contains spine batches,
182    /// we can figure out which legacy batch belongs in which spine batch by comparing the `desc`s.
183    /// Previously, we serialized a trace as just this list of batches. Keeping this data around
184    /// helps ensure backwards compatibility. In the near future, we may still keep some batches
185    /// here to help minimize the size of diffs -- rewriting all the hollow batches in a shard
186    /// can be prohibitively expensive. Eventually, we'd like to remove this in favour of the
187    /// collection below.
188    pub(crate) legacy_batches: BTreeMap<Arc<HollowBatch<T>>, ()>,
189    /// Hollow batches _with_ an associated ID. Spine batches can reference these hollow batches
190    /// by id directly.
191    pub(crate) hollow_batches: BTreeMap<SpineId, Arc<HollowBatch<T>>>,
192    /// Spine batches stored by ID. We reference hollow batches by ID, instead of inlining them,
193    /// to make differential updates smaller when two batches merge together. We also store the
194    /// level on the batch, instead of mapping from level to a list of batches... the level of a
195    /// spine batch doesn't change over time, but the list of batches at a particular level does.
196    pub(crate) spine_batches: BTreeMap<SpineId, ThinSpineBatch<T>>,
197    /// In-progress merges. We store this by spine id instead of level to prepare for some possible
198    /// generalizations to spine (merging N of M batches at a level). This is also a natural place
199    /// to store incremental merge progress in the future.
200    pub(crate) merges: BTreeMap<SpineId, ThinMerge<T>>,
201}
202
203impl<T: Timestamp + Lattice> Trace<T> {
204    pub(crate) fn flatten(&self) -> FlatTrace<T> {
205        let since = self.spine.since.clone();
206        let mut legacy_batches = BTreeMap::new();
207        let mut hollow_batches = BTreeMap::new();
208        let mut spine_batches = BTreeMap::new();
209        let mut merges = BTreeMap::new();
210
211        let mut push_spine_batch = |level: usize, batch: &SpineBatch<T>| {
212            let id = batch.id();
213            let desc = batch.desc.clone();
214            let mut parts = Vec::with_capacity(batch.parts.len());
215            let mut descs = Vec::with_capacity(batch.parts.len());
216            for IdHollowBatch { id, batch } in &batch.parts {
217                parts.push(*id);
218                descs.push(batch.desc.clone());
219                // Ideally, we'd like to put all batches in the hollow_batches collection, since
220                // tracking the spine id reduces ambiguity and makes diffing cheaper. However,
221                // we currently keep most batches in the legacy collection for backwards
222                // compatibility.
223                // As an exception, we add batches with empty time ranges to hollow_batches:
224                // they're otherwise not guaranteed to be unique, and since we only started writing
225                // them down recently there's no backwards compatibility risk.
226                if batch.desc.lower() == batch.desc.upper() {
227                    hollow_batches.insert(*id, Arc::clone(batch));
228                } else {
229                    legacy_batches.insert(Arc::clone(batch), ());
230                }
231            }
232
233            let spine_batch = ThinSpineBatch {
234                level,
235                desc,
236                parts,
237                descs,
238            };
239            spine_batches.insert(id, spine_batch);
240        };
241
242        for (level, state) in self.spine.merging.iter().enumerate() {
243            for batch in &state.batches {
244                push_spine_batch(level, batch);
245                if let Some(c) = &batch.active_compaction {
246                    let previous = merges.insert(batch.id, ThinMerge::fueled(batch));
247                    assert!(
248                        previous.is_none(),
249                        "recording a compaction for a batch that already exists! (level={level}, id={:?}, compaction={c:?})",
250                        batch.id,
251                    )
252                }
253            }
254            if let Some(IdFuelingMerge { id, merge }) = state.merge.as_ref() {
255                let previous = merges.insert(*id, ThinMerge::fueling(merge));
256                assert!(
257                    previous.is_none(),
258                    "fueling a merge for a batch that already exists! (level={level}, id={id:?}, merge={merge:?})"
259                )
260            }
261        }
262
263        if !self.roundtrip_structure {
264            assert!(hollow_batches.is_empty());
265            spine_batches.clear();
266            merges.clear();
267        }
268
269        FlatTrace {
270            since,
271            legacy_batches,
272            hollow_batches,
273            spine_batches,
274            merges,
275        }
276    }
277    pub(crate) fn unflatten(value: FlatTrace<T>) -> Result<Self, String> {
278        let FlatTrace {
279            since,
280            legacy_batches,
281            mut hollow_batches,
282            spine_batches,
283            mut merges,
284        } = value;
285
286        // If the flattened representation has spine batches (or is empty)
287        // we know to preserve the structure for this trace.
288        let roundtrip_structure = !spine_batches.is_empty() || legacy_batches.is_empty();
289
290        // The flattened trace is decoded from an untrusted blob: any invariant
291        // a crafted or corrupted value can violate must surface as a decode
292        // error here, never as a panic in the spine code below.
293        //
294        // Bound the total logical len of all batches. The spine's maintenance
295        // arithmetic (`len.next_power_of_two()`, summing lens of merged
296        // batches) overflows on absurd lens, and no real trace has anywhere
297        // near this many updates.
298        const MAX_TOTAL_LEN: usize = usize::MAX >> 3;
299        let mut total_len = 0usize;
300        for batch in legacy_batches.keys().chain(hollow_batches.values()) {
301            total_len = total_len
302                .checked_add(batch.len)
303                .filter(|len| *len <= MAX_TOTAL_LEN)
304                .ok_or_else(|| {
305                    format!("total len of batches exceeds the maximum trace size: {batch:?}")
306                })?;
307        }
308
309        // We need to look up legacy batches somehow, but we don't have a spine id for them.
310        // Instead, we rely on the fact that the spine must store them in antichain order.
311        // Our timestamp type may not be totally ordered, so we need to implement our own comparator
312        // here. Persist's invariants ensure that all the frontiers we're comparing are comparable,
313        // though.
314        let compare_chains = |left: &Antichain<T>, right: &Antichain<T>| {
315            if PartialOrder::less_than(left, right) {
316                Ordering::Less
317            } else if PartialOrder::less_than(right, left) {
318                Ordering::Greater
319            } else {
320                Ordering::Equal
321            }
322        };
323        let mut legacy_batches: Vec<_> = legacy_batches.into_iter().map(|(k, _)| k).collect();
324        legacy_batches.sort_by(|a, b| compare_chains(a.desc.lower(), b.desc.lower()).reverse());
325
326        let mut pop_batch = |id: SpineId,
327                             expected_desc: Option<&Description<T>>|
328         -> Result<_, String> {
329            if let Some(batch) = hollow_batches.remove(&id) {
330                if let Some(desc) = expected_desc {
331                    // We don't expect the desc's upper and lower to change for a given spine id.
332                    if desc.lower() != batch.desc.lower() || desc.upper() != batch.desc.upper() {
333                        return Err(format!(
334                            "hollow batch desc {:?} did not match the spine batch desc {:?} for {id:?}",
335                            batch.desc, desc
336                        ));
337                    }
338                    // Due to the way thin spine batches are diffed, the sinces can be out of sync.
339                    // This should be rare, and hopefully impossible once we change how diffs work.
340                    if desc.since() != batch.desc.since() {
341                        warn!(
342                            "unexpected since out of sync for spine batch: {:?} != {:?}",
343                            desc.since().elements(),
344                            batch.desc.since().elements()
345                        );
346                    }
347                }
348                return Ok(IdHollowBatch { id, batch });
349            }
350            let mut batch = legacy_batches
351                .pop()
352                .ok_or_else(|| format!("missing referenced hollow batch {id:?}"))?;
353
354            let Some(expected_desc) = expected_desc else {
355                return Ok(IdHollowBatch { id, batch });
356            };
357
358            if expected_desc.lower() != batch.desc.lower() {
359                return Err(format!(
360                    "hollow batch lower {:?} did not match expected lower {:?}",
361                    batch.desc.lower().elements(),
362                    expected_desc.lower().elements()
363                ));
364            }
365
366            // Empty legacy batches are not deterministic: different nodes may split them up
367            // in different ways. For now, we rearrange them such to match the spine data.
368            if batch.parts.is_empty() && batch.run_splits.is_empty() && batch.len == 0 {
369                let mut new_upper = batch.desc.upper().clone();
370
371                // While our current batch is too small, and there's another empty batch
372                // in the list, roll it in.
373                while PartialOrder::less_than(&new_upper, expected_desc.upper()) {
374                    let Some(next_batch) = legacy_batches.pop() else {
375                        break;
376                    };
377                    if next_batch.is_empty() {
378                        new_upper.clone_from(next_batch.desc.upper());
379                    } else {
380                        legacy_batches.push(next_batch);
381                        break;
382                    }
383                }
384
385                // If our current batch is too large, split it by the expected upper
386                // and preserve the remainder.
387                if PartialOrder::less_than(expected_desc.upper(), &new_upper) {
388                    legacy_batches.push(Arc::new(HollowBatch::empty(Description::new(
389                        expected_desc.upper().clone(),
390                        new_upper.clone(),
391                        batch.desc.since().clone(),
392                    ))));
393                    new_upper.clone_from(expected_desc.upper());
394                }
395                batch = Arc::new(HollowBatch::empty(Description::new(
396                    batch.desc.lower().clone(),
397                    new_upper,
398                    batch.desc.since().clone(),
399                )))
400            }
401
402            if expected_desc.upper() != batch.desc.upper() {
403                return Err(format!(
404                    "hollow batch upper {:?} did not match expected upper {:?}",
405                    batch.desc.upper().elements(),
406                    expected_desc.upper().elements()
407                ));
408            }
409
410            Ok(IdHollowBatch { id, batch })
411        };
412
413        let (upper, next_id) = if let Some((id, batch)) = spine_batches.last_key_value() {
414            (batch.desc.upper().clone(), id.1)
415        } else {
416            (Antichain::from_elem(T::minimum()), 0)
417        };
418        // Real spine levels are logarithmic in the total len of the trace, so
419        // this bound is far above any legitimate level while keeping the
420        // allocation below trivial.
421        const MAX_LEVELS: usize = 256;
422        let levels = spine_batches
423            .first_key_value()
424            .map(|(_, batch)| batch.level.saturating_add(1))
425            .unwrap_or(0);
426        if levels > MAX_LEVELS {
427            return Err(format!(
428                "spine level {} exceeds the maximum {MAX_LEVELS}",
429                levels - 1
430            ));
431        }
432        let mut merging = vec![MergeState::default(); levels];
433        for (id, batch) in spine_batches {
434            let level = batch.level;
435
436            if batch.descs.len() > batch.parts.len() {
437                return Err(format!(
438                    "spine batch {id:?} has more descs ({}) than parts ({})",
439                    batch.descs.len(),
440                    batch.parts.len()
441                ));
442            }
443            let descs = batch.descs.iter().map(Some).chain(std::iter::repeat_n(
444                None,
445                batch.parts.len() - batch.descs.len(),
446            ));
447            let parts = batch
448                .parts
449                .into_iter()
450                .zip_eq(descs)
451                .map(|(id, desc)| pop_batch(id, desc))
452                .collect::<Result<Vec<_>, _>>()?;
453            // A spine batch's parts tile its id range (`SpineBatch::id`
454            // `debug_assert`s the endpoints). Real batches always have at least
455            // one part: an empty batch still carries an empty hollow batch.
456            // Validate the full tiling, not just the endpoints: downstream
457            // maintenance (`fueled_merge_reqs_before_ms` -> `id_range` in
458            // compaction, `apply_merge_res_checked`) `assert_eq!`s that the
459            // collected part ids are contiguous, so non-adjacent parts that
460            // happen to hit the right endpoints would panic later instead of
461            // here.
462            if parts.first().map(|x| x.id.0) != Some(id.0)
463                || parts.last().map(|x| x.id.1) != Some(id.1)
464                || parts.windows(2).any(|w| w[0].id.1 != w[1].id.0)
465            {
466                return Err(format!(
467                    "spine batch {id:?} parts do not tile the batch's id range"
468                ));
469            }
470            let len = parts.iter().map(|p| (*p).batch.len).sum();
471            let active_compaction = merges.remove(&id).and_then(|m| m.active_compaction);
472            let batch = SpineBatch {
473                id,
474                desc: batch.desc,
475                parts,
476                active_compaction,
477                len,
478            };
479
480            let state = merging.get_mut(level).ok_or_else(|| {
481                format!("spine batch {id:?} level {level} out of bounds ({levels} levels)")
482            })?;
483
484            state.try_push_batch(batch)?;
485            if let Some(id) = state.id() {
486                if let Some(merge) = merges.remove(&id) {
487                    state.merge = Some(IdFuelingMerge {
488                        id,
489                        merge: FuelingMerge {
490                            since: merge.since,
491                            remaining_work: merge.remaining_work,
492                        },
493                    })
494                }
495            }
496        }
497
498        let mut trace = Trace {
499            spine: Spine {
500                effort: 1,
501                next_id,
502                since,
503                upper,
504                merging,
505            },
506            roundtrip_structure,
507        };
508
509        fn check_empty(name: &str, len: usize) -> Result<(), String> {
510            if len != 0 {
511                Err(format!("{len} {name} left after reconstructing spine"))
512            } else {
513                Ok(())
514            }
515        }
516
517        if roundtrip_structure {
518            check_empty("legacy batches", legacy_batches.len())?;
519        } else {
520            // If the structure wasn't actually serialized, we may have legacy batches left over.
521            for batch in legacy_batches.into_iter().rev() {
522                // `Spine::insert` asserts that pushed batches are non-empty
523                // and contiguous; check this here so that a corrupted batch
524                // results in a decode error instead of a panic.
525                if batch.desc.lower() == batch.desc.upper() {
526                    return Err(format!(
527                        "legacy batch has an empty time range: {:?}",
528                        batch.desc
529                    ));
530                }
531                if batch.desc.lower() != trace.upper() {
532                    return Err(format!(
533                        "legacy batch lower {:?} does not match the trace upper {:?}",
534                        batch.desc.lower().elements(),
535                        trace.upper().elements()
536                    ));
537                }
538                trace.push_batch_no_merge_reqs(Arc::unwrap_or_clone(batch));
539            }
540        }
541        check_empty("hollow batches", hollow_batches.len())?;
542        check_empty("merges", merges.len())?;
543
544        // The same check that's `debug_assert`ed when mutating a trace we
545        // built ourselves; for a trace reconstructed from untrusted data it
546        // must be a hard error, both to keep corrupted state from being used
547        // and because the write side would panic on it anyway (e.g. `Spine`'s
548        // batch invariants and the full-level/merge correspondence).
549        trace
550            .validate()
551            .map_err(|err| format!("reconstructed trace failed validation: {err}"))?;
552
553        Ok(trace)
554    }
555}
556
557#[derive(Clone, Debug, Default)]
558pub(crate) struct SpineMetrics {
559    pub compact_batches: u64,
560    pub compacting_batches: u64,
561    pub noncompact_batches: u64,
562}
563
564impl<T> Trace<T> {
565    pub fn since(&self) -> &Antichain<T> {
566        &self.spine.since
567    }
568
569    pub fn upper(&self) -> &Antichain<T> {
570        &self.spine.upper
571    }
572
573    pub fn map_batches<'a, F: FnMut(&'a HollowBatch<T>)>(&'a self, mut f: F) {
574        for batch in self.batches() {
575            f(batch);
576        }
577    }
578
579    pub fn batches(&self) -> impl Iterator<Item = &HollowBatch<T>> {
580        self.spine
581            .spine_batches()
582            .flat_map(|b| b.parts.as_slice())
583            .map(|b| &*b.batch)
584    }
585
586    pub fn num_spine_batches(&self) -> usize {
587        self.spine.spine_batches().count()
588    }
589
590    #[cfg(test)]
591    pub fn num_hollow_batches(&self) -> usize {
592        self.batches().count()
593    }
594
595    #[cfg(test)]
596    pub fn num_updates(&self) -> usize {
597        self.batches().map(|b| b.len).sum()
598    }
599}
600
601impl<T: Timestamp + Lattice> Trace<T> {
602    pub fn downgrade_since(&mut self, since: &Antichain<T>) {
603        self.spine.since.clone_from(since);
604    }
605
606    #[must_use]
607    pub fn push_batch(&mut self, batch: HollowBatch<T>) -> Vec<FueledMergeReq<T>> {
608        let mut merge_reqs = Vec::new();
609        self.spine.insert(
610            batch,
611            &mut SpineLog::Enabled {
612                merge_reqs: &mut merge_reqs,
613            },
614        );
615        debug_assert_eq!(self.spine.validate(), Ok(()), "{:?}", self);
616        // Spine::roll_up (internally used by insert) clears all batches out of
617        // levels below a target by walking up from level 0 and merging each
618        // level into the next (providing the necessary fuel). In practice, this
619        // means we'll get a series of requests like `(a, b), (a, b, c), ...`.
620        // It's a waste to do all of these (we'll throw away the results), so we
621        // filter out any that are entirely covered by some other request.
622        Self::remove_redundant_merge_reqs(merge_reqs)
623    }
624
625    pub fn claim_compaction(&mut self, id: SpineId, compaction: ActiveCompaction) {
626        // TODO: we ought to be able to look up the id for a batch by binary searching the levels.
627        // In the meantime, search backwards, since most compactions are for recent batches.
628        for batch in self.spine.spine_batches_mut().rev() {
629            if batch.id == id {
630                batch.active_compaction = Some(compaction);
631                break;
632            }
633        }
634    }
635
636    /// The same as [Self::push_batch] but without the `FueledMergeReq`s, which
637    /// account for a surprising amount of cpu in prod. database-issues#5411
638    pub(crate) fn push_batch_no_merge_reqs(&mut self, batch: HollowBatch<T>) {
639        self.spine.insert(batch, &mut SpineLog::Disabled);
640    }
641
642    /// Apply some amount of effort to trace maintenance.
643    ///
644    /// The units of effort are updates, and the method should be thought of as
645    /// analogous to inserting as many empty updates, where the trace is
646    /// permitted to perform proportionate work.
647    ///
648    /// Returns true if this did work and false if it left the spine unchanged.
649    #[must_use]
650    pub fn exert(&mut self, fuel: usize) -> (Vec<FueledMergeReq<T>>, bool) {
651        let mut merge_reqs = Vec::new();
652        let did_work = self.spine.exert(
653            fuel,
654            &mut SpineLog::Enabled {
655                merge_reqs: &mut merge_reqs,
656            },
657        );
658        debug_assert_eq!(self.spine.validate(), Ok(()), "{:?}", self);
659        // See the comment in [Self::push_batch].
660        let merge_reqs = Self::remove_redundant_merge_reqs(merge_reqs);
661        (merge_reqs, did_work)
662    }
663
664    /// Validates invariants.
665    ///
666    /// See `Spine::validate` for details.
667    pub fn validate(&self) -> Result<(), String> {
668        self.spine.validate()
669    }
670
671    /// Obtain all fueled merge reqs that either have no active compaction, or the previous
672    /// compaction was started at or before the threshold time, in order from oldest to newest.
673    pub(crate) fn fueled_merge_reqs_before_ms(
674        &self,
675        threshold_ms: u64,
676        threshold_writer: Option<WriterKey>,
677    ) -> impl Iterator<Item = FueledMergeReq<T>> + '_ {
678        self.spine
679            .spine_batches()
680            .filter(move |b| {
681                let noncompact = !b.is_compact();
682                let old_writer = threshold_writer.as_ref().map_or(false, |min_writer| {
683                    b.parts.iter().any(|b| {
684                        b.batch
685                            .parts
686                            .iter()
687                            .any(|p| p.writer_key().map_or(false, |writer| writer < *min_writer))
688                    })
689                });
690                noncompact || old_writer
691            })
692            .filter(move |b| {
693                // Either there's no active compaction, or the last active compaction
694                // is not after the timeout timestamp.
695                b.active_compaction
696                    .as_ref()
697                    .map_or(true, move |c| c.start_ms <= threshold_ms)
698            })
699            .map(|b| FueledMergeReq {
700                id: b.id,
701                desc: b.desc.clone(),
702                inputs: b.parts.clone(),
703            })
704    }
705
706    // This is only called with the results of one `insert` and so the length of
707    // `merge_reqs` is bounded by the number of levels in the spine (or possibly
708    // some small constant multiple?). The number of levels is logarithmic in the
709    // number of updates in the spine, so this number should stay very small. As
710    // a result, we simply use the naive O(n^2) algorithm here instead of doing
711    // anything fancy with e.g. interval trees.
712    fn remove_redundant_merge_reqs(
713        mut merge_reqs: Vec<FueledMergeReq<T>>,
714    ) -> Vec<FueledMergeReq<T>> {
715        // Returns true if b0 covers b1, false otherwise.
716        fn covers<T: PartialOrder>(b0: &FueledMergeReq<T>, b1: &FueledMergeReq<T>) -> bool {
717            // TODO: can we relax or remove this since check?
718            b0.id.covers(b1.id) && b0.desc.since() == b1.desc.since()
719        }
720
721        let mut ret = Vec::<FueledMergeReq<T>>::with_capacity(merge_reqs.len());
722        // In practice, merge_reqs will come in sorted such that the "large"
723        // requests are later. Take advantage of this by processing back to
724        // front.
725        while let Some(merge_req) = merge_reqs.pop() {
726            let covered = ret.iter().any(|r| covers(r, &merge_req));
727            if !covered {
728                // Now check if anything we've already staged is covered by this
729                // new req. In practice, the merge_reqs come in sorted and so
730                // this `retain` is a no-op.
731                ret.retain(|r| !covers(&merge_req, r));
732                ret.push(merge_req);
733            }
734        }
735        ret
736    }
737
738    pub fn spine_metrics(&self) -> SpineMetrics {
739        let mut metrics = SpineMetrics::default();
740        for batch in self.spine.spine_batches() {
741            if batch.is_compact() {
742                metrics.compact_batches += 1;
743            } else if batch.is_merging() {
744                metrics.compacting_batches += 1;
745            } else {
746                metrics.noncompact_batches += 1;
747            }
748        }
749        metrics
750    }
751}
752
753impl<T: Timestamp + Lattice + Codec64> Trace<T> {
754    pub fn apply_merge_res_checked<D: Codec64 + Monoid + PartialEq>(
755        &mut self,
756        res: &FueledMergeRes<T>,
757        metrics: &ColumnarMetrics,
758    ) -> ApplyMergeResult {
759        for batch in self.spine.spine_batches_mut().rev() {
760            let result = batch.maybe_replace_checked::<D>(res, metrics);
761            if result.matched() {
762                return result;
763            }
764        }
765        ApplyMergeResult::NotAppliedNoMatch
766    }
767
768    pub fn apply_merge_res_unchecked(&mut self, res: &FueledMergeRes<T>) -> ApplyMergeResult {
769        for batch in self.spine.spine_batches_mut().rev() {
770            let result = batch.maybe_replace_unchecked(res);
771            if result.matched() {
772                return result;
773            }
774        }
775        ApplyMergeResult::NotAppliedNoMatch
776    }
777
778    pub fn apply_tombstone_merge(&mut self, desc: &Description<T>) -> ApplyMergeResult {
779        for batch in self.spine.spine_batches_mut().rev() {
780            let result = batch.maybe_replace_with_tombstone(desc);
781            if result.matched() {
782                return result;
783            }
784        }
785        ApplyMergeResult::NotAppliedNoMatch
786    }
787}
788
789/// A log of what transitively happened during a Spine operation: e.g.
790/// FueledMergeReqs were generated.
791enum SpineLog<'a, T> {
792    Enabled {
793        merge_reqs: &'a mut Vec<FueledMergeReq<T>>,
794    },
795    Disabled,
796}
797
798#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
799pub enum CompactionInput {
800    /// We don't know what our inputs were; this should only be used for
801    /// unchecked legacy replacements.
802    Legacy,
803    /// This compaction output is a total replacement for all batches in this id range.
804    IdRange(SpineId),
805    /// This compaction output replaces the specified runs in this id range.
806    PartialBatch(SpineId, BTreeSet<RunId>),
807}
808
809#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
810pub struct SpineId(pub usize, pub usize);
811
812impl Display for SpineId {
813    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
814        write!(f, "[{}, {})", self.0, self.1)
815    }
816}
817
818impl Serialize for SpineId {
819    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
820    where
821        S: Serializer,
822    {
823        let SpineId(lo, hi) = self;
824        serializer.serialize_str(&format!("{lo}-{hi}"))
825    }
826}
827
828/// Creates a `SpineId` that covers the range of ids in the set.
829pub fn id_range(ids: BTreeSet<SpineId>) -> SpineId {
830    let mut id_iter = ids.iter().copied();
831    let Some(mut result) = id_iter.next() else {
832        panic!("at least one batch must be present")
833    };
834
835    for id in id_iter {
836        assert_eq!(
837            result.1, id.0,
838            "expected contiguous ids, but {result:?} is not adjacent to {id:?} in ids {ids:?}"
839        );
840        result.1 = id.1;
841    }
842    result
843}
844
845impl SpineId {
846    fn covers(self, other: SpineId) -> bool {
847        self.0 <= other.0 && other.1 <= self.1
848    }
849}
850
851#[derive(Debug, Clone, PartialEq)]
852pub struct IdHollowBatch<T> {
853    pub id: SpineId,
854    pub batch: Arc<HollowBatch<T>>,
855}
856
857#[derive(Debug, Clone, Eq, PartialEq, Serialize)]
858pub struct ActiveCompaction {
859    pub start_ms: u64,
860}
861
862#[derive(Debug, Clone, PartialEq)]
863struct SpineBatch<T> {
864    id: SpineId,
865    desc: Description<T>,
866    parts: Vec<IdHollowBatch<T>>,
867    active_compaction: Option<ActiveCompaction>,
868    // A cached version of parts.iter().map(|x| x.len).sum()
869    len: usize,
870}
871
872impl<T> SpineBatch<T> {
873    fn merged(batch: IdHollowBatch<T>) -> Self
874    where
875        T: Clone,
876    {
877        Self {
878            id: batch.id,
879            desc: batch.batch.desc.clone(),
880            len: batch.batch.len,
881            parts: vec![batch],
882            active_compaction: None,
883        }
884    }
885}
886
887#[derive(Debug, Copy, Clone)]
888pub enum ApplyMergeResult {
889    AppliedExact,
890    AppliedSubset,
891    NotAppliedNoMatch,
892    NotAppliedInvalidSince,
893    NotAppliedTooManyUpdates,
894}
895
896impl ApplyMergeResult {
897    pub fn applied(&self) -> bool {
898        match self {
899            ApplyMergeResult::AppliedExact | ApplyMergeResult::AppliedSubset => true,
900            _ => false,
901        }
902    }
903    pub fn matched(&self) -> bool {
904        match self {
905            ApplyMergeResult::AppliedExact
906            | ApplyMergeResult::AppliedSubset
907            | ApplyMergeResult::NotAppliedTooManyUpdates => true,
908            _ => false,
909        }
910    }
911}
912
913impl<T: Timestamp + Lattice> SpineBatch<T> {
914    pub fn lower(&self) -> &Antichain<T> {
915        self.desc().lower()
916    }
917
918    pub fn upper(&self) -> &Antichain<T> {
919        self.desc().upper()
920    }
921
922    fn id(&self) -> SpineId {
923        mz_ore::soft_assert_eq_no_log!(self.parts.first().map(|x| x.id.0), Some(self.id.0));
924        mz_ore::soft_assert_eq_no_log!(self.parts.last().map(|x| x.id.1), Some(self.id.1));
925        self.id
926    }
927
928    pub fn is_compact(&self) -> bool {
929        // A compact batch has at most one run.
930        // This check used to be if there was at most one hollow batch with at most one run,
931        // but that was a bit too strict since introducing incremental compaction.
932        // Incremental compaction can result in a batch with a single run, but multiple empty
933        // hollow batches, which we still consider compact. As levels are merged, we
934        // will eventually clean up the empty hollow batches.
935        self.parts
936            .iter()
937            .map(|p| p.batch.run_meta.len())
938            .sum::<usize>()
939            <= 1
940    }
941
942    pub fn is_merging(&self) -> bool {
943        self.active_compaction.is_some()
944    }
945
946    fn desc(&self) -> &Description<T> {
947        &self.desc
948    }
949
950    pub fn len(&self) -> usize {
951        // NB: This is an upper bound on len for a non-compact batch; we won't know for sure until
952        // we compact it.
953        debug_assert_eq!(
954            self.len,
955            self.parts.iter().map(|x| x.batch.len).sum::<usize>()
956        );
957        self.len
958    }
959
960    pub fn is_empty(&self) -> bool {
961        self.len() == 0
962    }
963
964    pub fn empty(
965        id: SpineId,
966        lower: Antichain<T>,
967        upper: Antichain<T>,
968        since: Antichain<T>,
969    ) -> Self {
970        SpineBatch::merged(IdHollowBatch {
971            id,
972            batch: Arc::new(HollowBatch::empty(Description::new(lower, upper, since))),
973        })
974    }
975
976    pub fn begin_merge(
977        bs: &[Self],
978        compaction_frontier: Option<AntichainRef<T>>,
979    ) -> Option<IdFuelingMerge<T>> {
980        let from = bs.first()?.id().0;
981        let until = bs.last()?.id().1;
982        let id = SpineId(from, until);
983        let mut sinces = bs.iter().map(|b| b.desc().since());
984        let mut since = sinces.next()?.clone();
985        for b in bs {
986            since.join_assign(b.desc().since())
987        }
988        if let Some(compaction_frontier) = compaction_frontier {
989            since.join_assign(&compaction_frontier.to_owned());
990        }
991        let remaining_work = bs.iter().map(|x| x.len()).sum();
992        Some(IdFuelingMerge {
993            id,
994            merge: FuelingMerge {
995                since,
996                remaining_work,
997            },
998        })
999    }
1000
1001    #[cfg(test)]
1002    fn describe(&self, extended: bool) -> String {
1003        let SpineBatch {
1004            id,
1005            parts,
1006            desc,
1007            active_compaction,
1008            len,
1009        } = self;
1010        let compaction = match active_compaction {
1011            None => "".to_owned(),
1012            Some(c) => format!(" (c@{})", c.start_ms),
1013        };
1014        match extended {
1015            false => format!(
1016                "[{}-{}]{:?}{:?}{}/{}{compaction}",
1017                id.0,
1018                id.1,
1019                desc.lower().elements(),
1020                desc.upper().elements(),
1021                parts.len(),
1022                len
1023            ),
1024            true => {
1025                format!(
1026                    "[{}-{}]{:?}{:?}{:?} {}/{}{}{compaction}",
1027                    id.0,
1028                    id.1,
1029                    desc.lower().elements(),
1030                    desc.upper().elements(),
1031                    desc.since().elements(),
1032                    parts.len(),
1033                    len,
1034                    parts
1035                        .iter()
1036                        .flat_map(|x| x.batch.parts.iter())
1037                        .map(|x| format!(" {}", x.printable_name()))
1038                        .collect::<Vec<_>>()
1039                        .join("")
1040                )
1041            }
1042        }
1043    }
1044}
1045
1046impl<T: Timestamp + Lattice + Codec64> SpineBatch<T> {
1047    fn diffs_sum<'a, D: Monoid + Codec64>(
1048        parts: impl IntoIterator<Item = &'a RunPart<T>>,
1049        metrics: &ColumnarMetrics,
1050    ) -> Option<D> {
1051        let mut sum = D::zero();
1052        for part in parts {
1053            sum.plus_equals(&part.diffs_sum::<D>(metrics)?);
1054        }
1055        Some(sum)
1056    }
1057
1058    /// Get the diff sum across all runs of the given batch.
1059    ///
1060    /// Returns `None` if any parts don't have statistics, or if any run may
1061    /// hold updates outside the batch's registered desc: the statistics count
1062    /// those updates but readers filter them out, so no sum computed from the
1063    /// statistics can be compared against data seen through a read.
1064    fn batch_diffs_sum<D: Monoid + Codec64>(
1065        batch: &HollowBatch<T>,
1066        metrics: &ColumnarMetrics,
1067    ) -> Option<D> {
1068        let mut sum = D::zero();
1069        for (meta, run) in batch.runs() {
1070            if meta.bounds_truncated() {
1071                return None;
1072            }
1073            sum.plus_equals(&Self::diffs_sum(run, metrics)?);
1074        }
1075        Some(sum)
1076    }
1077
1078    /// Get the diff sum from the given batch for the given runs.
1079    /// Returns `None` if the runs aren't present, any parts don't have
1080    /// statistics, or any of the runs may hold updates outside the batch's
1081    /// registered desc (see [Self::batch_diffs_sum]).
1082    fn diffs_sum_for_runs<D: Monoid + Codec64>(
1083        batch: &HollowBatch<T>,
1084        run_ids: &[RunId],
1085        metrics: &ColumnarMetrics,
1086    ) -> Option<D> {
1087        let mut run_ids = BTreeSet::from_iter(run_ids.iter().copied());
1088        let mut sum = D::zero();
1089
1090        for (meta, run) in batch.runs() {
1091            let id = meta.id?;
1092            if run_ids.remove(&id) {
1093                if meta.bounds_truncated() {
1094                    return None;
1095                }
1096                sum.plus_equals(&Self::diffs_sum(run, metrics)?);
1097            }
1098        }
1099
1100        run_ids.is_empty().then_some(sum)
1101    }
1102
1103    fn maybe_replace_with_tombstone(&mut self, desc: &Description<T>) -> ApplyMergeResult {
1104        let exact_match =
1105            desc.lower() == self.desc().lower() && desc.upper() == self.desc().upper();
1106
1107        let empty_batch = HollowBatch::empty(desc.clone());
1108        if exact_match {
1109            *self = SpineBatch::merged(IdHollowBatch {
1110                id: self.id(),
1111                batch: Arc::new(empty_batch),
1112            });
1113            return ApplyMergeResult::AppliedExact;
1114        }
1115
1116        if let Some((id, range)) = self.find_replacement_range(desc) {
1117            self.perform_subset_replacement(&empty_batch, id, range, None)
1118        } else {
1119            ApplyMergeResult::NotAppliedNoMatch
1120        }
1121    }
1122
1123    fn construct_batch_with_runs_replaced(
1124        original: &HollowBatch<T>,
1125        run_ids: &[RunId],
1126        replacement: &HollowBatch<T>,
1127    ) -> Result<HollowBatch<T>, ApplyMergeResult> {
1128        if run_ids.is_empty() {
1129            return Err(ApplyMergeResult::NotAppliedNoMatch);
1130        }
1131
1132        let orig_run_ids: BTreeSet<_> = original.runs().filter_map(|(meta, _)| meta.id).collect();
1133        let run_ids: BTreeSet<_> = run_ids.iter().cloned().collect();
1134        if !orig_run_ids.is_superset(&run_ids) {
1135            return Err(ApplyMergeResult::NotAppliedNoMatch);
1136        }
1137
1138        let runs: Vec<_> = original
1139            .runs()
1140            .filter(|(meta, _)| {
1141                !run_ids.contains(&meta.id.expect("id should be present at this point"))
1142            })
1143            .chain(replacement.runs())
1144            .collect();
1145
1146        let len = runs.iter().filter_map(|(meta, _)| meta.len).sum::<usize>();
1147
1148        let run_meta = runs
1149            .iter()
1150            .map(|(meta, _)| *meta)
1151            .cloned()
1152            .collect::<Vec<_>>();
1153
1154        let parts = runs
1155            .iter()
1156            .flat_map(|(_, parts)| *parts)
1157            .cloned()
1158            .collect::<Vec<_>>();
1159
1160        let run_splits = {
1161            let mut splits = Vec::with_capacity(run_meta.len().saturating_sub(1));
1162            let mut pointer = 0;
1163            for (i, (_, parts)) in runs.into_iter().enumerate() {
1164                if parts.is_empty() {
1165                    continue;
1166                }
1167                if i < run_meta.len() - 1 {
1168                    splits.push(pointer + parts.len());
1169                }
1170                pointer += parts.len();
1171            }
1172            splits
1173        };
1174
1175        Ok(HollowBatch::new(
1176            replacement.desc.clone(),
1177            parts,
1178            len,
1179            run_meta,
1180            run_splits,
1181        ))
1182    }
1183
1184    fn maybe_replace_checked<D>(
1185        &mut self,
1186        res: &FueledMergeRes<T>,
1187        metrics: &ColumnarMetrics,
1188    ) -> ApplyMergeResult
1189    where
1190        D: Monoid + Codec64 + PartialEq + Debug,
1191    {
1192        // The spine's and merge res's sinces don't need to match (which could occur if Spine
1193        // has been reloaded from state due to compare_and_set mismatch), but if so, the Spine
1194        // since must be in advance of the merge res since.
1195        if !PartialOrder::less_equal(res.output.desc.since(), self.desc().since()) {
1196            return ApplyMergeResult::NotAppliedInvalidSince;
1197        }
1198
1199        let new_diffs_sum = Self::diffs_sum(res.output.parts.iter(), metrics);
1200        let num_batches = self.parts.len();
1201
1202        let result = match &res.input {
1203            CompactionInput::IdRange(id) => {
1204                self.handle_id_range_replacement::<D>(res, id, new_diffs_sum, metrics)
1205            }
1206            CompactionInput::PartialBatch(id, runs) => {
1207                self.handle_partial_batch_replacement::<D>(res, *id, runs, new_diffs_sum, metrics)
1208            }
1209            CompactionInput::Legacy => self.maybe_replace_checked_classic::<D>(res, metrics),
1210        };
1211
1212        let num_batches_after = self.parts.len();
1213        assert!(
1214            num_batches_after <= num_batches,
1215            "replacing parts should not increase the number of batches"
1216        );
1217        result
1218    }
1219
1220    fn handle_id_range_replacement<D>(
1221        &mut self,
1222        res: &FueledMergeRes<T>,
1223        id: &SpineId,
1224        new_diffs_sum: Option<D>,
1225        metrics: &ColumnarMetrics,
1226    ) -> ApplyMergeResult
1227    where
1228        D: Monoid + Codec64 + PartialEq + Debug,
1229    {
1230        let range = self
1231            .parts
1232            .iter()
1233            .enumerate()
1234            .filter_map(|(i, p)| {
1235                if id.covers(p.id) {
1236                    Some((i, p.id))
1237                } else {
1238                    None
1239                }
1240            })
1241            .collect::<Vec<_>>();
1242
1243        let ids: BTreeSet<_> = range.iter().map(|(_, id)| *id).collect();
1244
1245        // If ids is empty, it means that we didn't find any parts that match the id range.
1246        // We also check that the id matches the range of ids we found.
1247        // At scale, sometimes regular compaction will race forced compaction,
1248        // for things like the catalog. In that case, we may have a
1249        // replacement that no longer lines up with the spine batches.
1250        // I think this is because forced compaction ignores the active_compaction
1251        // and just goes for it. This is slightly annoying but probably the right behavior
1252        // for a functions whose prefix is `force_`, so we just return
1253        // NotAppliedNoMatch here.
1254        if ids.is_empty() || id != &id_range(ids) {
1255            return ApplyMergeResult::NotAppliedNoMatch;
1256        }
1257
1258        // This is the range of hollow batches that we will replace.
1259        let (min, max) = match range.iter().map(|(i, _)| *i).minmax() {
1260            itertools::MinMaxResult::NoElements => return ApplyMergeResult::NotAppliedNoMatch,
1261            itertools::MinMaxResult::OneElement(elt) => (elt, elt),
1262            itertools::MinMaxResult::MinMax(min, max) => (min, max),
1263        };
1264        let replacement_range = min..max + 1;
1265
1266        // We need to replace a range of parts. Here we don't care about the run_indices
1267        // because we must be replacing the entire part(s)
1268        let old_diffs_sum =
1269            self.parts[replacement_range.clone()]
1270                .iter()
1271                .try_fold(D::zero(), |mut sum, p| {
1272                    sum.plus_equals(&Self::batch_diffs_sum::<D>(&p.batch, metrics)?);
1273                    Some(sum)
1274                });
1275
1276        Self::validate_diffs_sum_match(old_diffs_sum, new_diffs_sum, "id range replacement");
1277
1278        self.perform_subset_replacement(
1279            &res.output,
1280            *id,
1281            replacement_range,
1282            res.new_active_compaction.clone(),
1283        )
1284    }
1285
1286    fn handle_partial_batch_replacement<D>(
1287        &mut self,
1288        res: &FueledMergeRes<T>,
1289        id: SpineId,
1290        runs: &BTreeSet<RunId>,
1291        new_diffs_sum: Option<D>,
1292        metrics: &ColumnarMetrics,
1293    ) -> ApplyMergeResult
1294    where
1295        D: Monoid + Codec64 + PartialEq + Debug,
1296    {
1297        if runs.is_empty() {
1298            return ApplyMergeResult::NotAppliedNoMatch;
1299        }
1300
1301        let part = self.parts.iter().enumerate().find(|(_, p)| p.id == id);
1302        let Some((i, batch)) = part else {
1303            return ApplyMergeResult::NotAppliedNoMatch;
1304        };
1305        let replacement_range = i..(i + 1);
1306
1307        let replacement_desc = &res.output.desc;
1308        let existing_desc = &batch.batch.desc;
1309        assert_eq!(
1310            replacement_desc.lower(),
1311            existing_desc.lower(),
1312            "batch lower should match, but {:?} != {:?}",
1313            replacement_desc.lower(),
1314            existing_desc.lower()
1315        );
1316        assert_eq!(
1317            replacement_desc.upper(),
1318            existing_desc.upper(),
1319            "batch upper should match, but {:?} != {:?}",
1320            replacement_desc.upper(),
1321            existing_desc.upper()
1322        );
1323        if !PartialOrder::less_equal(existing_desc.since(), replacement_desc.since()) {
1324            error!(
1325                "batch since should advance, but {:?} !<= {:?}",
1326                existing_desc.since(),
1327                replacement_desc.since()
1328            );
1329            return ApplyMergeResult::NotAppliedInvalidSince;
1330        }
1331
1332        let batch = &batch.batch;
1333        let run_ids = runs.iter().cloned().collect::<Vec<_>>();
1334
1335        match Self::construct_batch_with_runs_replaced(batch, &run_ids, &res.output) {
1336            Ok(new_batch) => {
1337                let old_diffs_sum = Self::diffs_sum_for_runs::<D>(batch, &run_ids, metrics);
1338                Self::validate_diffs_sum_match(
1339                    old_diffs_sum,
1340                    new_diffs_sum,
1341                    "partial batch replacement",
1342                );
1343                let old_batch_diff_sum = Self::batch_diffs_sum::<D>(batch, metrics);
1344                let new_batch_diff_sum = Self::batch_diffs_sum::<D>(&new_batch, metrics);
1345                Self::validate_diffs_sum_match(
1346                    old_batch_diff_sum,
1347                    new_batch_diff_sum,
1348                    "sanity checking diffs sum for replaced runs",
1349                );
1350                self.perform_subset_replacement(
1351                    &new_batch,
1352                    id,
1353                    replacement_range,
1354                    res.new_active_compaction.clone(),
1355                )
1356            }
1357            Err(err) => err,
1358        }
1359    }
1360
1361    fn validate_diffs_sum_match<D>(
1362        old_diffs_sum: Option<D>,
1363        new_diffs_sum: Option<D>,
1364        context: &str,
1365    ) where
1366        D: Monoid + Codec64 + PartialEq + Debug,
1367    {
1368        let new_diffs_sum = new_diffs_sum.unwrap_or_else(D::zero);
1369        if let Some(old_diffs_sum) = old_diffs_sum {
1370            assert_eq!(
1371                old_diffs_sum, new_diffs_sum,
1372                "merge res diffs sum ({:?}) did not match spine batch diffs sum ({:?}) ({})",
1373                new_diffs_sum, old_diffs_sum, context
1374            )
1375        }
1376    }
1377
1378    /// This is the "legacy" way of replacing a spine batch with a merge result.
1379    /// It is used in moments when we don't have the full compaction input
1380    /// information.
1381    /// Eventually we should strive to roundtrip Spine IDs everywhere and
1382    /// deprecate this method.
1383    fn maybe_replace_checked_classic<D>(
1384        &mut self,
1385        res: &FueledMergeRes<T>,
1386        metrics: &ColumnarMetrics,
1387    ) -> ApplyMergeResult
1388    where
1389        D: Monoid + Codec64 + PartialEq + Debug,
1390    {
1391        // The spine's and merge res's sinces don't need to match (which could occur if Spine
1392        // has been reloaded from state due to compare_and_set mismatch), but if so, the Spine
1393        // since must be in advance of the merge res since.
1394        if !PartialOrder::less_equal(res.output.desc.since(), self.desc().since()) {
1395            return ApplyMergeResult::NotAppliedInvalidSince;
1396        }
1397
1398        let new_diffs_sum = Self::diffs_sum(res.output.parts.iter(), metrics);
1399
1400        // If our merge result exactly matches a spine batch, we can swap it in directly
1401        let exact_match = res.output.desc.lower() == self.desc().lower()
1402            && res.output.desc.upper() == self.desc().upper();
1403        if exact_match {
1404            let old_diffs_sum = self.parts.iter().try_fold(D::zero(), |mut sum, p| {
1405                sum.plus_equals(&Self::batch_diffs_sum::<D>(&p.batch, metrics)?);
1406                Some(sum)
1407            });
1408
1409            if let (Some(old_diffs_sum), Some(new_diffs_sum)) = (old_diffs_sum, new_diffs_sum) {
1410                assert_eq!(
1411                    old_diffs_sum, new_diffs_sum,
1412                    "merge res diffs sum ({:?}) did not match spine batch diffs sum ({:?})",
1413                    new_diffs_sum, old_diffs_sum
1414                );
1415            }
1416
1417            // Spine internally has an invariant about a batch being at some level
1418            // or higher based on the len. We could end up violating this invariant
1419            // if we increased the length of the batch.
1420            //
1421            // A res output with length greater than the existing spine batch implies
1422            // a compaction has already been applied to this range, and with a higher
1423            // rate of consolidation than this one. This could happen as a result of
1424            // compaction's memory bound limiting the amount of consolidation possible.
1425            if res.output.len > self.len() {
1426                return ApplyMergeResult::NotAppliedTooManyUpdates;
1427            }
1428            *self = SpineBatch::merged(IdHollowBatch {
1429                id: self.id(),
1430                batch: Arc::new(res.output.clone()),
1431            });
1432            return ApplyMergeResult::AppliedExact;
1433        }
1434
1435        // Try subset replacement
1436        if let Some((id, range)) = self.find_replacement_range(&res.output.desc) {
1437            let old_diffs_sum =
1438                self.parts[range.clone()]
1439                    .iter()
1440                    .try_fold(D::zero(), |mut sum, p| {
1441                        sum.plus_equals(&Self::batch_diffs_sum::<D>(&p.batch, metrics)?);
1442                        Some(sum)
1443                    });
1444
1445            if let (Some(old_diffs_sum), Some(new_diffs_sum)) = (old_diffs_sum, new_diffs_sum) {
1446                assert_eq!(
1447                    old_diffs_sum, new_diffs_sum,
1448                    "merge res diffs sum ({:?}) did not match spine batch diffs sum ({:?})",
1449                    new_diffs_sum, old_diffs_sum
1450                );
1451            }
1452
1453            self.perform_subset_replacement(
1454                &res.output,
1455                id,
1456                range,
1457                res.new_active_compaction.clone(),
1458            )
1459        } else {
1460            ApplyMergeResult::NotAppliedNoMatch
1461        }
1462    }
1463
1464    /// This is the even more legacy way of replacing a spine batch with a merge result.
1465    /// It is used in moments when we don't have the full compaction input
1466    /// information, and we don't have the diffs sum.
1467    /// Eventually we should strive to roundtrip Spine IDs and diffs sums everywhere and
1468    /// deprecate this method.
1469    fn maybe_replace_unchecked(&mut self, res: &FueledMergeRes<T>) -> ApplyMergeResult {
1470        // The spine's and merge res's sinces don't need to match (which could occur if Spine
1471        // has been reloaded from state due to compare_and_set mismatch), but if so, the Spine
1472        // since must be in advance of the merge res since.
1473        if !PartialOrder::less_equal(res.output.desc.since(), self.desc().since()) {
1474            return ApplyMergeResult::NotAppliedInvalidSince;
1475        }
1476
1477        // If our merge result exactly matches a spine batch, we can swap it in directly
1478        let exact_match = res.output.desc.lower() == self.desc().lower()
1479            && res.output.desc.upper() == self.desc().upper();
1480        if exact_match {
1481            // Spine internally has an invariant about a batch being at some level
1482            // or higher based on the len. We could end up violating this invariant
1483            // if we increased the length of the batch.
1484            //
1485            // A res output with length greater than the existing spine batch implies
1486            // a compaction has already been applied to this range, and with a higher
1487            // rate of consolidation than this one. This could happen as a result of
1488            // compaction's memory bound limiting the amount of consolidation possible.
1489            if res.output.len > self.len() {
1490                return ApplyMergeResult::NotAppliedTooManyUpdates;
1491            }
1492
1493            *self = SpineBatch::merged(IdHollowBatch {
1494                id: self.id(),
1495                batch: Arc::new(res.output.clone()),
1496            });
1497            return ApplyMergeResult::AppliedExact;
1498        }
1499
1500        // Try subset replacement
1501        if let Some((id, range)) = self.find_replacement_range(&res.output.desc) {
1502            self.perform_subset_replacement(
1503                &res.output,
1504                id,
1505                range,
1506                res.new_active_compaction.clone(),
1507            )
1508        } else {
1509            ApplyMergeResult::NotAppliedNoMatch
1510        }
1511    }
1512
1513    /// Find the range of parts that can be replaced by the merge result
1514    fn find_replacement_range(&self, desc: &Description<T>) -> Option<(SpineId, Range<usize>)> {
1515        // It is possible the structure of the spine has changed since the merge res
1516        // was created, such that it no longer exactly matches the description of a
1517        // spine batch. This can happen if another merge has happened in the interim,
1518        // or if spine needed to be rebuilt from state.
1519        //
1520        // When this occurs, we can still attempt to slot the merge res in to replace
1521        // the parts of a fueled merge. e.g. if the res is for `[1,3)` and the parts
1522        // are `[0,1),[1,2),[2,3),[3,4)`, we can swap out the middle two parts for res.
1523
1524        let mut lower = None;
1525        let mut upper = None;
1526
1527        for (i, batch) in self.parts.iter().enumerate() {
1528            if batch.batch.desc.lower() == desc.lower() {
1529                lower = Some((i, batch.id.0));
1530            }
1531            if batch.batch.desc.upper() == desc.upper() {
1532                upper = Some((i, batch.id.1));
1533            }
1534            if lower.is_some() && upper.is_some() {
1535                break;
1536            }
1537        }
1538
1539        match (lower, upper) {
1540            (Some((lower_idx, id_lower)), Some((upper_idx, id_upper))) => {
1541                Some((SpineId(id_lower, id_upper), lower_idx..(upper_idx + 1)))
1542            }
1543            _ => None,
1544        }
1545    }
1546
1547    /// Perform the actual subset replacement
1548    fn perform_subset_replacement(
1549        &mut self,
1550        res: &HollowBatch<T>,
1551        spine_id: SpineId,
1552        range: Range<usize>,
1553        new_active_compaction: Option<ActiveCompaction>,
1554    ) -> ApplyMergeResult {
1555        let SpineBatch {
1556            id,
1557            parts,
1558            desc,
1559            active_compaction: _,
1560            len: _,
1561        } = self;
1562
1563        let mut new_parts = vec![];
1564        new_parts.extend_from_slice(&parts[..range.start]);
1565        new_parts.push(IdHollowBatch {
1566            id: spine_id,
1567            batch: Arc::new(res.clone()),
1568        });
1569        new_parts.extend_from_slice(&parts[range.end..]);
1570
1571        let res = if range.len() == parts.len() {
1572            ApplyMergeResult::AppliedExact
1573        } else {
1574            ApplyMergeResult::AppliedSubset
1575        };
1576
1577        let new_spine_batch = SpineBatch {
1578            id: *id,
1579            desc: desc.to_owned(),
1580            len: new_parts.iter().map(|x| x.batch.len).sum(),
1581            parts: new_parts,
1582            active_compaction: new_active_compaction,
1583        };
1584
1585        if new_spine_batch.len() > self.len() {
1586            return ApplyMergeResult::NotAppliedTooManyUpdates;
1587        }
1588
1589        *self = new_spine_batch;
1590        res
1591    }
1592}
1593
1594#[derive(Debug, Clone, PartialEq, Serialize)]
1595pub struct FuelingMerge<T> {
1596    pub(crate) since: Antichain<T>,
1597    pub(crate) remaining_work: usize,
1598}
1599
1600#[derive(Debug, Clone, PartialEq, Serialize)]
1601pub struct IdFuelingMerge<T> {
1602    id: SpineId,
1603    merge: FuelingMerge<T>,
1604}
1605
1606impl<T: Timestamp + Lattice> FuelingMerge<T> {
1607    /// Perform some amount of work, decrementing `fuel`.
1608    ///
1609    /// If `fuel` is non-zero after the call, the merging is complete and one
1610    /// should call `done` to extract the merged results.
1611    fn work(&mut self, _: &[SpineBatch<T>], fuel: &mut isize) {
1612        // A negative `fuel` means a caller already overspent, so there is
1613        // nothing to spend here. Reading it as a `usize` instead would wrap to a
1614        // huge budget, spend `remaining_work` against it, and then underflow the
1615        // subtraction below.
1616        let available = usize::try_from(*fuel).unwrap_or(0);
1617        let used = std::cmp::min(available, self.remaining_work);
1618        self.remaining_work = self.remaining_work.saturating_sub(used);
1619        // `used <= available <= isize::MAX`, so neither conversion nor the
1620        // subtraction can overflow.
1621        *fuel -= isize::try_from(used).expect("used is bounded by fuel");
1622    }
1623
1624    /// Extracts merged results.
1625    ///
1626    /// This method should only be called after `work` has been called and has
1627    /// not brought `fuel` to zero. Otherwise, the merge is still in progress.
1628    fn done(
1629        self,
1630        bs: ArrayVec<SpineBatch<T>, BATCHES_PER_LEVEL>,
1631        log: &mut SpineLog<'_, T>,
1632    ) -> Option<SpineBatch<T>> {
1633        let first = bs.first()?;
1634        let last = bs.last()?;
1635        let id = SpineId(first.id().0, last.id().1);
1636        assert!(id.0 < id.1);
1637        let lower = first.desc().lower().clone();
1638        let upper = last.desc().upper().clone();
1639        let since = self.since;
1640
1641        // Special case empty batches.
1642        if bs.iter().all(SpineBatch::is_empty) {
1643            return Some(SpineBatch::empty(id, lower, upper, since));
1644        }
1645
1646        let desc = Description::new(lower, upper, since);
1647        let len = bs.iter().map(SpineBatch::len).sum();
1648
1649        // Pre-size the merged_parts Vec. Benchmarking has shown that, at least
1650        // in the worst case, the double iteration is absolutely worth having
1651        // merged_parts pre-sized.
1652        let mut merged_parts_len = 0;
1653        for b in &bs {
1654            merged_parts_len += b.parts.len();
1655        }
1656        let mut merged_parts = Vec::with_capacity(merged_parts_len);
1657        for b in bs {
1658            merged_parts.extend(b.parts)
1659        }
1660        // Sanity check the pre-size code.
1661        mz_ore::soft_assert_eq_no_log!(merged_parts.len(), merged_parts_len);
1662
1663        if let SpineLog::Enabled { merge_reqs } = log {
1664            merge_reqs.push(FueledMergeReq {
1665                id,
1666                desc: desc.clone(),
1667                inputs: merged_parts.clone(),
1668            });
1669        }
1670
1671        Some(SpineBatch {
1672            id,
1673            desc,
1674            len,
1675            parts: merged_parts,
1676            active_compaction: None,
1677        })
1678    }
1679}
1680
1681/// The maximum number of batches per level in the spine.
1682/// In practice, we probably want a larger max and a configurable soft cap, but using a
1683/// stack-friendly data structure and keeping this number low makes this safer during the
1684/// initial rollout.
1685const BATCHES_PER_LEVEL: usize = 2;
1686
1687/// An append-only collection of update batches.
1688///
1689/// The `Spine` is a general-purpose trace implementation based on collection
1690/// and merging immutable batches of updates. It is generic with respect to the
1691/// batch type, and can be instantiated for any implementor of `trace::Batch`.
1692///
1693/// ## Design
1694///
1695/// This spine is represented as a list of layers, where each element in the
1696/// list is either
1697///
1698///   1. MergeState::Vacant  empty
1699///   2. MergeState::Single  a single batch
1700///   3. MergeState::Double  a pair of batches
1701///
1702/// Each "batch" has the option to be `None`, indicating a non-batch that
1703/// nonetheless acts as a number of updates proportionate to the level at which
1704/// it exists (for bookkeeping).
1705///
1706/// Each of the batches at layer i contains at most 2^i elements. The sequence
1707/// of batches should have the upper bound of one match the lower bound of the
1708/// next. Batches may be logically empty, with matching upper and lower bounds,
1709/// as a bookkeeping mechanism.
1710///
1711/// Each batch at layer i is treated as if it contains exactly 2^i elements,
1712/// even though it may actually contain fewer elements. This allows us to
1713/// decouple the physical representation from logical amounts of effort invested
1714/// in each batch. It allows us to begin compaction and to reduce the number of
1715/// updates, without compromising our ability to continue to move updates along
1716/// the spine. We are explicitly making the trade-off that while some batches
1717/// might compact at lower levels, we want to treat them as if they contained
1718/// their full set of updates for accounting reasons (to apply work to higher
1719/// levels).
1720///
1721/// We maintain the invariant that for any in-progress merge at level k there
1722/// should be fewer than 2^k records at levels lower than k. That is, even if we
1723/// were to apply an unbounded amount of effort to those records, we would not
1724/// have enough records to prompt a merge into the in-progress merge. Ideally,
1725/// we maintain the extended invariant that for any in-progress merge at level
1726/// k, the remaining effort required (number of records minus applied effort) is
1727/// less than the number of records that would need to be added to reach 2^k
1728/// records in layers below.
1729///
1730/// ## Mathematics
1731///
1732/// When a merge is initiated, there should be a non-negative *deficit* of
1733/// updates before the layers below could plausibly produce a new batch for the
1734/// currently merging layer. We must determine a factor of proportionality, so
1735/// that newly arrived updates provide at least that amount of "fuel" towards
1736/// the merging layer, so that the merge completes before lower levels invade.
1737///
1738/// ### Deficit:
1739///
1740/// A new merge is initiated only in response to the completion of a prior
1741/// merge, or the introduction of new records from outside. The latter case is
1742/// special, and will maintain our invariant trivially, so we will focus on the
1743/// former case.
1744///
1745/// When a merge at level k completes, assuming we have maintained our invariant
1746/// then there should be fewer than 2^k records at lower levels. The newly
1747/// created merge at level k+1 will require up to 2^k+2 units of work, and
1748/// should not expect a new batch until strictly more than 2^k records are
1749/// added. This means that a factor of proportionality of four should be
1750/// sufficient to ensure that the merge completes before a new merge is
1751/// initiated.
1752///
1753/// When new records get introduced, we will need to roll up any batches at
1754/// lower levels, which we treat as the introduction of records. Each of these
1755/// virtual records introduced should either be accounted for the fuel it should
1756/// contribute, as it results in the promotion of batches closer to in-progress
1757/// merges.
1758///
1759/// ### Fuel sharing
1760///
1761/// We like the idea of applying fuel preferentially to merges at *lower*
1762/// levels, under the idea that they are easier to complete, and we benefit from
1763/// fewer total merges in progress. This does delay the completion of merges at
1764/// higher levels, and may not obviously be a total win. If we choose to do
1765/// this, we should make sure that we correctly account for completed merges at
1766/// low layers: they should still extract fuel from new updates even though they
1767/// have completed, at least until they have paid back any "debt" to higher
1768/// layers by continuing to provide fuel as updates arrive.
1769#[derive(Debug, Clone)]
1770struct Spine<T> {
1771    effort: usize,
1772    next_id: usize,
1773    since: Antichain<T>,
1774    upper: Antichain<T>,
1775    merging: Vec<MergeState<T>>,
1776}
1777
1778impl<T> Spine<T> {
1779    /// All batches in the spine, oldest to newest.
1780    pub fn spine_batches(&self) -> impl Iterator<Item = &SpineBatch<T>> {
1781        self.merging.iter().rev().flat_map(|m| &m.batches)
1782    }
1783
1784    /// All (mutable) batches in the spine, oldest to newest.
1785    pub fn spine_batches_mut(&mut self) -> impl DoubleEndedIterator<Item = &mut SpineBatch<T>> {
1786        self.merging.iter_mut().rev().flat_map(|m| &mut m.batches)
1787    }
1788}
1789
1790impl<T: Timestamp + Lattice> Spine<T> {
1791    /// Allocates a fueled `Spine`.
1792    ///
1793    /// This trace will merge batches progressively, with each inserted batch
1794    /// applying a multiple of the batch's length in effort to each merge. The
1795    /// `effort` parameter is that multiplier. This value should be at least one
1796    /// for the merging to happen; a value of zero is not helpful.
1797    pub fn new() -> Self {
1798        Spine {
1799            effort: 1,
1800            next_id: 0,
1801            since: Antichain::from_elem(T::minimum()),
1802            upper: Antichain::from_elem(T::minimum()),
1803            merging: Vec::new(),
1804        }
1805    }
1806
1807    /// Apply some amount of effort to trace maintenance.
1808    ///
1809    /// The units of effort are updates, and the method should be thought of as
1810    /// analogous to inserting as many empty updates, where the trace is
1811    /// permitted to perform proportionate work.
1812    ///
1813    /// Returns true if this did work and false if it left the spine unchanged.
1814    fn exert(&mut self, effort: usize, log: &mut SpineLog<'_, T>) -> bool {
1815        self.tidy_layers();
1816        if self.reduced() {
1817            return false;
1818        }
1819
1820        if self.merging.iter().any(|b| b.merge.is_some()) {
1821            let fuel = isize::try_from(effort).unwrap_or(isize::MAX);
1822            // If any merges exist, we can directly call `apply_fuel`.
1823            self.apply_fuel(&fuel, log);
1824        } else {
1825            // Otherwise, we'll need to introduce fake updates to move merges
1826            // along.
1827
1828            // Introduce an empty batch with roughly *effort number of virtual updates.
1829            let level = usize::cast_from(effort.next_power_of_two().trailing_zeros());
1830            let id = self.next_id();
1831            self.introduce_batch(
1832                SpineBatch::empty(
1833                    id,
1834                    self.upper.clone(),
1835                    self.upper.clone(),
1836                    self.since.clone(),
1837                ),
1838                level,
1839                log,
1840            );
1841        }
1842        true
1843    }
1844
1845    pub fn next_id(&mut self) -> SpineId {
1846        let id = self.next_id;
1847        self.next_id += 1;
1848        SpineId(id, self.next_id)
1849    }
1850
1851    // Ideally, this method acts as insertion of `batch`, even if we are not yet
1852    // able to begin merging the batch. This means it is a good time to perform
1853    // amortized work proportional to the size of batch.
1854    pub fn insert(&mut self, batch: HollowBatch<T>, log: &mut SpineLog<'_, T>) {
1855        assert!(batch.desc.lower() != batch.desc.upper());
1856        assert_eq!(batch.desc.lower(), &self.upper);
1857
1858        let id = self.next_id();
1859        let batch = SpineBatch::merged(IdHollowBatch {
1860            id,
1861            batch: Arc::new(batch),
1862        });
1863
1864        self.upper.clone_from(batch.upper());
1865
1866        // If `batch` and the most recently inserted batch are both empty,
1867        // we can just fuse them.
1868        if batch.is_empty() {
1869            if let Some(position) = self.merging.iter().position(|m| !m.is_vacant()) {
1870                if self.merging[position].is_single() && self.merging[position].is_empty() {
1871                    self.insert_at(batch, position);
1872                    // Since we just inserted a batch, we should always have work to complete...
1873                    // but otherwise we just leave this layer vacant.
1874                    if let Some(merged) = self.complete_at(position, log) {
1875                        self.merging[position] = MergeState::single(merged);
1876                    }
1877                    return;
1878                }
1879            }
1880        }
1881
1882        // Normal insertion for the batch.
1883        let index = batch.len().next_power_of_two();
1884        self.introduce_batch(batch, usize::cast_from(index.trailing_zeros()), log);
1885    }
1886
1887    /// Returns true when the trace is considered *structurally reduced*.
1888    ///
1889    /// Reduced == the total number of runs (across every
1890    /// `SpineBatch` and all of their inner hollow batches) is < 2. In other
1891    /// words, there are either zero runs (fully empty) or exactly one logical
1892    /// run of data remaining.
1893    fn reduced(&self) -> bool {
1894        self.spine_batches()
1895            .map(|b| {
1896                b.parts
1897                    .iter()
1898                    .map(|p| p.batch.run_meta.len())
1899                    .sum::<usize>()
1900            })
1901            .sum::<usize>()
1902            < 2
1903    }
1904
1905    /// Describes the merge progress of layers in the trace.
1906    ///
1907    /// Intended for diagnostics rather than public consumption.
1908    #[allow(dead_code)]
1909    fn describe(&self) -> Vec<(usize, usize)> {
1910        self.merging
1911            .iter()
1912            .map(|b| (b.batches.len(), b.len()))
1913            .collect()
1914    }
1915
1916    /// Introduces a batch at an indicated level.
1917    ///
1918    /// The level indication is often related to the size of the batch, but it
1919    /// can also be used to artificially fuel the computation by supplying empty
1920    /// batches at non-trivial indices, to move merges along.
1921    fn introduce_batch(
1922        &mut self,
1923        batch: SpineBatch<T>,
1924        batch_index: usize,
1925        log: &mut SpineLog<'_, T>,
1926    ) {
1927        // Step 0.  Determine an amount of fuel to use for the computation.
1928        //
1929        //          Fuel is used to drive maintenance of the data structure,
1930        //          and in particular are used to make progress through merges
1931        //          that are in progress. The amount of fuel to use should be
1932        //          proportional to the number of records introduced, so that
1933        //          we are guaranteed to complete all merges before they are
1934        //          required as arguments to merges again.
1935        //
1936        //          The fuel use policy is negotiable, in that we might aim
1937        //          to use relatively less when we can, so that we return
1938        //          control promptly, or we might account more work to larger
1939        //          batches. Not clear to me which are best, of if there
1940        //          should be a configuration knob controlling this.
1941
1942        // The amount of fuel to use is proportional to 2^batch_index, scaled by
1943        // a factor of self.effort which determines how eager we are in
1944        // performing maintenance work. We need to ensure that each merge in
1945        // progress receives fuel for each introduced batch, and so multiply by
1946        // that as well.
1947        if batch_index > 32 {
1948            println!("Large batch index: {}", batch_index);
1949        }
1950
1951        // We believe that eight units of fuel is sufficient for each introduced
1952        // record, accounted as four for each record, and a potential four more
1953        // for each virtual record associated with promoting existing smaller
1954        // batches. We could try and make this be less, or be scaled to merges
1955        // based on their deficit at time of instantiation. For now, we remain
1956        // conservative.
1957        //
1958        // `batch_index` is derived from a batch's `len`, which for a trace
1959        // reconstructed from an untrusted blob can be large enough that
1960        // `8 << batch_index` overflows. Saturate at `isize::MAX` rather than
1961        // wrapping: fuel is a budget, so more of it only completes merges sooner,
1962        // whereas a wrapped value can land negative and starve them. The result
1963        // is an `isize` so a fuel shortfall stays observable.
1964        let fuel = u32::try_from(batch_index)
1965            .ok()
1966            .and_then(|shift| 8usize.checked_shl(shift))
1967            // Scale up by the effort parameter, which is calibrated to one as the
1968            // minimum amount of effort.
1969            .and_then(|fuel| fuel.checked_mul(self.effort))
1970            .and_then(|fuel| isize::try_from(fuel).ok())
1971            .unwrap_or(isize::MAX);
1972
1973        // Step 1.  Apply fuel to each in-progress merge.
1974        //
1975        //          Before we can introduce new updates, we must apply any
1976        //          fuel to in-progress merges, as this fuel is what ensures
1977        //          that the merges will be complete by the time we insert
1978        //          the updates.
1979        self.apply_fuel(&fuel, log);
1980
1981        // Step 2.  We must ensure the invariant that adjacent layers do not
1982        //          contain two batches will be satisfied when we insert the
1983        //          batch. We forcibly completing all merges at layers lower
1984        //          than and including `batch_index`, so that the new batch is
1985        //          inserted into an empty layer.
1986        //
1987        //          We could relax this to "strictly less than `batch_index`"
1988        //          if the layer above has only a single batch in it, which
1989        //          seems not implausible if it has been the focus of effort.
1990        //
1991        //          This should be interpreted as the introduction of some
1992        //          volume of fake updates, and we will need to fuel merges
1993        //          by a proportional amount to ensure that they are not
1994        //          surprised later on. The number of fake updates should
1995        //          correspond to the deficit for the layer, which perhaps
1996        //          we should track explicitly.
1997        self.roll_up(batch_index, log);
1998
1999        // Step 3. This insertion should be into an empty layer. It is a logical
2000        //         error otherwise, as we may be violating our invariant, from
2001        //         which all wonderment derives.
2002        self.insert_at(batch, batch_index);
2003
2004        // Step 4. Tidy the largest layers.
2005        //
2006        //         It is important that we not tidy only smaller layers,
2007        //         as their ascension is what ensures the merging and
2008        //         eventual compaction of the largest layers.
2009        self.tidy_layers();
2010    }
2011
2012    /// Ensures that an insertion at layer `index` will succeed.
2013    ///
2014    /// This method is subject to the constraint that all existing batches
2015    /// should occur at higher levels, which requires it to "roll up" batches
2016    /// present at lower levels before the method is called. In doing this, we
2017    /// should not introduce more virtual records than 2^index, as that is the
2018    /// amount of excess fuel we have budgeted for completing merges.
2019    fn roll_up(&mut self, index: usize, log: &mut SpineLog<'_, T>) {
2020        // Ensure entries sufficient for `index`.
2021        while self.merging.len() <= index {
2022            self.merging.push(MergeState::default());
2023        }
2024
2025        // We only need to roll up if there are non-vacant layers.
2026        if self.merging[..index].iter().any(|m| !m.is_vacant()) {
2027            // Collect and merge all batches at layers up to but not including
2028            // `index`.
2029            let mut merged = None;
2030            for i in 0..index {
2031                if let Some(merged) = merged.take() {
2032                    self.insert_at(merged, i);
2033                }
2034                merged = self.complete_at(i, log);
2035            }
2036
2037            // The merged results should be introduced at level `index`, which
2038            // should be ready to absorb them (possibly creating a new merge at
2039            // the time).
2040            if let Some(merged) = merged {
2041                self.insert_at(merged, index);
2042            }
2043
2044            // If the insertion results in a merge, we should complete it to
2045            // ensure the upcoming insertion at `index` does not panic.
2046            if self.merging[index].is_full() {
2047                let merged = self.complete_at(index, log).expect("double batch");
2048                self.insert_at(merged, index + 1);
2049            }
2050        }
2051    }
2052
2053    /// Applies an amount of fuel to merges in progress.
2054    ///
2055    /// The supplied `fuel` is for each in progress merge, and if we want to
2056    /// spend the fuel non-uniformly (e.g. prioritizing merges at low layers) we
2057    /// could do so in order to maintain fewer batches on average (at the risk
2058    /// of completing merges of large batches later, but tbh probably not much
2059    /// later).
2060    pub fn apply_fuel(&mut self, fuel: &isize, log: &mut SpineLog<'_, T>) {
2061        // For the moment our strategy is to apply fuel independently to each
2062        // merge in progress, rather than prioritizing small merges. This sounds
2063        // like a great idea, but we need better accounting in place to ensure
2064        // that merges that borrow against later layers but then complete still
2065        // "acquire" fuel to pay back their debts.
2066        for index in 0..self.merging.len() {
2067            // Give each level independent fuel, for now.
2068            let mut fuel = *fuel;
2069            // Pass along various logging stuffs, in case we need to report
2070            // success.
2071            self.merging[index].work(&mut fuel);
2072            // `fuel` could have a deficit at this point, meaning we over-spent
2073            // when we took a merge step. We could ignore this, or maintain the
2074            // deficit and account future fuel against it before spending again.
2075            // It isn't clear why that would be especially helpful to do; we
2076            // might want to avoid overspends at multiple layers in the same
2077            // invocation (to limit latencies), but there is probably a rich
2078            // policy space here.
2079
2080            // If a merge completes, we can immediately merge it in to the next
2081            // level, which is "guaranteed" to be complete at this point, by our
2082            // fueling discipline.
2083            if self.merging[index].is_complete() {
2084                let complete = self.complete_at(index, log).expect("complete batch");
2085                self.insert_at(complete, index + 1);
2086            }
2087        }
2088    }
2089
2090    /// Inserts a batch at a specific location.
2091    ///
2092    /// This is a non-public internal method that can panic if we try and insert
2093    /// into a layer which already contains two batches (and is still in the
2094    /// process of merging).
2095    fn insert_at(&mut self, batch: SpineBatch<T>, index: usize) {
2096        // Ensure the spine is large enough.
2097        while self.merging.len() <= index {
2098            self.merging.push(MergeState::default());
2099        }
2100
2101        // Insert the batch at the location.
2102        let merging = &mut self.merging[index];
2103        merging.push_batch(batch);
2104        if merging.batches.is_full() {
2105            let compaction_frontier = Some(self.since.borrow());
2106            merging.merge = SpineBatch::begin_merge(&merging.batches[..], compaction_frontier)
2107        }
2108    }
2109
2110    /// Completes and extracts what ever is at layer `index`, leaving this layer vacant.
2111    fn complete_at(&mut self, index: usize, log: &mut SpineLog<'_, T>) -> Option<SpineBatch<T>> {
2112        self.merging[index].complete(log)
2113    }
2114
2115    /// Attempts to draw down large layers to size appropriate layers.
2116    fn tidy_layers(&mut self) {
2117        // If the largest layer is complete (not merging), we can attempt to
2118        // draw it down to the next layer. This is permitted if we can maintain
2119        // our invariant that below each merge there are at most half the
2120        // records that would be required to invade the merge.
2121        if !self.merging.is_empty() {
2122            let mut length = self.merging.len();
2123            if self.merging[length - 1].is_single() {
2124                // To move a batch down, we require that it contain few enough
2125                // records that the lower level is appropriate, and that moving
2126                // the batch would not create a merge violating our invariant.
2127                let appropriate_level = usize::cast_from(
2128                    self.merging[length - 1]
2129                        .len()
2130                        .next_power_of_two()
2131                        .trailing_zeros(),
2132                );
2133
2134                // Continue only as far as is appropriate
2135                while appropriate_level < length - 1 {
2136                    let current = &mut self.merging[length - 2];
2137                    if current.is_vacant() {
2138                        // Vacant batches can be absorbed.
2139                        self.merging.remove(length - 2);
2140                        length = self.merging.len();
2141                    } else {
2142                        if !current.is_full() {
2143                            // Single batches may initiate a merge, if sizes are
2144                            // within bounds, but terminate the loop either way.
2145
2146                            // Determine the number of records that might lead
2147                            // to a merge. Importantly, this is not the number
2148                            // of actual records, but the sum of upper bounds
2149                            // based on indices.
2150                            let mut smaller = 0;
2151                            for (index, batch) in self.merging[..(length - 2)].iter().enumerate() {
2152                                smaller += batch.batches.len() << index;
2153                            }
2154
2155                            if smaller <= (1 << length) / 8 {
2156                                // Remove the batch under consideration (shifting the deeper batches up a level),
2157                                // then merge in the single batch at the current level.
2158                                let state = self.merging.remove(length - 2);
2159                                assert_eq!(state.batches.len(), 1);
2160                                for batch in state.batches {
2161                                    self.insert_at(batch, length - 2);
2162                                }
2163                            }
2164                        }
2165                        break;
2166                    }
2167                }
2168            }
2169        }
2170    }
2171
2172    /// Checks invariants:
2173    /// - The lowers and uppers of all batches "line up".
2174    /// - The lower of the "minimum" batch is `antichain[T::minimum]`.
2175    /// - The upper of the "maximum" batch is `== self.upper`.
2176    /// - The since of each batch is `less_equal self.since`.
2177    /// - The `SpineIds` all "line up" and cover from `0` to `self.next_id`.
2178    /// - TODO: Verify fuel and level invariants.
2179    fn validate(&self) -> Result<(), String> {
2180        let mut id = SpineId(0, 0);
2181        let mut frontier = Antichain::from_elem(T::minimum());
2182        for x in self.merging.iter().rev() {
2183            if x.is_full() != x.merge.is_some() {
2184                return Err(format!(
2185                    "all (and only) full batches should have fueling merges (full={}, merge={:?})",
2186                    x.is_full(),
2187                    x.merge,
2188                ));
2189            }
2190
2191            if let Some(m) = &x.merge {
2192                if !x.is_full() {
2193                    return Err(format!(
2194                        "merge should only exist for full batches (len={:?}, merge={:?})",
2195                        x.batches.len(),
2196                        m.id,
2197                    ));
2198                }
2199                if x.id() != Some(m.id) {
2200                    return Err(format!(
2201                        "merge id should match the range of the batch ids (batch={:?}, merge={:?})",
2202                        x.id(),
2203                        m.id,
2204                    ));
2205                }
2206            }
2207
2208            // TODO: Anything we can validate about x.merge? It'd
2209            // be nice to assert that it's bigger than the len of the
2210            // two batches, but apply_merge_res might swap those lengths
2211            // out from under us.
2212            for batch in &x.batches {
2213                if batch.id().0 != id.1 {
2214                    return Err(format!(
2215                        "batch id {:?} does not match the previous id {:?}: {:?}",
2216                        batch.id(),
2217                        id,
2218                        self
2219                    ));
2220                }
2221                id = batch.id();
2222                if batch.desc().lower() != &frontier {
2223                    return Err(format!(
2224                        "batch lower {:?} does not match the previous upper {:?}: {:?}",
2225                        batch.desc().lower(),
2226                        frontier,
2227                        self
2228                    ));
2229                }
2230                frontier.clone_from(batch.desc().upper());
2231                if !PartialOrder::less_equal(batch.desc().since(), &self.since) {
2232                    return Err(format!(
2233                        "since of batch {:?} past the spine since {:?}: {:?}",
2234                        batch.desc().since(),
2235                        self.since,
2236                        self
2237                    ));
2238                }
2239            }
2240        }
2241        if self.next_id != id.1 {
2242            return Err(format!(
2243                "spine next_id {:?} does not match the last batch's id {:?}: {:?}",
2244                self.next_id, id, self
2245            ));
2246        }
2247        if self.upper != frontier {
2248            return Err(format!(
2249                "spine upper {:?} does not match the last batch's upper {:?}: {:?}",
2250                self.upper, frontier, self
2251            ));
2252        }
2253        Ok(())
2254    }
2255}
2256
2257/// Describes the state of a layer.
2258///
2259/// A layer can be empty, contain a single batch, or contain a pair of batches
2260/// that are in the process of merging into a batch for the next layer.
2261#[derive(Debug, Clone)]
2262struct MergeState<T> {
2263    batches: ArrayVec<SpineBatch<T>, BATCHES_PER_LEVEL>,
2264    merge: Option<IdFuelingMerge<T>>,
2265}
2266
2267impl<T> Default for MergeState<T> {
2268    fn default() -> Self {
2269        Self {
2270            batches: ArrayVec::new(),
2271            merge: None,
2272        }
2273    }
2274}
2275
2276impl<T: Timestamp + Lattice> MergeState<T> {
2277    /// An id that covers all the batches in the given merge state, assuming there are any.
2278    fn id(&self) -> Option<SpineId> {
2279        if let (Some(first), Some(last)) = (self.batches.first(), self.batches.last()) {
2280            Some(SpineId(first.id().0, last.id().1))
2281        } else {
2282            None
2283        }
2284    }
2285
2286    /// A new single-batch merge state.
2287    fn single(batch: SpineBatch<T>) -> Self {
2288        let mut state = Self::default();
2289        state.push_batch(batch);
2290        state
2291    }
2292
2293    /// Push a new batch at this level, checking invariants.
2294    fn push_batch(&mut self, batch: SpineBatch<T>) {
2295        self.try_push_batch(batch)
2296            .unwrap_or_else(|err| panic!("invalid batch push: {err}"));
2297    }
2298
2299    /// Fallible version of [Self::push_batch], for [Trace::unflatten], where
2300    /// the batches were decoded from an untrusted blob and a violated
2301    /// invariant must be a decode error rather than a panic.
2302    fn try_push_batch(&mut self, batch: SpineBatch<T>) -> Result<(), String> {
2303        if let Some(last) = self.batches.last() {
2304            if last.id().1 != batch.id().0 {
2305                return Err(format!(
2306                    "batch id {:?} does not chain with the previous id {:?}",
2307                    batch.id(),
2308                    last.id()
2309                ));
2310            }
2311            if last.upper() != batch.lower() {
2312                return Err(format!(
2313                    "batch lower {:?} does not match the previous upper {:?}",
2314                    batch.lower(),
2315                    last.upper()
2316                ));
2317            }
2318        }
2319        if self.merge.is_some() {
2320            return Err(format!(
2321                "attempted to insert batch into incomplete merge! (batch={:?}, batch_count={})",
2322                batch.id,
2323                self.batches.len(),
2324            ));
2325        }
2326        if self.batches.try_push(batch).is_err() {
2327            return Err("attempted to insert batch into full layer!".to_string());
2328        }
2329        Ok(())
2330    }
2331
2332    /// The number of actual updates contained in the level.
2333    fn len(&self) -> usize {
2334        self.batches.iter().map(SpineBatch::len).sum()
2335    }
2336
2337    /// True if this merge state contains no updates.
2338    fn is_empty(&self) -> bool {
2339        self.batches.iter().all(SpineBatch::is_empty)
2340    }
2341
2342    /// True if this level contains no batches.
2343    fn is_vacant(&self) -> bool {
2344        self.batches.is_empty()
2345    }
2346
2347    /// True only for a single-batch state.
2348    fn is_single(&self) -> bool {
2349        self.batches.len() == 1
2350    }
2351
2352    /// True if this merge cannot hold any more batches.
2353    /// (i.e. for a binary merge tree, true if this layer holds two batches.)
2354    fn is_full(&self) -> bool {
2355        self.batches.is_full()
2356    }
2357
2358    /// Immediately complete any merge.
2359    ///
2360    /// The result is either a batch, if there is a non-trivial batch to return
2361    /// or `None` if there is no meaningful batch to return.
2362    ///
2363    /// There is the additional option of input batches.
2364    fn complete(&mut self, log: &mut SpineLog<'_, T>) -> Option<SpineBatch<T>> {
2365        let mut this = mem::take(self);
2366        if this.batches.len() <= 1 {
2367            this.batches.pop()
2368        } else {
2369            // Merge the remaining batches, regardless of whether we have a fully fueled merge.
2370            let id_merge = this
2371                .merge
2372                .or_else(|| SpineBatch::begin_merge(&self.batches[..], None))?;
2373            id_merge.merge.done(this.batches, log)
2374        }
2375    }
2376
2377    /// True iff the layer is a complete merge, ready for extraction.
2378    fn is_complete(&self) -> bool {
2379        match &self.merge {
2380            Some(IdFuelingMerge { merge, .. }) => merge.remaining_work == 0,
2381            None => false,
2382        }
2383    }
2384
2385    /// Performs a bounded amount of work towards a merge.
2386    fn work(&mut self, fuel: &mut isize) {
2387        // We only perform work for merges in progress.
2388        if let Some(IdFuelingMerge { merge, .. }) = &mut self.merge {
2389            merge.work(&self.batches[..], fuel)
2390        }
2391    }
2392}
2393
2394#[cfg(test)]
2395pub mod datadriven {
2396    use mz_ore::fmt::FormatBuffer;
2397
2398    use crate::internal::datadriven::DirectiveArgs;
2399
2400    use super::*;
2401
2402    /// Shared state for a single [crate::internal::trace] [datadriven::TestFile].
2403    #[derive(Debug, Default)]
2404    pub struct TraceState {
2405        pub trace: Trace<u64>,
2406        pub merge_reqs: Vec<FueledMergeReq<u64>>,
2407    }
2408
2409    pub fn since_upper(
2410        datadriven: &TraceState,
2411        _args: DirectiveArgs,
2412    ) -> Result<String, anyhow::Error> {
2413        Ok(format!(
2414            "{:?}{:?}\n",
2415            datadriven.trace.since().elements(),
2416            datadriven.trace.upper().elements()
2417        ))
2418    }
2419
2420    pub fn batches(datadriven: &TraceState, _args: DirectiveArgs) -> Result<String, anyhow::Error> {
2421        let mut s = String::new();
2422        for b in datadriven.trace.spine.spine_batches() {
2423            s.push_str(b.describe(true).as_str());
2424            s.push('\n');
2425        }
2426        Ok(s)
2427    }
2428
2429    pub fn insert(
2430        datadriven: &mut TraceState,
2431        args: DirectiveArgs,
2432    ) -> Result<String, anyhow::Error> {
2433        for x in args
2434            .input
2435            .trim()
2436            .split('\n')
2437            .map(DirectiveArgs::parse_hollow_batch)
2438        {
2439            datadriven
2440                .merge_reqs
2441                .append(&mut datadriven.trace.push_batch(x));
2442        }
2443        Ok("ok\n".to_owned())
2444    }
2445
2446    pub fn downgrade_since(
2447        datadriven: &mut TraceState,
2448        args: DirectiveArgs,
2449    ) -> Result<String, anyhow::Error> {
2450        let since = args.expect("since");
2451        datadriven
2452            .trace
2453            .downgrade_since(&Antichain::from_elem(since));
2454        Ok("ok\n".to_owned())
2455    }
2456
2457    pub fn take_merge_req(
2458        datadriven: &mut TraceState,
2459        _args: DirectiveArgs,
2460    ) -> Result<String, anyhow::Error> {
2461        let mut s = String::new();
2462        for merge_req in std::mem::take(&mut datadriven.merge_reqs) {
2463            write!(
2464                s,
2465                "{:?}{:?}{:?} {}\n",
2466                merge_req.desc.lower().elements(),
2467                merge_req.desc.upper().elements(),
2468                merge_req.desc.since().elements(),
2469                merge_req
2470                    .inputs
2471                    .iter()
2472                    .flat_map(|x| x.batch.parts.iter())
2473                    .map(|x| x.printable_name())
2474                    .collect::<Vec<_>>()
2475                    .join(" ")
2476            );
2477        }
2478        Ok(s)
2479    }
2480
2481    pub fn apply_merge_res(
2482        datadriven: &mut TraceState,
2483        args: DirectiveArgs,
2484    ) -> Result<String, anyhow::Error> {
2485        let res = FueledMergeRes {
2486            output: DirectiveArgs::parse_hollow_batch(args.input),
2487            input: CompactionInput::Legacy,
2488            new_active_compaction: None,
2489        };
2490        match datadriven.trace.apply_merge_res_unchecked(&res) {
2491            ApplyMergeResult::AppliedExact => Ok("applied exact\n".into()),
2492            ApplyMergeResult::AppliedSubset => Ok("applied subset\n".into()),
2493            ApplyMergeResult::NotAppliedNoMatch => Ok("no-op\n".into()),
2494            ApplyMergeResult::NotAppliedInvalidSince => Ok("no-op invalid since\n".into()),
2495            ApplyMergeResult::NotAppliedTooManyUpdates => Ok("no-op too many updates\n".into()),
2496        }
2497    }
2498}
2499
2500#[cfg(test)]
2501pub(crate) mod tests {
2502    use std::ops::Range;
2503
2504    use proptest::prelude::*;
2505    use semver::Version;
2506
2507    use crate::internal::state::tests::{any_hollow_batch, any_hollow_batch_with_exact_runs};
2508
2509    use super::*;
2510
2511    pub fn any_trace<T: Arbitrary + Timestamp + Lattice>(
2512        num_batches: Range<usize>,
2513    ) -> impl Strategy<Value = Trace<T>> {
2514        Strategy::prop_map(
2515            (
2516                any::<Option<T>>(),
2517                proptest::collection::vec(any_hollow_batch::<T>(), num_batches),
2518                any::<bool>(),
2519                any::<u64>(),
2520            ),
2521            |(since, mut batches, roundtrip_structure, timeout_ms)| {
2522                let mut trace = Trace::<T>::default();
2523                trace.downgrade_since(&since.map_or_else(Antichain::new, Antichain::from_elem));
2524
2525                // Fix up the arbitrary HollowBatches so the lowers and uppers
2526                // align.
2527                batches.sort_by(|x, y| x.desc.upper().elements().cmp(y.desc.upper().elements()));
2528                let mut lower = Antichain::from_elem(T::minimum());
2529                for mut batch in batches {
2530                    // Overall trace since has to be past each batch's since.
2531                    if PartialOrder::less_than(trace.since(), batch.desc.since()) {
2532                        trace.downgrade_since(batch.desc.since());
2533                    }
2534                    batch.desc = Description::new(
2535                        lower.clone(),
2536                        batch.desc.upper().clone(),
2537                        batch.desc.since().clone(),
2538                    );
2539                    lower.clone_from(batch.desc.upper());
2540                    let _merge_req = trace.push_batch(batch);
2541                }
2542                let reqs: Vec<_> = trace
2543                    .fueled_merge_reqs_before_ms(timeout_ms, None)
2544                    .collect();
2545                for req in reqs {
2546                    trace.claim_compaction(req.id, ActiveCompaction { start_ms: 0 })
2547                }
2548                trace.roundtrip_structure = roundtrip_structure;
2549                trace
2550            },
2551        )
2552    }
2553
2554    #[mz_ore::test]
2555    #[cfg_attr(miri, ignore)] // proptest is too heavy for miri!
2556    fn test_roundtrips() {
2557        fn check(trace: Trace<i64>) {
2558            trace.validate().unwrap();
2559            let flat = trace.flatten();
2560            let unflat = Trace::unflatten(flat).unwrap();
2561            assert_eq!(trace, unflat);
2562        }
2563
2564        proptest!(|(trace in any_trace::<i64>(1..10))| { check(trace) })
2565    }
2566
2567    #[mz_ore::test]
2568    fn fueled_merge_reqs() {
2569        let mut trace: Trace<u64> = Trace::default();
2570        let fueled_reqs = trace.push_batch(crate::internal::state::tests::hollow(
2571            0,
2572            10,
2573            &["n0011500/p3122e2a1-a0c7-429f-87aa-1019bf4f5f86"],
2574            1000,
2575        ));
2576
2577        assert!(fueled_reqs.is_empty());
2578        assert_eq!(
2579            trace.fueled_merge_reqs_before_ms(u64::MAX, None).count(),
2580            0,
2581            "no merge reqs when not filtering by version"
2582        );
2583        assert_eq!(
2584            trace
2585                .fueled_merge_reqs_before_ms(
2586                    u64::MAX,
2587                    Some(WriterKey::for_version(&Version::new(0, 50, 0)))
2588                )
2589                .count(),
2590            0,
2591            "zero batches are older than a past version"
2592        );
2593        assert_eq!(
2594            trace
2595                .fueled_merge_reqs_before_ms(
2596                    u64::MAX,
2597                    Some(WriterKey::for_version(&Version::new(99, 99, 0)))
2598                )
2599                .count(),
2600            1,
2601            "one batch is older than a future version"
2602        );
2603    }
2604
2605    #[mz_ore::test]
2606    fn remove_redundant_merge_reqs() {
2607        fn req(lower: u64, upper: u64) -> FueledMergeReq<u64> {
2608            FueledMergeReq {
2609                id: SpineId(usize::cast_from(lower), usize::cast_from(upper)),
2610                desc: Description::new(
2611                    Antichain::from_elem(lower),
2612                    Antichain::from_elem(upper),
2613                    Antichain::new(),
2614                ),
2615                inputs: vec![],
2616            }
2617        }
2618
2619        // Empty
2620        assert_eq!(Trace::<u64>::remove_redundant_merge_reqs(vec![]), vec![]);
2621
2622        // Single
2623        assert_eq!(
2624            Trace::remove_redundant_merge_reqs(vec![req(0, 1)]),
2625            vec![req(0, 1)]
2626        );
2627
2628        // Duplicate
2629        assert_eq!(
2630            Trace::remove_redundant_merge_reqs(vec![req(0, 1), req(0, 1)]),
2631            vec![req(0, 1)]
2632        );
2633
2634        // Nothing covered
2635        assert_eq!(
2636            Trace::remove_redundant_merge_reqs(vec![req(0, 1), req(1, 2)]),
2637            vec![req(1, 2), req(0, 1)]
2638        );
2639
2640        // Covered
2641        assert_eq!(
2642            Trace::remove_redundant_merge_reqs(vec![req(1, 2), req(0, 3)]),
2643            vec![req(0, 3)]
2644        );
2645
2646        // Covered, lower equal
2647        assert_eq!(
2648            Trace::remove_redundant_merge_reqs(vec![req(0, 2), req(0, 3)]),
2649            vec![req(0, 3)]
2650        );
2651
2652        // Covered, upper equal
2653        assert_eq!(
2654            Trace::remove_redundant_merge_reqs(vec![req(1, 3), req(0, 3)]),
2655            vec![req(0, 3)]
2656        );
2657
2658        // Covered, unexpected order (doesn't happen in practice)
2659        assert_eq!(
2660            Trace::remove_redundant_merge_reqs(vec![req(0, 3), req(1, 2)]),
2661            vec![req(0, 3)]
2662        );
2663
2664        // Partially overlapping
2665        assert_eq!(
2666            Trace::remove_redundant_merge_reqs(vec![req(0, 2), req(1, 3)]),
2667            vec![req(1, 3), req(0, 2)]
2668        );
2669
2670        // Partially overlapping, the other order
2671        assert_eq!(
2672            Trace::remove_redundant_merge_reqs(vec![req(1, 3), req(0, 2)]),
2673            vec![req(0, 2), req(1, 3)]
2674        );
2675
2676        // Different sinces (doesn't happen in practice)
2677        let req015 = FueledMergeReq {
2678            id: SpineId(0, 1),
2679            desc: Description::new(
2680                Antichain::from_elem(0),
2681                Antichain::from_elem(1),
2682                Antichain::from_elem(5),
2683            ),
2684            inputs: vec![],
2685        };
2686        assert_eq!(
2687            Trace::remove_redundant_merge_reqs(vec![req(0, 1), req015.clone()]),
2688            vec![req015, req(0, 1)]
2689        );
2690    }
2691
2692    #[mz_ore::test]
2693    #[cfg_attr(miri, ignore)] // proptest is too heavy for miri!
2694    fn construct_batch_with_runs_replaced_test() {
2695        let batch_strategy = any_hollow_batch::<u64>();
2696        let to_replace_strategy = any_hollow_batch_with_exact_runs::<u64>(1);
2697
2698        let combined_strategy = (batch_strategy, to_replace_strategy)
2699            .prop_filter("non-empty batch", |(batch, _)| batch.run_meta.len() >= 1);
2700
2701        let final_strategy = combined_strategy.prop_flat_map(|(batch, to_replace)| {
2702            let batch_len = batch.run_meta.len();
2703            let batch_clone = batch.clone();
2704            let to_replace_clone = to_replace.clone();
2705
2706            proptest::collection::vec(any::<bool>(), batch_len)
2707                .prop_filter("at least one run selected", |mask| mask.iter().any(|&x| x))
2708                .prop_map(move |mask| {
2709                    let indices: Vec<usize> = mask
2710                        .iter()
2711                        .enumerate()
2712                        .filter_map(|(i, &selected)| if selected { Some(i) } else { None })
2713                        .collect();
2714                    (batch_clone.clone(), to_replace_clone.clone(), indices)
2715                })
2716        });
2717
2718        proptest!(|(
2719            (batch, to_replace, runs) in final_strategy
2720        )| {
2721            let original_run_ids: Vec<_> = batch.run_meta.iter().map(|x|
2722                x.id.unwrap().clone()
2723            ).collect();
2724
2725            let run_ids = runs.iter().map(|&i| original_run_ids[i].clone()).collect::<Vec<_>>();
2726
2727            let new_batch = SpineBatch::construct_batch_with_runs_replaced(
2728                &batch,
2729                &run_ids,
2730                &to_replace,
2731            ).unwrap();
2732
2733            let expected_len = batch.run_meta.len() - runs.len()
2734                + to_replace.run_meta.len();
2735            prop_assert!(new_batch.run_meta.len() == expected_len);
2736        });
2737    }
2738
2739    #[mz_ore::test]
2740    fn test_perform_subset_replacement() {
2741        let batch1 = crate::internal::state::tests::hollow::<u64>(0, 10, &["a"], 10);
2742        let batch2 = crate::internal::state::tests::hollow::<u64>(10, 20, &["b"], 10);
2743        let batch3 = crate::internal::state::tests::hollow::<u64>(20, 30, &["c"], 10);
2744
2745        let id_batch1 = IdHollowBatch {
2746            id: SpineId(0, 1),
2747            batch: Arc::new(batch1.clone()),
2748        };
2749        let id_batch2 = IdHollowBatch {
2750            id: SpineId(1, 2),
2751            batch: Arc::new(batch2.clone()),
2752        };
2753        let id_batch3 = IdHollowBatch {
2754            id: SpineId(2, 3),
2755            batch: Arc::new(batch3.clone()),
2756        };
2757
2758        let spine_batch = SpineBatch {
2759            id: SpineId(0, 3),
2760            desc: Description::new(
2761                Antichain::from_elem(0),
2762                Antichain::from_elem(30),
2763                Antichain::from_elem(0),
2764            ),
2765            parts: vec![id_batch1, id_batch2, id_batch3],
2766            active_compaction: None,
2767            len: 30,
2768        };
2769
2770        let res_exact = crate::internal::state::tests::hollow::<u64>(0, 30, &["d"], 30);
2771        let mut sb_exact = spine_batch.clone();
2772        let result = sb_exact.perform_subset_replacement(&res_exact, SpineId(0, 3), 0..3, None);
2773        assert!(matches!(result, ApplyMergeResult::AppliedExact));
2774        assert_eq!(sb_exact.parts.len(), 1);
2775        assert_eq!(sb_exact.len(), 30);
2776
2777        let res_subset = crate::internal::state::tests::hollow::<u64>(0, 20, &["e"], 20);
2778        let mut sb_subset = spine_batch.clone();
2779        let result = sb_subset.perform_subset_replacement(&res_subset, SpineId(0, 2), 0..2, None);
2780        assert!(matches!(result, ApplyMergeResult::AppliedSubset));
2781        assert_eq!(sb_subset.parts.len(), 2); // One new part + one old part
2782        assert_eq!(sb_subset.len(), 30);
2783
2784        let res_too_big = crate::internal::state::tests::hollow::<u64>(0, 30, &["f"], 31);
2785        let mut sb_too_big = spine_batch.clone();
2786        let result = sb_too_big.perform_subset_replacement(&res_too_big, SpineId(0, 3), 0..3, None);
2787        assert!(matches!(result, ApplyMergeResult::NotAppliedTooManyUpdates));
2788        assert_eq!(sb_too_big.parts.len(), 3);
2789        assert_eq!(sb_too_big.len(), 30);
2790    }
2791
2792    /// Inserting a batch whose `len` is large enough to saturate the fuel
2793    /// computation must not disturb an in-progress merge's accounting.
2794    ///
2795    /// `introduce_batch` derives its fuel from `8 << batch_index`, where
2796    /// `batch_index` is `len.next_power_of_two().trailing_zeros()`. A `len` near
2797    /// `2^60` pushes that past what an `isize` holds, and `Trace::unflatten`
2798    /// accepts such a `len` from an untrusted blob: its `MAX_TOTAL_LEN` guard
2799    /// only caps the total at `usize::MAX >> 3`.
2800    #[mz_ore::test]
2801    fn spine_fuel_isize_overflow() {
2802        let mut trace = Trace::<u64>::default();
2803        let mut push = |lower, upper, key: &str, len| {
2804            trace.push_batch_no_merge_reqs(crate::internal::state::tests::hollow::<u64>(
2805                lower,
2806                upper,
2807                &[key],
2808                len,
2809            ));
2810        };
2811        // Two same-size batches fill a level and begin a merge whose
2812        // `remaining_work` (the sum of their lens) far exceeds the `8 << 0` fuel
2813        // that a subsequent len-1 batch delivers, so the merge is still in
2814        // progress when the saturating batch arrives.
2815        push(0, 1, "a", 100);
2816        push(1, 2, "b", 100);
2817        push(2, 3, "c", 1);
2818        push(3, 4, "d", 1 << 60);
2819    }
2820}