[ENH]: Globalize dead letter queue in SysDB

tanujnay112 · tanujnay112 · commit e5ea3514525f · 2025-12-09T18:14:54.000-05:00
diff --git a/go/pkg/sysdb/coordinator/coordinator.go b/go/pkg/sysdb/coordinator/coordinator.go
@@ -331,3 +331,7 @@ func (s *Coordinator) FinishDatabaseDeletion(ctx context.Context, req *coordinat
 	}
 	return res, nil
 }
+
+func (s *Coordinator) IncrementCompactionFailureCount(ctx context.Context, collectionID types.UniqueID) (int32, error) {
+	return s.catalog.IncrementCompactionFailureCount(ctx, collectionID)
+}
diff --git a/go/pkg/sysdb/coordinator/model/collection.go b/go/pkg/sysdb/coordinator/model/collection.go
@@ -28,6 +28,7 @@ type Collection struct {
 	VersionFileName            string
 	CreatedAt                  time.Time
 	DatabaseId                 types.UniqueID
+	CompactionFailureCount     int32
 }
 
 type CollectionToGc struct {
diff --git a/go/pkg/sysdb/coordinator/model_db_convert.go b/go/pkg/sysdb/coordinator/model_db_convert.go
@@ -41,6 +41,7 @@ func convertCollectionToModel(collectionAndMetadataList []*dbmodel.CollectionAnd
 			CreatedAt:                  collectionAndMetadata.Collection.CreatedAt,
 			UpdatedAt:                  collectionAndMetadata.Collection.UpdatedAt.Unix(),
 			DatabaseId:                 types.MustParse(collectionAndMetadata.Collection.DatabaseID),
+			CompactionFailureCount:     collectionAndMetadata.Collection.CompactionFailureCount,
 		}
 		collection.Metadata = convertCollectionMetadataToModel(collectionAndMetadata.CollectionMetadata)
 		collections = append(collections, collection)
diff --git a/go/pkg/sysdb/coordinator/table_catalog.go b/go/pkg/sysdb/coordinator/table_catalog.go
@@ -2465,3 +2465,7 @@ func (tc *Catalog) GetVersionFileNamesForCollection(ctx context.Context, tenantI
 func (tc *Catalog) FinishDatabaseDeletion(ctx context.Context, cutoffTime time.Time) (uint64, error) {
 	return tc.metaDomain.DatabaseDb(ctx).FinishDatabaseDeletion(cutoffTime)
 }
+
+func (tc *Catalog) IncrementCompactionFailureCount(ctx context.Context, collectionID types.UniqueID) (int32, error) {
+	return tc.metaDomain.CollectionDb(ctx).IncrementCompactionFailureCount(collectionID.String())
+}
diff --git a/go/pkg/sysdb/grpc/collection_service.go b/go/pkg/sysdb/grpc/collection_service.go
@@ -767,3 +767,22 @@ func (s *Server) BatchGetCollectionSoftDeleteStatus(ctx context.Context, req *co
 	}
 	return res, nil
 }
+
+func (s *Server) IncrementCompactionFailureCount(ctx context.Context, req *coordinatorpb.IncrementCompactionFailureCountRequest) (*coordinatorpb.IncrementCompactionFailureCountResponse, error) {
+	collectionID := req.CollectionId
+	parsedCollectionID, err := types.ToUniqueID(&collectionID)
+	if err != nil {
+		log.Error("IncrementCompactionFailureCount failed. collection id format error", zap.Error(err), zap.String("collection_id", collectionID))
+		return nil, grpcutils.BuildInternalGrpcError(err.Error())
+	}
+
+	failureCount, err := s.coordinator.IncrementCompactionFailureCount(ctx, parsedCollectionID)
+	if err != nil {
+		log.Error("IncrementCompactionFailureCount failed", zap.Error(err), zap.String("collection_id", collectionID))
+		return nil, grpcutils.BuildInternalGrpcError(err.Error())
+	}
+
+	return &coordinatorpb.IncrementCompactionFailureCountResponse{
+		FailureCount: failureCount,
+	}, nil
+}
diff --git a/go/pkg/sysdb/grpc/proto_model_convert.go b/go/pkg/sysdb/grpc/proto_model_convert.go
@@ -62,7 +62,8 @@ func convertCollectionToProto(collection *model.Collection) *coordinatorpb.Colle
 			Seconds: collection.UpdatedAt,
 			Nanos:   0,
 		},
