[ENH]: move collection hard deletes to garbage collector

Author: codetheweb

rust/garbage_collector/src/construct_version_graph_orchestrator.rs

@@ -130,8 +130,10 @@ pub enum ConstructVersionGraphError {
     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
@@ -159,7 +161,8 @@ impl ChromaError for ConstructVersionGraphError {
             ConstructVersionGraphError::FetchVersionFilePaths(err) => err.code(),
             ConstructVersionGraphError::InvalidUuid(_) => ErrorCodes::Internal,
             ConstructVersionGraphError::InvalidTimestamp(_) => ErrorCodes::InvalidArgument,
-            ConstructVersionGraphError::ExpectedNodeNotFound => ErrorCodes::Internal,
+            ConstructVersionGraphError::ExpectedNodeNotFound(_, _) => ErrorCodes::Internal,
+            ConstructVersionGraphError::InvariantViolation(_) => ErrorCodes::Internal,
         }
     }
 }
@@ -263,6 +266,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;
@@ -279,23 +287,57 @@ impl ConstructVersionGraphOrchestrator {
                 }
             }
 
-            for dependency in self.version_dependencies.iter() {
-                let source_node = graph
-                    .node_indices()
-                    .find(|n| {
-                        let node = graph.node_weight(*n).expect("node index should exist");
-                        node.collection_id == dependency.source_collection_id
-                            && node.version == dependency.source_collection_version
-                    })
-                    .ok_or(ConstructVersionGraphError::ExpectedNodeNotFound)?;
+            let collection_id_soft_delete = self
+                .sysdb
+                .batch_get_collection_soft_delete_status(
+                    versions_by_collection_id.keys().cloned().collect(),
+                )
+                .await
+                .unwrap(); // todo
 
-                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)?;
+            for dependency in self.version_dependencies.iter() {
+                let source_node = match graph.node_indices().find(|n| {
+                    let node = graph.node_weight(*n).expect("node index should exist");
+                    node.collection_id == dependency.source_collection_id
+                        && node.version == dependency.source_collection_version
+                }) {
+                    Some(source_node) => source_node,
+                    None => {
+                        let is_source_soft_deleted = collection_id_soft_delete
+                            .get(&dependency.source_collection_id)
+                            .cloned()
+                            .unwrap_or(false);
+
+                        if !is_source_soft_deleted {
+                            return Err(ConstructVersionGraphError::ExpectedNodeNotFound(
+                                dependency.source_collection_id,
+                                Some(dependency.source_collection_version),
+                            ));
+                        }
+
+                        // Insert new node
+                        graph.add_node(VersionGraphNode {
+                            collection_id: dependency.source_collection_id,
+                            version: dependency.source_collection_version,
+                            status: VersionStatus::Deleted,
+                        })
+                    }
+                };
+
+                let target_node = match graph.node_indices().find(|n| {
+                    let node = graph.node_weight(*n).expect("node index should exist");
+                    node.collection_id == dependency.target_collection_id
+                }) {
+                    Some(target_node) => target_node,
+                    None => {
+                        // Insert new node
+                        graph.add_node(VersionGraphNode {
+                            collection_id: dependency.target_collection_id,
+                            version: 0,
+                            status: VersionStatus::Deleted,
+                        })
+                    }
+                };
 
                 graph.add_edge(source_node, target_node, ());
             }
