chore(epoch): Mergin two functions into one generic

ChainSafe · Sep 4, 2024 · 97ddc43 · 97ddc43
1 parent 50d4dab
commit 97ddc43
Showing 1 changed file with 8 additions and 27 deletions.
diff --git a/dot/state/epoch.go b/dot/state/epoch.go
@@ -1007,7 +1007,7 @@ func (s *EpochState) FinalizeBABENextEpochData(finalizedHeader *types.Header) er
 		if e <= nextEpoch {
 			delete(s.nextEpochData, e)
 			// remove the epoch data from the database
-			if err = s.deleteEpochDataFromDisk(e); err != nil {
+			if err = deleteDataFromDisk[types.NextEpochData](s.db, e, nextEpochDataPrefix); err != nil {
 				return fmt.Errorf("cannot delete next epoch data from the database: %w", err)
 			}
 		}
@@ -1016,13 +1016,15 @@ func (s *EpochState) FinalizeBABENextEpochData(finalizedHeader *types.Header) er
 	return nil
 }
 
-func (s *EpochState) deleteEpochDataFromDisk(epoch uint64) error {
-	configKeysToDelete, err := getDataKeysFromDisk[types.NextEpochData](s.db, nextEpochDataPrefix, epoch)
+// deleteDataFromDisk is a generic function that deletes all the nextEpochData or nextConfigData
+// for a given epoch from the database
+func deleteDataFromDisk[T types.NextEpochData | types.NextConfigDataV1](db database.Table, epoch uint64, prefix []byte) error {
+	keysToDelete, err := getDataKeysFromDisk[T](db, prefix, epoch)
 	if err != nil {
 		return fmt.Errorf("cannot get next config data keys from disk: %w", err)
 	}
-	batch := s.db.NewBatch()
-	for _, key := range configKeysToDelete {
+	batch := db.NewBatch()
+	for _, key := range keysToDelete {
 		err = batch.Del([]byte(key))
 		if err != nil {
 			return fmt.Errorf("cannot delete next config data from the database: %w", err)
@@ -1125,7 +1127,7 @@ func (s *EpochState) FinalizeBABENextConfigData(finalizedHeader *types.Header) e
 		if e <= nextEpoch {
 			delete(s.nextConfigData, e)
 			// remove the config data from the database
-			if err = s.deleteNextConfigDataFromDisk(e); err != nil {
+			if err = deleteDataFromDisk[types.NextConfigDataV1](s.db, e, nextConfigDataPrefix); err != nil {
 				return fmt.Errorf("cannot delete next config data from the database: %w", err)
 			}
 		}
@@ -1134,27 +1136,6 @@ func (s *EpochState) FinalizeBABENextConfigData(finalizedHeader *types.Header) e
 	return nil
 }
 
-func (s *EpochState) deleteNextConfigDataFromDisk(epoch uint64) error {
-	configKeysToDelete, err := getDataKeysFromDisk[types.NextConfigDataV1](s.db, nextConfigDataPrefix, epoch)
-	if err != nil {
-		return fmt.Errorf("cannot get next config data keys from disk: %w", err)
-	}
-
-	batch := s.db.NewBatch()
-	for _, key := range configKeysToDelete {
-		err = batch.Del([]byte(key))
-		if err != nil {
-			return fmt.Errorf("cannot delete next config data from the database: %w", err)
-		}
-	}
-
-	if err := batch.Flush(); err != nil {
-		return fmt.Errorf("cannot flush deletion batch: %w", err)
-	}
-
-	return nil
-}
-
 // findFinalizedHeaderForEpoch given a specific epoch (the key) will go through the hashes looking
 // for a database persisted hash (belonging to the finalized chain)
 // which contains the right configuration or data to be persisted and safely used