Persist: Add new API for run based spine replacement

src/persist-client/src/cli/admin.rs

@@ -40,7 +40,7 @@ use crate::internal::compact::{CompactConfig, CompactReq, Compactor};
 use crate::internal::encoding::Schemas;
 use crate::internal::gc::{GarbageCollector, GcReq};
 use crate::internal::machine::Machine;
-use crate::internal::trace::FueledMergeRes;
+use crate::internal::trace::{CompactionInput, FueledMergeRes};
 use crate::rpc::{NoopPubSubSender, PubSubSender};
 use crate::write::{WriteHandle, WriterId};
 use crate::{
@@ -509,7 +509,11 @@ where
                 start.elapsed(),
             );
             let (apply_res, maintenance) = machine
-                .merge_res(&FueledMergeRes { output: res.output })
+                .merge_res(&FueledMergeRes {
+                    output: res.output,
+                    input: CompactionInput::Legacy,
+                    new_active_compaction: None,
+                })
                 .await;
             if !maintenance.is_empty() {
                 info!("ignoring non-empty requested maintenance: {maintenance:?}")

src/persist-client/src/internal/compact.rs

@@ -49,7 +49,7 @@ use crate::internal::metrics::ShardMetrics;
 use crate::internal::state::{
     ENABLE_INCREMENTAL_COMPACTION, HollowBatch, RunMeta, RunOrder, RunPart,
 };
-use crate::internal::trace::{ApplyMergeResult, FueledMergeRes};
+use crate::internal::trace::{ApplyMergeResult, CompactionInput, FueledMergeRes};
 use crate::iter::{Consolidator, StructuredSort};
 use crate::{Metrics, PersistConfig, ShardId};
 
@@ -396,7 +396,11 @@ where
                     );
                     let res = Self::compact_all(stream, req.clone()).await?;
                     let maintenance = Self::apply(
-                        FueledMergeRes { output: res.output },
+                        FueledMergeRes {
+                            output: res.output,
+                            input: CompactionInput::Legacy,
+                            new_active_compaction: None,
+                        },
                         &metrics_clone,
                         &machine_clone,
                     )

src/persist-client/src/internal/machine.rs

