wip

Author: codetheweb

Cargo.lock

@@ -32,6 +32,7 @@ tracing = { workspace = true }
 thiserror = { workspace = true }
 humantime = { workspace = true }
 opentelemetry = { workspace = true }
+petgraph = "0.8.1"
 
 chroma-config = { workspace = true }
 chroma-error = { workspace = true }

rust/garbage_collector/Cargo.toml

@@ -36,9 +36,21 @@
 //!    - Permanently deletes marked versions from the system database
 //!    - Input: Version file, versions to delete, unused S3 files
 //!    - Output: Deletion confirmation
+//!
+//!
+//!
+//! ListFilesAtVersionOperator
+//! - input: version file & versions to check
+//! - output: full paths of all S3 files used by specified versions
+//!
 
+use std::collections::HashSet;
 use std::fmt::{Debug, Formatter};
+use std::str::FromStr;
 
+use crate::operators::fetch_lineage_file::{
+    FetchLineageFileError, FetchLineageFileInput, FetchLineageFileOperator, FetchLineageFileOutput,
+};
 use crate::types::CleanupMode;
 use async_trait::async_trait;
 use chroma_error::{ChromaError, ErrorCodes};
@@ -83,6 +95,7 @@ use prost::Message;
 
 pub struct GarbageCollectorOrchestrator {
     collection_id: CollectionUuid,
+    lineage_file_path: Option<String>,
     version_file_path: String,
     absolute_cutoff_time: DateTime<Utc>,
     sysdb_client: SysDb,
@@ -114,9 +127,11 @@ pub struct GarbageCollectorResponse {
 
 #[allow(clippy::too_many_arguments)]
 impl GarbageCollectorOrchestrator {
+    /// Lineage file path must be provided if this collection is part of a fork tree.
     pub fn new(
         collection_id: CollectionUuid,
         version_file_path: String,
+        lineage_file_path: Option<String>,
         absolute_cutoff_time: DateTime<Utc>,
         sysdb_client: SysDb,
         dispatcher: ComponentHandle<Dispatcher>,
@@ -126,6 +141,7 @@ impl GarbageCollectorOrchestrator {
         Self {
             collection_id,
             version_file_path,
+            lineage_file_path,
             absolute_cutoff_time,
             sysdb_client,
             dispatcher,
@@ -165,6 +181,8 @@ pub enum GarbageCollectorError {
     Aborted,
     #[error("DeleteUnusedFiles error: {0}")]
     DeleteUnusedFiles(#[from] DeleteUnusedFilesError),
+    #[error("FetchLineageFile error: {0}")]
+    FetchLineageFile(#[from] FetchLineageFileError),
 }
 
 impl ChromaError for GarbageCollectorError {
@@ -201,14 +219,15 @@ impl Orchestrator for GarbageCollectorOrchestrator {
             "Creating initial fetch version file task"
         );
 
-        vec![wrap(
-            Box::new(FetchVersionFileOperator {}),
-            FetchVersionFileInput {
-                version_file_path: self.version_file_path.clone(),
-                storage: self.storage.clone(),
-            },
-            ctx.receiver(),
-        )]
+        // vec![wrap(
+        //     Box::new(FetchVersionFileOperator {}),
+        //     FetchVersionFileInput {
+        //         version_file_path: self.version_file_path.clone(),
+        //         storage: self.storage.clone(),
+        //     },
+        //     ctx.receiver(),
+        // )]
+        vec![]
     }
 
     fn set_result_channel(
@@ -369,34 +388,76 @@ impl Handler<TaskResult<MarkVersionsAtSysDbOutput, MarkVersionsAtSysDbError>>
         message: TaskResult<MarkVersionsAtSysDbOutput, MarkVersionsAtSysDbError>,
         ctx: &ComponentContext<GarbageCollectorOrchestrator>,
     ) {
-        // Stage 3: After marking versions, compute unused files
         let output = match self.ok_or_terminate(message.into_inner(), ctx).await {
             Some(output) => output,
             None => return,
         };
 
-        let compute_task = wrap(
-            Box::new(ComputeUnusedFilesOperator::new(
-                self.collection_id.to_string(),
-                self.storage.clone(),
-                2, // min_versions_to_keep
-            )),
-            ComputeUnusedFilesInput {
-                version_file: output.version_file,
-                versions_to_delete: output.versions_to_delete,
-                oldest_version_to_keep: output.oldest_version_to_keep,
-            },
-            ctx.receiver(),
-        );
+        if let Some(lineage_file_path) = self.lineage_file_path.clone() {
+            let fetch_lineage_file_task = wrap(
+                Box::new(FetchLineageFileOperator::new()),
+                FetchLineageFileInput::new(self.storage.clone(), lineage_file_path),
+                ctx.receiver(),
+            );
 
-        if let Err(e) = self
-            .dispatcher()
-            .send(compute_task, Some(Span::current()))
-            .await
-        {
-            self.terminate_with_result(Err(GarbageCollectorError::Channel(e)), ctx)
-                .await;
-            return;
+            if let Err(e) = self
+                .dispatcher()
+                .send(fetch_lineage_file_task, Some(Span::current()))
+                .await
+            {
+                self.terminate_with_result(Err(GarbageCollectorError::Channel(e)), ctx)
+                    .await;
+                return;
+            }
+        }
+    }
+}
+
+#[async_trait]
+impl Handler<TaskResult<FetchLineageFileOutput, FetchLineageFileError>>
+    for GarbageCollectorOrchestrator
+{
+    type Result = ();
+
+    async fn handle(
+        &mut self,
+        message: TaskResult<FetchLineageFileOutput, FetchLineageFileError>,
+        ctx: &ComponentContext<GarbageCollectorOrchestrator>,
+    ) {
+        let output = match self.ok_or_terminate(message.into_inner(), ctx).await {
+            Some(output) => output,
+            None => return,
+        };
+
+        let mut collection_ids_in_tree = HashSet::new();
+        for dependency in output.0.dependencies {
+            // todo: no expect
+            let source_id = CollectionUuid::from_str(&dependency.source_collection_id)
+                .expect("Failed to parse source ID");
+            let target_id = CollectionUuid::from_str(&dependency.target_collection_id)
+                .expect("Failed to parse target ID");
+            collection_ids_in_tree.insert(source_id);
+            collection_ids_in_tree.insert(target_id);
+        }
+
+        // todo: need to remove self?
+
+        for collection_id in collection_ids_in_tree {
+            let fetch_version_file_task = wrap(
+                Box::new(FetchVersionFileOperator::new()),
+                FetchVersionFileInput::new(collection_id, self.storage.clone()),
+                ctx.receiver(),
+            );
+
+            if let Err(e) = self
+                .dispatcher()
+                .send(fetch_version_file_task, Some(Span::current()))
+                .await
+            {
+                self.terminate_with_result(Err(GarbageCollectorError::Channel(e)), ctx)
+                    .await;
+                return;
+            }
         }
     }
 }
@@ -871,6 +932,7 @@ mod tests {
         let orchestrator = GarbageCollectorOrchestrator::new(
             collection_id,
             collection_info.version_file_path.clone(),
+            None,                     // todo
             SystemTime::now().into(), //  immediately expire versions
             sysdb,
             dispatcher_handle,
@@ -1027,6 +1089,7 @@ mod tests {
         let orchestrator = GarbageCollectorOrchestrator::new(
             collection_id,
             collection_info.version_file_path.clone(),
+            None,                     // todo
             SystemTime::now().into(), //  immediately expire versions
             sysdb,
             dispatcher_handle,
@@ -1185,6 +1248,7 @@ mod tests {
         let orchestrator = GarbageCollectorOrchestrator::new(
             collection_id,
             collection_info.version_file_path.clone(),
+            None,                     // todo
             SystemTime::now().into(), //  immediately expire versions
             sysdb,
             dispatcher_handle,

rust/garbage_collector/src/garbage_collector_orchestrator.rs

@@ -5,12 +5,31 @@ use chroma_types::chroma_proto::CollectionLineageFile;
 use prost::Message;
 use thiserror::Error;
 
-struct FetchLineageFileInput {
+#[derive(Clone)]
+pub struct FetchLineageFileInput {
     storage: Storage,
     lineage_file_path: String,
 }
 
-struct FetchLineageFileOutput(CollectionLineageFile);
+impl std::fmt::Debug for FetchLineageFileInput {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("FetchLineageFileInput")
+            .field("lineage_file_path", &self.lineage_file_path)
+            .finish()
+    }
+}
+
+impl FetchLineageFileInput {
+    pub fn new(storage: Storage, lineage_file_path: String) -> Self {
+        Self {
+            storage,
+            lineage_file_path,
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct FetchLineageFileOutput(pub CollectionLineageFile);
 
 #[derive(Debug, Error)]
 pub enum FetchLineageFileError {
@@ -23,6 +42,12 @@ pub enum FetchLineageFileError {
 #[derive(Clone, Debug)]
 pub struct FetchLineageFileOperator {}
 
+impl FetchLineageFileOperator {
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
 #[async_trait]
 impl Operator<FetchLineageFileInput, FetchLineageFileOutput> for FetchLineageFileOperator {
     type Error = FetchLineageFileError;

rust/garbage_collector/src/operators/fetch_lineage_file.rs

@@ -14,21 +14,40 @@ use async_trait::async_trait;
 use chroma_error::{ChromaError, ErrorCodes};
 use chroma_storage::admissioncontrolleds3::StorageRequestPriority;
 use chroma_storage::{GetOptions, Storage, StorageError};
+use chroma_sysdb::SysDb;
 use chroma_system::{Operator, OperatorType};
+use chroma_types::CollectionUuid;
 use thiserror::Error;
 
 #[derive(Clone, Debug)]
 pub struct FetchVersionFileOperator {}
 
+impl FetchVersionFileOperator {
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
 pub struct FetchVersionFileInput {
-    pub version_file_path: String,
-    pub storage: Storage,
+    collection_id: CollectionUuid,
+    storage: Storage,
+    sysdb: SysDb,
+}
+
+impl FetchVersionFileInput {
+    pub fn new(collection_id: CollectionUuid, storage: Storage, sysdb: SysDb) -> Self {
+        Self {
+            collection_id,
+            storage,
+            sysdb,
+        }
+    }
 }
 
 impl Debug for FetchVersionFileInput {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("FetchVersionFileInput")
-            .field("version_file_path", &self.version_file_path)
+            .field("collection_id", &self.collection_id)
             .finish()
     }
 }
@@ -83,10 +102,10 @@ impl Operator<FetchVersionFileInput, FetchVersionFileOutput> for FetchVersionFil
         &self,
         input: &FetchVersionFileInput,
     ) -> Result<FetchVersionFileOutput, FetchVersionFileError> {
-        tracing::info!(
-            path = %input.version_file_path,
-            "Starting to fetch version file"
-        );
+        let collection =
+            input
+                .sysdb
+                .get_collections(collection_id, name, tenant, database, limit, offset);
 
         let content = input
             .storage