[ENH]: de-dupe block path prefixes during GC to reduce memory usage (#5766)

Author: codetheweb

rust/garbage_collector/src/garbage_collector_orchestrator_v2.rs

@@ -21,8 +21,8 @@ use crate::operators::list_files_at_version::{
     ListFilesAtVersionsOperator,
 };
 use crate::types::{
-    version_graph_to_collection_dependency_graph, CleanupMode, GarbageCollectorResponse,
-    VersionGraph, VersionGraphNode,
+    version_graph_to_collection_dependency_graph, CleanupMode, FilePathRefCountSet,
+    GarbageCollectorResponse, VersionGraph, VersionGraphNode,
 };
 use async_trait::async_trait;
 use chroma_blockstore::RootManager;
@@ -68,7 +68,7 @@ pub struct GarbageCollectorOrchestrator {
     versions_to_delete_output: Option<ComputeVersionsToDeleteOutput>,
     delete_unused_file_output: Option<DeleteUnusedFilesOutput>,
     delete_unused_log_output: Option<DeleteUnusedLogsOutput>,
-    file_ref_counts: HashMap<String, u32>,
+    file_ref_counts: FilePathRefCountSet,
     num_pending_tasks: usize,
     min_versions_to_keep: u32,
     enable_log_gc: bool,
@@ -119,7 +119,7 @@ impl GarbageCollectorOrchestrator {
             cleanup_mode,
             result_channel: None,
             version_files: HashMap::new(),
-            file_ref_counts: HashMap::new(),
+            file_ref_counts: FilePathRefCountSet::new(),
             versions_to_delete_output: None,
             delete_unused_file_output: None,
             delete_unused_log_output: None,
@@ -705,10 +705,8 @@ impl GarbageCollectorOrchestrator {
                     output.version
                 );
 
-                for file_path in output.file_paths {
-                    let count = self.file_ref_counts.entry(file_path).or_insert(0);
-                    *count += 1;
-                }
+                self.file_ref_counts
+                    .merge(FilePathRefCountSet::from_set(output.file_paths, 1));
             }
             CollectionVersionAction::Delete => {
                 tracing::debug!(
@@ -718,9 +716,8 @@ impl GarbageCollectorOrchestrator {
                     output.version
                 );
 
-                for file_path in output.file_paths {
-                    self.file_ref_counts.entry(file_path).or_insert(0);
-                }
+                self.file_ref_counts
+                    .merge(FilePathRefCountSet::from_set(output.file_paths, 0));
             }
         }
 
@@ -733,17 +730,8 @@ impl GarbageCollectorOrchestrator {
     ) -> Result<(), GarbageCollectorError> {
         // We now have results for all ListFilesAtVersionsOperator tasks that we spawned
         tracing::trace!("File ref counts: {:#?}", self.file_ref_counts);
-        let file_paths_to_delete = self
-            .file_ref_counts
-            .iter()
-            .filter_map(|(path, count)| {
-                if *count == 0 {
-                    Some(path.clone())
-                } else {
-                    None
-                }
-            })
-            .collect::<Vec<_>>();
+
+        let file_paths_to_delete = self.file_ref_counts.as_set(0);
 
         let delete_percentage =
             file_paths_to_delete.len() as f32 / self.file_ref_counts.len() as f32 * 100.0;
@@ -785,7 +773,6 @@ impl GarbageCollectorOrchestrator {
             )),
             DeleteUnusedFilesInput {
                 unused_s3_files: file_paths_to_delete,
-                hnsw_prefixes_for_deletion: vec![],
             },
             ctx.receiver(),
             self.context.task_cancellation_token.clone(),
@@ -822,7 +809,7 @@ impl GarbageCollectorOrchestrator {
             return Ok(());
         }
 
-        self.num_files_deleted += output.deleted_files.len() as u32;
+        self.num_files_deleted += output.num_files_deleted as u32;
 
         let versions_to_delete = self.versions_to_delete_output.as_ref().ok_or(
             GarbageCollectorError::InvariantViolation(

rust/garbage_collector/src/operators/delete_unused_files.rs

@@ -1,12 +1,13 @@
 use crate::types::CleanupMode;
+use crate::types::FilePathSet;
 use crate::types::RENAMED_FILE_PREFIX;
 use async_trait::async_trait;
 use chroma_error::{ChromaError, ErrorCodes};
 use chroma_storage::Storage;
 use chroma_storage::StorageError;
 use chroma_system::{Operator, OperatorType};
 use futures::stream::StreamExt;
-use std::collections::HashSet;
+use futures::TryStreamExt;
 use thiserror::Error;
 
 #[derive(Clone, Debug)]
@@ -39,13 +40,12 @@ impl DeleteUnusedFilesOperator {
 
 #[derive(Debug)]
 pub struct DeleteUnusedFilesInput {
-    pub unused_s3_files: Vec<String>,
-    pub hnsw_prefixes_for_deletion: Vec<String>,
+    pub unused_s3_files: FilePathSet,
 }
 
 #[derive(Debug)]
 pub struct DeleteUnusedFilesOutput {
-    pub deleted_files: HashSet<String>,
+    pub num_files_deleted: usize,
 }
 
 #[derive(Error, Debug)]
@@ -80,15 +80,10 @@ impl Operator<DeleteUnusedFilesInput, DeleteUnusedFilesOutput> for DeleteUnusedF
     ) -> Result<DeleteUnusedFilesOutput, DeleteUnusedFilesError> {
         tracing::debug!(
             files_count = input.unused_s3_files.len(),
-            hnsw_prefixes_count = input.hnsw_prefixes_for_deletion.len(),
             cleanup_mode = ?self.cleanup_mode,
             "Starting deletion of unused files"
         );
 
-        // Create a list that contains all files that will be deleted.
-        let mut all_files = input.unused_s3_files.clone();
-        all_files.extend(input.hnsw_prefixes_for_deletion.clone());
-
         // NOTE(rohit):
         // We don't want to fail the entire operation if one file fails to rename or delete.
         // It's possible that the file was already renamed/deleted in the last run that
@@ -97,8 +92,8 @@ impl Operator<DeleteUnusedFilesInput, DeleteUnusedFilesOutput> for DeleteUnusedF
             CleanupMode::DryRunV2 => {}
             CleanupMode::Rename => {
                 // Soft delete - rename the file
-                if !all_files.is_empty() {
-                    let mut rename_stream = futures::stream::iter(all_files.clone())
+                if !input.unused_s3_files.is_empty() {
+                    let mut rename_stream = futures::stream::iter(input.unused_s3_files.iter())
                         .map(move |file_path| {
                             let new_path = self.get_rename_path(&file_path);
                             self.rename_with_path(file_path, new_path)
@@ -123,18 +118,22 @@ impl Operator<DeleteUnusedFilesInput, DeleteUnusedFilesOutput> for DeleteUnusedF
             }
             CleanupMode::DeleteV2 => {
                 // Hard delete - remove the file
-                if !all_files.is_empty() {
-                    // The S3 DeleteObjects API allows up to 1000 objects per request
-                    for chunk in all_files.chunks(1000) {
-                        let result = self.storage.delete_many(chunk).await?;
-                        if !result.errors.is_empty() {
+                if !input.unused_s3_files.is_empty() {
+                    let mut delete_stream = futures::stream::iter(input.unused_s3_files.iter())
+                        // The S3 DeleteObjects API allows up to 1000 objects per request
+                        .chunks(1000)
+                        .then(|chunk| async move { self.storage.delete_many(chunk).await })
+                        .boxed();
+
+                    while let Some(delete_result) = delete_stream.try_next().await? {
+                        if !delete_result.errors.is_empty() {
                             // Log the errors but don't fail the operation
-                            for error in result.errors {
+                            for error in delete_result.errors {
                                 match error {
                                     StorageError::NotFound { path, source } => {
-                                        tracing::warn!("Delete file {path} not found: {source}")
+                                        tracing::info!("Delete file {path} not found: {source}")
                                     }
-                                    err => return Err(DeleteUnusedFilesError::StorageError(err)),
+                                    err => tracing::error!("Failed to delete: {err}"),
                                 }
                             }
                         }
@@ -144,13 +143,11 @@ impl Operator<DeleteUnusedFilesInput, DeleteUnusedFilesOutput> for DeleteUnusedF
         }
 
         Ok(DeleteUnusedFilesOutput {
-            deleted_files: all_files.into_iter().collect(),
+            num_files_deleted: input.unused_s3_files.len(),
         })
     }
 }
 
-impl DeleteUnusedFilesOperator {}
-
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -167,11 +164,14 @@ mod tests {
             .unwrap();
     }
 
-    async fn setup_test_files(storage: &Storage) -> (Vec<String>, Vec<String>) {
+    async fn setup_test_files(storage: &Storage) -> FilePathSet {
+        let mut set = FilePathSet::new();
+
         // Create regular test files
         let test_files = vec!["file1.txt", "file2.txt"];
         for file in &test_files {
             create_test_file(storage, file, b"test content").await;
+            set.insert_path(file.to_string());
         }
 
         // Create HNSW test files
@@ -181,40 +181,31 @@ mod tests {
             .collect::<Vec<String>>();
         for file in &hnsw_files {
             create_test_file(storage, file, b"test content").await;
+            set.insert_path(file.to_string());
         }
 
-        (
-            test_files.iter().map(|s| s.to_string()).collect(),
-            hnsw_files.iter().map(|s| s.to_string()).collect(),
-        )
+        set
     }
 
     #[tokio::test]
     async fn test_dry_run_mode() {
         let tmp_dir = TempDir::new().unwrap();
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
-        let (test_files, hnsw_files) = setup_test_files(&storage).await;
+        let unused_s3_files = setup_test_files(&storage).await;
 
         let operator = DeleteUnusedFilesOperator::new(
             storage.clone(),
             CleanupMode::DryRunV2,
             "test_tenant".to_string(),
         );
         let input = DeleteUnusedFilesInput {
-            unused_s3_files: test_files.clone(),
-            hnsw_prefixes_for_deletion: hnsw_files.clone(),
+            unused_s3_files: unused_s3_files.clone(),
         };
 
-        let result = operator.run(&input).await.unwrap();
+        operator.run(&input).await.unwrap();
 
         // Verify original files still exist
-        for file in &test_files {
-            assert!(result.deleted_files.contains(file));
-            assert!(Path::new(&tmp_dir.path().join(file)).exists());
-        }
-
-        for file in &hnsw_files {
-            assert!(result.deleted_files.contains(file));
+        for file in unused_s3_files.iter() {
             assert!(Path::new(&tmp_dir.path().join(file)).exists());
         }
     }
@@ -223,71 +214,50 @@ mod tests {
     async fn test_rename_mode() {
         let tmp_dir = TempDir::new().unwrap();
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
-        let (test_files, hnsw_files) = setup_test_files(&storage).await;
+        let unused_s3_files = setup_test_files(&storage).await;
 
         let operator = DeleteUnusedFilesOperator::new(
             storage.clone(),
             CleanupMode::Rename,
             "test_tenant".to_string(),
         );
         let input = DeleteUnusedFilesInput {
-            unused_s3_files: test_files.clone(),
-            hnsw_prefixes_for_deletion: hnsw_files.clone(),
+            unused_s3_files: unused_s3_files.clone(),
         };
 
-        let result = operator.run(&input).await.unwrap();
-
-        // Verify regular files were moved to deleted directory
-        for file in &test_files {
-            let original_path = tmp_dir.path().join(file);
-            let new_path = tmp_dir
-                .path()
-                .join(format!("{}{}/{}", RENAMED_FILE_PREFIX, "test_tenant", file));
-            assert!(!original_path.exists());
-            assert!(new_path.exists());
-            assert!(result.deleted_files.contains(file));
-        }
+        operator.run(&input).await.unwrap();
 
-        // Verify HNSW files were moved to deleted directory
-        for file in &hnsw_files {
-            let original_path = tmp_dir.path().join(file);
+        // Verify files were moved to deleted directory
+        for file in unused_s3_files.iter() {
+            let original_path = tmp_dir.path().join(&file);
             let new_path = tmp_dir
                 .path()
                 .join(format!("{}{}/{}", RENAMED_FILE_PREFIX, "test_tenant", file));
             assert!(!original_path.exists());
             assert!(new_path.exists());
-            assert!(result.deleted_files.contains(file));
         }
     }
 
     #[tokio::test]
     async fn test_delete_mode() {
         let tmp_dir = TempDir::new().unwrap();
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
-        let (test_files, hnsw_files) = setup_test_files(&storage).await;
+        let unused_s3_files = setup_test_files(&storage).await;
 
         let operator = DeleteUnusedFilesOperator::new(
             storage.clone(),
             CleanupMode::DeleteV2,
             "test_tenant".to_string(),
         );
         let input = DeleteUnusedFilesInput {
-            unused_s3_files: test_files.clone(),
-            hnsw_prefixes_for_deletion: hnsw_files.clone(),
+            unused_s3_files: unused_s3_files.clone(),
         };
 
-        let result = operator.run(&input).await.unwrap();
+        operator.run(&input).await.unwrap();
 
-        // Verify regular files were deleted
-        for file in &test_files {
-            assert!(!Path::new(&tmp_dir.path().join(file)).exists());
-            assert!(result.deleted_files.contains(file));
-        }
-
-        // Verify HNSW files were deleted
-        for file in &hnsw_files {
+        // Verify files were deleted
+        for file in unused_s3_files.iter() {
             assert!(!Path::new(&tmp_dir.path().join(file)).exists());
-            assert!(result.deleted_files.contains(file));
         }
     }
 
@@ -296,7 +266,8 @@ mod tests {
         let tmp_dir = TempDir::new().unwrap();
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
 
-        let unused_files = vec!["nonexistent.txt".to_string()];
+        let mut unused_files = FilePathSet::new();
+        unused_files.insert_path("nonexistent.txt".to_string());
 
         // Test Delete mode - should succeed but record the error in deletion list
         let delete_operator = DeleteUnusedFilesOperator::new(
@@ -307,7 +278,6 @@ mod tests {
         let result = delete_operator
             .run(&DeleteUnusedFilesInput {
                 unused_s3_files: unused_files.clone(),
-                hnsw_prefixes_for_deletion: vec![],
             })
             .await;
         assert!(result.is_ok());
@@ -321,7 +291,6 @@ mod tests {
         let result = rename_operator
             .run(&DeleteUnusedFilesInput {
                 unused_s3_files: unused_files.clone(),
-                hnsw_prefixes_for_deletion: vec![],
             })
             .await;
         assert!(result.is_ok());
@@ -335,7 +304,6 @@ mod tests {
         let result = list_operator
             .run(&DeleteUnusedFilesInput {
                 unused_s3_files: unused_files,
-                hnsw_prefixes_for_deletion: vec![],
             })
             .await;
         assert!(result.is_ok());