[staking] refactor statereader to support multiple contract indexers

action/protocol/staking/contractstake_indexer.go

@@ -6,19 +6,13 @@
 package staking
 
 import (
-	"context"
-	"math/big"
-
 	"github.com/iotexproject/iotex-address/address"
 )
 
 type (
 
 	// ContractStakingIndexer defines the interface of contract staking reader
 	ContractStakingIndexer interface {
-		// CandidateVotes returns the total staked votes of a candidate
-		// candidate identified by owner address
-		CandidateVotes(ctx context.Context, ownerAddr address.Address, height uint64) (*big.Int, error)
 		// Buckets returns active buckets
 		Buckets(height uint64) ([]*VoteBucket, error)
 		// BucketsByIndices returns active buckets by indices
@@ -27,6 +21,12 @@ type (
 		BucketsByCandidate(ownerAddr address.Address, height uint64) ([]*VoteBucket, error)
 		// TotalBucketCount returns the total number of buckets including burned buckets
 		TotalBucketCount(height uint64) (uint64, error)
+		// ContractAddress returns the contract address
+		ContractAddress() string
+	}
+	// ContractStakingIndexerWithBucketType defines the interface of contract staking reader with bucket type
+	ContractStakingIndexerWithBucketType interface {
+		ContractStakingIndexer
 		// BucketTypes returns the active bucket types
 		BucketTypes(height uint64) ([]*ContractStakingBucketType, error)
 	}

action/protocol/staking/protocol.go

@@ -80,7 +80,7 @@ type (
 		depositGas             DepositGas
 		config                 Configuration
 		candBucketsIndexer     *CandidatesBucketsIndexer
-		contractStakingIndexer ContractStakingIndexer
+		contractStakingIndexer ContractStakingIndexerWithBucketType
 		voteReviser            *VoteReviser
 		patch                  *PatchStore
 	}
@@ -121,7 +121,7 @@ func NewProtocol(
 	depositGas DepositGas,
 	cfg *BuilderConfig,
 	candBucketsIndexer *CandidatesBucketsIndexer,
-	contractStakingIndexer ContractStakingIndexer,
+	contractStakingIndexer ContractStakingIndexerWithBucketType,
 	correctCandsHeight uint64,
 	reviseHeights ...uint64,
 ) (*Protocol, error) {
@@ -518,18 +518,14 @@ func (p *Protocol) ActiveCandidates(ctx context.Context, sr protocol.StateReader
 	}
 	list := c.AllCandidates()
 	cand := make(CandidateList, 0, len(list))
-	featureCtx := protocol.MustGetFeatureCtx(ctx)
 	for i := range list {
-		if p.contractStakingIndexer != nil && featureCtx.AddContractStakingVotes {
-			// specifying the height param instead of query latest from indexer directly, aims to cause error when indexer falls behind
-			// currently there are two possible sr (i.e. factory or workingSet), it means the height could be chain height or current block height
-			// using height-1 will cover the two scenario while detect whether the indexer is lagging behind
-			contractVotes, err := p.contractStakingIndexer.CandidateVotes(ctx, list[i].GetIdentifier(), srHeight-1)
-			if err != nil {
-				return nil, errors.Wrap(err, "failed to get CandidateVotes from contractStakingIndexer")
-			}
-			list[i].Votes.Add(list[i].Votes, contractVotes)
+		// specifying the height param instead of query latest from indexer directly, aims to cause error when indexer falls behind.
+		// the reason of using srHeight-1 is contract indexer is not updated before the block is committed.
+		csVotes, err := p.contractStakingVotes(ctx, list[i].GetIdentifier(), srHeight-1)
+		if err != nil {
+			return nil, err
 		}
+		list[i].Votes.Add(list[i].Votes, csVotes)
 		active, err := p.isActiveCandidate(ctx, c, list[i])
 		if err != nil {
 			return nil, err
@@ -556,7 +552,7 @@ func (p *Protocol) ReadState(ctx context.Context, sr protocol.StateReader, metho
 	}
 
 	// stakeSR is the stake state reader including native and contract staking
-	stakeSR, err := newCompositeStakingStateReader(p.contractStakingIndexer, p.candBucketsIndexer, sr)
+	stakeSR, err := newCompositeStakingStateReader(p.candBucketsIndexer, sr, p.calculateVoteWeight, p.contractStakingIndexer)
 	if err != nil {
 		return nil, 0, err
 	}
@@ -696,6 +692,28 @@ func (p *Protocol) needToWriteCandsMap(ctx context.Context, height uint64) bool
 	return height >= p.config.PersistStakingPatchBlock && fCtx.CandCenterHasAlias(height)
 }
 
+func (p *Protocol) contractStakingVotes(ctx context.Context, candidate address.Address, height uint64) (*big.Int, error) {
+	featureCtx := protocol.MustGetFeatureCtx(ctx)
+	votes := big.NewInt(0)
+	if p.contractStakingIndexer != nil && featureCtx.AddContractStakingVotes {
+		btks, err := p.contractStakingIndexer.BucketsByCandidate(candidate, height)
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to get BucketsByCandidate from contractStakingIndexer")
+		}
+		for _, b := range btks {
+			if b.isUnstaked() {
+				continue
+			}
+			if featureCtx.FixContractStakingWeightedVotes {
+				votes.Add(votes, p.calculateVoteWeight(b, false))
+			} else {
+				votes.Add(votes, b.StakedAmount)
+			}
+		}
+	}
+	return votes, nil
+}
+
 func readCandCenterStateFromStateDB(sr protocol.StateReader) (CandidateList, CandidateList, CandidateList, error) {
 	var (
 		name, operator, owner CandidateList

action/protocol/staking/protocol_test.go

@@ -399,7 +399,7 @@ func TestProtocol_ActiveCandidates(t *testing.T) {
 	require := require.New(t)
 	ctrl := gomock.NewController(t)
 	sm := testdb.NewMockStateManagerWithoutHeightFunc(ctrl)
-	csIndexer := NewMockContractStakingIndexer(ctrl)
+	csIndexer := NewMockContractStakingIndexerWithBucketType(ctrl)
 
 	selfStake, _ := new(big.Int).SetString("1200000000000000000000000", 10)
 	cfg := genesis.Default.Staking
@@ -426,7 +426,9 @@ func TestProtocol_ActiveCandidates(t *testing.T) {
 		},
 	)
 	ctx = protocol.WithFeatureCtx(protocol.WithFeatureWithHeightCtx(ctx))
-	sm.EXPECT().Height().Return(blkHeight, nil).AnyTimes()
+	sm.EXPECT().Height().DoAndReturn(func() (uint64, error) {
+		return blkHeight, nil
+	}).AnyTimes()
 
 	v, err := p.Start(ctx, sm)
 	require.NoError(err)
@@ -436,20 +438,21 @@ func TestProtocol_ActiveCandidates(t *testing.T) {
 	require.NoError(err)
 
 	var csIndexerHeight, csVotes uint64
-	csIndexer.EXPECT().CandidateVotes(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, ownerAddr address.Address, height uint64) (*big.Int, error) {
+	csIndexer.EXPECT().BucketsByCandidate(gomock.Any(), gomock.Any()).DoAndReturn(func(ownerAddr address.Address, height uint64) ([]*VoteBucket, error) {
 		if height != csIndexerHeight {
-			return nil, errors.Errorf("invalid height")
+			return nil, errors.Errorf("invalid height %d", height)
 		}
-		return big.NewInt(int64(csVotes)), nil
+		return []*VoteBucket{
+			NewVoteBucket(identityset.Address(22), identityset.Address(22), big.NewInt(int64(csVotes)), 1, time.Now(), true),
+		}, nil
 	}).AnyTimes()
 
 	t.Run("contract staking indexer falls behind", func(t *testing.T) {
-		csIndexerHeight = 10
 		_, err := p.ActiveCandidates(ctx, sm, 0)
 		require.ErrorContains(err, "invalid height")
 	})
 
-	t.Run("contract staking indexer up to date", func(t *testing.T) {
+	t.Run("contract staking votes before Redsea", func(t *testing.T) {
 		csIndexerHeight = blkHeight - 1
 		csVotes = 0
 		cands, err := p.ActiveCandidates(ctx, sm, 0)
@@ -462,6 +465,27 @@ func TestProtocol_ActiveCandidates(t *testing.T) {
 		require.Len(cands, 1)
 		require.EqualValues(100, cands[0].Votes.Sub(cands[0].Votes, originCandVotes).Uint64())
 	})
+	t.Run("contract staking votes after Redsea", func(t *testing.T) {
+		blkHeight = genesis.Default.RedseaBlockHeight
+		ctx := protocol.WithBlockCtx(
+			genesis.WithGenesisContext(context.Background(), genesis.Default),
+			protocol.BlockCtx{
+				BlockHeight: blkHeight,
+			},
+		)
+		ctx = protocol.WithFeatureCtx(protocol.WithFeatureWithHeightCtx(ctx))
+		csIndexerHeight = blkHeight - 1
+		csVotes = 0
+		cands, err := p.ActiveCandidates(ctx, sm, 0)
+		require.NoError(err)
+		require.Len(cands, 1)
+		originCandVotes := cands[0].Votes
+		csVotes = 100
+		cands, err = p.ActiveCandidates(ctx, sm, 0)
+		require.NoError(err)
+		require.Len(cands, 1)
+		require.EqualValues(103, cands[0].Votes.Sub(cands[0].Votes, originCandVotes).Uint64())
+	})
 }
 
 func TestIsSelfStakeBucket(t *testing.T) {