fix(dot/digest): verify if next epoch already contains some definition

dot/digest/digest.go

@@ -8,7 +8,6 @@ import (
 	"errors"
 	"fmt"
 
-	"github.com/ChainSafe/gossamer/dot/state"
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/internal/log"
 	"github.com/ChainSafe/gossamer/lib/services"
@@ -17,8 +16,6 @@ import (
 
 var (
 	_ services.Service = &Handler{}
-
-	ErrDefineNextEpoch = errors.New("cannot define next epoch data and config")
 )
 
 // Handler is used to handle consensus messages and relevant authority updates to BABE and GRANDPA
@@ -242,9 +239,14 @@ func (h *Handler) handleBlockFinalisation(ctx context.Context) {
 				continue
 			}
 
-			err := h.persistBABEDigestsForNextEpoch(&info.Header)
+			err := h.epochState.FinalizeBABENextEpochData(&info.Header)
+			if err != nil {
+				h.logger.Errorf("failed to persist babe next epoch data: %s", err)
+			}
+
+			err = h.epochState.FinalizeBABENextConfigData(&info.Header)
 			if err != nil {
-				h.logger.Errorf("failed to store babe next epoch digest: %s", err)
+				h.logger.Errorf("failed to persist babe next epoch config: %s", err)
 			}
 
 			err = h.handleGrandpaChangesOnFinalization(info.Header.Number)
@@ -257,35 +259,6 @@ func (h *Handler) handleBlockFinalisation(ctx context.Context) {
 	}
 }
 
-// persistBABEDigestsForNextEpoch is called only when a block is finalised
-// and defines the correct next epoch data and next config data.
-func (h *Handler) persistBABEDigestsForNextEpoch(finalizedHeader *types.Header) error {
-	currEpoch, err := h.epochState.GetEpochForBlock(finalizedHeader)
-	if err != nil {
-		return fmt.Errorf("cannot get epoch for block %d (%s): %w",
-			finalizedHeader.Number, finalizedHeader.Hash(), err)
-	}
-
-	nextEpoch := currEpoch + 1
-	err = h.epochState.FinalizeBABENextEpochData(nextEpoch)
-	if err != nil && !errors.Is(err, state.ErrEpochNotInMemory) {
-		return fmt.Errorf("cannot finalize babe next epoch data for block number %d (%s): %w",
-			finalizedHeader.Number, finalizedHeader.Hash(), err)
-	}
-
-	err = h.epochState.FinalizeBABENextConfigData(nextEpoch)
-	if err == nil {
-		return nil
-	} else if errors.Is(err, state.ErrEpochNotInMemory) {
-		return fmt.Errorf("%w: %s", ErrDefineNextEpoch, err)
-	}
-
-	// the epoch state does not contains any information about the next epoch
-	return fmt.Errorf("cannot finalize babe next config data for block number %d (%s): %w",
-		finalizedHeader.Number, finalizedHeader.Hash(), err)
-
-}
-
 func (h *Handler) handleGrandpaChangesOnImport(num uint) error {
 	resume := h.grandpaResume
 	if resume != nil && num >= resume.atBlock {

dot/digest/interface.go

@@ -18,6 +18,8 @@ type BlockState interface {
 	FreeFinalisedNotifierChannel(ch chan *types.FinalisationInfo)
 }
 
+//go:generate mockgen -destination=mock_epoch_state_test.go -package $GOPACKAGE . EpochState
+
 // EpochState is the interface for state.EpochState
 type EpochState interface {
 	GetEpochForBlock(header *types.Header) (uint64, error)
@@ -26,8 +28,8 @@ type EpochState interface {
 
 	StoreBABENextEpochData(epoch uint64, hash common.Hash, nextEpochData types.NextEpochData)
 	StoreBABENextConfigData(epoch uint64, hash common.Hash, nextEpochData types.NextConfigData)
-	FinalizeBABENextEpochData(epoch uint64) error
-	FinalizeBABENextConfigData(epoch uint64) error
+	FinalizeBABENextEpochData(finalizedHeader *types.Header) error
+	FinalizeBABENextConfigData(finalizedHeader *types.Header) error
 }
 
 // GrandpaState is the interface for the state.GrandpaState

dot/state/epoch.go

@@ -540,11 +540,29 @@ func (s *EpochState) StoreBABENextConfigData(epoch uint64, hash common.Hash, nex
 // getting the set of hashes from the received epoch and for each hash
 // check if the header is in the database then it's been finalized and
 // thus we can also set the corresponding EpochData in the database
-func (s *EpochState) FinalizeBABENextEpochData(epoch uint64) error {
+func (s *EpochState) FinalizeBABENextEpochData(finalizedHeader *types.Header) error {
 	s.nextEpochDataLock.Lock()
 	defer s.nextEpochDataLock.Unlock()
 
-	finalizedNextEpochData, err := lookupForNextEpochPersistedHash(s.nextEpochData, s, epoch)
+	finalizedBlockEpoch, err := s.GetEpochForBlock(finalizedHeader)
+	if err != nil {
+		return fmt.Errorf("cannot get epoch for block %d (%s): %w",
+			finalizedHeader.Number, finalizedHeader.Hash(), err)
+	}
+
+	nextEpoch := finalizedBlockEpoch + 1
+
+	epochInDatabase, err := s.getEpochDataInDatabase(nextEpoch)
+	if err != nil && !errors.Is(err, chaindb.ErrKeyNotFound) {
+		return fmt.Errorf("cannot check if next epoch data is already defined for epoch %d: %w", nextEpoch, err)
+	}
+
+	// config already defined we don't need to lookup in the map
+	if epochInDatabase != nil {
+		return nil
+	}
+
+	finalizedNextEpochData, err := findFinalizedHeaderForEpoch(s.nextEpochData, s, nextEpoch)
 	if err != nil {
 		return fmt.Errorf("cannot find next epoch data: %w", err)
 	}
@@ -554,14 +572,14 @@ func (s *EpochState) FinalizeBABENextEpochData(epoch uint64) error {
 		return fmt.Errorf("cannot transform epoch data: %w", err)
 	}
 
-	err = s.SetEpochData(epoch, ed)
+	err = s.SetEpochData(nextEpoch, ed)
 	if err != nil {
 		return fmt.Errorf("cannot set epoch data: %w", err)
 	}
 
 	// remove previous epochs from the memory
 	for e := range s.nextEpochData {
-		if e <= epoch {
+		if e <= nextEpoch {
 			delete(s.nextEpochData, e)
 		}
 	}
@@ -573,35 +591,56 @@ func (s *EpochState) FinalizeBABENextEpochData(epoch uint64) error {
 // getting the set of hashes from the received epoch and for each hash
 // check if the header is in the database then it's been finalized and
 // thus we can also set the corresponding NextConfigData in the database
-func (s *EpochState) FinalizeBABENextConfigData(epoch uint64) error {
+func (s *EpochState) FinalizeBABENextConfigData(finalizedHeader *types.Header) error {
 	s.nextConfigDataLock.Lock()
 	defer s.nextConfigDataLock.Unlock()
 
-	finalizedNextConfigData, err := lookupForNextEpochPersistedHash(s.nextConfigData, s, epoch)
+	finalizedBlockEpoch, err := s.GetEpochForBlock(finalizedHeader)
 	if err != nil {
+		return fmt.Errorf("cannot get epoch for block %d (%s): %w",
+			finalizedHeader.Number, finalizedHeader.Hash(), err)
+	}
+
+	nextEpoch := finalizedBlockEpoch + 1
+
+	configInDatabase, err := s.getConfigDataInDatabase(nextEpoch)
+	if err != nil && !errors.Is(err, chaindb.ErrKeyNotFound) {
+		return fmt.Errorf("cannot check if next epoch config is already defined for epoch %d: %w", nextEpoch, err)
+	}
+
+	// config already defined we don't need to lookup in the map
+	if configInDatabase != nil {
+		return nil
+	}
+
+	// not every epoch will have `ConfigData`
+	finalizedNextConfigData, err := findFinalizedHeaderForEpoch(s.nextConfigData, s, nextEpoch)
+	if errors.Is(err, ErrEpochNotInMemory) {
+		return nil
+	} else if err != nil {
 		return fmt.Errorf("cannot find next config data: %w", err)
 	}
 
 	cd := finalizedNextConfigData.ToConfigData()
-	err = s.SetConfigData(epoch, cd)
+	err = s.SetConfigData(nextEpoch, cd)
 	if err != nil {
 		return fmt.Errorf("cannot set config data: %w", err)
 	}
 
 	// remove previous epochs from the memory
 	for e := range s.nextConfigData {
-		if e <= epoch {
+		if e <= nextEpoch {
 			delete(s.nextConfigData, e)
 		}
 	}
 
 	return nil
 }
 
-// lookupForNextEpochPersistedHash given a specific epoch (the key) will go through the hashes looking
+// 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 to be persisted and safely used
-func lookupForNextEpochPersistedHash[T types.NextConfigData | types.NextEpochData](
+func findFinalizedHeaderForEpoch[T types.NextConfigData | types.NextEpochData](
 	nextEpochMap map[uint64]map[common.Hash]T, es *EpochState, epoch uint64) (next *T, err error) {
 	hashes, has := nextEpochMap[epoch]
 	if !has {