[ENH]: sysdb changes to support moving collection hard deletes to garbage collector

Author: codetheweb

go/cmd/coordinator/cmd.go

@@ -44,13 +44,6 @@ func init() {
 	Cmd.Flags().IntVar(&conf.DBConfig.MaxOpenConns, "max-open-conns", 10, "MetaTable max open connections")
 	Cmd.Flags().StringVar(&conf.DBConfig.SslMode, "ssl-mode", "disable", "SSL mode for database connection")
 
-	// Soft deletes
-	Cmd.Flags().BoolVar(&conf.SoftDeleteEnabled, "soft-delete-enabled", true, "Enable soft deletes")
-	Cmd.Flags().DurationVar(&conf.SoftDeleteCleanupInterval, "soft-delete-cleanup-interval", 30*time.Second, "Soft delete cleanup interval")
-	Cmd.Flags().DurationVar(&conf.SoftDeleteMaxAge, "soft-delete-max-age", 72*time.Hour, "Soft delete max age")
-	Cmd.Flags().UintVar(&conf.SoftDeleteCleanupBatchSize, "soft-delete-cleanup-batch-size", 100, "Soft delete cleanup batch size")
-	// With above values, the soft delete cleaner can remove around 1000 collections in 5 minutes.
-
 	// Memberlist
 	Cmd.Flags().StringVar(&conf.KubernetesNamespace, "kubernetes-namespace", "chroma", "Kubernetes namespace")
 	Cmd.Flags().DurationVar(&conf.ReconcileInterval, "reconcile-interval", 100*time.Millisecond, "Reconcile interval")

go/pkg/common/errors.go

@@ -30,6 +30,7 @@ var (
 	ErrCollectionTooManyFork                 = errors.New("collection entry has too many forks")
 	ErrCollectionDeletedWithLocksHeld        = errors.New("collection got deleted concurrently even though select for update locks were held. Not possible unless corruption somehow")
 	ErrMissingLineageFileName                = errors.New("missing lineage file name in root collection entry")
+	ErrCollectionWasNotSoftDeleted           = errors.New("collection was not soft deleted")
 
 	// Collection metadata errors
 	ErrUnknownCollectionMetadataType = errors.New("collection metadata value type not supported")

go/pkg/sysdb/coordinator/coordinator.go

@@ -15,30 +15,18 @@ import (
 	"go.uber.org/zap"
 )
 
-// DeleteMode represents whether to perform a soft or hard delete
-type DeleteMode int
-
-const (
-	// SoftDelete marks records as deleted but keeps them in the database
-	SoftDelete DeleteMode = iota
-	// HardDelete permanently removes records from the database
-	HardDelete
-)
-
 // Coordinator is the top level component.
 // Currently, it only has the system catalog related APIs and will be extended to
 // support other functionalities such as membership managed and propagation.
 type Coordinator struct {
 	ctx         context.Context
 	catalog     Catalog
-	deleteMode  DeleteMode
 	objectStore *s3metastore.S3MetaStore
 }
 
-func NewCoordinator(ctx context.Context, deleteMode DeleteMode, objectStore *s3metastore.S3MetaStore, versionFileEnabled bool) (*Coordinator, error) {
+func NewCoordinator(ctx context.Context, objectStore *s3metastore.S3MetaStore, versionFileEnabled bool) (*Coordinator, error) {
 	s := &Coordinator{
 		ctx:         ctx,
-		deleteMode:  deleteMode,
 		objectStore: objectStore,
 	}
 
@@ -142,14 +130,11 @@ func (s *Coordinator) GetSoftDeletedCollections(ctx context.Context, collectionI
 	return s.catalog.GetSoftDeletedCollections(ctx, collectionID, tenantID, databaseName, limit)
 }
 
-func (s *Coordinator) DeleteCollection(ctx context.Context, deleteCollection *model.DeleteCollection) error {
-	if s.deleteMode == SoftDelete {
-		return s.catalog.DeleteCollection(ctx, deleteCollection, true)
-	}
-	return s.catalog.DeleteCollection(ctx, deleteCollection, false)
+func (s *Coordinator) SoftDeleteCollection(ctx context.Context, deleteCollection *model.DeleteCollection) error {
+	return s.catalog.DeleteCollection(ctx, deleteCollection, true)
 }
 
-func (s *Coordinator) CleanupSoftDeletedCollection(ctx context.Context, deleteCollection *model.DeleteCollection) error {
+func (s *Coordinator) FinishCollectionDeletion(ctx context.Context, deleteCollection *model.DeleteCollection) error {
 	return s.catalog.DeleteCollection(ctx, deleteCollection, false)
 }
 
@@ -267,7 +252,6 @@ func (s *Coordinator) BatchGetCollectionVersionFilePaths(ctx context.Context, re
 	return s.catalog.BatchGetCollectionVersionFilePaths(ctx, req.CollectionIds)
 }
 
-// SetDeleteMode sets the delete mode for testing
-func (c *Coordinator) SetDeleteMode(mode DeleteMode) {
-	c.deleteMode = mode
+func (s *Coordinator) BatchGetCollectionSoftDeleteStatus(ctx context.Context, req *coordinatorpb.BatchGetCollectionSoftDeleteStatusRequest) (*coordinatorpb.BatchGetCollectionSoftDeleteStatusResponse, error) {
+	return s.catalog.BatchGetCollectionSoftDeleteStatus(ctx, req.CollectionIds)
 }

go/pkg/sysdb/coordinator/coordinator_test.go

@@ -80,7 +80,7 @@ func (suite *APIsTestSuite) SetupTest() {
 		collection.Name = "collection_" + suite.T().Name() + strconv.Itoa(index)
 	}
 	ctx := context.Background()
-	c, err := NewCoordinator(ctx, SoftDelete, suite.s3MetaStore, true)
+	c, err := NewCoordinator(ctx, suite.s3MetaStore, true)
 	if err != nil {
 		suite.T().Fatalf("error creating coordinator: %v", err)
 	}
@@ -111,7 +111,7 @@ func (suite *APIsTestSuite) TearDownTest() {
 func testCollection(t *rapid.T) {
 	dbcore.ConfigDatabaseForTesting()
 	ctx := context.Background()
-	c, err := NewCoordinator(ctx, HardDelete, nil, false)
+	c, err := NewCoordinator(ctx, nil, false)
 	if err != nil {
 		t.Fatalf("error creating coordinator: %v", err)
 	}
@@ -164,7 +164,7 @@ func testCollection(t *rapid.T) {
 func testSegment(t *rapid.T) {
 	dbcore.ConfigDatabaseForTesting()
 	ctx := context.Background()
-	c, err := NewCoordinator(ctx, HardDelete, nil, false)
+	c, err := NewCoordinator(ctx, nil, false)
 	if err != nil {
 		t.Fatalf("error creating coordinator: %v", err)
 	}
@@ -486,7 +486,7 @@ func (suite *APIsTestSuite) TestCreateGetDeleteCollections() {
 		DatabaseName: suite.databaseName,
 		TenantID:     suite.tenantName,
 	}
-	err = suite.coordinator.DeleteCollection(ctx, deleteCollection)
+	err = suite.coordinator.SoftDeleteCollection(ctx, deleteCollection)
 	suite.NoError(err)
 
 	results, err = suite.coordinator.GetCollections(ctx, types.NilUniqueID(), nil, suite.tenantName, suite.databaseName, nil, nil)
@@ -508,7 +508,7 @@ func (suite *APIsTestSuite) TestCreateGetDeleteCollections() {
 	suite.Empty(byIDResult)
 
 	// Duplicate delete throws an exception
-	err = suite.coordinator.DeleteCollection(ctx, deleteCollection)
+	err = suite.coordinator.SoftDeleteCollection(ctx, deleteCollection)
 	suite.Error(err)
 
 	// Re-create the deleted collection
@@ -560,7 +560,7 @@ func (suite *APIsTestSuite) TestCreateGetDeleteCollections() {
 		DatabaseName: suite.databaseName,
 		TenantID:     suite.tenantName,
 	}
-	err = suite.coordinator.DeleteCollection(ctx, deleteCollection)
+	err = suite.coordinator.SoftDeleteCollection(ctx, deleteCollection)
 	suite.NoError(err)
 
 	// Verify collection and segment were deleted
@@ -572,68 +572,11 @@ func (suite *APIsTestSuite) TestCreateGetDeleteCollections() {
 	segments, err = suite.coordinator.GetSegments(ctx, segment.ID, nil, nil, createCollection.ID)
 	suite.NoError(err)
 	suite.NotEmpty(segments)
-	suite.coordinator.deleteMode = HardDelete
-	err = suite.coordinator.DeleteCollection(ctx, deleteCollection)
+	err = suite.coordinator.FinishCollectionDeletion(ctx, deleteCollection)
 	suite.NoError(err)
 	segments, err = suite.coordinator.GetSegments(ctx, segment.ID, nil, nil, createCollection.ID)
 	suite.NoError(err)
 	suite.Empty(segments)
-
-	// Check for forward and backward compatibility with soft and hard delete.
-	// 1. Create a collection with soft delete enabled.
-	// 2. Delete the collection (i.e. it will be marked as is_deleted)
-	// 3. Disable soft delete.
-	// 4. Query for the deleted collection. It should not be found.
-	// 5. Enable soft delete.
-	// 6. Query for the deleted collection. It should be found.
-
-	suite.coordinator.deleteMode = SoftDelete
-	collectionId := types.NewUniqueID()
-	suite.coordinator.CreateCollection(ctx, &model.CreateCollection{
-		ID:           collectionId,
-		Name:         "test_coll_fwd_bkwd_compat",
-		TenantID:     suite.tenantName,
-		DatabaseName: suite.databaseName,
-	})
-	suite.coordinator.DeleteCollection(ctx, &model.DeleteCollection{
-		ID:           collectionId,
-		DatabaseName: suite.databaseName,
-		TenantID:     suite.tenantName,
-	})
-	collection, err := suite.coordinator.GetCollections(ctx, collectionId, nil, suite.tenantName, suite.databaseName, nil, nil)
-	suite.NoError(err)
-	// Check that the collection is deleted
-	suite.Empty(collection)
-	// Toggle the mode.
-	suite.coordinator.deleteMode = HardDelete
-	collection, err = suite.coordinator.GetCollections(ctx, collectionId, nil, suite.tenantName, suite.databaseName, nil, nil)
-	suite.NoError(err)
-	// Check that the collection is still deleted
-	suite.Empty(collection)
-	// Create a collection and delete while being in HardDelete mode.
-	anotherCollectionId := types.NewUniqueID()
-	suite.coordinator.CreateCollection(ctx, &model.CreateCollection{
-		ID:           anotherCollectionId,
-		Name:         "test_coll_fwd_bkwd_compat",
-		TenantID:     suite.tenantName,
-		DatabaseName: suite.databaseName,
-	})
-	suite.coordinator.DeleteCollection(ctx, &model.DeleteCollection{
-		ID:           anotherCollectionId,
-		DatabaseName: suite.databaseName,
-		TenantID:     suite.tenantName,
-	})
-
-	// Toggle the mode.
-	suite.coordinator.deleteMode = SoftDelete
-	collection, err = suite.coordinator.GetCollections(ctx, collectionId, nil, suite.tenantName, suite.databaseName, nil, nil)
-	suite.NoError(err)
-	// Check that the collection is still deleted
-	suite.Empty(collection)
-	collection, err = suite.coordinator.GetCollections(ctx, anotherCollectionId, nil, suite.tenantName, suite.databaseName, nil, nil)
-	suite.NoError(err)
-	// Check that another collection is still deleted
-	suite.Empty(collection)
 }
 
 func (suite *APIsTestSuite) TestCollectionSize() {
@@ -1310,7 +1253,6 @@ func (suite *APIsTestSuite) TestSoftAndHardDeleteCollection() {
 	ctx := context.Background()
 
 	// Test Hard Delete scenario
-	suite.coordinator.deleteMode = HardDelete
 	// Create test collection
 	testCollection2 := &model.CreateCollection{
 		ID:           types.NewUniqueID(),
@@ -1324,7 +1266,13 @@ func (suite *APIsTestSuite) TestSoftAndHardDeleteCollection() {
 	suite.NoError(err)
 
 	// Hard delete the collection
-	err = suite.coordinator.DeleteCollection(ctx, &model.DeleteCollection{
+	err = suite.coordinator.SoftDeleteCollection(ctx, &model.DeleteCollection{
+		ID:           testCollection2.ID,
+		TenantID:     testCollection2.TenantID,
+		DatabaseName: testCollection2.DatabaseName,
+	})
+	suite.NoError(err)
+	err = suite.coordinator.FinishCollectionDeletion(ctx, &model.DeleteCollection{
 		ID:           testCollection2.ID,
 		TenantID:     testCollection2.TenantID,
 		DatabaseName: testCollection2.DatabaseName,
@@ -1343,7 +1291,6 @@ func (suite *APIsTestSuite) TestSoftAndHardDeleteCollection() {
 	suite.Empty(softDeletedResults)
 
 	// Test Soft Delete scenario
-	suite.coordinator.deleteMode = SoftDelete
 	// Create a test collection
 	testCollection := &model.CreateCollection{
 		ID:           types.NewUniqueID(),
@@ -1357,7 +1304,7 @@ func (suite *APIsTestSuite) TestSoftAndHardDeleteCollection() {
 	suite.NoError(err)
 
 	// Soft delete the collection
-	err = suite.coordinator.DeleteCollection(ctx, &model.DeleteCollection{
+	err = suite.coordinator.SoftDeleteCollection(ctx, &model.DeleteCollection{
 		ID:           testCollection.ID,
 		TenantID:     testCollection.TenantID,
 		DatabaseName: testCollection.DatabaseName,

go/pkg/sysdb/coordinator/table_catalog.go

@@ -564,6 +564,50 @@ func (tc *Catalog) hardDeleteCollection(ctx context.Context, deleteCollection *m
 			return common.ErrCollectionDeleteNonExistingCollection
 		}
 
+		if !collectionEntry.IsDeleted {
+			return common.ErrCollectionWasNotSoftDeleted
+		}
+
+		if collectionEntry.RootCollectionId != nil {
+			// This was a forked collection, so we need to update the lineage file
+			rootCollection, err := tc.metaDomain.CollectionDb(txCtx).GetCollectionWithoutMetadata(collectionEntry.RootCollectionId, nil, nil)
+			if err != nil {
+				return err
+			}
+			if rootCollection == nil {
+				return errors.New("root collection not found")
+			}
+
+			if rootCollection.LineageFileName == nil {
+				return errors.New("lineage file name is nil on root collection")
+			}
+
+			lineageFile, err := tc.getLineageFile(txCtx, rootCollection.LineageFileName)
+			if err != nil {
+				return err
+			}
+			// Remove collection being deleted from the dependencies
+			updatedDependencies := make([]*coordinatorpb.CollectionVersionDependency, 0)
+			for _, dependency := range lineageFile.Dependencies {
+				if dependency.TargetCollectionId != deleteCollection.ID.String() {
+					updatedDependencies = append(updatedDependencies, dependency)
+				}
+			}
+			lineageFile.Dependencies = updatedDependencies
+
+			newLineageFileId, err := uuid.NewV7()
+			if err != nil {
+				return err
+			}
+
+			newLineageFileFullName, err := tc.s3Store.PutLineageFile(collectionEntry.Tenant, collectionEntry.DatabaseID, rootCollection.ID, fmt.Sprintf("%s.binpb", newLineageFileId.String()), lineageFile)
+			if err != nil {
+				return err
+			}
+
+			tc.metaDomain.CollectionDb(txCtx).UpdateCollectionLineageFilePath(rootCollection.ID, rootCollection.LineageFileName, newLineageFileFullName)
+		}
+
 		// Delete collection and collection metadata.
 		collectionDeletedCount, err := tc.metaDomain.CollectionDb(txCtx).DeleteCollectionByID(collectionID.String())
 		if err != nil {
@@ -1053,7 +1097,12 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 			TargetCollectionId:      forkCollection.TargetCollectionID.String(),
 		})
 
-		newLineageFileBaseName := fmt.Sprintf("%s/%d/%s.binpb", sourceCollectionIDStr, sourceCollection.Version, forkCollection.TargetCollectionID)
+		newLineageFileId, err := uuid.NewV7()
+		if err != nil {
+			return err
+		}
+
+		newLineageFileBaseName := fmt.Sprintf("%s.binpb", newLineageFileId.String())
 		newLineageFileFullName, err = tc.s3Store.PutLineageFile(lineageFileTenantId, databaseID, rootCollectionIDStr, newLineageFileBaseName, lineageFile)
 		if err != nil {
 			return err
@@ -1964,7 +2013,7 @@ func (tc *Catalog) MarkVersionForDeletion(ctx context.Context, req *coordinatorp
 
 func (tc *Catalog) updateProtoRemoveVersionEntries(versionFilePb *coordinatorpb.CollectionVersionFile, versions []int64) error {
 	// Check if version history exists
-	if versionFilePb.GetVersionHistory() == nil || len(versionFilePb.GetVersionHistory().Versions) == 0 {
+	if versionFilePb.GetVersionHistory() == nil {
 		log.Error("version history not found")
 		return errors.New("version history not found")
 	}
@@ -2000,14 +2049,14 @@ func (tc *Catalog) getNumberOfActiveVersions(versionFilePb *coordinatorpb.Collec
 	return len(activeVersions)
 }
 
-func (tc *Catalog) getOldestVersionTs(versionFilePb *coordinatorpb.CollectionVersionFile) time.Time {
+func (tc *Catalog) getOldestVersionTs(versionFilePb *coordinatorpb.CollectionVersionFile) *time.Time {
 	if versionFilePb.GetVersionHistory() == nil || len(versionFilePb.GetVersionHistory().Versions) == 0 {
-		// Returning a zero timestamp that represents an unset value.
-		return time.Time{}
+		return nil
 	}
 	oldestVersionTs := versionFilePb.GetVersionHistory().Versions[0].CreatedAtSecs
 
-	return time.Unix(oldestVersionTs, 0)
+	ts := time.Unix(oldestVersionTs, 0)
+	return &ts
 }
 
 func (tc *Catalog) DeleteVersionEntriesForCollection(ctx context.Context, tenantID string, collectionID string, versions []int64) error {
@@ -2021,8 +2070,7 @@ func (tc *Catalog) DeleteVersionEntriesForCollection(ctx context.Context, tenant
 
 		// Read the existing version file
 		collectionIDPtr := &collectionID
-		isDeleted := false
-		collectionEntry, err := tc.metaDomain.CollectionDb(ctx).GetCollectionWithoutMetadata(collectionIDPtr, nil, &isDeleted)
+		collectionEntry, err := tc.metaDomain.CollectionDb(ctx).GetCollectionWithoutMetadata(collectionIDPtr, nil, nil)
 		if err != nil {
 			return err
 		}
@@ -2042,14 +2090,18 @@ func (tc *Catalog) DeleteVersionEntriesForCollection(ctx context.Context, tenant
 		}
 
 		numActiveVersions := tc.getNumberOfActiveVersions(versionFilePb)
-		if numActiveVersions < 1 {
+		if numActiveVersions < 1 && !collectionEntry.IsDeleted {
 			// No remaining valid versions after GC.
 			return errors.New("no valid versions after gc")
 		}
 
 		// Get the creation time of the oldest version.
 		oldestVersionTs := tc.getOldestVersionTs(versionFilePb)
-		if oldestVersionTs.IsZero() {
+		if oldestVersionTs == nil {
+			if !collectionEntry.IsDeleted {
+				return errors.New("oldest version timestamp is nil after GC, this should only happen if all versions are deleted")
+			}
+		} else if oldestVersionTs.IsZero() {
 			// This should never happen.
 			log.Error("oldest version timestamp is zero after GC.", zap.String("collection_id", collectionID))
 			// No versions to delete.
@@ -2069,7 +2121,7 @@ func (tc *Catalog) DeleteVersionEntriesForCollection(ctx context.Context, tenant
 		}
 
 		// Update the version file name in Postgres table as a CAS operation
-		rowsAffected, err := tc.metaDomain.CollectionDb(ctx).UpdateVersionRelatedFields(collectionID, existingVersionFileName, newVerFileFullPath, &oldestVersionTs, &numActiveVersions)
+		rowsAffected, err := tc.metaDomain.CollectionDb(ctx).UpdateVersionRelatedFields(collectionID, existingVersionFileName, newVerFileFullPath, oldestVersionTs, &numActiveVersions)
 		if err != nil {
 			// Delete the newly created version file from S3 since it is not needed
 			tc.s3Store.DeleteVersionFile(tenantID, collectionEntry.DatabaseID, collectionID, newVersionFileName)
@@ -2117,6 +2169,20 @@ func (tc *Catalog) BatchGetCollectionVersionFilePaths(ctx context.Context, colle
 	return &result, nil
 }
 
+func (tc *Catalog) BatchGetCollectionSoftDeleteStatus(ctx context.Context, collectionIds []string) (*coordinatorpb.BatchGetCollectionSoftDeleteStatusResponse, error) {
+	result := coordinatorpb.BatchGetCollectionSoftDeleteStatusResponse{
+		CollectionIdToIsSoftDeleted: make(map[string]bool),
+	}
+
+	status, err := tc.metaDomain.CollectionDb(ctx).BatchGetCollectionSoftDeleteStatus(collectionIds)
+	if err != nil {
+		return nil, err
+	}
+	result.CollectionIdToIsSoftDeleted = status
+
+	return &result, nil
+}
+
 func (tc *Catalog) GetVersionFileNamesForCollection(ctx context.Context, tenantID string, collectionID string) (string, error) {
 	collectionIDPtr := &collectionID
 	isDeleted := false