Wire up RFE

Author: Sicheng Pan

Diff for: rust/frontend/src/get_collection_with_segments_provider.rs

@@ -151,43 +151,55 @@ impl CollectionsWithSegmentsProvider {
                 return Ok(collection_and_segments_with_ttl.collection_and_segments);
             }
         }
-        // We acquire a lock to prevent the sysdb from experiencing a thundering herd.
-        // This can happen when a large number of threads try to get the same collection
-        // at the same time.
-        let _guard = self.sysdb_rpc_lock.lock(&collection_id).await;
-        // Double checked locking pattern to avoid lock contention in the
-        // happy path when the collection is already cached.
-        if let Some(collection_and_segments_with_ttl) = self
-            .collections_with_segments_cache
-            .get(&collection_id)
-            .await?
-        {
-            if collection_and_segments_with_ttl.expires_at
-                > SystemTime::now()
-                    .duration_since(UNIX_EPOCH)
-                    .expect("Do not deploy before UNIX epoch")
+
+        let collection_and_segments_sysdb = {
+            // We acquire a lock to prevent the sysdb from experiencing a thundering herd.
+            // This can happen when a large number of threads try to get the same collection
+            // at the same time.
+            let _guard = self.sysdb_rpc_lock.lock(&collection_id).await;
+            // Double checked locking pattern to avoid lock contention in the
+            // happy path when the collection is already cached.
+            if let Some(collection_and_segments_with_ttl) = self
+                .collections_with_segments_cache
+                .get(&collection_id)
+                .await?
             {
-                return Ok(collection_and_segments_with_ttl.collection_and_segments);
+                if collection_and_segments_with_ttl.expires_at
+                    > SystemTime::now()
+                        .duration_since(UNIX_EPOCH)
+                        .expect("Do not deploy before UNIX epoch")
+                {
+                    return Ok(collection_and_segments_with_ttl.collection_and_segments);
+                }
             }
-        }
-        tracing::info!("Cache miss for collection {}", collection_id);
-        let collection_and_segments_sysdb = self
-            .sysdb_client
-            .get_collection_with_segments(collection_id)
-            .await?;
-        let collection_and_segments_sysdb_with_ttl = CollectionAndSegmentsWithTtl {
-            collection_and_segments: collection_and_segments_sysdb.clone(),
-            expires_at: SystemTime::now()
-                .duration_since(UNIX_EPOCH)
-                .expect("Do not deploy before UNIX epoch")
-                + Duration::from_secs(self.cache_ttl_secs as u64), // Cache for 1 minute
+            tracing::info!("Cache miss for collection {}", collection_id);
+            self.sysdb_client
+                .get_collection_with_segments(collection_id)
+                .await?
         };
+
+        self.set_collection_with_segments(collection_and_segments_sysdb.clone())
+            .await;
+        Ok(collection_and_segments_sysdb)
+    }
+
+    pub(crate) async fn set_collection_with_segments(
+        &mut self,
+        collection_and_segments: CollectionAndSegments,
+    ) {
         // Insert only if the collection dimension is set.
-        if collection_and_segments_sysdb.collection.dimension.is_some() {
+        if collection_and_segments.collection.dimension.is_some() {
+            let collection_id = collection_and_segments.collection.collection_id;
+            let collection_and_segments_with_ttl = CollectionAndSegmentsWithTtl {
+                collection_and_segments,
+                expires_at: SystemTime::now()
+                    .duration_since(UNIX_EPOCH)
+                    .expect("Do not deploy before UNIX epoch")
+                    + Duration::from_secs(self.cache_ttl_secs as u64), // Cache for 1 minute
+            };
             self.collections_with_segments_cache
-                .insert(collection_id, collection_and_segments_sysdb_with_ttl)
+                .insert(collection_id, collection_and_segments_with_ttl)
                 .await;
         }
-        Ok(collection_and_segments_sysdb)
     }
 }

Diff for: rust/frontend/src/impls/service_based_frontend.rs

