EVM-846 Stake manager post epoch problem (#1935)

0xPolygon · Sep 28, 2023 · 02e2d47 · 02e2d47
1 parent e700006
commit 02e2d47
Show file tree

Hide file tree

Showing 7 changed files with 123 additions and 163 deletions.
diff --git a/consensus/polybft/consensus_runtime.go b/consensus/polybft/consensus_runtime.go
@@ -235,6 +235,7 @@ func (c *consensusRuntime) initStakeManager(logger hcf.Logger) error {
 		contracts.ValidatorSetContract,
 		c.config.PolyBFTConfig.Bridge.CustomSupernetManagerAddr,
 		c.config.blockchain,
+		c.config.polybftBackend,
 		int(c.config.PolyBFTConfig.MaxValidatorSetSize),
 	)
 
@@ -480,10 +481,6 @@ func (c *consensusRuntime) restartEpoch(header *types.Header) (*epochMetadata, e
 		return nil, err
 	}
 
-	if err := c.stakeManager.PostEpoch(reqObj); err != nil {
-		return nil, err
-	}
-
 	return &epochMetadata{
 		Number:            epochNumber,
 		Validators:        validatorSet,

diff --git a/consensus/polybft/stake_manager.go b/consensus/polybft/stake_manager.go
@@ -3,6 +3,7 @@ package polybft
 import (
 	"bytes"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"math/big"
 	"sort"
@@ -29,7 +30,6 @@ var (
 // and updating validator set based on changed stake
 type StakeManager interface {
 	PostBlock(req *PostBlockRequest) error
-	PostEpoch(req *PostEpochRequest) error
 	UpdateValidatorSet(epoch uint64, currentValidatorSet validator.AccountSet) (*validator.ValidatorSetDelta, error)
 }
 
@@ -38,7 +38,6 @@ type StakeManager interface {
 type dummyStakeManager struct{}
 
 func (d *dummyStakeManager) PostBlock(req *PostBlockRequest) error { return nil }
-func (d *dummyStakeManager) PostEpoch(req *PostEpochRequest) error { return nil }
 func (d *dummyStakeManager) UpdateValidatorSet(epoch uint64,
 	currentValidatorSet validator.AccountSet) (*validator.ValidatorSetDelta, error) {
 	return &validator.ValidatorSetDelta{}, nil
@@ -56,6 +55,7 @@ type stakeManager struct {
 	supernetManagerContract types.Address
 	maxValidatorSetSize     int
 	eventsGetter            *eventsGetter[*contractsapi.TransferEvent]
+	polybftBackend          polybftBackend
 }
 
 // newStakeManager returns a new instance of stake manager
@@ -66,6 +66,7 @@ func newStakeManager(
 	key ethgo.Key,
 	validatorSetAddr, supernetManagerAddr types.Address,
 	blockchain blockchainBackend,
+	polybftBackend polybftBackend,
 	maxValidatorSetSize int,
 ) (*stakeManager, error) {
 	eventsGetter := &eventsGetter[*contractsapi.TransferEvent]{
@@ -89,6 +90,7 @@ func newStakeManager(
 		supernetManagerContract: supernetManagerAddr,
 		maxValidatorSetSize:     maxValidatorSetSize,
 		eventsGetter:            eventsGetter,
+		polybftBackend:          polybftBackend,
 	}
 
 	if err := sm.init(blockchain); err != nil {
@@ -98,25 +100,10 @@ func newStakeManager(
 	return sm, nil
 }
 
-// PostEpoch saves the initial validator set to db
-func (s *stakeManager) PostEpoch(req *PostEpochRequest) error {
-	if req.NewEpochID != 1 {
-		return nil
-	}
-
-	// save initial validator set as full validator set in db
-	return s.state.StakeStore.insertFullValidatorSet(validatorSetState{
-		BlockNumber:          0,
-		EpochID:              0,
-		UpdatedAtBlockNumber: 0,
-		Validators:           newValidatorStakeMap(req.ValidatorSet.Accounts()),
-	})
-}
-
 // PostBlock is called on every insert of finalized block (either from consensus or syncer)
 // It will read any transfer event that happened in block and update full validator set in db
 func (s *stakeManager) PostBlock(req *PostBlockRequest) error {
-	fullValidatorSet, err := s.state.StakeStore.getFullValidatorSet()
+	fullValidatorSet, err := s.getOrInitValidatorSet()
 	if err != nil {
 		return err
 	}
@@ -150,12 +137,7 @@ func (s *stakeManager) init(blockchain blockchainBackend) error {
 	currentHeader := blockchain.CurrentHeader()
 	currentBlockNumber := currentHeader.Number
 
-	// early exit if this is genesis block
-	if currentBlockNumber == 0 {
-		return nil
-	}
-
-	validatorSet, err := s.state.StakeStore.getFullValidatorSet()
+	validatorSet, err := s.getOrInitValidatorSet()
 	if err != nil {
 		return err
 	}
@@ -193,6 +175,33 @@ func (s *stakeManager) init(blockchain blockchainBackend) error {
 	return s.state.StakeStore.insertFullValidatorSet(validatorSet)
 }
 
+func (s *stakeManager) getOrInitValidatorSet() (validatorSetState, error) {
+	validatorSet, err := s.state.StakeStore.getFullValidatorSet()
+	if err != nil {
+		if !errors.Is(err, errNoFullValidatorSet) {
+			return validatorSetState{}, err
+		}
+
+		validators, err := s.polybftBackend.GetValidators(0, nil)
+		if err != nil {
+			return validatorSetState{}, err
+		}
+
+		validatorSet = validatorSetState{
+			BlockNumber:          0,
+			EpochID:              0,
+			UpdatedAtBlockNumber: 0,
+			Validators:           newValidatorStakeMap(validators),
+		}
+
+		if err = s.state.StakeStore.insertFullValidatorSet(validatorSet); err != nil {
+			return validatorSetState{}, err
+		}
+	}
+
+	return validatorSet, nil
+}
+
 func (s *stakeManager) updateWithReceipts(
 	fullValidatorSet *validatorSetState,
 	transferEvents []*contractsapi.TransferEvent,

diff --git a/consensus/polybft/stake_manager_fuzz_test.go b/consensus/polybft/stake_manager_fuzz_test.go
@@ -34,66 +34,6 @@ type updateValidatorSetF struct {
 	VotingPower int64
 }
 
-func FuzzTestStakeManagerPostEpoch(f *testing.F) {
-	state := newTestState(f)
-
-	seeds := []epochIDValidatorsF{
-		{
-			EpochID:    0,
-			Validators: validator.NewTestValidators(f, 6).GetPublicIdentities(),
-		},
-		{
-			EpochID:    1,
-			Validators: validator.NewTestValidators(f, 42).GetPublicIdentities(),
-		},
-		{
-			EpochID:    42,
-			Validators: validator.NewTestValidators(f, 6).GetPublicIdentities(),
-		},
-	}
-
-	for _, seed := range seeds {
-		data, err := json.Marshal(seed)
-		if err != nil {
-			return
-		}
-
-		f.Add(data)
-	}
-
-	f.Fuzz(func(t *testing.T, input []byte) {
-		stakeManager := &stakeManager{
-			logger:              hclog.NewNullLogger(),
-			state:               state,
-			maxValidatorSetSize: 10,
-		}
-
-		var data epochIDValidatorsF
-		if err := json.Unmarshal(input, &data); err != nil {
-			t.Skip(err)
-		}
-
-		invalidDataFormat := false
-		for _, v := range data.Validators {
-			if err := ValidateStruct(*v); err != nil {
-				invalidDataFormat = true
-			}
-		}
-		if invalidDataFormat {
-			t.Skip()
-		}
-
-		err := stakeManager.PostEpoch(&PostEpochRequest{
-			NewEpochID: data.EpochID,
-			ValidatorSet: validator.NewValidatorSet(
-				data.Validators,
-				stakeManager.logger,
-			),
-		})
-		require.NoError(t, err)
-	})
-}
-
 func FuzzTestStakeManagerPostBlock(f *testing.F) {
 	var (
 		allAliases        = []string{"A", "B", "C", "D", "E", "F"}
@@ -163,6 +103,11 @@ func FuzzTestStakeManagerPostBlock(f *testing.F) {
 			bcMock.On("GetReceiptsByHash", mock.Anything).Return([]*types.Receipt{{}}, error(nil)).Once()
 		}
 
+		// insert initial full validator set
+		require.NoError(t, state.StakeStore.insertFullValidatorSet(validatorSetState{
+			Validators: newValidatorStakeMap(validators.GetPublicIdentities(initialSetAliases...)),
+		}))
+
 		stakeManager, err := newStakeManager(
 			hclog.NewNullLogger(),
 			state,
@@ -171,15 +116,11 @@ func FuzzTestStakeManagerPostBlock(f *testing.F) {
 			validatorSetAddr,
 			types.StringToAddress("0x0002"),
 			bcMock,
+			nil,
 			5,
 		)
 		require.NoError(t, err)
 
-		// insert initial full validator set
-		require.NoError(t, state.StakeStore.insertFullValidatorSet(validatorSetState{
-			Validators: newValidatorStakeMap(validators.GetPublicIdentities(initialSetAliases...)),
-		}))
-
 		receipt := &types.Receipt{
 			Logs: []*types.Log{
 				createTestLogForTransferEvent(
@@ -213,13 +154,18 @@ func FuzzTestStakeManagerUpdateValidatorSet(f *testing.F) {
 	bcMock := new(blockchainMock)
 	bcMock.On("CurrentHeader").Return(&types.Header{Number: 0})
 
+	err := state.StakeStore.insertFullValidatorSet(validatorSetState{
+		Validators: newValidatorStakeMap(validators.GetPublicIdentities())})
+	require.NoError(f, err)
+
 	stakeManager, err := newStakeManager(
 		hclog.NewNullLogger(),
 		state,
 		nil,
 		wallet.NewEcdsaSigner(validators.GetValidator("A").Key()),
 		types.StringToAddress("0x0001"), types.StringToAddress("0x0002"),
 		bcMock,
+		nil,
 		10,
 	)
 	require.NoError(f, err)