-		DatabaseId: &dbId,
+		DatabaseId:             &dbId,
+		CompactionFailureCount: collection.CompactionFailureCount,
 	}
 
 	if collection.RootCollectionID != nil {
diff --git a/go/pkg/sysdb/metastore/db/dao/collection.go b/go/pkg/sysdb/metastore/db/dao/collection.go
@@ -550,6 +550,7 @@ func (s *collectionDb) UpdateLogPositionAndVersionInfo(
 		"size_bytes_post_compaction":    sizeBytesPostCompaction,
 		"last_compaction_time_secs":     lastCompactionTimeSecs,
 		"num_versions":                  numVersions,
+		"compaction_failure_count":      0, // Reset on successful compaction
 	}
 
 	if schemaStr != nil {
@@ -592,9 +593,9 @@ func (s *collectionDb) UpdateLogPositionVersionTotalRecordsAndLogicalSize(collec
 	version := currentCollectionVersion + 1
 	// only writing if schemaStr is not nil to avoid overwriting the schemaStr
 	if schemaStr != nil {
-		err = s.db.Model(&dbmodel.Collection{}).Where("id = ?", collectionID).Updates(map[string]interface{}{"log_position": logPosition, "version": version, "total_records_post_compaction": totalRecordsPostCompaction, "size_bytes_post_compaction": sizeBytesPostCompaction, "last_compaction_time_secs": lastCompactionTimeSecs, "tenant": tenant, "schema_str": schemaStr}).Error
+		err = s.db.Model(&dbmodel.Collection{}).Where("id = ?", collectionID).Updates(map[string]interface{}{"log_position": logPosition, "version": version, "total_records_post_compaction": totalRecordsPostCompaction, "size_bytes_post_compaction": sizeBytesPostCompaction, "last_compaction_time_secs": lastCompactionTimeSecs, "tenant": tenant, "schema_str": schemaStr, "compaction_failure_count": 0}).Error
 	} else {
-		err = s.db.Model(&dbmodel.Collection{}).Where("id = ?", collectionID).Updates(map[string]interface{}{"log_position": logPosition, "version": version, "total_records_post_compaction": totalRecordsPostCompaction, "size_bytes_post_compaction": sizeBytesPostCompaction, "last_compaction_time_secs": lastCompactionTimeSecs, "tenant": tenant}).Error
+		err = s.db.Model(&dbmodel.Collection{}).Where("id = ?", collectionID).Updates(map[string]interface{}{"log_position": logPosition, "version": version, "total_records_post_compaction": totalRecordsPostCompaction, "size_bytes_post_compaction": sizeBytesPostCompaction, "last_compaction_time_secs": lastCompactionTimeSecs, "tenant": tenant, "compaction_failure_count": 0}).Error
 	}
 	if err != nil {
 		return 0, err
@@ -714,3 +715,16 @@ func (s *collectionDb) BatchGetCollectionSoftDeleteStatus(collectionIDs []string
 	}
 	return result, nil
 }
+
+func (s *collectionDb) IncrementCompactionFailureCount(collectionID string) (int32, error) {
+	var collection dbmodel.Collection
+	err := s.db.Model(&dbmodel.Collection{}).
+		Where("id = ?", collectionID).
+		UpdateColumn("compaction_failure_count", gorm.Expr("compaction_failure_count + 1")).
+		First(&collection).Error
+	if err != nil {
+		log.Error("IncrementCompactionFailureCount failed", zap.Error(err), zap.String("collectionID", collectionID))
+		return 0, err
+	}
+	return collection.CompactionFailureCount, nil
+}
diff --git a/go/pkg/sysdb/metastore/db/dbmodel/collection.go b/go/pkg/sysdb/metastore/db/dbmodel/collection.go
@@ -28,6 +28,7 @@ type Collection struct {
 	NumVersions                uint32          `gorm:"num_versions;type:integer;default:0"`
 	OldestVersionTs            time.Time       `gorm:"oldest_version_ts;type:timestamp"`
 	Tenant                     string          `gorm:"tenant"`
+	CompactionFailureCount     int32           `gorm:"compaction_failure_count;default:0"`
 }
 
 type CollectionToGc struct {
@@ -73,4 +74,5 @@ type ICollectionDb interface {
 	UpdateCollectionLineageFilePath(collectionID string, currentLineageFilePath *string, newLineageFilePath string) error
 	BatchGetCollectionVersionFilePaths(collectionIDs []string) (map[string]string, error)
 	BatchGetCollectionSoftDeleteStatus(collectionIDs []string) (map[string]bool, error)
+	IncrementCompactionFailureCount(collectionID string) (int32, error)
 }
diff --git a/go/pkg/sysdb/metastore/db/migrations/atlas.sum b/go/pkg/sysdb/metastore/db/migrations/atlas.sum
@@ -1,4 +1,4 @@
-h1:TI+7Y8LdtEPUUsspBjfrE9KdjS8XS8SLfKUoSJhmhEI=
+h1:Qu6wEENNffsixNyG6bzBEGihu8sC5xa6dtQB37QF7Gg=
 20240313233558.sql h1:Gv0TiSYsqGoOZ2T2IWvX4BOasauxool8PrBOIjmmIdg=
 20240321194713.sql h1:kVkNpqSFhrXGVGFFvL7JdK3Bw31twFcEhI6A0oCFCkg=
 20240327075032.sql h1:nlr2J74XRU8erzHnKJgMr/tKqJxw9+R6RiiEBuvuzgo=
@@ -26,3 +26,4 @@ h1:TI+7Y8LdtEPUUsspBjfrE9KdjS8XS8SLfKUoSJhmhEI=
 20251114125442.sql h1:oRHN+AO+xYnYa3aF0QzSa3T/TRi8ETCydp2sDT/nSnI=
 20251114134400.sql h1:N30qnVNjR+d4RoArJ11YrixyIsNODxqXFpgiREEhczs=
 20251116154842.sql h1:G0qy4MPDayH+Y9/Dm9PS2xwvyLt+nmIw90uJTzZSJUM=
+20251209143000.sql h1:r7locVeKMfzp+/olpFQHx1Shy7juhR5Y+xWZw3vel5M=
diff --git a/idl/chromadb/proto/chroma.proto b/idl/chromadb/proto/chroma.proto
@@ -74,6 +74,9 @@ message Collection {
   // This is the database id of the collection.
   optional string database_id = 17;
   optional string schema_str = 18;
+  // Number of consecutive compaction failures for this collection
+  // Defaults to 0.
+  int32 compaction_failure_count = 19;
 }
 
 message Database {
diff --git a/idl/chromadb/proto/coordinator.proto b/idl/chromadb/proto/coordinator.proto
@@ -666,6 +666,14 @@ message FinishAttachedFunctionDeletionRequest {
 
 message FinishAttachedFunctionDeletionResponse {}
 
+message IncrementCompactionFailureCountRequest {
+  string collection_id = 1;
+}
+
+message IncrementCompactionFailureCountResponse {
+  int32 failure_count = 1;
+}
+
 service SysDB {
   rpc CreateDatabase(CreateDatabaseRequest) returns (CreateDatabaseResponse) {}
   rpc GetDatabase(GetDatabaseRequest) returns (GetDatabaseResponse) {}
@@ -714,4 +722,5 @@ service SysDB {
   rpc GetSoftDeletedAttachedFunctions(GetSoftDeletedAttachedFunctionsRequest) returns (GetSoftDeletedAttachedFunctionsResponse) {}
   rpc FinishAttachedFunctionDeletion(FinishAttachedFunctionDeletionRequest) returns (FinishAttachedFunctionDeletionResponse) {}
   rpc FlushCollectionCompactionAndAttachedFunction(FlushCollectionCompactionAndAttachedFunctionRequest) returns (FlushCollectionCompactionAndAttachedFunctionResponse) {}
+  rpc IncrementCompactionFailureCount(IncrementCompactionFailureCountRequest) returns (IncrementCompactionFailureCountResponse) {}
 }
diff --git a/rust/sysdb/src/sqlite.rs b/rust/sysdb/src/sqlite.rs
@@ -379,6 +379,7 @@ impl SqliteSysDb {
             lineage_file_path: None,
             updated_at: SystemTime::UNIX_EPOCH,
             database_id: database_uuid,
+            compaction_failure_count: 0,
         })
     }
 
@@ -860,6 +861,7 @@ impl SqliteSysDb {
                     lineage_file_path: None,
                     updated_at: SystemTime::UNIX_EPOCH,
                     database_id,
+                    compaction_failure_count: 0,
                 }))
             })
             .collect::<Result<Vec<_>, GetCollectionsError>>()?;
diff --git a/rust/sysdb/src/sysdb.rs b/rust/sysdb/src/sysdb.rs
@@ -346,6 +346,7 @@ impl SysDb {
                     lineage_file_path: None,
                     updated_at: SystemTime::now(),
                     database_id: DatabaseUuid::new(),
+                    compaction_failure_count: 0,
                 };
 
                 test_sysdb.add_collection(collection.clone());
@@ -694,6 +695,19 @@ impl SysDb {
         }
     }
 
