[ENH] Implement log forking

Author: Sicheng Pan

Diff for: go/pkg/log/repository/log.go

@@ -126,6 +126,47 @@ func (r *LogRepository) PullRecords(ctx context.Context, collectionId string, of
 	return
 }
 
+func (r *LogRepository) ForkRecords(ctx context.Context, sourceCollectionID string, targetCollectionID string) (err error) {
+	var tx pgx.Tx
+	tx, err = r.conn.BeginTx(ctx, pgx.TxOptions{})
+	if err != nil {
+		trace_log.Error("Error in begin transaction for forking logs in log service", zap.Error(err), zap.String("sourceCollectionID", sourceCollectionID))
+		return
+	}
+	queriesWithTx := r.queries.WithTx(tx)
+	defer func() {
+		if err != nil {
+			tx.Rollback(ctx)
+		} else {
+			err = tx.Commit(ctx)
+		}
+	}()
+
+	err = queriesWithTx.LockCollection(ctx, sourceCollectionID)
+	if err != nil {
+		trace_log.Error("Error locking collection for fork", zap.String("sourceCollectionID", sourceCollectionID))
+		return
+	}
+	err = queriesWithTx.ForkCollectionOffset(ctx, log.ForkCollectionOffsetParams{
+		ID:   sourceCollectionID,
+		ID_2: targetCollectionID,
+	})
+	if err != nil {
+		trace_log.Error("Error forking log offset", zap.String("sourceCollectionID", sourceCollectionID))
+		return
+	}
+	err = queriesWithTx.ForkCollectionRecord(ctx, log.ForkCollectionRecordParams{
+		CollectionID:   sourceCollectionID,
+		CollectionID_2: targetCollectionID,
+	})
+	if err != nil {
+		trace_log.Error("Error forking log record", zap.String("sourceCollectionID", sourceCollectionID))
+		return
+	}
+
+	return
+}
+
 func (r *LogRepository) GetAllCollectionInfoToCompact(ctx context.Context, minCompactionSize uint64) (collectionToCompact []log.GetAllCollectionsToCompactRow, err error) {
 	collectionToCompact, err = r.queries.GetAllCollectionsToCompact(ctx, int64(minCompactionSize))
 	if collectionToCompact == nil {

Diff for: go/pkg/log/server/server.go

@@ -56,7 +56,7 @@ func (s *logServer) ScoutLogs(ctx context.Context, req *logservicepb.ScoutLogsRe
 		return
 	}
 	// +1 to convert from the (] bound to a [) bound.
-	res = &logservicepb.ScoutLogsResponse {
+	res = &logservicepb.ScoutLogsResponse{
 		FirstUninsertedRecordOffset: int64(limit + 1),
 	}
 	return
@@ -90,6 +90,27 @@ func (s *logServer) PullLogs(ctx context.Context, req *logservicepb.PullLogsRequ
 	return
 }
 
+func (s *logServer) ForkLogs(ctx context.Context, req *logservicepb.ForkLogsRequest) (res *logservicepb.ForkLogsResponse, err error) {
+	var sourceCollectionID types.UniqueID
+	var targetCollectionID types.UniqueID
+	sourceCollectionID, err = types.ToUniqueID(&req.SourceCollectionId)
+	if err != nil {
+		return
+	}
+	targetCollectionID, err = types.ToUniqueID(&req.TargetCollectionId)
+	if err != nil {
+		return
+	}
+
+	err = s.lr.ForkRecords(ctx, sourceCollectionID.String(), targetCollectionID.String())
+	if err != nil {
+		return
+	}
+
+	res = &logservicepb.ForkLogsResponse{}
+	return
+}
+
 func (s *logServer) GetAllCollectionInfoToCompact(ctx context.Context, req *logservicepb.GetAllCollectionInfoToCompactRequest) (res *logservicepb.GetAllCollectionInfoToCompactResponse, err error) {
 	var collectionToCompact []log.GetAllCollectionsToCompactRow
 	collectionToCompact, err = s.lr.GetAllCollectionInfoToCompact(ctx, req.MinCompactionSize)

Diff for: go/pkg/log/store/db/queries.sql.go

@@ -58,3 +58,18 @@ SELECT id FROM collection;
 
 -- name: GetLastCompactedOffset :one
 SELECT record_compaction_offset_position FROM collection c WHERE c.id = $1;
+
+-- name: LockCollection :exec
+SELECT * FROM collection WHERE id = $1 FOR UPDATE;
+
+-- name: ForkCollectionOffset :exec
+INSERT INTO collection (id, record_compaction_offset_position, record_enumeration_offset_position)
+    SELECT $2, collection.record_compaction_offset_position, collection.record_enumeration_offset_position
+    FROM collection
+    WHERE collection.id = $1;
+
+-- name: ForkCollectionRecord :exec
+INSERT INTO record_log ("offset", collection_id, timestamp, record)
+    SELECT record_log.offset, $2, record_log.timestamp, record_log.record
+    FROM record_log
+    WHERE record_log.collection_id = $1;

Diff for: go/pkg/log/store/queries/queries.sql

@@ -1416,11 +1416,9 @@ func (suite *APIsTestSuite) TestForkCollection() {
 	}
 
 	sourceFlushCollectionCompaction := &model.FlushCollectionCompaction{
-		ID:       sourceCreateCollection.ID,
-		TenantID: sourceCreateCollection.TenantID,
-		// TODO: Inherit log position after log fork is implemented
-		// LogPosition:              1000,
-		LogPosition:              0,
+		ID:                       sourceCreateCollection.ID,
+		TenantID:                 sourceCreateCollection.TenantID,
+		LogPosition:              1000,
 		CurrentCollectionVersion: 0,
 		FlushSegmentCompactions: []*model.FlushSegmentCompaction{
 			sourceFlushMetadataSegment,

Diff for: go/pkg/sysdb/coordinator/coordinator_test.go

@@ -868,18 +868,16 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 		}
 
 		createCollection := &model.CreateCollection{
-			ID:                   forkCollection.TargetCollectionID,
-			Name:                 forkCollection.TargetCollectionName,
-			ConfigurationJsonStr: sourceCollection.ConfigurationJsonStr,
-			Dimension:            sourceCollection.Dimension,
-			Metadata:             sourceCollection.Metadata,
-			GetOrCreate:          false,
-			TenantID:             sourceCollection.TenantID,
-			DatabaseName:         sourceCollection.DatabaseName,
-			Ts:                   ts.Unix(),
-			// TODO: Inherit log position after log fork is implemented
-			// LogPosition:                sourceCollection.LogPosition,
-			LogPosition:                0,
+			ID:                         forkCollection.TargetCollectionID,
+			Name:                       forkCollection.TargetCollectionName,
+			ConfigurationJsonStr:       sourceCollection.ConfigurationJsonStr,
+			Dimension:                  sourceCollection.Dimension,
+			Metadata:                   sourceCollection.Metadata,
+			GetOrCreate:                false,
+			TenantID:                   sourceCollection.TenantID,
+			DatabaseName:               sourceCollection.DatabaseName,
+			Ts:                         ts.Unix(),
+			LogPosition:                sourceCollection.LogPosition,
 			RootCollectionId:           rootCollectionIDStr,
 			TotalRecordsPostCompaction: sourceCollection.TotalRecordsPostCompaction,
 			SizeBytesPostCompaction:    sourceCollection.SizeBytesPostCompaction,

Diff for: go/pkg/sysdb/coordinator/table_catalog.go

@@ -42,6 +42,15 @@ message PullLogsResponse {
   repeated LogRecord records = 1;
 }
 
+message ForkLogsRequest {
+  string source_collection_id = 1;
+  string target_collection_id = 2;
+}
+
+message ForkLogsResponse {
+  // Empty  
+}
+
 message CollectionInfo {
   string collection_id = 1;
   // The log offset of the first log entry of the collection that needs to be compacted
@@ -81,6 +90,7 @@ service LogService {
   rpc PushLogs(PushLogsRequest) returns (PushLogsResponse) {}
   rpc ScoutLogs(ScoutLogsRequest) returns (ScoutLogsResponse) {}
   rpc PullLogs(PullLogsRequest) returns (PullLogsResponse) {}
+  rpc ForkLogs(ForkLogsRequest) returns (ForkLogsResponse) {}
   rpc GetAllCollectionInfoToCompact(GetAllCollectionInfoToCompactRequest) returns (GetAllCollectionInfoToCompactResponse) {}
   rpc UpdateCollectionLogOffset(UpdateCollectionLogOffsetRequest) returns (UpdateCollectionLogOffsetResponse) {}
   rpc PurgeDirtyForCollection(PurgeDirtyForCollectionRequest) returns (PurgeDirtyForCollectionResponse) {}

Diff for: idl/chromadb/proto/logservice.proto

@@ -552,6 +552,9 @@ impl ServiceBasedFrontend {
         }: ForkCollectionRequest,
     ) -> Result<ForkCollectionResponse, ForkCollectionError> {
         let target_collection_id = CollectionUuid::new();
+        self.log_client
+            .fork_logs(source_collection_id, target_collection_id)
+            .await?;
         let collection_and_segments = self
             .sysdb_client
             .fork_collection(

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

@@ -18,6 +18,7 @@ use chroma_types::chroma_proto::{
     PushLogsRequest, PushLogsResponse, ScoutLogsRequest, ScoutLogsResponse,
     UpdateCollectionLogOffsetRequest, UpdateCollectionLogOffsetResponse,
 };
+use chroma_types::chroma_proto::{ForkLogsRequest, ForkLogsResponse};
 use chroma_types::CollectionUuid;
 use figment::providers::{Env, Format, Yaml};
 use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
@@ -763,6 +764,13 @@ impl LogService for LogServer {
         .await
     }
 
+    async fn fork_logs(
+        &self,
+        _request: Request<ForkLogsRequest>,
+    ) -> Result<Response<ForkLogsResponse>, Status> {
+        unimplemented!("Log forking is unimplemented for WAL3 for now")
+    }
+
     #[tracing::instrument(info, skip(self, request), err(Display))]
     async fn get_all_collection_info_to_compact(
         &self,

Diff for: rust/log-service/src/lib.rs

@@ -57,6 +57,23 @@ impl ChromaError for GrpcPushLogsError {
     }
 }
 
+#[derive(Error, Debug)]
+pub enum GrpcForkLogsError {
+    #[error("Please backoff exponentially and retry")]
+    Backoff,
+    #[error("Failed to push logs")]
+    FailedToForkLogs(#[from] tonic::Status),
+}
+
+impl ChromaError for GrpcForkLogsError {
+    fn code(&self) -> ErrorCodes {
+        match self {
+            GrpcForkLogsError::Backoff => ErrorCodes::Unavailable,
+            GrpcForkLogsError::FailedToForkLogs(_) => ErrorCodes::Internal,
+        }
+    }
+}
+
 #[derive(Error, Debug)]
 pub enum GrpcGetCollectionsWithNewDataError {
     #[error("Failed to fetch")]
@@ -306,6 +323,24 @@ impl GrpcLog {
         Ok(())
     }
 
+    pub(super) async fn fork_logs(
+        &mut self,
+        source_collection_id: CollectionUuid,
+        target_collection_id: CollectionUuid,
+    ) -> Result<(), GrpcForkLogsError> {
+        self.client_for(source_collection_id)
+            .fork_logs(chroma_proto::ForkLogsRequest {
+                source_collection_id: source_collection_id.to_string(),
+                target_collection_id: target_collection_id.to_string(),
+            })
+            .await
+            .map(|_| ())
+            .map_err(|err| match err.code() {
+                tonic::Code::Unavailable => GrpcForkLogsError::Backoff,
+                _ => err.into(),
+            })
+    }
+
     pub(crate) async fn get_collections_with_new_data(
         &mut self,
         min_compaction_size: u64,

Diff for: rust/log/src/grpc_log.rs

@@ -3,7 +3,9 @@ use crate::in_memory_log::InMemoryLog;
 use crate::sqlite_log::SqliteLog;
 use crate::types::CollectionInfo;
 use chroma_error::ChromaError;
-use chroma_types::{CollectionUuid, LogRecord, OperationRecord, ResetError, ResetResponse};
+use chroma_types::{
+    CollectionUuid, ForkCollectionError, LogRecord, OperationRecord, ResetError, ResetResponse,
+};
 use std::fmt::Debug;
 
 #[derive(Clone, Debug)]
@@ -92,6 +94,22 @@ impl Log {
         }
     }
 
+    #[tracing::instrument(skip(self))]
+    pub async fn fork_logs(
+        &mut self,
+        source_collection_id: CollectionUuid,
+        target_collection_id: CollectionUuid,
+    ) -> Result<(), ForkCollectionError> {
+        match self {
+            Log::Sqlite(_) => Err(ForkCollectionError::Local),
+            Log::Grpc(log) => log
+                .fork_logs(source_collection_id, target_collection_id)
+                .await
+                .map_err(|err| err.boxed().into()),
+            Log::InMemory(_) => Err(ForkCollectionError::Local),
+        }
+    }
+
     #[tracing::instrument(skip(self))]
     pub async fn get_collections_with_new_data(
         &mut self,