@@ -308,6 +350,14 @@ impl ConstructVersionGraphOrchestrator {
 
             tracing::trace!("Version files: {:#?}", self.version_files);
 
+            let components = petgraph::algo::connected_components(&graph);
+            if components != 1 {
+                return Err(ConstructVersionGraphError::InvariantViolation(format!(
+                    "Graph is not fully connected, found {} components",
+                    components
+                )));
+            }
+
             self.terminate_with_result(
                 Ok(ConstructVersionGraphResponse {
                     graph,

rust/garbage_collector/src/garbage_collector_orchestrator_v2.rs

@@ -32,7 +32,7 @@ use chroma_system::{
     PanicError, System, TaskError, TaskResult,
 };
 use chroma_types::chroma_proto::{CollectionVersionFile, VersionListForCollection};
-use chroma_types::CollectionUuid;
+use chroma_types::{CollectionUuid, DeleteCollectionError};
 use chrono::{DateTime, Utc};
 use std::collections::{HashMap, HashSet};
 use std::str::FromStr;
@@ -61,6 +61,9 @@ pub struct GarbageCollectorOrchestrator {
     num_pending_tasks: usize,
     min_versions_to_keep: u32,
     graph: Option<VersionGraph>,
+    soft_deleted_collections_to_gc: HashSet<CollectionUuid>,
+    tenant: Option<String>,
+    database_name: Option<String>,
 
     num_files_deleted: u32,
     num_versions_deleted: u32,
@@ -107,6 +110,9 @@ impl GarbageCollectorOrchestrator {
             num_pending_tasks: 0,
             min_versions_to_keep,
             graph: None,
+            soft_deleted_collections_to_gc: HashSet::new(),
+            tenant: None,
+            database_name: None,
 
             num_files_deleted: 0,
             num_versions_deleted: 0,
@@ -146,6 +152,8 @@ pub enum GarbageCollectorError {
     InvariantViolation(String),
     #[error("Could not parse UUID: {0}")]
     UnparsableUuid(#[from] uuid::Error),
+    #[error("Collection deletion failed: {0}")]
+    CollectionDeletionFailed(#[from] DeleteCollectionError),
 }
 
 impl ChromaError for GarbageCollectorError {
@@ -216,13 +224,35 @@ impl GarbageCollectorOrchestrator {
             self.lineage_file_path.clone(),
         );
         let output = orchestrator.run(self.system.clone()).await?;
+
+        let collection_ids = output.version_files.keys().cloned().collect::<Vec<_>>();
+        let soft_delete_statuses = self
+            .sysdb_client
+            .batch_get_collection_soft_delete_status(collection_ids)
+            .await
+            .unwrap(); // todo
+                       // todo: filter out root collection?
+        self.soft_deleted_collections_to_gc = soft_delete_statuses
+            .iter()
+            .filter_map(
+                |(collection_id, status)| {
+                    if *status {
+                        Some(*collection_id)
+                    } else {
+                        None
+                    }
+                },
+            )
+            .collect();
+
         self.version_files = output.version_files;
         self.graph = Some(output.graph.clone());
 
         let task = wrap(
             Box::new(ComputeVersionsToDeleteOperator {}),
             ComputeVersionsToDeleteInput {
                 graph: output.graph,
+                soft_deleted_collections: self.soft_deleted_collections_to_gc.clone(),
                 cutoff_time: self.absolute_cutoff_time,
                 min_versions_to_keep: self.min_versions_to_keep,
             },
@@ -262,6 +292,13 @@ impl GarbageCollectorOrchestrator {
             })
             .collect();
 
+        self.pending_mark_versions_at_sysdb_tasks = output
+            .versions
+            .keys()
+            .filter(|collection_id| !self.soft_deleted_collections_to_gc.contains(collection_id))
+            .cloned()
+            .collect();
+
         for (collection_id, versions) in &output.versions {
             let version_file = self
                 .version_files
@@ -398,7 +435,7 @@ impl GarbageCollectorOrchestrator {
                     output.collection_id, output.version
                 )))?;
 
-            let root = graph
+            let root_index = graph
                 .node_indices()
                 .find(|&n| {
                     graph
@@ -409,20 +446,26 @@ impl GarbageCollectorOrchestrator {
                 .ok_or(GarbageCollectorError::InvariantViolation(
                     "Expected to find root node".to_string(),
                 ))?;
-
-            let versions_from_root_to_this_node =
-                petgraph::algo::astar(graph, root, |finish| finish == this_node, |_| 1, |_| 0)
-                    .ok_or(GarbageCollectorError::InvariantViolation(format!(
-                        "Expected to find path from root to node for {}@v{}",
-                        output.collection_id, output.version
-                    )))?
-                    .1
-                    .into_iter()
-                    .map(|i| {
-                        let node = graph.node_weight(i).expect("Node should exist");
-                        node.version
-                    })
-                    .collect::<Vec<_>>();
+            let root = graph.node_weight(root_index).expect("Node should exist");
+
+            let versions_from_root_to_this_node = petgraph::algo::astar(
+                graph,
+                root_index,
+                |finish| finish == this_node,
+                |_| 1,
+                |_| 0,
+            )
+            .ok_or(GarbageCollectorError::InvariantViolation(format!(
+                "Expected to find path from root ({}@v{}) to node for {}@v{}",
+                root.collection_id, root.version, output.collection_id, output.version
+            )))?
+            .1
+            .into_iter()
+            .map(|i| {
+                let node = graph.node_weight(i).expect("Node should exist");
+                node.version
+            })
+            .collect::<Vec<_>>();
             let are_all_versions_v0 = versions_from_root_to_this_node
                 .iter()
                 .all(|&version| version == 0);
@@ -523,14 +566,15 @@ impl GarbageCollectorOrchestrator {
                     "Expected there to be at least one version file".to_string(),
                 ))?;
         // Assumes that all collections in a fork tree are under the same tenant
-        let tenant_id = version_file
-            .collection_info_immutable
-            .as_ref()
-            .ok_or(GarbageCollectorError::InvariantViolation(
+        let collection_info = version_file.collection_info_immutable.as_ref().ok_or(
+            GarbageCollectorError::InvariantViolation(
                 "Expected collection_info_immutable to be set".to_string(),
-            ))?
-            .tenant_id
-            .clone();
+            ),
+        )?;
+        let tenant_id = collection_info.tenant_id.clone();
+        self.tenant = Some(tenant_id.clone());
+        let database_name = collection_info.database_name.clone();
+        self.database_name = Some(database_name.clone());
 
         let task = wrap(
             Box::new(DeleteUnusedFilesOperator::new(
@@ -659,6 +703,91 @@ impl GarbageCollectorOrchestrator {
 
         Ok(())
     }
+
+    async fn handle_delete_versions_output(
+        &mut self,
+        output: DeleteVersionsAtSysDbOutput,
+        ctx: &ComponentContext<Self>,
+    ) -> Result<(), GarbageCollectorError> {
+        tracing::trace!("Received DeleteVersionsAtSysDbOutput: {:#?}", output);
+        self.num_versions_deleted += output.versions_to_delete.versions.len() as u32;
+
+        self.num_pending_tasks -= 1;
+        if self.num_pending_tasks == 0 {
+            for collection_id in self.soft_deleted_collections_to_gc.drain() {
+                let graph =
+                    self.graph
+                        .as_ref()
+                        .ok_or(GarbageCollectorError::InvariantViolation(
+                            "Expected graph to be set".to_string(),
+                        ))?;
+
+                let first_collection_node = graph
+                    .node_indices()
+                    .filter(|&n| {
+                        let node = graph.node_weight(n).expect("Node should exist");
+                        node.collection_id == collection_id
+                    })
+                    .max_by(|a, b| {
+                        // todo: does this sort in the right direction?
+                        let a_node = graph.node_weight(*a).expect("Node should exist");
+                        let b_node = graph.node_weight(*b).expect("Node should exist");
+                        b_node.version.cmp(&a_node.version)
+                    })
+                    .ok_or(GarbageCollectorError::InvariantViolation(format!(
+                        "Expected to find node for collection {}",
+                        collection_id
+                    )))?;
+
+                let first_node = graph
+                    .node_weight(first_collection_node)
+                    .expect("Node should exist");
+                tracing::debug!(
+                    "First node for collection {}: {:#?}",
+                    collection_id,
+                    first_node
+                );
+
+                let mut dfs = petgraph::visit::Dfs::new(graph, first_collection_node);
+                let mut seen_collection_ids: HashSet<CollectionUuid> = HashSet::new();
+
+                while let Some(nx) = dfs.next(&graph) {
+                    let node = graph.node_weight(nx).expect("Node should exist");
+                    seen_collection_ids.insert(node.collection_id);
+                }
+
+                let is_leaf_node_in_fork_tree =
+                    seen_collection_ids.len() == 1 && seen_collection_ids.contains(&collection_id);
+
+                if is_leaf_node_in_fork_tree {
+                    self.sysdb_client
+                        .finish_collection_deletion(
+                            self.tenant.clone().ok_or(
+                                GarbageCollectorError::InvariantViolation(
+                                    "Expected tenant to be set".to_string(),
+                                ),
+                            )?,
+                            self.database_name.clone().ok_or(
+                                GarbageCollectorError::InvariantViolation(
+                                    "Expected database to be set".to_string(),
+                                ),
+                            )?,
+                            collection_id,
+                        )
+                        .await?;
+                }
+            }
+
+            let response = GarbageCollectorResponse {
+                num_files_deleted: self.num_files_deleted,
+                num_versions_deleted: self.num_versions_deleted,
+            };
+
+            self.terminate_with_result(Ok(response), ctx).await;
+        }
+
+        Ok(())
+    }
 }
 
 #[async_trait]
@@ -774,18 +903,9 @@ impl Handler<TaskResult<DeleteVersionsAtSysDbOutput, DeleteVersionsAtSysDbError>
             Some(output) => output,
             None => return,
         };
-        tracing::trace!("Received DeleteVersionsAtSysDbOutput: {:#?}", output);
-        self.num_versions_deleted += output.versions_to_delete.versions.len() as u32;
-
-        self.num_pending_tasks -= 1;
-        if self.num_pending_tasks == 0 {
-            let response = GarbageCollectorResponse {
-                num_files_deleted: self.num_files_deleted,
-                num_versions_deleted: self.num_versions_deleted,
-            };
 
-            self.terminate_with_result(Ok(response), ctx).await;
-        }
+        let res = self.handle_delete_versions_output(output, ctx).await;
+        self.ok_or_terminate(res, ctx).await;
     }
 }