+    /// Increment the compaction failure count for a collection.
+    /// Returns the new failure count.
+    pub async fn increment_compaction_failure_count(
+        &mut self,
+        collection_id: CollectionUuid,
+    ) -> Result<i32, IncrementCompactionFailureCountError> {
+        match self {
+            SysDb::Grpc(grpc) => grpc.increment_compaction_failure_count(collection_id).await,
+            SysDb::Test(test) => Ok(test.increment_compaction_failure_count(collection_id)),
+            SysDb::Sqlite(_) => todo!(),
+        }
+    }
+
     pub async fn reset(&mut self) -> Result<ResetResponse, ResetError> {
         match self {
             SysDb::Grpc(grpc) => grpc.reset().await,
@@ -1722,6 +1736,18 @@ impl GrpcSysDb {
         Ok(res.into_inner().collection_id_to_success)
     }
 
+    async fn increment_compaction_failure_count(
+        &mut self,
+        collection_id: CollectionUuid,
+    ) -> Result<i32, IncrementCompactionFailureCountError> {
+        let req = chroma_proto::IncrementCompactionFailureCountRequest {
+            collection_id: collection_id.0.to_string(),
+        };
+
+        let res = self.client.increment_compaction_failure_count(req).await?;
+        Ok(res.into_inner().failure_count)
+    }
+
     async fn update_tenant(
         &mut self,
         tenant_id: String,
@@ -2194,6 +2220,20 @@ impl ChromaError for DeleteCollectionVersionError {
     }
 }
 
+#[derive(Error, Debug)]
+pub enum IncrementCompactionFailureCountError {
+    #[error("Failed to increment compaction failure count: {0}")]
+    FailedToIncrement(#[from] tonic::Status),
+}
+
+impl ChromaError for IncrementCompactionFailureCountError {
+    fn code(&self) -> ErrorCodes {
+        match self {
+            IncrementCompactionFailureCountError::FailedToIncrement(_) => ErrorCodes::Internal,
+        }
+    }
+}
+
 //////////////////////////  Attached Function Operations //////////////////////////
 
 impl SysDb {
diff --git a/rust/sysdb/src/test_sysdb.rs b/rust/sysdb/src/test_sysdb.rs
@@ -693,6 +693,18 @@ impl TestSysDb {
         inner.tenant_resource_names.insert(tenant_id, resource_name);
         Ok(UpdateTenantResponse {})
     }
+
+    /// Increment the compaction failure count for a collection.
+    /// Returns the new failure count.
+    pub fn increment_compaction_failure_count(&mut self, collection_id: CollectionUuid) -> i32 {
+        let mut inner = self.inner.lock();
+        if let Some(collection) = inner.collections.get_mut(&collection_id) {
+            collection.compaction_failure_count += 1;
+            collection.compaction_failure_count
+        } else {
+            0
+        }
+    }
 }
 
 fn attached_function_to_proto(
diff --git a/rust/types/src/collection.rs b/rust/types/src/collection.rs
@@ -132,6 +132,10 @@ pub struct Collection {
     pub updated_at: SystemTime,
     #[serde(skip)]
     pub database_id: DatabaseUuid,
+    /// Number of consecutive compaction failures for this collection.
+    /// Used by the scheduler to track and skip collections that repeatedly fail compaction.
+    #[serde(skip)]
+    pub compaction_failure_count: i32,
 }
 
 impl Default for Collection {
@@ -155,6 +159,7 @@ impl Default for Collection {
             lineage_file_path: None,
             updated_at: SystemTime::now(),
             database_id: DatabaseUuid::new(),
+            compaction_failure_count: 0,
         }
     }
 }
@@ -333,6 +338,7 @@ impl TryFrom<chroma_proto::Collection> for Collection {
             lineage_file_path: proto_collection.lineage_file_path,
             updated_at,
             database_id,
+            compaction_failure_count: proto_collection.compaction_failure_count,
         })
     }
 }
