[ENH]: garbage collector v2 orchestrator (supports forking)

Author: codetheweb

Cargo.lock

@@ -97,8 +97,8 @@ func (s *Coordinator) GetTenant(ctx context.Context, getTenant *model.GetTenant)
 	return tenant, nil
 }
 
-func (s *Coordinator) CreateCollectionAndSegments(ctx context.Context, createCollection *model.CreateCollection, createSegments []*model.CreateSegment) (*model.Collection, bool, error) {
-	collection, created, err := s.catalog.CreateCollectionAndSegments(ctx, createCollection, createSegments, createCollection.Ts)
+func (s *Coordinator) CreateCollectionAndSegments(ctx context.Context, createCollection *model.CreateCollection, createSegments []*model.CreateSegment, initialFilePaths []*model.FlushSegmentCompaction) (*model.Collection, bool, error) {
+	collection, created, err := s.catalog.CreateCollectionAndSegments(ctx, createCollection, createSegments, initialFilePaths, createCollection.Ts)
 	if err != nil {
 		return nil, false, err
 	}
@@ -251,6 +251,10 @@ func (s *Coordinator) ListCollectionsToGc(ctx context.Context, cutoffTimeSecs *u
 	return s.catalog.ListCollectionsToGc(ctx, cutoffTimeSecs, limit, tenantID)
 }
 
+func (s *Coordinator) GetCollectionToGcInfo(ctx context.Context, collectionID types.UniqueID, tenantID string) (*model.CollectionToGc, error) {
+	return s.catalog.GetCollectionToGcInfo(ctx, collectionID, tenantID)
+}
+
 func (s *Coordinator) ListCollectionVersions(ctx context.Context, collectionID types.UniqueID, tenantID string, maxCount *int64, versionsBefore *int64, versionsAtOrAfter *int64, includeMarkedForDeletion bool) ([]*coordinatorpb.CollectionVersionInfo, error) {
 	return s.catalog.ListCollectionVersions(ctx, collectionID, tenantID, maxCount, versionsBefore, versionsAtOrAfter, includeMarkedForDeletion)
 }

go/pkg/sysdb/coordinator/coordinator.go

@@ -41,22 +41,14 @@ func convertCollectionToModel(collectionAndMetadataList []*dbmodel.CollectionAnd
 	return collections
 }
 
-func convertCollectionToGcToModel(collectionToGc []*dbmodel.CollectionToGc) []*model.CollectionToGc {
-	if collectionToGc == nil {
-		return nil
-	}
-	collections := make([]*model.CollectionToGc, 0, len(collectionToGc))
-	for _, collectionInfo := range collectionToGc {
-		collection := model.CollectionToGc{
-			ID:              types.MustParse(collectionInfo.ID),
-			Name:            collectionInfo.Name,
-			VersionFilePath: collectionInfo.VersionFileName,
-			TenantID:        collectionInfo.TenantID,
-			LineageFilePath: collectionInfo.LineageFileName,
-		}
-		collections = append(collections, &collection)
+func convertCollectionToGcToModel(collectionToGc dbmodel.CollectionToGc) model.CollectionToGc {
+	return model.CollectionToGc{
+		ID:              types.MustParse(collectionToGc.ID),
+		Name:            collectionToGc.Name,
+		VersionFilePath: collectionToGc.VersionFileName,
+		TenantID:        collectionToGc.TenantID,
+		LineageFilePath: collectionToGc.LineageFileName,
 	}
-	return collections
 }
 
 func convertCollectionMetadataToModel(collectionMetadataList []*dbmodel.CollectionMetadata) *model.CollectionMetadata[model.CollectionMetadataValueType] {

go/pkg/sysdb/coordinator/model_db_convert.go

@@ -456,10 +456,38 @@ func (tc *Catalog) ListCollectionsToGc(ctx context.Context, cutoffTimeSecs *uint
 	if err != nil {
 		return nil, err
 	}
-	collections := convertCollectionToGcToModel(collectionsToGc)
+
+	if collectionsToGc == nil {
+		return nil, nil
+	}
+
+	collections := make([]*model.CollectionToGc, 0, len(collectionsToGc))
+	for _, collectionInfo := range collectionsToGc {
+		collection := convertCollectionToGcToModel(*collectionInfo)
+		collections = append(collections, &collection)
+	}
+
 	return collections, nil
 }
 
+func (tc *Catalog) GetCollectionToGcInfo(ctx context.Context, collectionID types.UniqueID, tenantID string) (*model.CollectionToGc, error) {
+	tracer := otel.Tracer
+	if tracer != nil {
+		_, span := tracer.Start(ctx, "Catalog.GetCollectionToGcInfo")
+		defer span.End()
+	}
+
+	collectionToGc, err := tc.metaDomain.CollectionDb(ctx).GetCollectionToGcInfo(collectionID.String(), tenantID)
+	if err != nil {
+		return nil, err
+	}
+	if collectionToGc == nil {
+		return nil, common.ErrCollectionNotFound
+	}
+	result := convertCollectionToGcToModel(*collectionToGc)
+	return &result, nil
+}
+
 func (tc *Catalog) GetCollectionWithSegments(ctx context.Context, collectionID types.UniqueID) (*model.Collection, []*model.Segment, error) {
 	tracer := otel.Tracer
 	if tracer != nil {
@@ -952,7 +980,7 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 			flushFilePaths = append(flushFilePaths, flushFilePath)
 		}
 
-		_, _, err = tc.CreateCollectionAndSegments(txCtx, createCollection, createSegments, ts.Unix())
+		_, _, err = tc.CreateCollectionAndSegments(txCtx, createCollection, createSegments, flushFilePaths, ts.Unix())
 		if err != nil {
 			return err
 		}
@@ -1090,7 +1118,22 @@ func (tc *Catalog) createSegmentImpl(txCtx context.Context, createSegment *model
 	return result, nil
 }
 
-func (tc *Catalog) createFirstVersionFile(ctx context.Context, databaseID string, createCollection *model.CreateCollection, createSegments []*model.CreateSegment, ts types.Timestamp) (string, error) {
+func (tc *Catalog) createFirstVersionFile(ctx context.Context, databaseID string, createCollection *model.CreateCollection, createSegments []*model.CreateSegment, initialFilePaths []*model.FlushSegmentCompaction, ts types.Timestamp) (string, error) {
+	segmentCompactionInfos := make([]*coordinatorpb.FlushSegmentCompactionInfo, 0, len(createSegments))
+	for _, compaction := range initialFilePaths {
+		// Convert map[string][]string to map[string]*coordinatorpb.FilePaths
+		convertedPaths := make(map[string]*coordinatorpb.FilePaths)
+		for k, v := range compaction.FilePaths {
+			convertedPaths[k] = &coordinatorpb.FilePaths{Paths: v}
+		}
+
+		info := &coordinatorpb.FlushSegmentCompactionInfo{
+			SegmentId: compaction.ID.String(),
+			FilePaths: convertedPaths,
+		}
+		segmentCompactionInfos = append(segmentCompactionInfos, info)
+	}
+
 	collectionVersionFilePb := &coordinatorpb.CollectionVersionFile{
 		CollectionInfoImmutable: &coordinatorpb.CollectionInfoImmutable{
 			TenantId:               createCollection.TenantID,
@@ -1104,6 +1147,9 @@ func (tc *Catalog) createFirstVersionFile(ctx context.Context, databaseID string
 				{
 					Version:       0,
 					CreatedAtSecs: int64(ts),
+					SegmentInfo: &coordinatorpb.CollectionSegmentInfo{
+						SegmentCompactionInfo: segmentCompactionInfos,
+					},
 				},
 			},
 		},
@@ -1117,7 +1163,7 @@ func (tc *Catalog) createFirstVersionFile(ctx context.Context, databaseID string
 	return fullFilePath, nil
 }
 
-func (tc *Catalog) CreateCollectionAndSegments(ctx context.Context, createCollection *model.CreateCollection, createSegments []*model.CreateSegment, ts types.Timestamp) (*model.Collection, bool, error) {
+func (tc *Catalog) CreateCollectionAndSegments(ctx context.Context, createCollection *model.CreateCollection, createSegments []*model.CreateSegment, initialFilePaths []*model.FlushSegmentCompaction, ts types.Timestamp) (*model.Collection, bool, error) {
 	var resultCollection *model.Collection
 	created := false
 
@@ -1139,7 +1185,7 @@ func (tc *Catalog) CreateCollectionAndSegments(ctx context.Context, createCollec
 			return nil, false, common.ErrDatabaseNotFound
 		}
 
-		versionFileName, err = tc.createFirstVersionFile(ctx, databases[0].ID, createCollection, createSegments, ts)
+		versionFileName, err = tc.createFirstVersionFile(ctx, databases[0].ID, createCollection, createSegments, initialFilePaths, ts)
 		if err != nil {
 			return nil, false, err
 		}

go/pkg/sysdb/coordinator/table_catalog.go

@@ -78,6 +78,7 @@ func (s *Server) CreateCollection(ctx context.Context, req *coordinatorpb.Create
 
 	// Convert the request segments to create segment models
 	createSegments := []*model.CreateSegment{}
+	initialFilePaths := []*model.FlushSegmentCompaction{}
 	for _, segment := range req.Segments {
 		createSegment, err := convertSegmentToModel(segment)
 		if err != nil {
@@ -91,10 +92,19 @@ func (s *Server) CreateCollection(ctx context.Context, req *coordinatorpb.Create
 			return res, grpcutils.BuildInternalGrpcError(err.Error())
 		}
 		createSegments = append(createSegments, createSegment)
+		filePaths := make(map[string][]string)
+		for key, filePath := range segment.FilePaths {
+			filePaths[key] = filePath.Paths
+		}
+		flushSegmentCompaction := &model.FlushSegmentCompaction{
+			ID:        createSegment.ID,
+			FilePaths: filePaths,
+		}
+		initialFilePaths = append(initialFilePaths, flushSegmentCompaction)
 	}
 
 	// Create the collection and segments
-	collection, created, err := s.coordinator.CreateCollectionAndSegments(ctx, createCollection, createSegments)
+	collection, created, err := s.coordinator.CreateCollectionAndSegments(ctx, createCollection, createSegments, initialFilePaths)
 	if err != nil {
 		log.Error("CreateCollection failed. error creating collection", zap.Error(err), zap.String("collection_id", req.Id), zap.String("collection_name", req.Name))
 		res.Collection = &coordinatorpb.Collection{
@@ -513,6 +523,29 @@ func (s *Server) ListCollectionsToGc(ctx context.Context, req *coordinatorpb.Lis
 	return res, nil
 }
 
+func (s *Server) GetCollectionToGcInfo(ctx context.Context, req *coordinatorpb.GetCollectionToGcInfoRequest) (*coordinatorpb.GetCollectionToGcInfoResponse, error) {
+	collectionID, err := types.ToUniqueID(&req.CollectionId)
+	if err != nil {
+		log.Error("GetCollectionInfoToGc failed. error parsing collection id", zap.Error(err), zap.String("collection_id", req.CollectionId))
+		return nil, grpcutils.BuildInternalGrpcError(err.Error())
+	}
+	collection, err := s.coordinator.GetCollectionToGcInfo(ctx, collectionID, req.TenantId)
+	if err != nil {
+		log.Error("GetCollectionInfoToGc failed", zap.Error(err), zap.String("collection_id", req.CollectionId))
+		return nil, grpcutils.BuildInternalGrpcError(err.Error())
+	}
+	res := &coordinatorpb.GetCollectionToGcInfoResponse{
+		Collection: &coordinatorpb.CollectionToGcInfo{
+			Id:              collection.ID.String(),
+			Name:            collection.Name,
+			VersionFilePath: collection.VersionFilePath,
+			TenantId:        collection.TenantID,
+			LineageFilePath: collection.LineageFilePath,
+		},
+	}
+	return res, nil
+}
+
 // Mark the versions for deletion.
 // GC minics a 2PC protocol.
 // 1. Mark the versions for deletion by calling MarkVersionForDeletion.

go/pkg/sysdb/grpc/collection_service.go

@@ -101,6 +101,20 @@ func (s *collectionDb) ListCollectionsToGc(cutoffTimeSecs *uint64, limit *uint64
 	return collections, nil
 }
 
+func (s *collectionDb) GetCollectionToGcInfo(collectionID string, tenantID string) (*dbmodel.CollectionToGc, error) {
+	var collection dbmodel.CollectionToGc
+	err := s.read_db.Table("collections").
+		Select("collections.id, collections.name, collections.version_file_name, collections.oldest_version_ts, databases.tenant_id, COALESCE(NULLIF(root.lineage_file_name, ''), NULLIF(collections.lineage_file_name, '')) AS lineage_file_name").
+		Joins("INNER JOIN databases ON collections.database_id = databases.id").
+		Joins("LEFT JOIN collections AS root ON collections.root_collection_id = root.id").
+		Where("collections.id = ? AND databases.tenant_id = ?", collectionID, tenantID).
+		First(&collection).Error
+	if err != nil {
+		return nil, err
+	}
+	return &collection, nil
+}
+
 func (s *collectionDb) getCollections(id *string, name *string, tenantID string, databaseName string, limit *int32, offset *int32, is_deleted bool) (collectionWithMetdata []*dbmodel.CollectionAndMetadata, err error) {
 	type Result struct {
 		// Collection fields

go/pkg/sysdb/metastore/db/dao/collection.go

@@ -65,6 +65,7 @@ type ICollectionDb interface {
 	GetCollectionEntry(collectionID *string, databaseName *string) (*Collection, error)
 	GetCollectionSize(collectionID string) (uint64, error)
 	ListCollectionsToGc(cutoffTimeSecs *uint64, limit *uint64, tenantID *string) ([]*CollectionToGc, error)
+	GetCollectionToGcInfo(collectionID string, tenantID string) (*CollectionToGc, error)
 	UpdateVersionRelatedFields(collectionID, existingVersionFileName, newVersionFileName string, oldestVersionTs *time.Time, numActiveVersions *int) (int64, error)
 	LockCollection(collectionID string) error
 	UpdateCollectionLineageFilePath(collectionID string, currentLineageFilePath string, newLineageFilePath string) error

go/pkg/sysdb/metastore/db/dbmodel/collection.go

@@ -449,6 +449,15 @@ message ListCollectionsToGcResponse {
   repeated CollectionToGcInfo collections = 1;
 }
 
+message GetCollectionToGcInfoRequest {
+  string collection_id = 1;
+  string tenant_id = 2;
+}
+
+message GetCollectionToGcInfoResponse {
+  CollectionToGcInfo collection = 1;
+}
+
 message MarkVersionForDeletionRequest {
   int64 epoch_id = 1;
   repeated VersionListForCollection versions = 2;
@@ -504,6 +513,7 @@ service SysDB {
   rpc ListCollectionVersions(ListCollectionVersionsRequest) returns (ListCollectionVersionsResponse) {}
   rpc GetCollectionSize(GetCollectionSizeRequest) returns (GetCollectionSizeResponse) {}
   rpc ListCollectionsToGc(ListCollectionsToGcRequest) returns (ListCollectionsToGcResponse) {}
+  rpc GetCollectionToGcInfo(GetCollectionToGcInfoRequest) returns (GetCollectionToGcInfoResponse) {}
   rpc MarkVersionForDeletion(MarkVersionForDeletionRequest) returns (MarkVersionForDeletionResponse) {}
   rpc DeleteCollectionVersion(DeleteCollectionVersionRequest) returns (DeleteCollectionVersionResponse) {}
   rpc BatchGetCollectionVersionFilePaths(BatchGetCollectionVersionFilePathsRequest) returns (BatchGetCollectionVersionFilePathsResponse) {}

idl/chromadb/proto/coordinator.proto

@@ -21,6 +21,7 @@ use uuid::Uuid;
 /// * `Result<Uuid, Box<dyn ChromaError>>` - The UUID of the created sparse index file
 pub async fn create_test_sparse_index(
     storage: &Storage,
+    root_id: Uuid,
     block_ids: Vec<Uuid>,
     prefix: Option<String>,
 ) -> Result<Uuid, Box<dyn ChromaError>> {
@@ -43,7 +44,6 @@ pub async fn create_test_sparse_index(
     sparse_index.set_count(block_ids[0], 1)?;
 
     // Create and save the sparse index file
-    let root_id = Uuid::new_v4();
     let root_writer = RootWriter::new(Version::V1_1, root_id, sparse_index);
     let root_manager = RootManager::new(storage.clone(), Box::new(NopCache));
     root_manager.flush::<&str>(&root_writer).await?;
@@ -87,7 +87,8 @@ mod tests {
         let storage = Storage::Local(LocalStorage::new(temp_dir.path().to_str().unwrap()));
 
         let block_ids = vec![Uuid::new_v4(), Uuid::new_v4(), Uuid::new_v4()];
-        let result = create_test_sparse_index(&storage, block_ids.clone(), None).await;
+        let result =
+            create_test_sparse_index(&storage, Uuid::new_v4(), block_ids.clone(), None).await;
         assert!(result.is_ok());
 
         // Verify the sparse index was created by trying to read it
@@ -107,7 +108,7 @@ mod tests {
         let temp_dir = TempDir::new().unwrap();
         let storage = Storage::Local(LocalStorage::new(temp_dir.path().to_str().unwrap()));
 
-        let result = create_test_sparse_index(&storage, vec![], None).await;
+        let result = create_test_sparse_index(&storage, Uuid::new_v4(), vec![], None).await;
         assert!(matches!(
             result,
             Err(e) if e.to_string().contains("Cannot create sparse index with empty block IDs")
@@ -121,7 +122,7 @@ mod tests {
 
         let block_ids = vec![Uuid::new_v4(), Uuid::new_v4()];
         let prefix = Some("custom".to_string());
-        let result = create_test_sparse_index(&storage, block_ids, prefix).await;
+        let result = create_test_sparse_index(&storage, Uuid::new_v4(), block_ids, prefix).await;
         assert!(result.is_ok());
     }
 }

rust/blockstore/src/test_utils/sparse_index_test_utils.rs

@@ -57,3 +57,4 @@ tokio-test = "0.4"
 rand = { workspace = true }
 itertools = { workspace = true }
 tracing-test = { version = "0.2.5" }
+tracing-subscriber = { workspace = true }