fix(dot/digest): BABE NextEpochData and NextConfigData should be set on finalization

dot/core/service.go

@@ -216,9 +216,6 @@ func (s *Service) handleBlock(block *types.Block, state *rtstorage.TrieState) er
 	logger.Debugf("imported block %s and stored state trie with root %s",
 		block.Header.Hash(), state.MustRoot())
 
-	// handle consensus digests
-	s.digestHandler.HandleDigests(&block.Header)
-
 	rt, err := s.blockState.GetRuntime(&block.Header.ParentHash)
 	if err != nil {
 		return err

dot/digest/digest.go

@@ -6,6 +6,7 @@ package digest
 import (
 	"context"
 	"errors"
+	"fmt"
 
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/internal/log"
@@ -118,12 +119,14 @@ func (h *Handler) NextGrandpaAuthorityChange() (next uint) {
 func (h *Handler) HandleDigests(header *types.Header) {
 	for i, d := range header.Digest.Types {
 		val, ok := d.Value().(types.ConsensusDigest)
-		if ok {
-			err := h.handleConsensusDigest(&val, header)
-			if err != nil {
-				h.logger.Errorf("cannot handle digests for block number %d, index %d, digest %s: %s",
-					header.Number, i, d.Value(), err)
-			}
+		if !ok {
+			continue
+		}
+
+		err := h.handleConsensusDigest(&val, header)
+		if err != nil {
+			h.logger.Errorf("cannot handle digests for block number %d, index %d, digest %s: %s",
+				header.Number, i, d.Value(), err)
 		}
 	}
 }
@@ -175,15 +178,28 @@ func (h *Handler) handleGrandpaConsensusDigest(digest scale.VaryingDataType, hea
 }
 
 func (h *Handler) handleBabeConsensusDigest(digest scale.VaryingDataType, header *types.Header) error {
+	currEpoch, err := h.epochState.GetEpochForBlock(header)
+	if err != nil {
+		return fmt.Errorf("cannot get epoch for block %d (%s): %w",
+			header.Number, header.Hash(), err)
+	}
+
+	nextEpoch := currEpoch + 1
+	headerHash := header.Hash()
+
 	switch val := digest.Value().(type) {
 	case types.NextEpochData:
-		h.logger.Debugf("handling BABENextEpochData data: %v", digest)
-		return h.handleNextEpochData(val, header)
+		h.logger.Debugf("stored BABENextEpochData data: %v for hash: %s to epoch: %d", digest, headerHash, nextEpoch)
+		h.epochState.StoreBABENextEpochData(nextEpoch, headerHash, val)
+		return nil
+
 	case types.BABEOnDisabled:
 		return h.handleBABEOnDisabled(val, header)
+
 	case types.NextConfigData:
-		h.logger.Debugf("handling BABENextConfigData data: %v", digest)
-		return h.handleNextConfigData(val, header)
+		h.logger.Debugf("stored BABENextConfigData data: %v for hash: %s to epoch: %d", digest, headerHash, nextEpoch)
+		h.epochState.StoreBABENextConfigData(nextEpoch, headerHash, val)
+		return nil
 	}
 
 	return errors.New("invalid consensus digest data")
@@ -197,6 +213,7 @@ func (h *Handler) handleBlockImport(ctx context.Context) {
 				continue
 			}
 
+			h.HandleDigests(&block.Header)
 			err := h.handleGrandpaChangesOnImport(block.Header.Number)
 			if err != nil {
 				h.logger.Errorf("failed to handle grandpa changes on block import: %s", err)
@@ -215,7 +232,12 @@ func (h *Handler) handleBlockFinalisation(ctx context.Context) {
 				continue
 			}
 
-			err := h.handleGrandpaChangesOnFinalization(info.Header.Number)
+			err := h.setBABEDigestsOnFinalization(&info.Header)
+			if err != nil {
+				h.logger.Errorf("failed to store babe next epoch digest: %s", err)
+			}
+
+			err = h.handleGrandpaChangesOnFinalization(info.Header.Number)
 			if err != nil {
 				h.logger.Errorf("failed to handle grandpa changes on block finalisation: %s", err)
 			}
@@ -225,6 +247,34 @@ func (h *Handler) handleBlockFinalisation(ctx context.Context) {
 	}
 }
 
+// setBABEDigestsOnFinalization is called only when a block is finalised
+// and defines the correct next epoch data and next config data
+func (h *Handler) setBABEDigestsOnFinalization(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 {
+		h.logger.Debugf("failed to store BABENextEpochData for block number %d (%s): %w",
+			finalizedHeader.Number, finalizedHeader.Hash(), err)
+		return err
+	}
+
+	err = h.epochState.FinalizeBABENextConfigDataToFinalize(nextEpoch)
+	if err != nil {
+		h.logger.Debugf("failed to store BABENextConfigData for block number %d (%s): %w",
+			finalizedHeader.Number, finalizedHeader.Hash(), err)
+		return err
+	}
+
+	return nil
+}
+
 func (h *Handler) handleGrandpaChangesOnImport(num uint) error {
 	resume := h.grandpaResume
 	if resume != nil && num >= resume.atBlock {
@@ -233,17 +283,12 @@ func (h *Handler) handleGrandpaChangesOnImport(num uint) error {
 
 	fc := h.grandpaForcedChange
 	if fc != nil && num >= fc.atBlock {
-		err := h.grandpaState.IncrementSetID()
+		curr, err := h.grandpaState.IncrementSetID()
 		if err != nil {
 			return err
 		}
 
 		h.grandpaForcedChange = nil
-		curr, err := h.grandpaState.GetCurrentSetID()
-		if err != nil {
-			return err
-		}
-
 		h.logger.Debugf("incremented grandpa set id %d", curr)
 	}
 
@@ -258,17 +303,12 @@ func (h *Handler) handleGrandpaChangesOnFinalization(num uint) error {
 
 	sc := h.grandpaScheduledChange
 	if sc != nil && num >= sc.atBlock {
-		err := h.grandpaState.IncrementSetID()
+		curr, err := h.grandpaState.IncrementSetID()
 		if err != nil {
 			return err
 		}
 
 		h.grandpaScheduledChange = nil
-		curr, err := h.grandpaState.GetCurrentSetID()
-		if err != nil {
-			return err
-		}
-
 		h.logger.Debugf("incremented grandpa set id %d", curr)
 	}
 
@@ -381,32 +421,3 @@ func (h *Handler) handleBABEOnDisabled(_ types.BABEOnDisabled, _ *types.Header)
 	h.logger.Debug("handling BABEOnDisabled")
 	return nil
 }
-
-func (h *Handler) handleNextEpochData(act types.NextEpochData, header *types.Header) error {
-	currEpoch, err := h.epochState.GetEpochForBlock(header)
-	if err != nil {
-		return err
-	}
-
-	// set EpochState epoch data for upcoming epoch
-	data, err := act.ToEpochData()
-	if err != nil {
-		return err
-	}
-
-	h.logger.Debugf("setting data for block number %d and epoch %d with data: %v",
-		header.Number, currEpoch+1, data)
-	return h.epochState.SetEpochData(currEpoch+1, data)
-}
-
-func (h *Handler) handleNextConfigData(config types.NextConfigData, header *types.Header) error {
-	currEpoch, err := h.epochState.GetEpochForBlock(header)
-	if err != nil {
-		return err
-	}
-
-	h.logger.Debugf("setting BABE config data for block number %d and epoch %d with data: %v",
-		header.Number, currEpoch+1, config.ToConfigData())
-	// set EpochState config data for upcoming epoch
-	return h.epochState.SetConfigData(currEpoch+1, config.ToConfigData())
-}

dot/digest/digest_test.go

@@ -7,6 +7,8 @@ import (
 	"testing"
 	"time"
 
+	"context"
+
 	"github.com/ChainSafe/gossamer/dot/state"
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/internal/log"
@@ -377,16 +379,13 @@ func createHeaderWithPreDigest(t *testing.T, slotNumber uint64) *types.Header {
 	require.NoError(t, err)
 
 	return &types.Header{
+		Number: 1,
 		Digest: digest,
 	}
 }
 
 func TestHandler_HandleNextEpochData(t *testing.T) {
-	expData := common.MustHexToBytes("0x0108d43593c715fdd31c61141abd04a99fd6822c8558854ccde39a5684e7a56da27d01000000000000008eaf04151687736326c9fea17e25fc5287613693c912909cb226aa4794f26a4801000000000000004d58630000000000000000000000000000000000000000000000000000000000") //nolint:lll
-
-	handler := newTestHandler(t)
-	handler.Start()
-	defer handler.Stop()
+	expectedDigestBytes := common.MustHexToBytes("0x0108d43593c715fdd31c61141abd04a99fd6822c8558854ccde39a5684e7a56da27d01000000000000008eaf04151687736326c9fea17e25fc5287613693c912909cb226aa4794f26a4801000000000000004d58630000000000000000000000000000000000000000000000000000000000") //nolint:lll
 
 	keyring, err := keystore.NewSr25519Keyring()
 	require.NoError(t, err)
@@ -411,7 +410,7 @@ func TestHandler_HandleNextEpochData(t *testing.T) {
 	data, err := scale.Marshal(digest)
 	require.NoError(t, err)
 
-	require.Equal(t, expData, data)
+	require.Equal(t, expectedDigestBytes, data)
 
 	d := &types.ConsensusDigest{
 		ConsensusEngineID: types.BabeEngineID,
@@ -420,9 +419,34 @@ func TestHandler_HandleNextEpochData(t *testing.T) {
 
 	header := createHeaderWithPreDigest(t, 10)
 
+	finalisedCh := make(chan *types.FinalisationInfo)
+
+	handler := newTestHandler(t)
+	handler.finalised = finalisedCh
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	doneCh := make(chan struct{})
+
+	go func() {
+		defer close(doneCh)
+		handler.handleBlockFinalisation(ctx)
+	}()
+
 	err = handler.handleConsensusDigest(d, header)
 	require.NoError(t, err)
 
+	finalisedCh <- &types.FinalisationInfo{
+		Header: *header,
+		Round:  1,
+		SetID:  1,
+	}
+
+	// Before check the epoch data was stored
+	// we need to wait for both handle functions to finish
+	cancel()
+	<-doneCh
+
 	stored, err := handler.epochState.(*state.EpochState).GetEpochData(1)
 	require.NoError(t, err)
 
@@ -437,10 +461,6 @@ func TestHandler_HandleNextEpochData(t *testing.T) {
 }
 
 func TestHandler_HandleNextConfigData(t *testing.T) {
-	handler := newTestHandler(t)
-	handler.Start()
-	defer handler.Stop()
-
 	var digest = types.NewBabeConsensusDigest()
 	err := digest.Set(types.NextConfigData{
 		C1:             1,
@@ -459,9 +479,33 @@ func TestHandler_HandleNextConfigData(t *testing.T) {
 
 	header := createHeaderWithPreDigest(t, 10)
 
+	finalisedCh := make(chan *types.FinalisationInfo)
+
+	handler := newTestHandler(t)
+	handler.finalised = finalisedCh
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	doneCh := make(chan struct{})
+	go func() {
+		defer close(doneCh)
+		handler.handleBlockFinalisation(ctx)
+	}()
+
 	err = handler.handleConsensusDigest(d, header)
 	require.NoError(t, err)
 
+	finalisedCh <- &types.FinalisationInfo{
+		Header: *header,
+		Round:  1,
+		SetID:  1,
+	}
+
+	// Before check the config data was stored
+	// we need to wait for both handle functions finish
+	cancel()
+	<-doneCh
+
 	act, ok := digest.Value().(types.NextConfigData)
 	if !ok {
 		t.Fatal()

dot/digest/interface.go

@@ -5,6 +5,7 @@ package digest
 
 import (
 	"github.com/ChainSafe/gossamer/dot/types"
+	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/ChainSafe/gossamer/lib/grandpa"
 )
 
@@ -22,12 +23,17 @@ type EpochState interface {
 	GetEpochForBlock(header *types.Header) (uint64, error)
 	SetEpochData(epoch uint64, info *types.EpochData) error
 	SetConfigData(epoch uint64, info *types.ConfigData) error
+
+	StoreBABENextEpochData(epoch uint64, hash common.Hash, nextEpochData types.NextEpochData)
+	StoreBABENextConfigData(epoch uint64, hash common.Hash, nextEpochData types.NextConfigData)
+	FinalizeBABENextEpochData(epoch uint64) error
+	FinalizeBABENextConfigDataToFinalize(epoch uint64) error
 }
 
 // GrandpaState is the interface for the state.GrandpaState
 type GrandpaState interface {
 	SetNextChange(authorities []grandpa.Voter, number uint) error
-	IncrementSetID() error
+	IncrementSetID() (newSetID uint64, err error)
 	SetNextPause(number uint) error
 	SetNextResume(number uint) error
 	GetCurrentSetID() (uint64, error)

dot/state/block.go

@@ -183,7 +183,8 @@ func (bs *BlockState) GenesisHash() common.Hash {
 	return bs.genesisHash
 }
 
-// HasHeader returns if the db contains a header with the given hash
+// HasHeader returns if the hash are stored in the memory unfinalised blocks or
+// persisted in database
 func (bs *BlockState) HasHeader(hash common.Hash) (bool, error) {
 	if bs.unfinalisedBlocks.getBlock(hash) != nil {
 		return true, nil
@@ -192,6 +193,11 @@ func (bs *BlockState) HasHeader(hash common.Hash) (bool, error) {
 	return bs.db.Has(headerKey(hash))
 }
 
+// HasHeaderInDatabase returns if only database contains a header with the given hash
+func (bs *BlockState) HasHeaderInDatabase(hash common.Hash) (bool, error) {
+	return bs.db.Has(headerKey(hash))
+}
+
 // GetHeader returns a BlockHeader for a given hash
 func (bs *BlockState) GetHeader(hash common.Hash) (header *types.Header, err error) {
 	header = bs.unfinalisedBlocks.getBlockHeader(hash)