@@ -1438,6 +1438,7 @@ pub mod datadriven {
     use crate::internal::paths::{BlobKey, BlobKeyPrefix, PartialBlobKey};
     use crate::internal::state::{BatchPart, RunOrder, RunPart};
     use crate::internal::state_versions::EncodedRollup;
+    use crate::internal::trace::CompactionInput;
     use crate::read::{Listen, ListenEvent, READER_LEASE_DURATION};
     use crate::rpc::NoopPubSubSender;
     use crate::tests::new_test_client;
@@ -2464,7 +2465,11 @@ pub mod datadriven {
             .clone();
         let (merge_res, maintenance) = datadriven
             .machine
-            .merge_res(&FueledMergeRes { output: batch })
+            .merge_res(&FueledMergeRes {
+                output: batch,
+                input: CompactionInput::Legacy,
+                new_active_compaction: None,
+            })
             .await;
         datadriven.routine.push(maintenance);
         Ok(format!(

src/persist-client/src/internal/state.rs

@@ -1053,6 +1053,29 @@ impl<T> HollowBatch<T> {
         }
     }
 
+    #[cfg(test)]
+    pub(crate) fn new_run_for_test(
+        desc: Description<T>,
+        parts: Vec<RunPart<T>>,
+        len: usize,
+        run_id: RunId,
+    ) -> Self {
+        let run_meta = if parts.is_empty() {
+            vec![]
+        } else {
+            let mut meta = RunMeta::default();
+            meta.id = Some(run_id);
+            vec![meta]
+        };
+        Self {
+            desc,
+            len,
+            parts,
+            run_splits: vec![],
+            run_meta,
+        }
+    }
+
     /// An empty hollow batch, representing no updates over the given desc.
     pub(crate) fn empty(desc: Description<T>) -> Self {
         Self {
@@ -1810,7 +1833,9 @@ where
             return Break(NoOpStateTransition(ApplyMergeResult::NotAppliedNoMatch));
         }
 
-        let apply_merge_result = self.trace.apply_merge_res_checked::<D>(res, metrics);
+        let apply_merge_result = self
+            .trace
+            .apply_merge_res_checked_classic::<D>(res, metrics);
         Continue(apply_merge_result)
     }
 
@@ -2162,10 +2187,7 @@ where
             // We have a nonempty batch: replace it with an empty batch and return.
             // This should not produce an excessively large diff: if it did, we wouldn't have been
             // able to append that batch in the first place.
-            let fake_merge = FueledMergeRes {
-                output: HollowBatch::empty(desc),
-            };
-            let result = self.trace.apply_tombstone_merge(&fake_merge);
+            let result = self.trace.apply_tombstone_merge(&desc);
             assert!(
                 result.matched(),
                 "merge with a matching desc should always match"
@@ -2833,34 +2855,92 @@ pub(crate) mod tests {
         }
     }
 
+    pub fn any_hollow_batch_with_exact_runs<T: Arbitrary + Timestamp>(
+        num_runs: usize,
+    ) -> impl Strategy<Value = HollowBatch<T>> {
+        (
+            any::<T>(),
+            any::<T>(),
+            any::<T>(),
+            proptest::collection::vec(any_run_part::<T>(), num_runs + 1..20),
+            any::<usize>(),
+        )
+            .prop_map(move |(t0, t1, since, parts, len)| {
+                let (lower, upper) = if t0 <= t1 {
+                    (Antichain::from_elem(t0), Antichain::from_elem(t1))
+                } else {
+                    (Antichain::from_elem(t1), Antichain::from_elem(t0))
+                };
+                let since = Antichain::from_elem(since);
+
+                let run_splits = (1..num_runs)
+                    .map(|i| i * parts.len() / num_runs)
+                    .collect::<Vec<_>>();
+
+                let run_meta = (0..num_runs)
+                    .map(|_| {
+                        let mut meta = RunMeta::default();
+                        meta.id = Some(RunId::new());
+                        meta
+                    })
+                    .collect::<Vec<_>>();
+
+                HollowBatch::new(
+                    Description::new(lower, upper, since),
+                    parts,
+                    len % 10,
+                    run_meta,
+                    run_splits,
+                )
+            })
+    }
+
     pub fn any_hollow_batch<T: Arbitrary + Timestamp>() -> impl Strategy<Value = HollowBatch<T>> {
         Strategy::prop_map(
             (
                 any::<T>(),
                 any::<T>(),
                 any::<T>(),
-                proptest::collection::vec(any_run_part::<T>(), 0..3),
+                proptest::collection::vec(any_run_part::<T>(), 0..20),
                 any::<usize>(),
-                any::<bool>(),
+                0..=10usize,
+                proptest::collection::vec(any::<RunId>(), 10),
             ),
-            |(t0, t1, since, parts, len, runs)| {
+            |(t0, t1, since, parts, len, num_runs, run_ids)| {
                 let (lower, upper) = if t0 <= t1 {
                     (Antichain::from_elem(t0), Antichain::from_elem(t1))
                 } else {
                     (Antichain::from_elem(t1), Antichain::from_elem(t0))
                 };
                 let since = Antichain::from_elem(since);
-                if runs && parts.len() > 2 {
-                    let split_at = parts.len() / 2;
+                if num_runs > 0 && parts.len() > 2 && num_runs < parts.len() {
+                    let run_splits = (1..num_runs)
+                        .map(|i| i * parts.len() / num_runs)
+                        .collect::<Vec<_>>();
+
+                    let run_meta = (0..num_runs)
+                        .enumerate()
+                        .map(|(i, _)| {
+                            let mut meta = RunMeta::default();
+                            meta.id = Some(run_ids[i]);
+                            meta
+                        })
+                        .collect::<Vec<_>>();
+
                     HollowBatch::new(
                         Description::new(lower, upper, since),
                         parts,
                         len % 10,
-                        vec![RunMeta::default(), RunMeta::default()],
-                        vec![split_at],
+                        run_meta,
+                        run_splits,
                     )
                 } else {
-                    HollowBatch::new_run(Description::new(lower, upper, since), parts, len % 10)
+                    HollowBatch::new_run_for_test(
+                        Description::new(lower, upper, since),
+                        parts,
+                        len % 10,
+                        run_ids[0],
+                    )
                 }
             },
         )

src/persist-client/src/internal/state_diff.rs

@@ -32,6 +32,7 @@ use crate::internal::state::{
     LeasedReaderState, ProtoStateField, ProtoStateFieldDiffType, ProtoStateFieldDiffs, RunPart,
     State, StateCollections, WriterState,
 };
+use crate::internal::trace::CompactionInput;
 use crate::internal::trace::{FueledMergeRes, SpineId, ThinMerge, ThinSpineBatch, Trace};
 use crate::read::LeasedReaderId;
 use crate::write::WriterId;
@@ -877,7 +878,11 @@ fn apply_diffs_spine<T: Timestamp + Lattice + Codec64>(
 
     // Fast-path: compaction
     if let Some((_inputs, output)) = sniff_compaction(&diffs) {
-        let res = FueledMergeRes { output };
+        let res = FueledMergeRes {
+            output,
+            input: CompactionInput::Legacy,
+            new_active_compaction: None,
+        };
         // We can't predict how spine will arrange the batches when it's
         // hydrated. This means that something that is maintaining a Spine
         // starting at some seqno may not exactly match something else
@@ -1444,7 +1449,11 @@ mod tests {
                             leader
                                 .collections
                                 .trace
-                                .apply_merge_res_unchecked(&FueledMergeRes { output });
+                                .apply_merge_res_unchecked(&FueledMergeRes {
+                                    output,
+                                    input: CompactionInput::Legacy,
+                                    new_active_compaction: None,
+                                });
                         }
                     }
                 }