Plumb incremental compaction all the way through

Author: DAlperin

src/persist-client/src/batch.rs

@@ -17,6 +17,21 @@ use std::mem;
 use std::sync::Arc;
 use std::time::Instant;
 
+use crate::async_runtime::IsolatedRuntime;
+use crate::cfg::{BATCH_BUILDER_MAX_OUTSTANDING_PARTS, MiB};
+use crate::error::InvalidUsage;
+use crate::internal::compact::{CompactConfig, Compactor};
+use crate::internal::encoding::{LazyInlineBatchPart, LazyPartStats, LazyProto, Schemas};
+use crate::internal::machine::{Machine, retry_external};
+use crate::internal::merge::{MergeTree, Pending};
+use crate::internal::metrics::{BatchWriteMetrics, Metrics, RetryMetrics, ShardMetrics};
+use crate::internal::paths::{PartId, PartialBatchKey, WriterKey};
+use crate::internal::state::{
+    BatchPart, HollowBatch, HollowBatchPart, HollowRun, HollowRunRef, ProtoInlineBatchPart,
+    RunMeta, RunOrder, RunPart,
+};
+use crate::stats::{STATS_BUDGET_BYTES, STATS_COLLECTION_ENABLED, untrimmable_columns};
+use crate::{PersistConfig, ShardId};
 use arrow::array::{Array, Int64Array};
 use bytes::Bytes;
 use differential_dataflow::difference::Semigroup;
@@ -46,21 +61,6 @@ use timely::PartialOrder;
 use timely::order::TotalOrder;
 use timely::progress::{Antichain, Timestamp};
 use tracing::{Instrument, debug_span, trace_span, warn};
-use crate::async_runtime::IsolatedRuntime;
-use crate::cfg::{BATCH_BUILDER_MAX_OUTSTANDING_PARTS, MiB};
-use crate::error::InvalidUsage;
-use crate::internal::compact::{CompactConfig, Compactor};
-use crate::internal::encoding::{LazyInlineBatchPart, LazyPartStats, LazyProto, Schemas};
-use crate::internal::machine::retry_external;
-use crate::internal::merge::{MergeTree, Pending};
-use crate::internal::metrics::{BatchWriteMetrics, Metrics, RetryMetrics, ShardMetrics};
-use crate::internal::paths::{PartId, PartialBatchKey, WriterKey};
-use crate::internal::state::{
-    BatchPart, HollowBatch, HollowBatchPart, HollowRun, HollowRunRef, ProtoInlineBatchPart,
-    RunMeta, RunOrder, RunPart,
-};
-use crate::stats::{STATS_BUDGET_BYTES, STATS_COLLECTION_ENABLED, untrimmable_columns};
-use crate::{PersistConfig, ShardId};
 
 include!(concat!(env!("OUT_DIR"), "/mz_persist_client.batch.rs"));
 
@@ -681,7 +681,9 @@ where
     ) -> Option<HollowBatch<T>> {
         let runs = self.parts.finish_completed_runs();
 
-        if runs.is_empty() { return None; }
+        if runs.is_empty() {
+            return None;
+        }
 
         let mut run_parts = vec![];
         let mut run_splits = vec![];
@@ -701,8 +703,6 @@ where
             });
             run_parts.extend(parts);
         }
-        //TODO(dov): should we fetch the last part and constrain the upper bound
-        // to whatever we have actually flushed?
         let desc = registered_desc;
 
         let batch = HollowBatch::new(desc, run_parts, self.num_updates, run_meta, run_splits);
@@ -877,8 +877,8 @@ impl<T: Timestamp + Codec64> BatchParts<T> {
                             shard_metrics,
                             isolated_runtime,
                             write_schemas,
+                            &None,
                             None,
-                            None
                         )
                         .await
                         .expect("successful compaction");
@@ -977,7 +977,6 @@ impl<T: Timestamp + Codec64> BatchParts<T> {
         }
     }
 
