[ENH] Fork the wal3 log. (chroma-core#4416)

Author: warpbuild-benchmark-bot[bot]

rust/log-service/src/lib.rs

@@ -766,9 +766,79 @@ impl LogService for LogServer {
 
     async fn fork_logs(
         &self,
-        _request: Request<ForkLogsRequest>,
+        request: Request<ForkLogsRequest>,
     ) -> Result<Response<ForkLogsResponse>, Status> {
-        unimplemented!("Log forking is unimplemented for WAL3 for now")
+        let request = request.into_inner();
+        let source_collection_id = Uuid::parse_str(&request.source_collection_id)
+            .map(CollectionUuid)
+            .map_err(|_| Status::invalid_argument("Failed to parse collection id"))?;
+        let target_collection_id = Uuid::parse_str(&request.target_collection_id)
+            .map(CollectionUuid)
+            .map_err(|_| Status::invalid_argument("Failed to parse collection id"))?;
+        let source_prefix = storage_prefix_for_log(source_collection_id);
+        let target_prefix = storage_prefix_for_log(target_collection_id);
+        let span = tracing::info_span!(
+            "fork_logs",
+            source_collection_id = source_collection_id.to_string(),
+            target_collection_id = target_collection_id.to_string(),
+        );
+        let storage = Arc::clone(&self.storage);
+        let options = self.config.writer.clone();
+
+        async move {
+            let log_reader = LogReader::new(
+                self.config.reader.clone(),
+                Arc::clone(&storage),
+                source_prefix.clone(),
+            );
+            let cursors = CursorStore::new(
+                CursorStoreOptions::default(),
+                Arc::clone(&storage),
+                source_prefix,
+                "copy task".to_string(),
+            );
+            let cursor_name = &COMPACTION;
+            let witness = cursors.load(cursor_name).await.map_err(|err| {
+                Status::new(err.code().into(), format!("Failed to load cursor: {}", err))
+            })?;
+            // This is the existing compaction_offset, which is the last record that was compacted.
+            let offset = witness
+                .map(|x| x.1.position)
+                .unwrap_or(LogPosition::from_offset(0));
+            wal3::copy(
+                &storage,
+                &options,
+                &log_reader,
+                // + 1 to get to the first uncompacted record.
+                offset + 1u64,
+                target_prefix.clone(),
+            )
+            .await
+            .map_err(|err| {
+                Status::new(err.code().into(), format!("Failed to copy log: {}", err))
+            })?;
+            let log_reader = LogReader::new(
+                self.config.reader.clone(),
+                Arc::clone(&storage),
+                target_prefix,
+            );
+            // This is the next record to insert, so we'll have to adjust downwards.
+            let max_offset = log_reader.maximum_log_position().await.map_err(|err| {
+                Status::new(err.code().into(), format!("Failed to copy log: {}", err))
+            })?;
+            if max_offset < offset {
+                return Err(Status::new(
+                    chroma_error::ErrorCodes::Internal.into(),
+                    format!("max_offset={:?} < offset={:?}", max_offset, offset),
+                ));
+            }
+            Ok(Response::new(ForkLogsResponse {
+                compaction_offset: offset.offset(),
+                enumeration_offset: (max_offset - 1u64).offset(),
+            }))
+        }
+        .instrument(span)
+        .await
     }
 
     #[tracing::instrument(info, skip(self, request), err(Display))]

rust/wal3/src/copy.rs

@@ -0,0 +1,129 @@
+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())
+}
+
+pub async fn copy(
+    storage: &Storage,
+    options: &LogWriterOptions,
+    reader: &LogReader,
+    offset: LogPosition,
+    target: String,
+) -> Result<(), Error> {
+    let Some(manifest) = reader.manifest().await? else {
+        return Err(Error::UninitializedLog);
+    };
+    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 setsum = snapshots
+        .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>();
+    let manifest = Manifest {
+        setsum,
+        acc_bytes,
+        writer: "copy task".to_string(),
+        snapshots,
+        fragments,
+    };
+    Manifest::initialize_from_manifest(options, storage, &target, manifest).await?;
+    Ok(())
+}

rust/wal3/src/lib.rs

@@ -7,6 +7,7 @@ use setsum::Setsum;
 
 mod backoff;
 mod batch_manager;
+mod copy;
 mod cursors;
 mod manifest;
 mod manifest_manager;
@@ -15,6 +16,7 @@ mod writer;
 
 pub use backoff::ExponentialBackoff;
 pub use batch_manager::BatchManager;
+pub use copy::copy;
 pub use cursors::{Cursor, CursorName, CursorStore, Witness};
 pub use manifest::{Manifest, Snapshot, SnapshotPointer};
 pub use manifest_manager::ManifestManager;
@@ -219,6 +221,16 @@ impl std::ops::Add<usize> for LogPosition {
     }
 }
 