@@ -377,6 +383,7 @@ impl TryFrom<Collection> for chroma_proto::Collection {
             lineage_file_path: value.lineage_file_path,
             updated_at: Some(value.updated_at.into()),
             database_id: Some(value.database_id.0.to_string()),
+            compaction_failure_count: value.compaction_failure_count,
         })
     }
 }
@@ -460,6 +467,7 @@ mod test {
                 nanos: 1,
             }),
             database_id: Some("00000000-0000-0000-0000-000000000000".to_string()),
+            compaction_failure_count: 0,
         };
         let converted_collection: Collection = proto_collection.try_into().unwrap();
         assert_eq!(
diff --git a/rust/worker/src/compactor/compaction_manager.rs b/rust/worker/src/compactor/compaction_manager.rs
diff --git a/rust/worker/src/compactor/config.rs b/rust/worker/src/compactor/config.rs
diff --git a/rust/worker/src/compactor/scheduler.rs b/rust/worker/src/compactor/scheduler.rs

Original file line number	Diff line number	Diff line change
`@@ -331,3 +331,7 @@ func (s Coordinator) FinishDatabaseDeletion(ctx context.Context, req coordinat`
`331`	`331`	`}`
`332`	`332`	`return res, nil`
`333`	`333`	`}`
	`334`	`+`
	`335`	`+func (s *Coordinator) IncrementCompactionFailureCount(ctx context.Context, collectionID types.UniqueID) (int32, error) {`
	`336`	`+ return s.catalog.IncrementCompactionFailureCount(ctx, collectionID)`
	`337`	`+}`
Original file line number	Diff line number	Diff line change
`@@ -28,6 +28,7 @@ type Collection struct {`
`28`	`28`	`VersionFileName string`
`29`	`29`	`CreatedAt time.Time`
`30`	`30`	`DatabaseId types.UniqueID`
	`31`	`+ CompactionFailureCount int32`
`31`	`32`	`}`
`32`	`33`
`33`	`34`	`type CollectionToGc struct {`
Original file line number	Diff line number	Diff line change
`@@ -41,6 +41,7 @@ func convertCollectionToModel(collectionAndMetadataList []*dbmodel.CollectionAnd`
`41`	`41`	`CreatedAt: collectionAndMetadata.Collection.CreatedAt,`
`42`	`42`	`UpdatedAt: collectionAndMetadata.Collection.UpdatedAt.Unix(),`
`43`	`43`	`DatabaseId: types.MustParse(collectionAndMetadata.Collection.DatabaseID),`
	`44`	`+ CompactionFailureCount: collectionAndMetadata.Collection.CompactionFailureCount,`
`44`	`45`	`}`
`45`	`46`	`collection.Metadata = convertCollectionMetadataToModel(collectionAndMetadata.CollectionMetadata)`
`46`	`47`	`collections = append(collections, collection)`
Original file line number	Diff line number	Diff line change
`@@ -2465,3 +2465,7 @@ func (tc *Catalog) GetVersionFileNamesForCollection(ctx context.Context, tenantI`
`2465`	`2465`	`func (tc *Catalog) FinishDatabaseDeletion(ctx context.Context, cutoffTime time.Time) (uint64, error) {`
`2466`	`2466`	`return tc.metaDomain.DatabaseDb(ctx).FinishDatabaseDeletion(cutoffTime)`
`2467`	`2467`	`}`
	`2468`	`+`
	`2469`	`+func (tc *Catalog) IncrementCompactionFailureCount(ctx context.Context, collectionID types.UniqueID) (int32, error) {`
	`2470`	`+ return tc.metaDomain.CollectionDb(ctx).IncrementCompactionFailureCount(collectionID.String())`
	`2471`	`+}`
Original file line number	Diff line number	Diff line change
`@@ -62,7 +62,8 @@ func convertCollectionToProto(collection model.Collection) coordinatorpb.Colle`
`62`	`62`	`Seconds: collection.UpdatedAt,`
`63`	`63`	`Nanos: 0,`
`64`	`64`	`},`
`65`		`- DatabaseId: &dbId,`
	`65`	`+ DatabaseId: &dbId,`
	`66`	`+ CompactionFailureCount: collection.CompactionFailureCount,`
`66`	`67`	`}`
`67`	`68`
`68`	`69`	`if collection.RootCollectionID != nil {`
Original file line number	Diff line number	Diff line change
`@@ -74,6 +74,9 @@ message Collection {`
`74`	`74`	`// This is the database id of the collection.`
`75`	`75`	`optional string database_id = 17;`
`76`	`76`	`optional string schema_str = 18;`
	`77`	`+ // Number of consecutive compaction failures for this collection`
	`78`	`+ // Defaults to 0.`
	`79`	`+ int32 compaction_failure_count = 19;`
`77`	`80`	`}`
`78`	`81`
`79`	`82`	`message Database {`