-
     pub(crate) async fn write<K: Codec, V: Codec, D: Codec64>(
         &mut self,
         write_schemas: &Schemas<K, V>,
@@ -1268,18 +1267,22 @@ impl<T: Timestamp + Codec64> BatchParts<T> {
 
     pub(crate) fn finish_completed_runs(&self) -> Vec<(RunOrder, Vec<RunPart<T>>)> {
         match &self.writing_runs {
-            WritingRuns::Ordered(order, tree) => tree.iter()
-                .take_while(|part| matches!(part, Pending::Finished(_))).map(|part| match part {
-                Pending::Finished(p) => (order.clone(), vec![p.clone()]),
-                _ => (order.clone(), vec![]),
-            }).collect(),
-            WritingRuns::Compacting(tree) => tree.iter()
+            WritingRuns::Ordered(order, tree) => tree
+                .iter()
+                .take_while(|part| matches!(part, Pending::Finished(_)))
+                .map(|part| match part {
+                    Pending::Finished(p) => (order.clone(), vec![p.clone()]),
+                    _ => (order.clone(), vec![]),
+                })
+                .collect(),
+            WritingRuns::Compacting(tree) => tree
+                .iter()
                 .take_while(|(_, run)| matches!(run, Pending::Finished(_)))
                 .map(|(order, run)| match run {
                     Pending::Finished(parts) => (order.clone(), parts.clone()),
                     _ => (order.clone(), vec![]),
                 })
-                .collect()
+                .collect(),
         }
     }
 

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

@@ -500,6 +500,7 @@ where
                 Arc::new(IsolatedRuntime::default()),
                 req.clone(),
                 schemas,
+                &machine,
             );
             pin_mut!(stream);
 

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

@@ -397,6 +397,7 @@ where
                             Arc::clone(&machine_clone.isolated_runtime),
                             req.clone(),
                             compaction_schema,
+                            &machine_clone,
                         );
 
                         let maintenance =
@@ -570,6 +571,7 @@ where
         isolated_runtime: Arc<IsolatedRuntime>,
         req: CompactReq<T>,
         write_schemas: Schemas<K, V>,
+        machine: &Machine<K, V, T, D>,
     ) -> impl Stream<Item = Result<FueledMergeRes<T>, anyhow::Error>> {
         async_stream::stream! {
             let _ = Self::validate_req(&req)?;
@@ -605,9 +607,6 @@ where
             let run_reserved_memory_bytes =
                 cfg.compaction_memory_bound_bytes - in_progress_part_reserved_memory_bytes;
 
-            //if there is an active compaction
-            //We should call compact_runs with the same
-
             // Flatten the input batches into a single list of runs
             let ordered_runs =
                 Self::order_runs(&req, cfg.batch.preferred_order, &*blob, &*metrics).await?;
@@ -616,6 +615,8 @@ where
             let chunked_runs =
                 Self::chunk_runs(&ordered_runs, &cfg, &*metrics, run_reserved_memory_bytes);
 
+            let (incremental_tx, mut incremental_rx) = mpsc::channel(0);
+
             let total_chunked_runs = chunked_runs.len();
             let mut applied = 0;
             for (runs, run_chunk_max_memory_usage) in chunked_runs {
@@ -643,11 +644,16 @@ where
                     Arc::clone(&shard_metrics),
                     Arc::clone(&isolated_runtime),
                     write_schemas.clone(),
-                    None,
-                    None
+                    &req.prev_batch,
+                    Some(incremental_tx.clone())
                 )
                 .await?;
 
+                while let Some(res) = incremental_rx.recv().await {
+                    let now = SYSTEM_TIME.clone();
+                    machine.checkpoint_compaction_progress(&res, now()).await;
+                }
+
                 let (parts, run_splits, run_meta, updates) =
                     (batch.parts, batch.run_splits, batch.run_meta, batch.len);
 
@@ -690,6 +696,7 @@ where
 
                 applied += 1;
             }
+            drop(incremental_tx);
         }
     }
 
@@ -827,6 +834,28 @@ where
         Ok(ordered_runs)
     }
 
