[ENH]: ListCollectionsToGc should return lineage file path, group by fork tree

Author: codetheweb

go/pkg/sysdb/coordinator/coordinator.go

@@ -247,8 +247,8 @@ func (s *Coordinator) FlushCollectionCompaction(ctx context.Context, flushCollec
 	return s.catalog.FlushCollectionCompaction(ctx, flushCollectionCompaction)
 }
 
-func (s *Coordinator) ListCollectionsToGc(ctx context.Context, cutoffTimeSecs *uint64, limit *uint64) ([]*model.CollectionToGc, error) {
-	return s.catalog.ListCollectionsToGc(ctx, cutoffTimeSecs, limit)
+func (s *Coordinator) ListCollectionsToGc(ctx context.Context, cutoffTimeSecs *uint64, limit *uint64, tenantID *string) ([]*model.CollectionToGc, error) {
+	return s.catalog.ListCollectionsToGc(ctx, cutoffTimeSecs, limit, tenantID)
 }
 
 func (s *Coordinator) ListCollectionVersions(ctx context.Context, collectionID types.UniqueID, tenantID string, maxCount *int64, versionsBefore *int64, versionsAtOrAfter *int64, includeMarkedForDeletion bool) ([]*coordinatorpb.CollectionVersionInfo, error) {

go/pkg/sysdb/coordinator/coordinator_test.go

@@ -79,7 +79,7 @@ func (suite *APIsTestSuite) SetupTest() {
 		collection.Name = "collection_" + suite.T().Name() + strconv.Itoa(index)
 	}
 	ctx := context.Background()
-	c, err := NewCoordinator(ctx, SoftDelete, suite.s3MetaStore, false)
+	c, err := NewCoordinator(ctx, SoftDelete, suite.s3MetaStore, true)
 	if err != nil {
 		suite.T().Fatalf("error creating coordinator: %v", err)
 	}
@@ -1482,6 +1482,17 @@ func (suite *APIsTestSuite) TestForkCollection() {
 	collections, err := suite.coordinator.GetCollections(ctx, forkCollectionWithSameName.TargetCollectionID, nil, suite.tenantName, suite.databaseName, nil, nil)
 	suite.NoError(err)
 	suite.Empty(collections)
+
+	res, err := suite.coordinator.ListCollectionsToGc(ctx, nil, nil, nil)
+	suite.NoError(err)
+	suite.NotEmpty(res)
+	suite.Equal(1, len(res))
+	// ListCollectionsToGc groups by fork trees and should always return the root of the tree
+	suite.Equal(forkCollectionWithSameName.SourceCollectionID, res[0].ID)
+
+	exists, err := suite.s3MetaStore.HasObjectWithPrefix(ctx, *res[0].LineageFilePath)
+	suite.NoError(err)
+	suite.True(exists, "Lineage file should exist in S3")
 }
 
 func (suite *APIsTestSuite) TestCountForks() {

go/pkg/sysdb/coordinator/model/collection.go

@@ -29,6 +29,7 @@ type CollectionToGc struct {
 	Name            string
 	VersionFilePath string
 	LatestVersion   int64
+	LineageFilePath *string
 }
 
 type CreateCollection struct {

go/pkg/sysdb/coordinator/model_db_convert.go

@@ -53,6 +53,7 @@ func convertCollectionToGcToModel(collectionToGc []*dbmodel.CollectionToGc) []*m
 			VersionFilePath: collectionInfo.VersionFileName,
 			LatestVersion:   int64(collectionInfo.Version),
 			TenantID:        collectionInfo.TenantID,
+			LineageFilePath: collectionInfo.LineageFileName,
 		}
 		collections = append(collections, &collection)
 	}

go/pkg/sysdb/coordinator/table_catalog.go

@@ -444,14 +444,14 @@ func (tc *Catalog) GetCollectionSize(ctx context.Context, collectionID types.Uni
 	return total_records_post_compaction, nil
 }
 
-func (tc *Catalog) ListCollectionsToGc(ctx context.Context, cutoffTimeSecs *uint64, limit *uint64) ([]*model.CollectionToGc, error) {
+func (tc *Catalog) ListCollectionsToGc(ctx context.Context, cutoffTimeSecs *uint64, limit *uint64, tenantID *string) ([]*model.CollectionToGc, error) {
 	tracer := otel.Tracer
 	if tracer != nil {
 		_, span := tracer.Start(ctx, "Catalog.ListCollectionsToGc")
 		defer span.End()
 	}
 
-	collectionsToGc, err := tc.metaDomain.CollectionDb(ctx).ListCollectionsToGc(cutoffTimeSecs, limit)
+	collectionsToGc, err := tc.metaDomain.CollectionDb(ctx).ListCollectionsToGc(cutoffTimeSecs, limit, tenantID)
 
 	if err != nil {
 		return nil, err

go/pkg/sysdb/coordinator/table_catalog_test.go

@@ -676,6 +676,7 @@ func TestCatalog_ListCollectionsToGc(t *testing.T) {
 	limit := uint64(10)
 
 	// Mock collections to return
+	lineageFileName := "lineage_file_1"
 	collectionsToGc := []*dbmodel.CollectionToGc{
 		{
 			ID:              "00000000-0000-0000-0000-000000000001",
@@ -692,15 +693,16 @@ func TestCatalog_ListCollectionsToGc(t *testing.T) {
 			VersionFileName: "2_existing_version",
 			OldestVersionTs: time.Now().Add(-36 * time.Hour), // 36 hours ago
 			NumVersions:     2,
+			LineageFileName: &lineageFileName,
 		},
 	}
 
 	// Setup mock behaviors
 	mockMetaDomain.On("CollectionDb", mock.Anything).Return(mockCollectionDb)
-	mockCollectionDb.On("ListCollectionsToGc", &cutoffTimeSecs, &limit).Return(collectionsToGc, nil)
+	mockCollectionDb.On("ListCollectionsToGc", &cutoffTimeSecs, &limit, (*string)(nil)).Return(collectionsToGc, nil)
 
 	// Execute test
-	result, err := catalog.ListCollectionsToGc(context.Background(), &cutoffTimeSecs, &limit)
+	result, err := catalog.ListCollectionsToGc(context.Background(), &cutoffTimeSecs, &limit, nil)
 
 	// Verify results
 	assert.NoError(t, err)
@@ -718,6 +720,7 @@ func TestCatalog_ListCollectionsToGc(t *testing.T) {
 	assert.Equal(t, "collection2", result[1].Name)
 	assert.Equal(t, int64(2), result[1].LatestVersion)
 	assert.Equal(t, "2_existing_version", result[1].VersionFilePath)
+	assert.Equal(t, "lineage_file_1", *result[1].LineageFilePath)
 
 	// Verify mock expectations
 	mockMetaDomain.AssertExpectations(t)
@@ -748,10 +751,10 @@ func TestCatalog_ListCollectionsToGc_NilParameters(t *testing.T) {
 
 	// Setup mock behaviors
 	mockMetaDomain.On("CollectionDb", mock.Anything).Return(mockCollectionDb)
-	mockCollectionDb.On("ListCollectionsToGc", (*uint64)(nil), (*uint64)(nil)).Return(collectionsToGc, nil)
+	mockCollectionDb.On("ListCollectionsToGc", (*uint64)(nil), (*uint64)(nil), (*string)(nil)).Return(collectionsToGc, nil)
 
 	// Execute test with nil parameters
-	result, err := catalog.ListCollectionsToGc(context.Background(), nil, nil)
+	result, err := catalog.ListCollectionsToGc(context.Background(), nil, nil, nil)
 
 	// Verify results
 	assert.NoError(t, err)

go/pkg/sysdb/grpc/collection_service.go

@@ -495,7 +495,7 @@ func (s *Server) ListCollectionsToGc(ctx context.Context, req *coordinatorpb.Lis
 		absoluteCutoffTimeSecs = &cutoffTime
 	}
 
-	collectionsToGc, err := s.coordinator.ListCollectionsToGc(ctx, absoluteCutoffTimeSecs, req.Limit)
+	collectionsToGc, err := s.coordinator.ListCollectionsToGc(ctx, absoluteCutoffTimeSecs, req.Limit, req.TenantId)
 	if err != nil {
 		log.Error("ListCollectionsToGc failed", zap.Error(err))
 		return nil, grpcutils.BuildInternalGrpcError(err.Error())
@@ -507,6 +507,7 @@ func (s *Server) ListCollectionsToGc(ctx context.Context, req *coordinatorpb.Lis
 			Name:            collectionToGc.Name,
 			VersionFilePath: collectionToGc.VersionFilePath,
 			TenantId:        collectionToGc.TenantID,
+			LineageFilePath: collectionToGc.LineageFilePath,
 		})
 	}
 	return res, nil

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

@@ -56,17 +56,29 @@ func (s *collectionDb) GetCollections(id *string, name *string, tenantID string,
 	return s.getCollections(id, name, tenantID, databaseName, limit, offset, false)
 }
 
-func (s *collectionDb) ListCollectionsToGc(cutoffTimeSecs *uint64, limit *uint64) ([]*dbmodel.CollectionToGc, error) {
-	var collections []*dbmodel.CollectionToGc
-	// Use the read replica for this so as to not overwhelm the writer.
+func (s *collectionDb) ListCollectionsToGc(cutoffTimeSecs *uint64, limit *uint64, tenantID *string) ([]*dbmodel.CollectionToGc, error) {
+	// There are three types of collections:
+	// 1. Regular: a collection created by a normal call to create_collection(). Does not have a root_collection_id or a lineage_file_name.
+	// 2. Root of fork tree: a collection created by a call to create_collection() which was later the source of a fork with fork(). Has a lineage_file_name.
+	// 3. Fork of a root: a collection created by a call to fork(). Has a root_collection_id.
+	//
+	// For the purposes of this method, we group by fork "trees". A fork tree is a root collection and all its forks (or, in the case of regular collections, a single collection). For every fork tree, we check if at least one collection in the tree meets the GC requirements. If so, we return the root collection of the tree. We ignore forks in the response as the garbage collector will GC forks when run on the root collection.
+
+	sub := s.read_db.Table("collections").
+		Select("COALESCE(NULLIF(root_collection_id, ''), id) AS id, MAX(version) AS max_version, MIN(oldest_version_ts) AS min_oldest_version_ts, SUM(num_versions) AS sum_num_versions").
+		Group("COALESCE(NULLIF(root_collection_id, ''), id)").
+		Where("version_file_name IS NOT NULL").
+		Where("version_file_name != ''")
+
+	if tenantID != nil {
+		sub = sub.Where("tenant = ?", *tenantID)
+	}
+
 	query := s.read_db.Table("collections").
-		Select("collections.id, collections.name, collections.version, collections.version_file_name, collections.oldest_version_ts, collections.num_versions, databases.tenant_id").
+		Select("collections.id, collections.name, collections.version, collections.version_file_name, sub.min_oldest_version_ts AS oldest_version_ts, sub.sum_num_versions AS num_versions, databases.tenant_id, collections.lineage_file_name").
 		Joins("INNER JOIN databases ON collections.database_id = databases.id").
-		Where("version > 0").
-		Where("version_file_name IS NOT NULL").
-		Where("version_file_name != ''").
-		Where("root_collection_id IS NULL OR root_collection_id = ''").
-		Where("lineage_file_name IS NULL OR lineage_file_name = ''")
+		Joins("INNER JOIN (?) AS sub ON collections.id = sub.id", sub)
+
 	// Apply cutoff time filter only if provided
 	if cutoffTimeSecs != nil {
 		cutoffTime := time.Unix(int64(*cutoffTimeSecs), 0)
@@ -80,6 +92,7 @@ func (s *collectionDb) ListCollectionsToGc(cutoffTimeSecs *uint64, limit *uint64
 		query = query.Limit(int(*limit))
 	}
 
+	var collections []*dbmodel.CollectionToGc
 	err := query.Find(&collections).Error
 	if err != nil {
 		return nil, err

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

@@ -37,6 +37,7 @@ type CollectionToGc struct {
 	VersionFileName string    `gorm:"version_file_name"`
 	OldestVersionTs time.Time `gorm:"oldest_version_ts;type:timestamp"`
 	NumVersions     uint32    `gorm:"num_versions;default:0"`
+	LineageFileName *string   `gorm:"lineage_file_name"`
 }
 
 func (v Collection) TableName() string {
@@ -65,7 +66,7 @@ type ICollectionDb interface {
 		sizeBytesPostCompaction uint64, lastCompactionTimeSecs uint64) (int64, error)
 	GetCollectionEntry(collectionID *string, databaseName *string) (*Collection, error)
 	GetCollectionSize(collectionID string) (uint64, error)
-	ListCollectionsToGc(cutoffTimeSecs *uint64, limit *uint64) ([]*CollectionToGc, error)
+	ListCollectionsToGc(cutoffTimeSecs *uint64, limit *uint64, 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/mocks/ICollectionDb.go

@@ -429,6 +429,8 @@ message ListCollectionsToGcRequest {
   // This also allows for a cheap and stateless pagination without using offsets.
   optional uint64 limit = 2;
 
+  optional string tenant_id = 3;
+
   // Design NOTE: When GC calls DeleteCollectionVersion, sysdb will update the
   // time associated with the oldest version of the collection. This allows
   // sysdb to return the collections that have not been GCed for a long time.
@@ -440,6 +442,7 @@ message CollectionToGcInfo {
   string version_file_path = 3;
   int64 latest_version = 4;
   string tenant_id = 5;
+  optional string lineage_file_path = 6;
 }
 
 message ListCollectionsToGcResponse {