[HOTFIX] applying PR #4867 to release/2025-06-13 (#4875)

This PR cherry-picks the commit 6e5ec8f
onto release/2025-06-13. If there are unresolved conflicts, please
resolve them manually.

Co-authored-by: Max Isom <[email protected]>

Author: chroma-droid

rust/garbage_collector/src/construct_version_graph_orchestrator.rs

@@ -30,6 +30,7 @@ use petgraph::{dot::Dot, graph::DiGraph};
 use std::{
     collections::{HashMap, HashSet},
     str::FromStr,
+    sync::Arc,
 };
 use thiserror::Error;
 use tokio::sync::oneshot::{error::RecvError, Sender};
@@ -55,7 +56,7 @@ pub struct ConstructVersionGraphOrchestrator {
     lineage_file_path: Option<String>,
 
     version_dependencies: Vec<VersionDependency>,
-    version_files: HashMap<CollectionUuid, CollectionVersionFile>,
+    version_files: HashMap<CollectionUuid, Arc<CollectionVersionFile>>,
     num_pending_tasks: usize,
 }
 
@@ -88,7 +89,7 @@ impl ConstructVersionGraphOrchestrator {
 #[derive(Debug)]
 #[allow(dead_code)]
 pub struct ConstructVersionGraphResponse {
-    pub version_files: HashMap<CollectionUuid, CollectionVersionFile>,
+    pub version_files: HashMap<CollectionUuid, Arc<CollectionVersionFile>>,
     pub graph: VersionGraph,
 }
 

rust/garbage_collector/src/garbage_collector_orchestrator.rs

@@ -37,8 +37,6 @@
 //!    - Input: Version file, versions to delete, unused S3 files
 //!    - Output: Deletion confirmation
 
-use std::fmt::{Debug, Formatter};
-
 use crate::types::{CleanupMode, GarbageCollectorResponse};
 use async_trait::async_trait;
 use chroma_error::{ChromaError, ErrorCodes};
@@ -51,6 +49,8 @@ use chroma_system::{
 use chroma_types::chroma_proto::CollectionVersionFile;
 use chroma_types::CollectionUuid;
 use chrono::{DateTime, Utc};
+use std::fmt::{Debug, Formatter};
+use std::sync::Arc;
 use thiserror::Error;
 use tokio::sync::oneshot::{error::RecvError, Sender};
 use tracing::Span;
@@ -87,7 +87,7 @@ pub struct GarbageCollectorOrchestrator {
     dispatcher: ComponentHandle<Dispatcher>,
     storage: Storage,
     result_channel: Option<Sender<Result<GarbageCollectorResponse, GarbageCollectorError>>>,
-    pending_version_file: Option<CollectionVersionFile>,
+    pending_version_file: Option<Arc<CollectionVersionFile>>,
     pending_versions_to_delete: Option<chroma_types::chroma_proto::VersionListForCollection>,
     pending_epoch_id: Option<i64>,
     num_versions_deleted: u32,

rust/garbage_collector/src/garbage_collector_orchestrator_v2.rs

@@ -40,6 +40,7 @@ use chrono::{DateTime, Utc};
 use petgraph::algo::toposort;
 use std::collections::{HashMap, HashSet};
 use std::str::FromStr;
+use std::sync::Arc;
 use std::time::SystemTime;
 use thiserror::Error;
 use tokio::sync::oneshot::{error::RecvError, Sender};
@@ -59,7 +60,7 @@ pub struct GarbageCollectorOrchestrator {
     root_manager: RootManager,
     result_channel: Option<Sender<Result<GarbageCollectorResponse, GarbageCollectorError>>>,
     cleanup_mode: CleanupMode,
-    version_files: HashMap<CollectionUuid, CollectionVersionFile>,
+    version_files: HashMap<CollectionUuid, Arc<CollectionVersionFile>>,
     versions_to_delete_output: Option<ComputeVersionsToDeleteOutput>,
     pending_mark_versions_at_sysdb_tasks: HashSet<CollectionUuid>,
     pending_list_files_at_version_tasks: HashSet<(CollectionUuid, i64)>,

rust/garbage_collector/src/operators/compute_unused_files.rs

@@ -8,7 +8,10 @@ use chroma_types::{
     chroma_proto::{CollectionSegmentInfo, CollectionVersionFile, VersionListForCollection},
     HNSW_PATH,
 };
-use std::collections::{HashMap, HashSet};
+use std::{
+    collections::{HashMap, HashSet},
+    sync::Arc,
+};
 use thiserror::Error;
 use uuid::Uuid;
 
@@ -171,7 +174,7 @@ impl ComputeUnusedFilesOperator {
 
 #[derive(Clone)]
 pub struct ComputeUnusedFilesInput {
-    pub version_file: CollectionVersionFile,
+    pub version_file: Arc<CollectionVersionFile>,
     pub versions_to_delete: VersionListForCollection,
     pub oldest_version_to_keep: i64,
 }
@@ -239,7 +242,9 @@ impl Operator<ComputeUnusedFilesInput, ComputeUnusedFilesOutput> for ComputeUnus
         // Build a map to version to segment_info
         let mut version_to_segment_info = HashMap::new();
         let version_history = version_file
+            .as_ref()
             .version_history
+            .as_ref()
             .ok_or(ComputeUnusedFilesError::MissingVersionHistory)?;
 
         // Check if version history is empty
@@ -465,7 +470,7 @@ mod tests {
 
         let input = ComputeUnusedFilesInput {
             oldest_version_to_keep: 3,
-            version_file: chroma_proto::CollectionVersionFile {
+            version_file: Arc::new(chroma_proto::CollectionVersionFile {
                 collection_info_immutable: None,
                 version_history: Some(CollectionVersionHistory {
                     versions: vec![
@@ -520,7 +525,7 @@ mod tests {
                         },
                     ],
                 }),
-            },
+            }),
             versions_to_delete: chroma_proto::VersionListForCollection {
                 versions: vec![1, 2],
                 collection_id: "test_collection".to_string(),
@@ -570,12 +575,12 @@ mod tests {
 
         // Create input with missing version info
         let input = ComputeUnusedFilesInput {
-            version_file: chroma_proto::CollectionVersionFile {
+            version_file: Arc::new(chroma_proto::CollectionVersionFile {
                 collection_info_immutable: None,
                 version_history: Some(CollectionVersionHistory {
                     versions: vec![], // Empty version history wrapped in Some
                 }),
-            },
+            }),
             versions_to_delete: chroma_proto::VersionListForCollection {
                 versions: vec![1, 2], // Versions that don't exist in history
                 collection_id: "test_collection".to_string(),
@@ -606,7 +611,7 @@ mod tests {
         );
 
         let input = ComputeUnusedFilesInput {
-            version_file: chroma_proto::CollectionVersionFile {
+            version_file: Arc::new(chroma_proto::CollectionVersionFile {
                 collection_info_immutable: None,
                 version_history: Some(CollectionVersionHistory {
                     versions: vec![
@@ -639,7 +644,7 @@ mod tests {
                         },
                     ],
                 }),
-            },
+            }),
             versions_to_delete: chroma_proto::VersionListForCollection {
                 versions: vec![1, 2], // Try to delete 2 versions
                 collection_id: "test_collection".to_string(),

rust/garbage_collector/src/operators/compute_versions_to_delete.rs

@@ -4,22 +4,22 @@ use chroma_system::{Operator, OperatorType};
 use chroma_types::chroma_proto::{CollectionVersionFile, VersionListForCollection};
 use chrono::{DateTime, Utc};
 use humantime::format_duration;
-use std::time::Duration;
+use std::{sync::Arc, time::Duration};
 use thiserror::Error;
 
 #[derive(Clone, Debug)]
 pub struct ComputeVersionsToDeleteOperator {}
 
 #[derive(Debug)]
 pub struct ComputeVersionsToDeleteInput {
-    pub version_file: CollectionVersionFile,
+    pub version_file: Arc<CollectionVersionFile>,
     pub cutoff_time: DateTime<Utc>,
     pub min_versions_to_keep: u32,
 }
 
 #[derive(Debug)]
 pub struct ComputeVersionsToDeleteOutput {
-    pub version_file: CollectionVersionFile,
+    pub version_file: Arc<CollectionVersionFile>,
     pub versions_to_delete: VersionListForCollection,
     pub oldest_version_to_keep: i64,
 }
@@ -54,7 +54,7 @@ impl Operator<ComputeVersionsToDeleteInput, ComputeVersionsToDeleteOutput>
         &self,
         input: &ComputeVersionsToDeleteInput,
     ) -> Result<ComputeVersionsToDeleteOutput, ComputeVersionsToDeleteError> {
-        let mut version_file = input.version_file.clone();
+        let version_file = input.version_file.clone();
         let collection_info = version_file
             .collection_info_immutable
             .as_ref()
@@ -73,6 +73,7 @@ impl Operator<ComputeVersionsToDeleteInput, ComputeVersionsToDeleteOutput>
         let mut marked_versions = Vec::new();
         let mut oldest_version_to_keep = 0;
 
+        let mut version_file = version_file.as_ref().clone();
         if let Some(ref mut version_history) = version_file.version_history {
             let mut unique_versions_seen = 0;
             let mut last_version = None;
@@ -149,7 +150,7 @@ impl Operator<ComputeVersionsToDeleteInput, ComputeVersionsToDeleteOutput>
         );
 
         Ok(ComputeVersionsToDeleteOutput {
-            version_file,
+            version_file: Arc::new(version_file),
             versions_to_delete,
             oldest_version_to_keep,
         })
@@ -198,7 +199,7 @@ mod tests {
             ],
         };
 
-        let version_file = CollectionVersionFile {
+        let version_file = Arc::new(CollectionVersionFile {
             version_history: Some(version_history),
             collection_info_immutable: Some(CollectionInfoImmutable {
                 tenant_id: "test_tenant".to_string(),
@@ -207,7 +208,7 @@ mod tests {
                 dimension: 0,
                 ..Default::default()
             }),
-        };
+        });
 
         let input = ComputeVersionsToDeleteInput {
             version_file,
@@ -221,7 +222,12 @@ mod tests {
             .unwrap();
 
         // Verify the results.
-        let versions = &result.version_file.version_history.unwrap().versions;
+        let versions = &result
+            .version_file
+            .version_history
+            .as_ref()
+            .unwrap()
+            .versions;
         assert!(versions[0].marked_for_deletion);
         assert!(versions[1].marked_for_deletion);
         assert!(!versions[2].marked_for_deletion); // Version 2 should be kept.

rust/garbage_collector/src/operators/delete_versions_at_sysdb.rs

@@ -5,7 +5,7 @@ use chroma_sysdb::SysDb;
 use chroma_system::{Operator, OperatorType};
 use chroma_types::chroma_proto::{CollectionVersionFile, VersionListForCollection};
 use futures::stream::StreamExt;
-use std::collections::HashSet;
+use std::{collections::HashSet, sync::Arc};
 use thiserror::Error;
 
 #[derive(Clone)]
@@ -22,7 +22,7 @@ impl std::fmt::Debug for DeleteVersionsAtSysDbOperator {
 
 #[derive(Debug)]
 pub struct DeleteVersionsAtSysDbInput {
-    pub version_file: CollectionVersionFile,
+    pub version_file: Arc<CollectionVersionFile>,
     pub epoch_id: i64,
     pub sysdb_client: SysDb,
     pub versions_to_delete: VersionListForCollection,
@@ -31,7 +31,7 @@ pub struct DeleteVersionsAtSysDbInput {
 
 #[derive(Debug)]
 pub struct DeleteVersionsAtSysDbOutput {
-    pub version_file: CollectionVersionFile,
+    pub version_file: Arc<CollectionVersionFile>,
     pub versions_to_delete: VersionListForCollection,
     pub unused_s3_files: HashSet<String>,
 }
@@ -209,10 +209,10 @@ mod tests {
         let sysdb = SysDb::Test(TestSysDb::new());
 
         // Create a version file with actual version history
-        let version_file = CollectionVersionFile {
+        let version_file = Arc::new(CollectionVersionFile {
             version_history: Some(chroma_proto::CollectionVersionHistory { versions: vec![] }),
             ..Default::default()
-        };
+        });
 
         let versions_to_delete = VersionListForCollection {
             collection_id: "test_collection".to_string(),
@@ -245,7 +245,7 @@ mod tests {
         let tmp_dir = TempDir::new().unwrap();
         let storage = Storage::Local(LocalStorage::new(tmp_dir.path().to_str().unwrap()));
         let sysdb = SysDb::Test(TestSysDb::new());
-        let version_file = CollectionVersionFile::default();
+        let version_file = Arc::new(CollectionVersionFile::default());
         let versions_to_delete = VersionListForCollection {
             collection_id: "test_collection".to_string(),
             database_id: "default".to_string(),
@@ -368,7 +368,7 @@ mod tests {
         }
 
         // Create version file with history
-        let version_file = CollectionVersionFile {
+        let version_file = Arc::new(CollectionVersionFile {
             version_history: Some(chroma_proto::CollectionVersionHistory {
                 versions: vec![
                     chroma_proto::CollectionVersionInfo {
@@ -386,7 +386,7 @@ mod tests {
                 ],
             }),
             ..Default::default()
-        };
+        });
 
         let versions_to_delete = VersionListForCollection {
             collection_id: "test_collection".to_string(),
@@ -439,7 +439,7 @@ mod tests {
         }
 
         // Create version file with history
-        let version_file = CollectionVersionFile {
+        let version_file = Arc::new(CollectionVersionFile {
             version_history: Some(chroma_proto::CollectionVersionHistory {
                 versions: vec![
                     chroma_proto::CollectionVersionInfo {
@@ -457,7 +457,7 @@ mod tests {
                 ],
             }),
             ..Default::default()
-        };
+        });
 
         let versions_to_delete = VersionListForCollection {
             collection_id: "test_collection".to_string(),

rust/garbage_collector/src/operators/fetch_version_file.rs

@@ -17,6 +17,7 @@ use chroma_types::CollectionUuid;
 use prost::Message;
 use std::fmt::{Debug, Formatter};
 use std::str::FromStr;
+use std::sync::Arc;
 use thiserror::Error;
 
 #[derive(Clone, Debug, Default)]
@@ -53,7 +54,7 @@ impl Debug for FetchVersionFileInput {
 #[allow(dead_code)]
 #[derive(Debug)]
 pub struct FetchVersionFileOutput {
-    pub file: CollectionVersionFile,
+    pub file: Arc<CollectionVersionFile>,
     pub collection_id: CollectionUuid,
 }
 
@@ -128,7 +129,7 @@ impl Operator<FetchVersionFileInput, FetchVersionFileOutput> for FetchVersionFil
         .map_err(FetchVersionFileError::InvalidUuid)?;
 
         Ok(FetchVersionFileOutput {
-            file: version_file,
+            file: Arc::new(version_file),
             collection_id,
         })
     }
@@ -212,7 +213,7 @@ mod tests {
         let result = operator.run(&input).await.expect("Failed to run operator");
 
         // Verify the content
-        assert_eq!(result.file, test_file);
+        assert_eq!(result.file, test_file.into());
 
         // Cleanup - Note: object_store doesn't have a delete method,
         // but the test bucket should be cleaned up between test runs

rust/garbage_collector/src/operators/list_files_at_version.rs

@@ -3,22 +3,22 @@ use chroma_blockstore::{arrow::provider::RootManagerError, RootManager};
 use chroma_storage::StorageError;
 use chroma_system::{Operator, OperatorType};
 use chroma_types::{chroma_proto::CollectionVersionFile, CollectionUuid, HNSW_PATH};
-use std::{collections::HashSet, str::FromStr};
+use std::{collections::HashSet, str::FromStr, sync::Arc};
 use thiserror::Error;
 use tokio::task::{JoinError, JoinSet};
 use uuid::Uuid;
 
 #[derive(Debug)]
 pub struct ListFilesAtVersionInput {
     root_manager: RootManager,
-    version_file: CollectionVersionFile,
+    version_file: Arc<CollectionVersionFile>,
     version: i64,
 }
 
 impl ListFilesAtVersionInput {
     pub fn new(
         root_manager: RootManager,
-        version_file: CollectionVersionFile,
+        version_file: Arc<CollectionVersionFile>,
         version: i64,
     ) -> Self {
         Self {