+    fn combine_hollow_batch_with_previous(
+        previous_batch: &HollowBatch<T>,
+        batch: &HollowBatch<T>,
+    ) -> HollowBatch<T> {
+        // Simplifying assumption: you can't combine batches with different descriptions.
+        assert_eq!(previous_batch.desc, batch.desc);
+        let len = previous_batch.len + batch.len;
+        let mut parts = Vec::with_capacity(previous_batch.parts.len() + batch.parts.len());
+        parts.extend(previous_batch.parts.clone());
+        parts.extend(batch.parts.clone());
+        assert!(previous_batch.run_splits.is_empty());
+        assert!(batch.run_splits.is_empty());
+        assert_eq!(previous_batch.run_meta, batch.run_meta);
+        HollowBatch::new(
+            previous_batch.desc.clone(),
+            parts,
+            len,
+            previous_batch.run_meta.clone(),
+            previous_batch.run_splits.clone(),
+        )
+    }
+
     /// Compacts runs together. If the input runs are sorted, a single run will be created as output.
     ///
     /// Maximum possible memory usage is `(# runs + 2) * [crate::PersistConfig::blob_target_size]`
@@ -840,7 +869,7 @@ where
         shard_metrics: Arc<ShardMetrics>,
         isolated_runtime: Arc<IsolatedRuntime>,
         write_schemas: Schemas<K, V>,
-        batch_so_far: Option<HollowBatch<T>>,
+        batch_so_far: &Option<HollowBatch<T>>,
         incremental_tx: Option<Sender<HollowBatch<T>>>,
     ) -> Result<HollowBatch<T>, anyhow::Error> {
         // TODO: Figure out a more principled way to allocate our memory budget.
@@ -864,7 +893,7 @@ where
         // the config allows BatchBuilder to do its normal pipelining of writes.
         batch_cfg.inline_writes_single_max_bytes = 0;
 
-        if let Some(batch_so_far) = batch_so_far {
+        if let Some(batch_so_far) = batch_so_far.as_ref() {
             let last_part = batch_so_far
                 .last_part(shard_id.clone(), &*blob, &metrics)
                 .await;
@@ -994,12 +1023,20 @@ where
                 break;
             };
 
+            batch.flush_part(desc.clone(), updates).await;
+
             if let Some(tx) = incremental_tx.as_ref() {
                 // This is where we record whatever parts were successfully flushed
                 // to blob. That way we can resume an interrupted compaction later.
                 let partial_batch = batch.batch_with_finished_parts(desc.clone());
+
                 if let Some(partial_batch) = partial_batch {
-                    match tx.send(partial_batch).await {
+                    let hollow_batch = if let Some(batch_so_far) = batch_so_far.as_ref() {
+                        Self::combine_hollow_batch_with_previous(batch_so_far, &partial_batch)
+                    } else {
+                        partial_batch
+                    };
+                    match tx.send(hollow_batch).await {
                         Ok(_) => {
                             // metrics.compaction.incremental_batch_sent.inc();
                         }
@@ -1010,11 +1047,6 @@ where
                     };
                 }
             }
-
-            // part
-            // upper bound kvt
-            // find the last part where the lower bound < recorded upper bound
-            batch.flush_part(desc.clone(), updates).await;
         }
         let mut batch = batch.finish(desc.clone()).await?;
 
@@ -1036,8 +1068,14 @@ where
                 .await;
         }
 
+        let hollow_batch = if let Some(batch_so_far) = batch_so_far.as_ref() {
+            Self::combine_hollow_batch_with_previous(batch_so_far, &batch.into_hollow_batch())
+        } else {
+            batch.into_hollow_batch()
+        };
+
         timings.record(&metrics);
-        Ok(batch.into_hollow_batch())
+        Ok(hollow_batch)
     }
 
     fn validate_req(req: &CompactReq<T>) -> Result<(), anyhow::Error> {
@@ -1163,6 +1201,7 @@ mod tests {
             Arc::new(IsolatedRuntime::default()),
             req.clone(),
             schemas.clone(),
+            &write.machine,
         );
 
         let res = Compactor::<String, String, u64, i64>::compact_all(stream, req.clone())
@@ -1243,6 +1282,7 @@ mod tests {
             Arc::new(IsolatedRuntime::default()),
             req.clone(),
             schemas.clone(),
+            &write.machine,
         );
 
         let res =

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

@@ -501,7 +501,7 @@ where
                                 .into_iter()
                                 .map(|b| Arc::unwrap_or_clone(b.batch))
                                 .collect(),
-                            prev_batch: None,
+                            prev_batch: req.active_compaction,
                         };
                         compact_reqs.push(req);
                     }
