impl: fp reference counting logic

proto/babylon/incentive/rewards.proto

@@ -22,16 +22,18 @@ message FinalityProviderHistoricalRewards {
     (gogoproto.nullable) = false,
     (gogoproto.castrepeated) = "github.com/cosmos/cosmos-sdk/types.Coins"
   ];
-  // TODO(rafilx): add reference count for state prunning
+  // The reference count indicates the number of objects
+  // which might need to reference this historical entry at any point.
   // https://github.com/cosmos/cosmos-sdk/blob/d9c53bfefc1e75a3c6b09065ea8b3a836cda0d18/x/distribution/types/distribution.pb.go#L98
+  uint32 reference_count = 2;
 }
 
 // FinalityProviderCurrentRewards represents the current rewards of the pool of
 // BTC delegations that delegated for this finality provider is entitled to.
 // Note: This rewards are for the BTC delegators that delegated to this FP
 // the FP itself is not the owner or can withdraw this rewards.
 // If a slash event happens with this finality provider, all the delegations
-// need to withdraw to the RewardGauge and the related scrutures should be
+// need to withdraw to the RewardGauge and the related structures should be
 // deleted. Key: Prefix + Finality provider bech32 address.
 message FinalityProviderCurrentRewards {
   // CurrentRewards is the current rewards that the finality provider have and

testutil/datagen/incentive.go

@@ -160,7 +160,7 @@ func GenRandomBTCDelegationRewardsTracker(r *rand.Rand) itypes.BTCDelegationRewa
 
 func GenRandomFPHistRwd(r *rand.Rand) itypes.FinalityProviderHistoricalRewards {
 	rwd := GenRandomCoins(r)
-	return itypes.NewFinalityProviderHistoricalRewards(rwd)
+	return itypes.NewFinalityProviderHistoricalRewards(rwd, uint32(1))
 }
 
 func GenRandomFPHistRwdWithDecimals(r *rand.Rand) itypes.FinalityProviderHistoricalRewards {

x/incentive/keeper/grpc_query_test.go

@@ -139,7 +139,7 @@ func FuzzDelegationRewardsQuery(f *testing.F) {
 
 		// set start historical rewards corresponding to btcRwd.StartPeriodCumulativeReward
 		amtRwdInHistStart := fpCurrentRwd.CurrentRewards.MulInt(types.DecimalRewards).QuoInt(math.NewInt(2))
-		startHist := types.NewFinalityProviderHistoricalRewards(amtRwdInHistStart)
+		startHist := types.NewFinalityProviderHistoricalRewards(amtRwdInHistStart, uint32(1))
 
 		sfphrKey := collections.Join(fp.Bytes(), btcRwd.StartPeriodCumulativeReward)
 
@@ -153,7 +153,7 @@ func FuzzDelegationRewardsQuery(f *testing.F) {
 		// set end period historical rewards
 		// end period for calculation is fpCurrentRwd.Period-1
 		amtRwdInHistEnd := amtRwdInHistStart.Add(fpCurrentRwd.CurrentRewards.MulInt(types.DecimalRewards).QuoInt(fpCurrentRwd.TotalActiveSat)...)
-		endHist := types.NewFinalityProviderHistoricalRewards(amtRwdInHistEnd)
+		endHist := types.NewFinalityProviderHistoricalRewards(amtRwdInHistEnd, uint32(1))
 
 		efphrKey := collections.Join(fp.Bytes(), fpCurrentRwd.Period-1)
 

x/incentive/keeper/reward_tracker.go

@@ -40,7 +40,7 @@ func (k Keeper) BtcDelegationActivated(ctx context.Context, fp, del sdk.AccAddre
 // BtcDelegationUnbonded it modifies the total amount of satoshi staked
 // for the delegation (fp, del) and for the finality provider by subtracting.
 // Since it modifies the active amount it triggers the increment of fp period,
-// creationg of new historical reward, withdraw of rewards to gauge
+// creation of new historical reward, withdraw of rewards to gauge
 // and initialization of a new delegation.
 // It errors out if the unbond amount is higher than the total amount staked.
 func (k Keeper) BtcDelegationUnbonded(ctx context.Context, fp, del sdk.AccAddress, sat sdkmath.Int) error {
@@ -92,7 +92,25 @@ func (k Keeper) btcDelegationModified(
 	ctx context.Context,
 	fp, del sdk.AccAddress,
 ) error {
-	return k.btcDelegationModifiedWithPreInitDel(ctx, fp, del, func(ctx context.Context, fp, del sdk.AccAddress) error { return nil })
+	endedPeriod, err := k.IncrementFinalityProviderPeriod(ctx, fp)
+	if err != nil {
+		return err
+	}
+
+	if err = k.CalculateBTCDelegationRewardsAndSendToGauge(ctx, fp, del, endedPeriod); err != nil {
+		return err
+	}
+
+	btcDelRwdTracker, err := k.GetBTCDelegationRewardsTracker(ctx, fp, del)
+	if err != nil {
+		return err
+	}
+
+	if err := k.decrementReferenceCount(ctx, fp, btcDelRwdTracker.StartPeriodCumulativeReward); err != nil {
+		return err
+	}
+
+	return k.initializeBTCDelegation(ctx, fp, del)
 }
 
 // btcDelegationModifiedWithPreInitDel does the procedure when a BTC delegation has
@@ -112,11 +130,11 @@ func (k Keeper) btcDelegationModifiedWithPreInitDel(
 		return err
 	}
 
-	if err := k.CalculateBTCDelegationRewardsAndSendToGauge(ctx, fp, del, endedPeriod); err != nil {
+	if err = k.CalculateBTCDelegationRewardsAndSendToGauge(ctx, fp, del, endedPeriod); err != nil {
 		return err
 	}
 
-	if err := preInitializeDelegation(ctx, fp, del); err != nil {
+	if err = preInitializeDelegation(ctx, fp, del); err != nil {
 		return err
 	}
 
@@ -240,7 +258,13 @@ func (k Keeper) IncrementFinalityProviderPeriod(ctx context.Context, fp sdk.AccA
 		return 0, err
 	}
 
-	newFpHistoricalRwd := types.NewFinalityProviderHistoricalRewards(fpHistoricalRwd.CumulativeRewardsPerSat.Add(currentRewardsPerSat...))
+	// decrement reference count
+	err = k.decrementReferenceCount(ctx, fp, fpCurrentRwd.Period-1)
+	if err != nil {
+		return 0, err
+	}
+
+	newFpHistoricalRwd := types.NewFinalityProviderHistoricalRewards(fpHistoricalRwd.CumulativeRewardsPerSat.Add(currentRewardsPerSat...), uint32(1))
 	if err := k.setFinalityProviderHistoricalRewards(ctx, fp, fpCurrentRwd.Period, newFpHistoricalRwd); err != nil {
 		return 0, err
 	}
@@ -254,12 +278,43 @@ func (k Keeper) IncrementFinalityProviderPeriod(ctx context.Context, fp sdk.AccA
 	return fpCurrentRwd.Period, nil
 }
 
+func (k Keeper) incrementReferenceCount(ctx context.Context, fp sdk.AccAddress, period uint64) error {
+	historical, err := k.GetFinalityProviderHistoricalRewards(ctx, fp, period)
+	if err != nil {
+		return err
+	}
+
+	historical.ReferenceCount++
+	if historical.ReferenceCount > 2 {
+		panic("reference count should never be greater than 2")
+	}
+
+	return k.setFinalityProviderHistoricalRewards(ctx, fp, period, historical)
+}
+
+func (k Keeper) decrementReferenceCount(ctx context.Context, fp sdk.AccAddress, period uint64) error {
+	historical, err := k.GetFinalityProviderHistoricalRewards(ctx, fp, period)
+	if err != nil {
+		return err
+	}
+
+	if historical.ReferenceCount == 0 {
+		panic("cannot set negative reference count")
+	}
+	historical.ReferenceCount--
+	if historical.ReferenceCount == 0 {
+		k.deleteFinalityProviderHistoricalRewards(ctx, fp, period)
+	}
+
+	return k.setFinalityProviderHistoricalRewards(ctx, fp, period, historical)
+}
+
 // initializeFinalityProvider initializes a new finality provider current rewards at period 1, empty rewards and zero sats
 // and also creates a new historical rewards at period 0 and zero rewards as well.
 // It does not verifies if it exists prior to overwrite, who calls it needs to verify.
 func (k Keeper) initializeFinalityProvider(ctx context.Context, fp sdk.AccAddress) (types.FinalityProviderCurrentRewards, error) {
 	// historical rewards starts at the period 0
-	err := k.setFinalityProviderHistoricalRewards(ctx, fp, 0, types.NewFinalityProviderHistoricalRewards(sdk.NewCoins()))
+	err := k.setFinalityProviderHistoricalRewards(ctx, fp, 0, types.NewFinalityProviderHistoricalRewards(sdk.NewCoins(), uint32(1)))
 	if err != nil {
 		return types.FinalityProviderCurrentRewards{}, err
 	}
@@ -276,16 +331,21 @@ func (k Keeper) initializeFinalityProvider(ctx context.Context, fp sdk.AccAddres
 // modification to the amount of satoshi staked from this btc delegator to
 // this finality provider (activivation or unbonding) of BTC delegations, it
 // should withdraw all rewards (send to gauge) and initialize a new BTCDelegationRewardsTracker.
-// TODO: add reference count to keep track of possible prunning state of val rewards
 func (k Keeper) initializeBTCDelegation(ctx context.Context, fp, del sdk.AccAddress) error {
 	// period has already been incremented prior to call this function
 	// it is needed to store the period ended by this delegation action
 	// as a starting point of the delegation rewards calculation
-	valCurrentRewards, err := k.GetFinalityProviderCurrentRewards(ctx, fp)
+	fpCurrentRwd, err := k.GetFinalityProviderCurrentRewards(ctx, fp)
+	if err != nil {
+		return err
+	}
+	previousPeriod := fpCurrentRwd.Period - 1
+
+	// increment reference count for the period we're going to track
+	err = k.incrementReferenceCount(ctx, fp, previousPeriod)
 	if err != nil {
 		return err
 	}
-	previousPeriod := valCurrentRewards.Period - 1
 
 	btcDelRwdTracker, err := k.GetBTCDelegationRewardsTracker(ctx, fp, del)
 	if err != nil {

x/incentive/keeper/reward_tracker_store.go

@@ -213,6 +213,13 @@ func (k Keeper) setFinalityProviderHistoricalRewards(ctx context.Context, fp sdk
 	return k.finalityProviderHistoricalRewards.Set(ctx, collections.Join(fp.Bytes(), period), rwd)
 }
 
+// deleteFinalityProviderHistoricalRewards deletes the historical rewards of FP based on the (fp, period) pair
+func (k Keeper) deleteFinalityProviderHistoricalRewards(ctx context.Context, fp sdk.AccAddress, period uint64) {
+	if err := k.finalityProviderHistoricalRewards.Remove(ctx, collections.Join(fp.Bytes(), period)); err != nil {
+		k.Logger(sdk.UnwrapSDKContext(ctx)).Error("error deleting FinalityProviderHistoricalRewards", "error", err)
+	}
+}
+
 // subDelegationSat subtracts an amount of active stake from the BTCDelegationRewardsTracker
 // and the FinalityProviderCurrentRewards.
 // There is no need to check if the fp or delegation exists, because they should exist

x/incentive/keeper/reward_tracker_test.go

@@ -355,7 +355,7 @@ func FuzzCheckCalculateDelegationRewardsBetween(f *testing.F) {
 		endingPeriod := btcRwd.StartPeriodCumulativeReward + 1
 
 		// creates a bad historical ending period that has less rewards than the starting one
-		err = k.setFinalityProviderHistoricalRewards(ctx, fp, endingPeriod, types.NewFinalityProviderHistoricalRewards(historicalStartPeriod.CumulativeRewardsPerSat.QuoInt(math.NewInt(2))))
+		err = k.setFinalityProviderHistoricalRewards(ctx, fp, endingPeriod, types.NewFinalityProviderHistoricalRewards(historicalStartPeriod.CumulativeRewardsPerSat.QuoInt(math.NewInt(2)), uint32(1)))
 		require.NoError(t, err)
 		require.Panics(t, func() {
 			_, _ = k.calculateDelegationRewardsBetween(ctx, fp, btcRwd, endingPeriod)
@@ -432,7 +432,7 @@ func FuzzCheckIncrementFinalityProviderPeriod(f *testing.F) {
 		require.NoError(t, err)
 
 		amtRwdInHistorical := fpCurrentRwd.CurrentRewards.MulInt(types.DecimalRewards).QuoInt(math.NewInt(2))
-		err = k.setFinalityProviderHistoricalRewards(ctx, fp, fpCurrentRwd.Period-1, types.NewFinalityProviderHistoricalRewards(amtRwdInHistorical))
+		err = k.setFinalityProviderHistoricalRewards(ctx, fp, fpCurrentRwd.Period-1, types.NewFinalityProviderHistoricalRewards(amtRwdInHistorical, uint32(1)))
 		require.NoError(t, err)
 
 		endedPeriod, err = k.IncrementFinalityProviderPeriod(ctx, fp)
@@ -469,19 +469,32 @@ func FuzzCheckInitializeBTCDelegation(f *testing.F) {
 		err = k.SetFinalityProviderCurrentRewards(ctx, fp, fpCurrentRwd)
 		require.NoError(t, err)
 
+		err = k.initializeBTCDelegation(ctx, fp, del)
+		require.EqualError(t, err, types.ErrFPHistoricalRewardsNotFound.Error())
+
+		fpHistoricalRwd := datagen.GenRandomFPHistRwd(r)
+		err = k.setFinalityProviderHistoricalRewards(ctx, fp, fpCurrentRwd.Period-1, fpHistoricalRwd)
+		require.NoError(t, err)
+
 		err = k.initializeBTCDelegation(ctx, fp, del)
 		require.EqualError(t, err, types.ErrBTCDelegationRewardsTrackerNotFound.Error())
 
 		delBtcRwdTrackerBeforeInitialize := datagen.GenRandomBTCDelegationRewardsTracker(r)
 		err = k.setBTCDelegationRewardsTracker(ctx, fp, del, delBtcRwdTrackerBeforeInitialize)
 		require.NoError(t, err)
 
+		// increment period since reference count is already increased in the previous
+		// initializeBTCDelegation
+		endedPeriod, err := k.IncrementFinalityProviderPeriod(ctx, fp)
+		require.NoError(t, err)
+		require.Equal(t, endedPeriod, fpCurrentRwd.Period)
+
 		err = k.initializeBTCDelegation(ctx, fp, del)
 		require.NoError(t, err)
 
 		actBtcDelRwdTracker, err := k.GetBTCDelegationRewardsTracker(ctx, fp, del)
 		require.NoError(t, err)
-		require.Equal(t, fpCurrentRwd.Period-1, actBtcDelRwdTracker.StartPeriodCumulativeReward)
+		require.Equal(t, fpCurrentRwd.Period, actBtcDelRwdTracker.StartPeriodCumulativeReward)
 		require.Equal(t, delBtcRwdTrackerBeforeInitialize.TotalActiveSat, actBtcDelRwdTracker.TotalActiveSat)
 	})
 }
@@ -528,14 +541,14 @@ func TestIncrementFinalityProviderPeriod(t *testing.T) {
 	require.Equal(t, fp1EndedPeriod, uint64(1))
 
 	checkFpCurrentRwd(t, ctx, k, fp1, fp1EndedPeriod, sdk.NewCoins(), math.NewInt(0))
-	checkFpHistoricalRwd(t, ctx, k, fp1, 0, sdk.NewCoins())
+	checkFpHistoricalRwdWithRefCount(t, ctx, k, fp1, 0, sdk.NewCoins(), uint32(1))
 
 	rwdAddedToPeriod1 := newBaseCoins(2_000000) // 2bbn
 	err = k.AddFinalityProviderRewardsForBtcDelegations(ctx, fp1, rwdAddedToPeriod1)
 	require.NoError(t, err)
 
 	// historical should not modify the rewards for the period already created
-	checkFpHistoricalRwd(t, ctx, k, fp1, 0, sdk.NewCoins())
+	checkFpHistoricalRwdWithRefCount(t, ctx, k, fp1, 0, sdk.NewCoins(), uint32(1))
 	checkFpCurrentRwd(t, ctx, k, fp1, fp1EndedPeriod, rwdAddedToPeriod1, math.NewInt(0))
 
 	// needs to add some voting power so it can calculate the amount of rewards per share
@@ -548,18 +561,19 @@ func TestIncrementFinalityProviderPeriod(t *testing.T) {
 	require.Equal(t, fp1EndedPeriod, uint64(1))
 
 	// now the historical that just ended should have as cumulative rewards 4000ubbn 2_000000ubbn/500sats
-	checkFpHistoricalRwd(t, ctx, k, fp1, fp1EndedPeriod, newBaseCoins(4000).MulInt(types.DecimalRewards))
+	checkFpHistoricalRwdWithRefCount(t, ctx, k, fp1, fp1EndedPeriod, newBaseCoins(4000).MulInt(types.DecimalRewards), uint32(1))
 	checkFpCurrentRwd(t, ctx, k, fp1, fp1EndedPeriod+1, sdk.NewCoins(), satsDelegated)
 
 	fp2EndedPeriod, err := k.IncrementFinalityProviderPeriod(ctx, fp2)
 	require.NoError(t, err)
 	require.Equal(t, fp2EndedPeriod, uint64(1))
 }
 
-func checkFpHistoricalRwd(t *testing.T, ctx sdk.Context, k *Keeper, fp sdk.AccAddress, period uint64, expectedRwd sdk.Coins) {
+func checkFpHistoricalRwdWithRefCount(t *testing.T, ctx sdk.Context, k *Keeper, fp sdk.AccAddress, period uint64, expectedRwd sdk.Coins, expectedRefCount uint32) {
 	historical, err := k.GetFinalityProviderHistoricalRewards(ctx, fp, period)
 	require.NoError(t, err)
 	require.Equal(t, historical.CumulativeRewardsPerSat.String(), expectedRwd.String())
+	require.Equal(t, historical.ReferenceCount, expectedRefCount)
 }
 
 func checkFpCurrentRwd(t *testing.T, ctx sdk.Context, k *Keeper, fp sdk.AccAddress, expectedPeriod uint64, expectedRwd sdk.Coins, totalActiveSat math.Int) {

x/incentive/types/rewards.go

@@ -58,9 +58,10 @@ func NewFinalityProviderCurrentRewards(currentRewards sdk.Coins, period uint64,
 	}
 }
 
-func NewFinalityProviderHistoricalRewards(cumulativeRewardsPerSat sdk.Coins) FinalityProviderHistoricalRewards {
+func NewFinalityProviderHistoricalRewards(cumulativeRewardsPerSat sdk.Coins, referenceCount uint32) FinalityProviderHistoricalRewards {
 	return FinalityProviderHistoricalRewards{
 		CumulativeRewardsPerSat: cumulativeRewardsPerSat,
+		ReferenceCount:          referenceCount,
 	}
 }