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