refactor rewarding

x/finality/abci.go

@@ -36,8 +36,11 @@ func EndBlocker(ctx context.Context, k keeper.Keeper) ([]abci.ValidatorUpdate, e
 		// bit in a bit array of size params.SignedBlocksWindow)
 		// once this height is judged as `missed`, the judgement is irreversible
 		heightToExamine := sdk.UnwrapSDKContext(ctx).HeaderInfo().Height - k.GetParams(ctx).FinalitySigTimeout
+
 		if heightToExamine >= 1 {
 			k.HandleLiveness(ctx, heightToExamine)
+
+			k.HandleRewarding(ctx, heightToExamine)
 		}
 	}
 

x/finality/keeper/gov_test.go

@@ -57,7 +57,7 @@ func TestHandleResumeFinalityProposal(t *testing.T) {
 	}
 
 	// tally blocks and none of them should be finalised
-	iKeeper.EXPECT().RewardBTCStaking(gomock.Any(), gomock.Any(), gomock.Any()).Return().AnyTimes()
+	iKeeper.EXPECT().RewardBTCStaking(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return().AnyTimes()
 	ctx = datagen.WithCtxHeight(ctx, currentHeight)
 	fKeeper.TallyBlocks(ctx)
 	for i := haltingHeight; i < currentHeight; i++ {

x/finality/keeper/rewarding.go

@@ -0,0 +1,83 @@
+package keeper
+
+import (
+	"context"
+	"fmt"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	"github.com/babylonlabs-io/babylon/x/finality/types"
+)
+
+func (k Keeper) HandleRewarding(ctx context.Context, targetHeight int64) {
+	// rewarding is executed in a range of [nextHeightToReward, heightToExamine]
+	// this is we don't know when a block will be finalized and we need ensure
+	// every finalized block will be processed to reward
+	nextHeightToReward := k.GetNextHeightToReward(ctx)
+	if nextHeightToReward == 0 {
+		// first time to call reward, set it to activated height
+		activatedHeight, err := k.GetBTCStakingActivatedHeight(ctx)
+		if err != nil {
+			panic(err)
+		}
+		nextHeightToReward = activatedHeight
+	}
+	copiedNextHeightToReward := nextHeightToReward
+
+	for height := nextHeightToReward; height <= uint64(targetHeight); height++ {
+		block, err := k.GetBlock(ctx, height)
+		if err != nil {
+			panic(err)
+		}
+		if !block.Finalized {
+			break
+		}
+		k.rewardBTCStaking(ctx, height)
+		nextHeightToReward = height + 1
+	}
+
+	if nextHeightToReward != copiedNextHeightToReward {
+		k.SetNextHeightToReward(ctx, nextHeightToReward)
+	}
+}
+
+func (k Keeper) rewardBTCStaking(ctx context.Context, height uint64) {
+	// distribute rewards to BTC staking stakeholders w.r.t. the voting power distribution cache
+	dc := k.GetVotingPowerDistCache(ctx, height)
+	if dc == nil {
+		// failing to get a voting power distribution cache before distributing reward is a programming error
+		panic(fmt.Errorf("voting power distribution cache not found at height %d", height))
+	}
+
+	// get all the voters for the height
+	voterBTCPKs := k.GetVoters(ctx, height)
+
+	// reward active finality providers
+	k.IncentiveKeeper.RewardBTCStaking(ctx, height, dc, voterBTCPKs)
+
+	// remove reward distribution cache afterwards
+	k.RemoveVotingPowerDistCache(ctx, height)
+}
+
+// SetNextHeightToReward sets the next height to reward as the given height
+func (k Keeper) SetNextHeightToReward(ctx context.Context, height uint64) {
+	store := k.storeService.OpenKVStore(ctx)
+	heightBytes := sdk.Uint64ToBigEndian(height)
+	if err := store.Set(types.NextHeightToRewardKey, heightBytes); err != nil {
+		panic(err)
+	}
+}
+
+// GetNextHeightToReward gets the next height to reward
+func (k Keeper) GetNextHeightToReward(ctx context.Context) uint64 {
+	store := k.storeService.OpenKVStore(ctx)
+	bz, err := store.Get(types.NextHeightToRewardKey)
+	if err != nil {
+		panic(err)
+	}
+	if bz == nil {
+		return 0
+	}
+	height := sdk.BigEndianToUint64(bz)
+	return height
+}

x/finality/keeper/tallying.go

@@ -84,16 +84,6 @@ func (k Keeper) finalizeBlock(ctx context.Context, block *types.IndexedBlock) {
 	k.SetBlock(ctx, block)
 	// set next height to finalise as height+1
 	k.setNextHeightToFinalize(ctx, block.Height+1)
-	// distribute rewards to BTC staking stakeholders w.r.t. the voting power distribution cache
-	dc := k.GetVotingPowerDistCache(ctx, block.Height)
-	if dc == nil {
-		// failing to get a voting power distribution cache before distributing reward is a programming error
-		panic(fmt.Errorf("voting power distribution cache not found at height %d", block.Height))
-	}
-	// reward active finality providers
-	k.IncentiveKeeper.RewardBTCStaking(ctx, block.Height, dc)
-	// remove reward distribution cache afterwards
-	k.RemoveVotingPowerDistCache(ctx, block.Height)
 	// record the last finalized height metric
 	types.RecordLastFinalizedHeight(block.Height)
 }

x/finality/keeper/tallying_bench_test.go

@@ -42,7 +42,7 @@ func benchmarkTallyBlocks(b *testing.B, numFPs int) {
 	}
 
 	// TODO: test incentive
-	iKeeper.EXPECT().RewardBTCStaking(gomock.Any(), gomock.Any(), gomock.Any()).Return().AnyTimes()
+	iKeeper.EXPECT().RewardBTCStaking(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return().AnyTimes()
 	// Start the CPU profiler
 	cpuProfileFile := fmt.Sprintf("/tmp/finality-tally-blocks-%d-cpu.pprof", numFPs)
 	f, err := os.Create(cpuProfileFile)

x/finality/keeper/tallying_test.go

@@ -93,8 +93,6 @@ func FuzzTallying_FinalizingSomeBlocks(f *testing.F) {
 				require.NoError(t, err)
 			}
 		}
-		// we don't test incentive in this function
-		iKeeper.EXPECT().RewardBTCStaking(gomock.Any(), gomock.Any(), gomock.Any()).Return().Times(int(numWithQCs))
 		// tally blocks and none of them should be finalised
 		ctx = datagen.WithCtxHeight(ctx, activatedHeight+10-1)
 		fKeeper.TallyBlocks(ctx)

x/finality/types/expected_keepers.go

@@ -32,6 +32,6 @@ type CheckpointingKeeper interface {
 // IncentiveKeeper defines the expected interface needed for distributing rewards
 // and refund transaction fee for finality signatures
 type IncentiveKeeper interface {
-	RewardBTCStaking(ctx context.Context, height uint64, filteredDc *VotingPowerDistCache)
+	RewardBTCStaking(ctx context.Context, height uint64, filteredDc *VotingPowerDistCache, voters map[string]struct{})
 	IndexRefundableMsg(ctx context.Context, msg sdk.Msg)
 }

x/finality/types/keys.go

@@ -49,4 +49,5 @@ var (
 	FinalityProviderMissedBlockBitmapKeyPrefix = collections.NewPrefix(9) // key prefix for missed block bitmap
 	VotingPowerKey                             = []byte{0x10}             // key prefix for the voting power
 	VotingPowerDistCacheKey                    = []byte{0x11}             // key prefix for voting power distribution cache
+	NextHeightToRewardKey                      = []byte{0x012}            // key prefix for next height to reward
 )

x/incentive/keeper/btc_staking_gauge.go

@@ -4,16 +4,17 @@ import (
 	"context"
 
 	"cosmossdk.io/store/prefix"
-	ftypes "github.com/babylonlabs-io/babylon/x/finality/types"
-	"github.com/babylonlabs-io/babylon/x/incentive/types"
 	"github.com/cosmos/cosmos-sdk/runtime"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	ftypes "github.com/babylonlabs-io/babylon/x/finality/types"
+	"github.com/babylonlabs-io/babylon/x/incentive/types"
 )
 
 // RewardBTCStaking distributes rewards to finality providers/delegations at a given height according
 // to the filtered reward distribution cache (that only contains voted finality providers)
 // (adapted from https://github.com/cosmos/cosmos-sdk/blob/release/v0.47.x/x/distribution/keeper/allocation.go#L12-L64)
-func (k Keeper) RewardBTCStaking(ctx context.Context, height uint64, dc *ftypes.VotingPowerDistCache) {
+func (k Keeper) RewardBTCStaking(ctx context.Context, height uint64, dc *ftypes.VotingPowerDistCache, voters map[string]struct{}) {
 	gauge := k.GetBTCStakingGauge(ctx, height)
 	if gauge == nil {
 		// failing to get a reward gauge at previous height is a programming error
@@ -28,6 +29,12 @@ func (k Keeper) RewardBTCStaking(ctx context.Context, height uint64, dc *ftypes.
 		if i >= int(dc.NumActiveFps) {
 			break
 		}
+
+		// skip if finality provider didn't vote
+		if _, ok := voters[fp.BtcPk.MarshalHex()]; !ok {
+			continue
+		}
+
 		// get coins that will be allocated to the finality provider and its BTC delegations
 		fpPortion := dc.GetFinalityProviderPortion(fp)
 		coinsForFpsAndDels := gauge.GetCoinsPortion(fpPortion)

x/incentive/keeper/btc_staking_gauge_test.go

@@ -4,12 +4,13 @@ import (
 	"math/rand"
 	"testing"
 
-	"github.com/babylonlabs-io/babylon/testutil/datagen"
-	testkeeper "github.com/babylonlabs-io/babylon/testutil/keeper"
-	"github.com/babylonlabs-io/babylon/x/incentive/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/require"
+
+	"github.com/babylonlabs-io/babylon/testutil/datagen"
+	testkeeper "github.com/babylonlabs-io/babylon/testutil/keeper"
+	"github.com/babylonlabs-io/babylon/x/incentive/types"
 )
 
 func FuzzRewardBTCStaking(f *testing.F) {
@@ -60,8 +61,14 @@ func FuzzRewardBTCStaking(f *testing.F) {
 			}
 		}
 
+		// create voter map from the voting power cache
+		voterMap := make(map[string]struct{})
+		for _, fp := range dc.FinalityProviders {
+			voterMap[fp.BtcPk.MarshalHex()] = struct{}{}
+		}
+
 		// distribute rewards in the gauge to finality providers/delegations
-		keeper.RewardBTCStaking(ctx, height, dc)
+		keeper.RewardBTCStaking(ctx, height, dc, voterMap)
 
 		// assert consistency between reward map and reward gauge
 		for addrStr, reward := range fpRewardMap {