@@ -21,13 +21,14 @@ use chroma_types::{
     CreateTenantError, CreateTenantRequest, CreateTenantResponse, DeleteCollectionError,
     DeleteCollectionRecordsError, DeleteCollectionRecordsRequest, DeleteCollectionRecordsResponse,
     DeleteCollectionRequest, DeleteDatabaseError, DeleteDatabaseRequest, DeleteDatabaseResponse,
-    GetCollectionError, GetCollectionRequest, GetCollectionResponse, GetCollectionsError,
-    GetDatabaseError, GetDatabaseRequest, GetDatabaseResponse, GetRequest, GetResponse,
-    GetTenantError, GetTenantRequest, GetTenantResponse, HealthCheckResponse, HeartbeatError,
-    HeartbeatResponse, Include, KnnIndex, ListCollectionsRequest, ListCollectionsResponse,
-    ListDatabasesError, ListDatabasesRequest, ListDatabasesResponse, Operation, OperationRecord,
-    QueryError, QueryRequest, QueryResponse, ResetError, ResetResponse, Segment, SegmentScope,
-    SegmentType, SegmentUuid, UpdateCollectionError, UpdateCollectionRecordsError,
+    ForkCollectionError, ForkCollectionRequest, ForkCollectionResponse, GetCollectionError,
+    GetCollectionRequest, GetCollectionResponse, GetCollectionsError, GetDatabaseError,
+    GetDatabaseRequest, GetDatabaseResponse, GetRequest, GetResponse, GetTenantError,
+    GetTenantRequest, GetTenantResponse, HealthCheckResponse, HeartbeatError, HeartbeatResponse,
+    Include, KnnIndex, ListCollectionsRequest, ListCollectionsResponse, ListDatabasesError,
+    ListDatabasesRequest, ListDatabasesResponse, Operation, OperationRecord, QueryError,
+    QueryRequest, QueryResponse, ResetError, ResetResponse, Segment, SegmentScope, SegmentType,
+    SegmentUuid, UpdateCollectionError, UpdateCollectionRecordsError,
     UpdateCollectionRecordsRequest, UpdateCollectionRecordsResponse, UpdateCollectionRequest,
     UpdateCollectionResponse, UpsertCollectionRecordsError, UpsertCollectionRecordsRequest,
     UpsertCollectionRecordsResponse, VectorIndexConfiguration, Where,
@@ -542,6 +543,33 @@ impl ServiceBasedFrontend {
         Ok(DeleteCollectionRecordsResponse {})
     }
 
+    pub async fn fork_collection(
+        &mut self,
+        ForkCollectionRequest {
+            source_collection_id,
+            target_collection_name,
+            ..
+        }: ForkCollectionRequest,
+    ) -> Result<ForkCollectionResponse, ForkCollectionError> {
+        let target_collection_id = CollectionUuid::new();
+        let collection_and_segments = self
+            .sysdb_client
+            .fork_collection(
+                source_collection_id,
+                target_collection_id,
+                target_collection_name,
+            )
+            .await?;
+        let collection = collection_and_segments.collection.clone();
+
+        // Update the cache.
+        self.collections_with_segments_provider
+            .set_collection_with_segments(collection_and_segments)
+            .await;
+
+        Ok(collection)
+    }
+
     pub async fn add(
         &mut self,
         AddCollectionRecordsRequest {

Diff for: rust/types/src/api_types.rs

@@ -767,6 +767,35 @@ impl ChromaError for DeleteCollectionError {
     }
 }
 
+#[non_exhaustive]
+#[derive(Clone, Validate, Serialize, ToSchema)]
+pub struct ForkCollectionRequest {
+    pub tenant_id: String,
+    pub database_name: String,
+    pub source_collection_id: CollectionUuid,
+    pub target_collection_name: String,
+}
+
+impl ForkCollectionRequest {
+    pub fn try_new(
+        tenant_id: String,
+        database_name: String,
+        source_collection_id: CollectionUuid,
+        target_collection_name: String,
+    ) -> Result<Self, ChromaValidationError> {
+        let request = Self {
+            tenant_id,
+            database_name,
+            source_collection_id,
+            target_collection_name,
+        };
+        request.validate().map_err(ChromaValidationError::from)?;
+        Ok(request)
+    }
+}
+
+pub type ForkCollectionResponse = Collection;
+
 #[derive(Error, Debug)]
 pub enum ForkCollectionError {
     #[error("Collection [{0}] already exists")]