chore(adr-39): Refactor rewarding (#360)

Implements babylonlabs-io/pm#145 including:
- Rewards are calculated with a timeout `finality_sig_timeout`
- Rewards should not be assigned to finality providers that have not
voted for the block when calculating rewards.

CHANGELOG.md

@@ -63,6 +63,7 @@ for rewards
 
 - [#341](https://github.com/babylonlabs-io/babylon/pull/341) Select parameters
 for pre-approval flow based on BTC LC tip height
+- [#360](https://github.com/babylonlabs-io/babylon/pull/360) Refactor rewarding
 
 ## v0.18.2
 

test/e2e/btc_staking_e2e_test.go

@@ -298,6 +298,9 @@ func (s *BTCStakingTestSuite) Test3CommitPublicRandomnessAndSubmitFinalitySignat
 		s.T().Logf("the block %d is finalized", activatedHeight)
 	}, true)
 
+	finalityParams := nonValidatorNode.QueryFinalityParams()
+	nonValidatorNode.WaitForNextBlocks(uint64(finalityParams.FinalitySigTimeout))
+
 	// ensure finality provider has received rewards after the block is finalised
 	fpRewardGauges, err := nonValidatorNode.QueryRewardGauge(fpBabylonAddr)
 	s.NoError(err)

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/rewarding_test.go

@@ -0,0 +1,110 @@
+package keeper_test
+
+import (
+	"math/rand"
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/require"
+
+	"github.com/babylonlabs-io/babylon/testutil/datagen"
+	keepertest "github.com/babylonlabs-io/babylon/testutil/keeper"
+	"github.com/babylonlabs-io/babylon/x/finality/types"
+)
+
+func FuzzHandleRewarding(f *testing.F) {
+	datagen.AddRandomSeedsToFuzzer(f, 10)
+
+	f.Fuzz(func(t *testing.T, seed int64) {
+		r := rand.New(rand.NewSource(seed))
+		ctrl := gomock.NewController(t)
+		defer ctrl.Finish()
+
+		// Setup keepers
+		bsKeeper := types.NewMockBTCStakingKeeper(ctrl)
+		iKeeper := types.NewMockIncentiveKeeper(ctrl)
+		cKeeper := types.NewMockCheckpointingKeeper(ctrl)
+		fKeeper, ctx := keepertest.FinalityKeeper(t, bsKeeper, iKeeper, cKeeper)
+
+		// Activate BTC staking protocol at a random height
+		activatedHeight := datagen.RandomInt(r, 10) + 1
+		fpPK, err := datagen.GenRandomBIP340PubKey(r)
+		require.NoError(t, err)
+		fKeeper.SetVotingPower(ctx, fpPK.MustMarshal(), activatedHeight, 1)
+
+		totalBlocks := uint64(10)
+		targetHeight := activatedHeight + totalBlocks - 1
+
+		// First phase: Index blocks with none finalized
+		for i := activatedHeight; i <= targetHeight; i++ {
+			fKeeper.SetBlock(ctx, &types.IndexedBlock{
+				Height:    i,
+				AppHash:   datagen.GenRandomByteArray(r, 32),
+				Finalized: false,
+			})
+
+			// Set voting power distribution cache for each height
+			dc := types.NewVotingPowerDistCache()
+			dc.AddFinalityProviderDistInfo(&types.FinalityProviderDistInfo{
+				BtcPk:          fpPK,
+				TotalBondedSat: 1,
+			})
+			fKeeper.SetVotingPowerDistCache(ctx, i, dc)
+		}
+
+		// First call to HandleRewarding - expect no rewards
+		ctx = datagen.WithCtxHeight(ctx, targetHeight)
+		fKeeper.HandleRewarding(ctx, int64(targetHeight))
+
+		nextHeight := fKeeper.GetNextHeightToReward(ctx)
+		require.Equal(t, uint64(0), nextHeight,
+			"next height is not updated when no blocks finalized")
+
+		// Second phase: Finalize some blocks
+		firstBatchFinalized := datagen.RandomInt(r, 5) + 1
+		for i := activatedHeight; i < activatedHeight+firstBatchFinalized; i++ {
+			block, err := fKeeper.GetBlock(ctx, i)
+			require.NoError(t, err)
+			block.Finalized = true
+			fKeeper.SetBlock(ctx, block)
+		}
+
+		// Expect rewards for first batch of finalized blocks
+		iKeeper.EXPECT().
+			RewardBTCStaking(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+			Return().
+			Times(int(firstBatchFinalized))
+
+		// Second call to HandleRewarding
+		fKeeper.HandleRewarding(ctx, int64(targetHeight))
+
+		nextHeight = fKeeper.GetNextHeightToReward(ctx)
+		expectedNextHeight := activatedHeight + firstBatchFinalized
+		require.Equal(t, expectedNextHeight, nextHeight,
+			"next height should be after first batch of finalized blocks")
+
+		// Third phase: Finalize more blocks
+		secondBatchFinalized := datagen.RandomInt(r, int(totalBlocks-firstBatchFinalized)) + 1
+		for i := expectedNextHeight; i < expectedNextHeight+secondBatchFinalized; i++ {
+			block, err := fKeeper.GetBlock(ctx, i)
+			require.NoError(t, err)
+			block.Finalized = true
+			fKeeper.SetBlock(ctx, block)
+		}
+
+		// Expect rewards for second batch of finalized blocks
+		iKeeper.EXPECT().
+			RewardBTCStaking(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+			Return().
+			Times(int(secondBatchFinalized))
+
+		// Final call to HandleRewarding
+		fKeeper.HandleRewarding(ctx, int64(targetHeight))
+
+		// Verify final state
+		finalNextHeight := fKeeper.GetNextHeightToReward(ctx)
+		expectedFinalHeight := expectedNextHeight + secondBatchFinalized
+		require.Equal(t, expectedFinalHeight, finalNextHeight,
+			"next height should be after second batch of finalized blocks")
+	})
+}

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,7 +32,7 @@ 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)
 	BtcDelegationActivated(ctx context.Context, fp, del sdk.AccAddress, sat uint64) error
 	BtcDelegationUnbonded(ctx context.Context, fp, del sdk.AccAddress, sat uint64) error

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/finality/types/power_table.go

@@ -4,8 +4,9 @@ import (
 	"sort"
 
 	sdkmath "cosmossdk.io/math"
-	bstypes "github.com/babylonlabs-io/babylon/x/btcstaking/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	bstypes "github.com/babylonlabs-io/babylon/x/btcstaking/types"
 )
 
 func NewVotingPowerDistCache() *VotingPowerDistCache {
@@ -135,11 +136,6 @@ func (dc *VotingPowerDistCache) GetInactiveFinalityProviderSet() map[string]*Fin
 	return inactiveFps
 }
 
-// GetFinalityProviderPortion returns the portion of a finality provider's voting power out of the total voting power
-func (dc *VotingPowerDistCache) GetFinalityProviderPortion(v *FinalityProviderDistInfo) sdkmath.LegacyDec {
-	return sdkmath.LegacyNewDec(int64(v.TotalBondedSat)).QuoTruncate(sdkmath.LegacyNewDec(int64(dc.TotalVotingPower)))
-}
-
 func NewFinalityProviderDistInfo(fp *bstypes.FinalityProvider) *FinalityProviderDistInfo {
 	return &FinalityProviderDistInfo{
 		BtcPk:          fp.BtcPk,

x/incentive/keeper/btc_staking_gauge.go

@@ -3,35 +3,51 @@ package keeper
 import (
 	"context"
 
+	sdkmath "cosmossdk.io/math"
 	"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
 		panic("failed to get a reward gauge at previous height")
 	}
 
-	// reward each of the finality provider and its BTC delegations in proportion
+	// calculate total voting power of voters
+	var totalVotingPowerOfVoters uint64
+	for i, fp := range dc.FinalityProviders {
+		if i >= int(dc.NumActiveFps) {
+			break
+		}
+		if _, ok := voters[fp.BtcPk.MarshalHex()]; ok {
+			totalVotingPowerOfVoters += fp.TotalBondedSat
+		}
+	}
+
+	// distribute rewards according to voting power portions for voters
 	for i, fp := range dc.FinalityProviders {
-		// only reward the first NumActiveFps finality providers
-		// note that ApplyActiveFinalityProviders is called before saving `dc`
-		// in DB so that the top dc.NumActiveFps ones in dc.FinalityProviders
-		// are the active finality providers
 		if i >= int(dc.NumActiveFps) {
 			break
 		}
-		// get coins that will be allocated to the finality provider and its BTC delegations
-		fpPortion := dc.GetFinalityProviderPortion(fp)
+
+		if _, ok := voters[fp.BtcPk.MarshalHex()]; !ok {
+			continue
+		}
+
+		// calculate the portion of a finality provider's voting power out of the total voting power of the voters
+		fpPortion := sdkmath.LegacyNewDec(int64(fp.TotalBondedSat)).
+			QuoTruncate(sdkmath.LegacyNewDec(int64(totalVotingPowerOfVoters)))
 		coinsForFpsAndDels := gauge.GetCoinsPortion(fpPortion)
+
 		// reward the finality provider with commission
 		coinsForCommission := types.GetCoinsPortion(coinsForFpsAndDels, *fp.Commission)
 		k.accumulateRewardGauge(ctx, types.FinalityProviderType, fp.GetAddress(), coinsForCommission)