[ENH]: move collection hard deletes to garbage collector

Author: codetheweb

rust/garbage_collector/src/construct_version_graph_orchestrator.rs

@@ -21,7 +21,9 @@ use chroma_system::{
     wrap, ChannelError, ComponentContext, ComponentHandle, Dispatcher, Handler, Orchestrator,
     PanicError, TaskError, TaskMessage, TaskResult,
 };
-use chroma_types::{chroma_proto::CollectionVersionFile, CollectionUuid};
+use chroma_types::{
+    chroma_proto::CollectionVersionFile, BatchGetCollectionSoftDeleteStatusError, CollectionUuid,
+};
 use chrono::DateTime;
 use petgraph::{dot::Dot, graph::DiGraph};
 use std::{
@@ -125,13 +127,17 @@ pub enum ConstructVersionGraphError {
     FetchLineageFile(#[from] FetchLineageFileError),
     #[error("Error fetching version file paths: {0}")]
     FetchVersionFilePaths(#[from] GetVersionFilePathsError),
+    #[error("Could not get collection soft delete status: {0}")]
+    BatchGetCollectionSoftDeleteStatus(#[from] BatchGetCollectionSoftDeleteStatusError),
 
     #[error("Invalid UUID: {0}")]
     InvalidUuid(#[from] uuid::Error),
     #[error("Invalid timestamp: {0}")]
     InvalidTimestamp(i64),
-    #[error("Expected node not found while constructing graph")]
-    ExpectedNodeNotFound,
+    #[error("Expected node not found while constructing graph (collection {0}@v{1:?})")]
+    ExpectedNodeNotFound(CollectionUuid, Option<i64>),
+    #[error("Invariant violation: {0}")]
+    InvariantViolation(String),
 }
 
 impl<E> From<TaskError<E>> for ConstructVersionGraphError
@@ -157,9 +163,11 @@ impl ChromaError for ConstructVersionGraphError {
             ConstructVersionGraphError::FetchVersionFile(err) => err.code(),
             ConstructVersionGraphError::FetchLineageFile(err) => err.code(),
             ConstructVersionGraphError::FetchVersionFilePaths(err) => err.code(),
+            ConstructVersionGraphError::BatchGetCollectionSoftDeleteStatus(err) => err.code(),
             ConstructVersionGraphError::InvalidUuid(_) => ErrorCodes::Internal,
             ConstructVersionGraphError::InvalidTimestamp(_) => ErrorCodes::InvalidArgument,
-            ConstructVersionGraphError::ExpectedNodeNotFound => ErrorCodes::Internal,
+            ConstructVersionGraphError::ExpectedNodeNotFound(_, _) => ErrorCodes::Internal,
+            ConstructVersionGraphError::InvariantViolation(_) => ErrorCodes::Internal,
         }
     }
 }
@@ -215,9 +223,11 @@ impl ConstructVersionGraphOrchestrator {
         ctx: &ComponentContext<ConstructVersionGraphOrchestrator>,
     ) -> Result<(), ConstructVersionGraphError> {
         if self.num_pending_tasks == 0 {
+            // This map will be used as a basis for building the graph
             let mut versions_by_collection_id: HashMap<CollectionUuid, Vec<(i64, VersionStatus)>> =
                 HashMap::new();
 
+            // Add all known versions from version files to map
             for (collection_id, version_file) in self.version_files.iter() {
                 if let Some(versions) = &version_file.version_history {
                     for version in versions.versions.iter() {
@@ -242,6 +252,7 @@ impl ConstructVersionGraphOrchestrator {
                 }
             }
 
+            // If any version appears as a version dependency (from the lineage file) but does not already exist in the map from the version files, the version must have been deleted.
             for dependency in self.version_dependencies.iter() {
                 let source_collection_id = dependency.source_collection_id;
                 let source_collection_version = dependency.source_collection_version;
@@ -263,6 +274,11 @@ impl ConstructVersionGraphOrchestrator {
                 versions.sort_unstable_by_key(|v| v.0);
             }
 
+            tracing::trace!(
+                "Versions by collection ID: {:#?}",
+                versions_by_collection_id
+            );
+
             let mut graph = DiGraph::new();
             for (collection_id, versions) in versions_by_collection_id.iter() {
                 let mut prev_node = None;
@@ -273,12 +289,14 @@ impl ConstructVersionGraphOrchestrator {
                         status: *status,
                     });
                     if let Some(prev) = prev_node {
+                        // Add edge between each successive pair of collection versions
                         graph.add_edge(prev, node, ());
                     }
                     prev_node = Some(node);
                 }
             }
 
+            // Add edges for forked collections
             for dependency in self.version_dependencies.iter() {
                 let source_node = graph
                     .node_indices()
@@ -287,15 +305,25 @@ impl ConstructVersionGraphOrchestrator {
                         node.collection_id == dependency.source_collection_id
                             && node.version == dependency.source_collection_version
                     })
-                    .ok_or(ConstructVersionGraphError::ExpectedNodeNotFound)?;
+                    .ok_or_else(|| {
+                        ConstructVersionGraphError::ExpectedNodeNotFound(
+                            dependency.source_collection_id,
+                            Some(dependency.source_collection_version),
+                        )
+                    })?;
 
                 let target_node = graph
                     .node_indices()
                     .find(|n| {
                         let node = graph.node_weight(*n).expect("node index should exist");
                         node.collection_id == dependency.target_collection_id
                     })
-                    .ok_or(ConstructVersionGraphError::ExpectedNodeNotFound)?;
+                    .ok_or_else(|| {
+                        ConstructVersionGraphError::ExpectedNodeNotFound(
+                            dependency.target_collection_id,
+                            None,
+                        )
+                    })?;
 
                 graph.add_edge(source_node, target_node, ());
             }
@@ -308,6 +336,15 @@ impl ConstructVersionGraphOrchestrator {
 
             tracing::trace!("Version files: {:#?}", self.version_files);
 
+            let components = petgraph::algo::connected_components(&graph);
+            if components != 1 {
+                // This is a defensive check, it should never happen
+                return Err(ConstructVersionGraphError::InvariantViolation(format!(
+                    "Graph is not fully connected, found {} components",
+                    components
+                )));
+            }
+
             self.terminate_with_result(
                 Ok(ConstructVersionGraphResponse {
                     graph,