@@ -611,13 +611,15 @@ where
         let metrics = Arc::clone(&self.applier.metrics);
 
         //TODO(dov): new metric
-        let _ = self.apply_unbatched_idempotent_cmd(&metrics.cmds.merge_res, |_, _, state| {
-            let ret = state.apply_compaction_progress(batch_so_far, current_ts);
-            if let Continue(_) = ret {
-                // metrics.state.compaction_progress_applied.inc();
-            }
-            ret
-        }).await;
+        let _ = self
+            .apply_unbatched_idempotent_cmd(&metrics.cmds.merge_res, |_, _, state| {
+                let ret = state.apply_compaction_progress(batch_so_far, current_ts);
+                if let Continue(_) = ret {
+                    // metrics.state.compaction_progress_applied.inc();
+                }
+                ret
+            })
+            .await;
     }
 
     pub async fn merge_res(
@@ -2033,6 +2035,7 @@ pub mod datadriven {
             Arc::clone(&datadriven.client.isolated_runtime),
             req_clone,
             SCHEMAS.clone(),
+            &datadriven.machine,
         );
 
         let res = Compactor::<String, (), u64, i64>::compact_all(stream, req.clone()).await?;

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

@@ -16,7 +16,7 @@ use std::collections::BTreeMap;
 use std::fmt::{Debug, Formatter};
 use std::marker::PhantomData;
 use std::ops::ControlFlow::{self, Break, Continue};
-use std::ops::{Deref, DerefMut};
+use std::ops::{Add, Deref, DerefMut};
 use std::time::Duration;
 
 use arrow::array::{Array, ArrayData, make_array};
@@ -647,10 +647,16 @@ impl<T: Timestamp + Codec64 + Sync> RunPart<T> {
         match self {
             RunPart::Single(p) => Ok(Box::new(p.clone())),
             RunPart::Many(r) => {
-                let fetched = r.get(shard_id, blob, metrics).await.ok_or_else(|| MissingBlob(r.key.complete(&shard_id)))?;
-                let last_part = fetched.parts.last().ok_or_else(|| MissingBlob(r.key.complete(&shard_id)))?;
+                let fetched = r
+                    .get(shard_id, blob, metrics)
+                    .await
+                    .ok_or_else(|| MissingBlob(r.key.complete(&shard_id)))?;
+                let last_part = fetched
+                    .parts
+                    .last()
+                    .ok_or_else(|| MissingBlob(r.key.complete(&shard_id)))?;
                 Ok(Box::pin(last_part.last_part(shard_id, blob, metrics)).await?)
-            },
+            }
         }
     }
 }
@@ -1687,7 +1693,7 @@ where
                 req.id,
                 ActiveCompaction {
                     start_ms: heartbeat_timestamp_ms,
-                    batch_so_far: None
+                    batch_so_far: None,
                 },
             )
         }
@@ -1727,10 +1733,11 @@ where
 
         let new_active_compaction = ActiveCompaction {
             start_ms: new_ts,
-            batch_so_far: Some(batch_so_far.clone())
+            batch_so_far: Some(batch_so_far.clone()),
         };
 
-        self.trace.apply_incremental_compaction(&new_active_compaction);
+        self.trace
+            .apply_incremental_compaction(&new_active_compaction);
 
         Continue(())
     }