[ENH] wal3::copy implemented using scan/AWS copy (#4803)

rescrv · web-flow · commit bf2b0c06df2a · 2025-06-10T15:22:54.000-07:00
## Description of changes

This PR uses the `LogReader`'s standard `scan` method to select
fragments for copying.  It then issues a copy in parallel and writes
them to the manifest as one list, without accounting for snapshots.

Alternatives considered:
- Potentially walk the snapshots and copy/paste snapshots.  This is what
  was done before, but it could only prune to boundaries of snapshot
  pointers in the root/manifest of the tree.
- Use the GC logic to prune.  This was originally my intent, but it had
  a hidden downside:  The GC would walk all old snapshots to determine
  what data it had to delete.  Implementing a switch to skip this
  behavior essentially made two distinct functions in one, obviating the
  advantage of reusing the GC code.

## Test plan

Integration tests cover copy.

- [X] Tests pass locally with `pytest` for python, `yarn test` for js,
`cargo test` for rust

## Documentation Changes

N/A
diff --git a/rust/wal3/src/copy.rs b/rust/wal3/src/copy.rs
@@ -1,93 +1,10 @@
+use std::sync::Arc;
+
 use chroma_storage::Storage;
 use setsum::Setsum;
 
-use crate::manifest::{unprefixed_snapshot_path, Manifest, Snapshot};
 use crate::reader::LogReader;
-use crate::writer::copy_parquet;
-use crate::{Error, Fragment, LogPosition, LogWriterOptions, SnapshotPointer};
-
-pub async fn copy_snapshot(
-    storage: &Storage,
-    options: &LogWriterOptions,
-    reader: &LogReader,
-    root: &SnapshotPointer,
-    offset: LogPosition,
-    target: &str,
-) -> Result<SnapshotPointer, Error> {
-    let Some(snapshot) =
-        Snapshot::load(&options.throttle_manifest, storage, &reader.prefix, root).await?
-    else {
-        return Err(Error::CorruptManifest(format!(
-            "snapshot {} does not exist",
-            root.setsum.hexdigest(),
-        )));
-    };
-    let mut dropped = vec![];
-    let mut snapshots = vec![];
-    for snapshot in &snapshot.snapshots {
-        if snapshot.limit > offset {
-            snapshots.push(
-                Box::pin(copy_snapshot(
-                    storage, options, reader, snapshot, offset, target,
-                ))
-                .await?,
-            );
-        } else {
-            dropped.push(snapshot.setsum);
-        }
-    }
-    let mut fragments = vec![];
-    for fragment in &snapshot.fragments {
-        if fragment.limit > offset {
-            fragments.push(copy_fragment(storage, options, reader, fragment, target).await?);
-        } else {
-            dropped.push(fragment.setsum);
-        }
-    }
-    let dropped = dropped.iter().fold(Setsum::default(), |x, y| x + *y);
-    let kept_snapshots = snapshots
-        .iter()
-        .fold(Setsum::default(), |x, y| x + y.setsum);
-    let kept_fragments = fragments
-        .iter()
-        .fold(Setsum::default(), |x, y| x + y.setsum);
-    if dropped + kept_snapshots + kept_fragments != root.setsum {
-        // NOTE(rescrv):  If you see this error you have to figure out where data is lost.  This
-        // will require writing a test case rather than trying to deduce it from the setsums.
-        return Err(Error::CorruptManifest(
-            "Copying failed because the setsum was not balanced".to_string(),
-        ));
-    }
-    let depth = snapshots.iter().map(|x| x.depth + 1).max().unwrap_or(0);
-    let snapshot = Snapshot {
-        path: unprefixed_snapshot_path(kept_snapshots + kept_fragments),
-        depth,
-        setsum: kept_snapshots + kept_fragments,
-        writer: "copy task".to_string(),
-        snapshots,
-        fragments,
-    };
-    snapshot
-        .install(&options.throttle_manifest, storage, target)
-        .await
-}
-
-pub async fn copy_fragment(
-    storage: &Storage,
-    options: &LogWriterOptions,
-    reader: &LogReader,
-    frag: &Fragment,
-    target: &str,
-) -> Result<Fragment, Error> {
-    copy_parquet(
-        options,
-        storage,
-        &format!("{}/{}", reader.prefix, frag.path),
-        &format!("{}/{}", target, frag.path),
-    )
-    .await?;
-    Ok(frag.clone())
-}
+use crate::{prefixed_fragment_path, Error, Limits, LogPosition, LogWriterOptions, Manifest};
 
 pub async fn copy(
     storage: &Storage,
@@ -96,35 +13,36 @@ pub async fn copy(
     offset: LogPosition,
     target: String,
 ) -> Result<(), Error> {
-    let manifest = reader
-        .manifest()
-        .await?
-        .unwrap_or(Manifest::new_empty("copy task"));
-    let mut snapshots = vec![];
-    for snapshot in &manifest.snapshots {
-        snapshots.push(copy_snapshot(storage, options, reader, snapshot, offset, &target).await?);
-    }
-    let mut fragments = vec![];
-    for fragment in &manifest.fragments {
-        fragments.push(copy_fragment(storage, options, reader, fragment, &target).await?);
+    let fragments = reader.scan(offset, Limits::UNLIMITED).await?;
+    let mut futures = vec![];
+    for fragment in fragments.into_iter() {
+        let target = &target;
+        futures.push(async move {
+            storage
+                .copy(
+                    &prefixed_fragment_path(&reader.prefix, fragment.seq_no),
+                    &prefixed_fragment_path(target, fragment.seq_no),
+                )
+                .await
+                .map(|_| fragment)
+        });
     }
-    let setsum = snapshots
+    let fragments = futures::future::try_join_all(futures)
+        .await
+        .map_err(Arc::new)?;
+    let setsum = fragments
         .iter()
         .map(|x| x.setsum)
-        .fold(Setsum::default(), |x, y| x + y)
-        + fragments
-            .iter()
-            .map(|x| x.setsum)
-            .fold(Setsum::default(), |x, y| x + y);
-    let acc_bytes = snapshots.iter().map(|x| x.num_bytes).sum::<u64>()
-        + fragments.iter().map(|x| x.num_bytes).sum::<u64>();
+        .fold(Setsum::default(), |x, y| x + y);
+    let acc_bytes = fragments.iter().map(|x| x.num_bytes).sum::<u64>();
+    let initial_offset = Some(fragments.iter().map(|f| f.start).min().unwrap_or(offset));
     let manifest = Manifest {
         setsum,
         acc_bytes,
         writer: "copy task".to_string(),
-        snapshots,
+        snapshots: vec![],
         fragments,
-        initial_offset: manifest.initial_offset,
+        initial_offset,
     };
     Manifest::initialize_from_manifest(options, storage, &target, manifest).await?;
     Ok(())
diff --git a/rust/wal3/src/lib.rs b/rust/wal3/src/lib.rs
@@ -502,6 +502,15 @@ where
 
 ////////////////////////////////////////// Fragment Paths //////////////////////////////////////////
 
+pub fn prefixed_fragment_path(prefix: &str, fragment_seq_no: FragmentSeqNo) -> String {
+    format!(
+        "{}/log/Bucket={:016x}/FragmentSeqNo={:016x}.parquet",
+        prefix,
+        fragment_seq_no.bucket(),
+        fragment_seq_no.0,
+    )
+}
+
 pub fn unprefixed_fragment_path(fragment_seq_no: FragmentSeqNo) -> String {
     format!(
         "log/Bucket={:016x}/FragmentSeqNo={:016x}.parquet",
diff --git a/rust/wal3/src/manifest.rs b/rust/wal3/src/manifest.rs
@@ -118,6 +118,7 @@ impl Snapshot {
         }
         for frag in self.fragments.iter() {
             calculated_setsum += frag.setsum;
+            bytes_read += frag.num_bytes;
         }
         if calculated_setsum != self.setsum {
             return Err(ScrubError::CorruptManifest{
diff --git a/rust/wal3/src/reader.rs b/rust/wal3/src/reader.rs
@@ -22,13 +22,21 @@ fn ranges_overlap(lhs: (LogPosition, LogPosition), rhs: (LogPosition, LogPositio
 }
 
 /// Limits allows encoding things like offset, timestamp, and byte size limits for the read.
-#[derive(Clone, Debug, Default, serde::Deserialize, serde::Serialize)]
+#[derive(Copy, Clone, Debug, Default, serde::Deserialize, serde::Serialize)]
 pub struct Limits {
     pub max_files: Option<u64>,
     pub max_bytes: Option<u64>,
     pub max_records: Option<u64>,
 }
 
+impl Limits {
+    pub const UNLIMITED: Limits = Limits {
+        max_files: None,
+        max_bytes: None,
+        max_records: None,
+    };
+}
+
 /// LogReader is a reader for the log.
 pub struct LogReader {
     options: LogReaderOptions,
@@ -757,7 +765,7 @@ mod tests {
             max_bytes: None,
             max_records: Some(100), // Would need data up to exactly offset 200
         };
-        let result = LogReader::scan_from_manifest(&manifest, from, limits.clone());
+        let result = LogReader::scan_from_manifest(&manifest, from, limits);
         assert!(
             result.is_some(),
             "Should succeed when request stays within manifest coverage"
@@ -836,11 +844,8 @@ mod tests {
             fragments: vec![],
             initial_offset: None,
         };
-        let result_empty = LogReader::scan_from_manifest(
-            &empty_manifest,
-            LogPosition::from_offset(0),
-            limits.clone(),
-        );
+        let result_empty =
+            LogReader::scan_from_manifest(&empty_manifest, LogPosition::from_offset(0), limits);
         assert!(
             result_empty.is_none(),
             "Should return None for empty manifest"
diff --git a/rust/wal3/tests/test_k8s_integration_82_copy_then_update_dst.rs b/rust/wal3/tests/test_k8s_integration_82_copy_then_update_dst.rs
@@ -15,7 +15,7 @@ async fn test_k8s_integration_82_copy_then_update_dst() {
     Manifest::initialize(
         &LogWriterOptions::default(),
         &storage,
-        "test_k8s_integration_80_copy_source",
+        "test_k8s_integration_82_copy_then_update_dst_source",
         "init",
     )
     .await
@@ -29,7 +29,7 @@ async fn test_k8s_integration_82_copy_then_update_dst() {
             ..LogWriterOptions::default()
         },
         Arc::clone(&storage),
-        "test_k8s_integration_80_copy_source",
+        "test_k8s_integration_82_copy_then_update_dst_source",
         "load and scrub writer",
         (),
     )
@@ -45,7 +45,7 @@ async fn test_k8s_integration_82_copy_then_update_dst() {
     let reader = LogReader::open(
         LogReaderOptions::default(),
         Arc::clone(&storage),
-        "test_k8s_integration_80_copy_source".to_string(),
+        "test_k8s_integration_82_copy_then_update_dst_source".to_string(),
     )
     .await
     .unwrap();
@@ -55,15 +55,15 @@ async fn test_k8s_integration_82_copy_then_update_dst() {
         &LogWriterOptions::default(),
         &reader,
         LogPosition::default(),
-        "test_k8s_integration_80_copy_target".to_string(),
+        "test_k8s_integration_82_copy_then_update_dst_target".to_string(),
     )
     .await
     .unwrap();
     // Scrub the copy.
     let copied = LogReader::open(
         LogReaderOptions::default(),
         Arc::clone(&storage),
-        "test_k8s_integration_80_copy_target".to_string(),
+        "test_k8s_integration_82_copy_then_update_dst_target".to_string(),
     )
     .await
     .unwrap();
@@ -82,7 +82,7 @@ async fn test_k8s_integration_82_copy_then_update_dst() {
             ..LogWriterOptions::default()
         },
         Arc::clone(&storage),
-        "test_k8s_integration_80_copy_target",
+        "test_k8s_integration_82_copy_then_update_dst_target",
         "load and scrub writer",
         (),
     )

Original file line number	Diff line number	Diff line change
`@@ -118,6 +118,7 @@ impl Snapshot {`
`118`	`118`	`}`
`119`	`119`	`for frag in self.fragments.iter() {`
`120`	`120`	`calculated_setsum += frag.setsum;`
	`121`	`+ bytes_read += frag.num_bytes;`
`121`	`122`	`}`
`122`	`123`	`if calculated_setsum != self.setsum {`
`123`	`124`	`return Err(ScrubError::CorruptManifest{`