chore(epoch): Adding correct use of batch and updating deletion test

dot/state/epoch.go

@@ -70,7 +70,7 @@ type GenesisEpochDescriptor struct {
 
 // EpochState tracks information related to each epoch
 type EpochState struct {
-	db           GetterPutterNewBatcherPrefixIter
+	db           database.Table
 	baseState    *BaseState
 	blockState   *BlockState
 	epochLength  uint64 // measured in slots
@@ -599,7 +599,7 @@ type prefixedKeyBuilder func(epoch uint64) []byte
 // updateEpochDefinitionKey updates the informations from database
 // by querying the raw bytes from prefix + oldEpoch and inserting
 // at prefix + newEpoch and return the values stored at prefix + oldEpoch
-func updateEpochDefinitionKey(db GetterPutterNewBatcherPrefixIter,
+func updateEpochDefinitionKey(db GetterPutterNewBatcher,
 	oldEpoch, newEpoch uint64, usePrefix prefixedKeyBuilder) ([]byte, error) {
 	rawBytes, err := db.Get(usePrefix(oldEpoch))
 	if err != nil {
@@ -631,7 +631,7 @@ func updateEpochDefinitionKey(db GetterPutterNewBatcherPrefixIter,
 }
 
 func getAndUpdateEpochDefinitionKey[T types.ConfigData | types.EpochDataRaw](
-	db GetterPutterNewBatcherPrefixIter, oldEpoch, newEpoch uint64, usePrefix prefixedKeyBuilder) (*T, error) {
+	db GetterPutterNewBatcher, oldEpoch, newEpoch uint64, usePrefix prefixedKeyBuilder) (*T, error) {
 	rawBytes, err := updateEpochDefinitionKey(db, oldEpoch, newEpoch, usePrefix)
 	if err != nil {
 		return nil, fmt.Errorf("updating epoch key definition: %w", err)
@@ -1021,20 +1021,23 @@ func (s *EpochState) deleteEpochDataFromDisk(epoch uint64) error {
 	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 = s.db.NewBatch().Del([]byte(key))
+		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
 }
 
 // getDataKeysFromDisk is a generic function that returns all the nextEpochData or  nextConfigData keys
 // for a given epoch from the database
 func getDataKeysFromDisk[T types.NextEpochData | types.NextConfigDataV1](
-	db GetterPutterNewBatcherPrefixIter, prefix []byte, currentEpoch uint64) (
+	db database.Table, prefix []byte, currentEpoch uint64) (
 	[]string, error) {
 
 	dataKeys := []string{}
@@ -1047,14 +1050,17 @@ func getDataKeysFromDisk[T types.NextEpochData | types.NextConfigDataV1](
 	defer iter.Release()
 
 	for iter.First(); iter.Valid(); iter.Next() {
-		_, epoch, fork, err := getNextEpochOrConfigData[T](iter)
+		key := string(iter.Key())
+
+		keyWithoutPrefix := strings.Split(key, "-")[1]
+		epochPart := strings.Split(keyWithoutPrefix, ":")[0]
+		epoch, err := strconv.ParseUint(epochPart, 10, 64)
 		if err != nil {
 			return dataKeys, err
 		}
-		// if the epoch is the current epoch, then we append the key to the dataKeys
+
 		if epoch == currentEpoch {
-			key := append(prefix, []byte(fmt.Sprintf("-%d:%s", epoch, fork))...)
-			dataKeys = append(dataKeys, string(key))
+			dataKeys = append(dataKeys, key)
 		}
 
 	}
@@ -1134,12 +1140,18 @@ func (s *EpochState) deleteNextConfigDataFromDisk(epoch uint64) error {
 		return fmt.Errorf("cannot get next config data keys from disk: %w", err)
 	}
 
+	batch := s.db.NewBatch()
 	for _, key := range configKeysToDelete {
-		err = s.db.NewBatch().Del([]byte(key))
+		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
 }
 

dot/state/epoch_test.go

@@ -1138,11 +1138,10 @@ func TestDeleteNextEpochDataAndConfig(t *testing.T) {
 	err = epochState.FinalizeBABENextEpochData(expectedHeader)
 	require.NoError(t, err)
 
-	// Check if iterators are invalid, making sure that the data was deleted
-	iter, err := epochState.db.NewPrefixIterator(nextConfigDataPrefix)
-	require.NoError(t, err)
-	require.Equal(t, iter.Valid(), false)
-	epochState.db.NewPrefixIterator(nextEpochDataPrefix)
+	// Check if the next epoch data and config data are not stored in the database
+	// after finalization
+	epochState, err = NewEpochState(db, epochState.blockState, config.BABEConfigurationTestDefault)
 	require.NoError(t, err)
-	require.Equal(t, iter.Valid(), false)
+	require.Equal(t, 0, len(epochState.nextEpochData))
+	require.Equal(t, 0, len(epochState.nextConfigData))
 }

dot/state/inmemory_storage.go

@@ -34,7 +34,7 @@ type InmemoryStorageState struct {
 	blockState *BlockState
 	tries      *Tries
 
-	db GetterPutterNewBatcherPrefixIter
+	db GetterPutterNewBatcher
 	sync.RWMutex
 
 	// change notifiers

dot/state/interfaces.go

@@ -31,11 +31,10 @@ type GetPutter interface {
 
 // GetterPutterNewBatcher has methods to get values and create a
 // new batch.
-type GetterPutterNewBatcherPrefixIter interface {
+type GetterPutterNewBatcher interface {
 	Getter
 	Putter
 	NewBatcher
-	NewPrefixIterator
 }
 
 // Getter gets a value corresponding to the given key.
@@ -72,7 +71,3 @@ type BabeConfigurer interface {
 type Telemetry interface {
 	SendMessage(msg json.Marshaler)
 }
-
-type NewPrefixIterator interface {
-	NewPrefixIterator(prefix []byte) (database.Iterator, error)
-}