+impl std::ops::Sub<u64> for LogPosition {
+    type Output = LogPosition;
+
+    fn sub(self, rhs: u64) -> Self::Output {
+        LogPosition {
+            offset: self.offset.wrapping_sub(rhs),
+        }
+    }
+}
+
 impl std::ops::Sub<LogPosition> for LogPosition {
     type Output = u64;
 

rust/wal3/src/manifest.rs

@@ -21,15 +21,15 @@ use crate::{
 
 /////////////////////////////////////////////// paths //////////////////////////////////////////////
 
-fn manifest_path(prefix: &str) -> String {
+pub fn manifest_path(prefix: &str) -> String {
     format!("{prefix}/manifest/MANIFEST")
 }
 
-fn unprefixed_snapshot_path(setsum: Setsum) -> String {
+pub fn unprefixed_snapshot_path(setsum: Setsum) -> String {
     format!("snapshot/SNAPSHOT.{}", setsum.hexdigest())
 }
 
-fn snapshot_setsum(path: &str) -> Result<Setsum, Error> {
+pub fn snapshot_setsum(path: &str) -> Result<Setsum, Error> {
     let setsum = path
         .strip_prefix("snapshot/SNAPSHOT.")
         .ok_or_else(|| Error::CorruptManifest(format!("unparseable snapshot path: {}", path,)))?;
@@ -199,7 +199,7 @@ impl Snapshot {
         options: &ThrottleOptions,
         storage: &Storage,
         prefix: &str,
-    ) -> Result<(), Error> {
+    ) -> Result<SnapshotPointer, Error> {
         let exp_backoff = crate::backoff::ExponentialBackoff::new(
             options.throughput as f64,
             options.headroom as f64,
@@ -214,15 +214,15 @@ impl Snapshot {
             let options = PutOptions::if_not_exists(StorageRequestPriority::P0);
             match storage.put_bytes(&path, payload, options).await {
                 Ok(_) => {
-                    return Ok(());
+                    return Ok(self.to_pointer());
                 }
                 Err(StorageError::Precondition { path: _, source: _ }) => {
                     // NOTE(rescrv):  This is something of a lie.  We know that someone put the
                     // file before us, and we know the setsum of the file is embedded in the path.
                     // Because the setsum is only calculable if you have the file and we assume
                     // non-malicious code, anyone who puts the same setsum as us has, in all
                     // likelihood, put something referencing the same content as us.
-                    return Ok(());
+                    return Ok(self.to_pointer());
                 }
                 Err(e) => {
                     tracing::error!("error uploading manifest: {e:?}");
@@ -279,6 +279,17 @@ impl Snapshot {
             (None, None) => LogPosition::default(),
         }
     }
+
+    pub fn to_pointer(&self) -> SnapshotPointer {
+        SnapshotPointer {
+            setsum: self.setsum,
+            path_to_snapshot: self.path.clone(),
+            depth: self.depth,
+            start: self.minimum_log_position(),
+            limit: self.maximum_log_position(),
+            num_bytes: self.num_bytes(),
+        }
+    }
 }
 
 ///////////////////////////////////////////// Manifest /////////////////////////////////////////////
@@ -533,7 +544,7 @@ impl Manifest {
     /// Initialize the log with an empty manifest.
     #[tracing::instrument(skip(storage), err(Display))]
     pub async fn initialize(
-        _: &LogWriterOptions,
+        options: &LogWriterOptions,
         storage: &Storage,
         prefix: &str,
         writer: &str,
@@ -546,6 +557,17 @@ impl Manifest {
             snapshots: vec![],
             fragments: vec![],
         };
+        Self::initialize_from_manifest(options, storage, prefix, initial).await
+    }
+
+    /// Initialize the log with an empty manifest.
+    #[tracing::instrument(skip(storage), err(Display))]
+    pub async fn initialize_from_manifest(
+        _: &LogWriterOptions,
+        storage: &Storage,
+        prefix: &str,
+        initial: Manifest,
+    ) -> Result<(), Error> {
         let payload = serde_json::to_string(&initial)
             .map_err(|e| Error::CorruptManifest(format!("could not encode JSON manifest: {e:?}")))?
             .into_bytes();

rust/wal3/src/reader.rs

@@ -28,7 +28,7 @@ pub struct Limits {
 pub struct LogReader {
     options: LogReaderOptions,
     storage: Arc<Storage>,
-    prefix: String,
+    pub(crate) prefix: String,
 }
 
 impl LogReader {

rust/wal3/src/writer.rs

@@ -4,7 +4,7 @@ use std::time::{Duration, Instant, SystemTime};
 
 use arrow::array::{ArrayRef, BinaryArray, RecordBatch, UInt64Array};
 use chroma_storage::admissioncontrolleds3::StorageRequestPriority;
-use chroma_storage::{PutOptions, Storage, StorageError};
+use chroma_storage::{GetOptions, PutOptions, Storage, StorageError};
 use parquet::arrow::ArrowWriter;
 use parquet::basic::Compression;
 use parquet::file::properties::WriterProperties;
@@ -520,3 +520,41 @@ pub async fn upload_parquet(
         }
     }
 }
+
+#[tracing::instrument(skip(options, storage))]
+pub async fn copy_parquet(
+    options: &LogWriterOptions,
+    storage: &Storage,
+    source: &str,
+    target: &str,
+) -> Result<(), Error> {
+    let parquet = storage
+        .get(source, GetOptions::new(StorageRequestPriority::P0))
+        .await
+        .map_err(Arc::new)?;
+    let exp_backoff: ExponentialBackoff = options.throttle_fragment.into();
+    let start = Instant::now();
+    loop {
+        match storage
+            .put_bytes(
+                target,
+                parquet.to_vec(),
+                PutOptions::if_not_exists(StorageRequestPriority::P0),
+            )
+            .await
+        {
+            Ok(_) => return Ok(()),
+            Err(StorageError::Precondition { path: _, source: _ }) => return Ok(()),
+            Err(err) => {
+                if start.elapsed() > Duration::from_secs(60) {
+                    return Err(Error::StorageError(Arc::new(err)));
+                }
+                let mut backoff = exp_backoff.next();
+                if backoff > Duration::from_secs(3_600) {
+                    backoff = Duration::from_secs(3_600);
+                }
+                tokio::time::sleep(backoff).await;
+            }
+        }
+    }
+}