CNS-866: make pool methods support all denoms (lavanet#1200)

x/rewards/keeper/grpc_query_block_reward.go

@@ -23,7 +23,7 @@ func (k Keeper) BlockReward(goCtx context.Context, req *types.QueryBlockRewardRe
 	blocksToNextTimerExpiry := k.BlocksToNextTimerExpiry(ctx)
 
 	// get validator block pool balance
-	blockPoolBalance := k.TotalPoolTokens(ctx, types.ValidatorsRewardsDistributionPoolName)
+	blockPoolBalance := k.TotalPoolTokens(ctx, types.ValidatorsRewardsDistributionPoolName, k.stakingKeeper.BondDenom(ctx))
 	if blocksToNextTimerExpiry == 0 {
 		return nil, utils.LavaFormatWarning("blocksToNextTimerExpiry is zero", fmt.Errorf("critical: Attempt to divide by zero"),
 			utils.LogAttr("blocksToNextTimerExpiry", blocksToNextTimerExpiry),

x/rewards/keeper/grpc_query_pools.go

@@ -19,19 +19,19 @@ func (k Keeper) Pools(goCtx context.Context, req *types.QueryPoolsRequest) (*typ
 	pools := []types.PoolInfo{
 		{
 			Name:    string(types.ValidatorsRewardsDistributionPoolName),
-			Balance: sdk.NewCoin(k.stakingKeeper.BondDenom(ctx), k.TotalPoolTokens(ctx, types.ValidatorsRewardsDistributionPoolName)),
+			Balance: sdk.NewCoin(k.stakingKeeper.BondDenom(ctx), k.TotalPoolTokens(ctx, types.ValidatorsRewardsDistributionPoolName, k.stakingKeeper.BondDenom(ctx))),
 		},
 		{
 			Name:    string(types.ValidatorsRewardsAllocationPoolName),
-			Balance: sdk.NewCoin(k.stakingKeeper.BondDenom(ctx), k.TotalPoolTokens(ctx, types.ValidatorsRewardsAllocationPoolName)),
+			Balance: sdk.NewCoin(k.stakingKeeper.BondDenom(ctx), k.TotalPoolTokens(ctx, types.ValidatorsRewardsAllocationPoolName, k.stakingKeeper.BondDenom(ctx))),
 		},
 		{
 			Name:    string(types.ProviderRewardsDistributionPool),
-			Balance: sdk.NewCoin(k.stakingKeeper.BondDenom(ctx), k.TotalPoolTokens(ctx, types.ProviderRewardsDistributionPool)),
+			Balance: sdk.NewCoin(k.stakingKeeper.BondDenom(ctx), k.TotalPoolTokens(ctx, types.ProviderRewardsDistributionPool, k.stakingKeeper.BondDenom(ctx))),
 		},
 		{
 			Name:    string(types.ProvidersRewardsAllocationPool),
-			Balance: sdk.NewCoin(k.stakingKeeper.BondDenom(ctx), k.TotalPoolTokens(ctx, types.ProvidersRewardsAllocationPool)),
+			Balance: sdk.NewCoin(k.stakingKeeper.BondDenom(ctx), k.TotalPoolTokens(ctx, types.ProvidersRewardsAllocationPool, k.stakingKeeper.BondDenom(ctx))),
 		},
 	}
 

x/rewards/keeper/helpers_test.go

@@ -39,7 +39,7 @@ func newTester(t *testing.T, addValidator bool) *tester {
 	}
 
 	ts.plan = common.CreateMockPlan()
-	monthlyProvidersPool := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, rewardsTypes.ProviderRewardsDistributionPool)
+	monthlyProvidersPool := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, rewardsTypes.ProviderRewardsDistributionPool, ts.BondDenom())
 	ts.plan.Price.Amount = monthlyProvidersPool.QuoRaw(5).AddRaw(5)
 	ts.plan.PlanPolicy.EpochCuLimit = monthlyProvidersPool.Uint64() * 5
 	ts.plan.PlanPolicy.TotalCuLimit = monthlyProvidersPool.Uint64() * 5

x/rewards/keeper/pool.go

@@ -8,19 +8,19 @@ import (
 )
 
 // TotalPoolTokens gets the total tokens supply from a pool
-func (k Keeper) TotalPoolTokens(ctx sdk.Context, pool types.Pool) math.Int {
+func (k Keeper) TotalPoolTokens(ctx sdk.Context, pool types.Pool, denom string) math.Int {
 	poolAddr := k.accountKeeper.GetModuleAddress(string(pool))
-	return k.bankKeeper.GetBalance(ctx, poolAddr, k.stakingKeeper.BondDenom(ctx)).Amount
+	return k.bankKeeper.GetBalance(ctx, poolAddr, denom).Amount
 }
 
 // BurnPoolTokens removes coins from a pool module account
-func (k Keeper) BurnPoolTokens(ctx sdk.Context, pool types.Pool, amt math.Int) error {
+func (k Keeper) BurnPoolTokens(ctx sdk.Context, pool types.Pool, amt math.Int, denom string) error {
 	if !amt.IsPositive() {
 		// skip as no coins need to be burned
 		return nil
 	}
 
-	coins := sdk.NewCoins(sdk.NewCoin(k.stakingKeeper.BondDenom(ctx), amt))
+	coins := sdk.NewCoins(sdk.NewCoin(denom, amt))
 
 	return k.bankKeeper.BurnCoins(ctx, string(pool), coins)
 }

x/rewards/keeper/pool_test.go

@@ -78,7 +78,7 @@ func TestRewardsModuleSetup(t *testing.T) {
 func TestBurnRateParam(t *testing.T) {
 	ts := newTester(t, true)
 	lifetime := int64(types.RewardsAllocationPoolsLifetime)
-	allocPoolBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsAllocationPoolName).Int64()
+	allocPoolBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsAllocationPoolName, ts.BondDenom()).Int64()
 
 	// advance a month to trigger monthly pool refill callback
 	// to see why these 3 are called, see general note 2
@@ -90,7 +90,7 @@ func TestBurnRateParam(t *testing.T) {
 	// default burn rate = 1, distribution pool's old balance should be wiped
 	// current balance should be exactly the expected monthly quota minus block reward
 	expectedMonthlyQuota := allocPoolBalance / (lifetime - 1)
-	distPoolBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsDistributionPoolName).Int64()
+	distPoolBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsDistributionPoolName, ts.BondDenom()).Int64()
 	require.Equal(t, expectedMonthlyQuota, distPoolBalance)
 
 	// change the burn rate param to be zero
@@ -105,12 +105,12 @@ func TestBurnRateParam(t *testing.T) {
 	resp, err = ts.QueryRewardsPools()
 	require.NoError(t, err)
 	ts.AdvanceBlock(time.Duration(resp.TimeToRefill) * time.Second)
-	prevDistPoolBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsDistributionPoolName).Int64()
+	prevDistPoolBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsDistributionPoolName, ts.BondDenom()).Int64()
 	testkeeper.EndBlock(ts.Ctx, ts.Keepers)
 
 	// burn rate = 0, distribution pool's old balance should not be wiped
 	// current balance should be previous balance (minus block reward) plus new quota
-	distPoolBalance = ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsDistributionPoolName).Int64()
+	distPoolBalance = ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsDistributionPoolName, ts.BondDenom()).Int64()
 	require.Equal(t, prevDistPoolBalance+expectedMonthlyQuota, distPoolBalance)
 }
 
@@ -125,11 +125,11 @@ func TestAllocationPoolMonthlyQuota(t *testing.T) {
 	// calc expectedMonthlyQuota. Check that it was subtracted from the allocation pool and added
 	// to the distribution pool (its balance should be the monthly quota minus the fee collector's balance)
 	expectedMonthlyQuota := allocationPoolBalance / lifetime
-	currentAllocPoolBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsAllocationPoolName)
+	currentAllocPoolBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsAllocationPoolName, ts.BondDenom())
 	require.Equal(t, expectedMonthlyQuota, allocationPoolBalance-currentAllocPoolBalance.Int64())
 
 	feeCollectorBalance := ts.GetBalance(ts.feeCollector())
-	currentDistPoolBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsDistributionPoolName)
+	currentDistPoolBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsDistributionPoolName, ts.BondDenom())
 	require.Equal(t, expectedMonthlyQuota, feeCollectorBalance+currentDistPoolBalance.Int64())
 
 	// check the monthly quota is as expected with advancement of months
@@ -141,7 +141,7 @@ func TestAllocationPoolMonthlyQuota(t *testing.T) {
 
 		monthsLeft := ts.Keepers.Rewards.AllocationPoolMonthsLeft(ts.Ctx)
 		prevAllocPoolBalance := currentAllocPoolBalance
-		currentAllocPoolBalance = ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsAllocationPoolName)
+		currentAllocPoolBalance = ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsAllocationPoolName, ts.BondDenom())
 
 		var monthlyQuota int64
 		if monthsLeft != 0 {
@@ -160,7 +160,7 @@ func TestAllocationPoolMonthlyQuota(t *testing.T) {
 	ts.AdvanceMonths(1)
 	ts.AdvanceBlock()
 	testkeeper.EndBlock(ts.Ctx, ts.Keepers)
-	currentAllocPoolBalance = ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsAllocationPoolName)
+	currentAllocPoolBalance = ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsAllocationPoolName, ts.BondDenom())
 	require.True(t, currentAllocPoolBalance.IsZero())
 
 	// advance another month to distribute the last quota to the fee collector
@@ -174,7 +174,7 @@ func TestAllocationPoolMonthlyQuota(t *testing.T) {
 		ts.AdvanceBlock()
 		testkeeper.EndBlock(ts.Ctx, ts.Keepers)
 
-		currentAllocPoolBalance = ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsAllocationPoolName)
+		currentAllocPoolBalance = ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsAllocationPoolName, ts.BondDenom())
 		require.True(t, currentAllocPoolBalance.IsZero())
 
 		currentFeeCollectorBalance := ts.GetBalance(ts.feeCollector())
@@ -212,7 +212,7 @@ func TestValidatorBlockRewards(t *testing.T) {
 	res, err := ts.QueryRewardsBlockReward()
 	require.NoError(t, err)
 	blockReward := res.Reward.Amount
-	distPoolBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsDistributionPoolName)
+	distPoolBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsDistributionPoolName, ts.BondDenom())
 	blocksToNextExpiry := ts.Keepers.Rewards.BlocksToNextTimerExpiry(ts.Ctx)
 	bondedTargetFactor := sdk.OneDec().MulInt(blockReward).MulInt64(blocksToNextExpiry).QuoInt(distPoolBalance).TruncateInt()
 	require.True(t, bondedTargetFactor.Equal(expectedBondedTargetFactor))
@@ -228,7 +228,7 @@ func TestValidatorBlockRewards(t *testing.T) {
 	blockReward = res.Reward.Amount
 
 	// transfer half of the total distribution pool balance to the allocation pool
-	distPoolBalance = ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsDistributionPoolName)
+	distPoolBalance = ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsDistributionPoolName, ts.BondDenom())
 	err = ts.Keepers.BankKeeper.SendCoinsFromModuleToModule(
 		ts.Ctx,
 		string(types.ValidatorsRewardsDistributionPoolName),
@@ -263,7 +263,7 @@ func TestValidatorBlockRewards(t *testing.T) {
 	require.NoError(t, err)
 	blockReward = res.Reward.Amount
 	bondedTargetFactor = ts.Keepers.Rewards.BondedTargetFactor(ts.Ctx).TruncateInt()
-	distPoolBalance = ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsDistributionPoolName)
+	distPoolBalance = ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ValidatorsRewardsDistributionPoolName, ts.BondDenom())
 	blocksToNextExpiry = bondedTargetFactor.Mul(distPoolBalance).Quo(blockReward).Int64()
 	require.Equal(t, refBlocksToExpiry, blocksToNextExpiry)
 }

x/rewards/keeper/providers.go

@@ -30,7 +30,7 @@ func (k Keeper) AggregateRewards(ctx sdk.Context, provider, chainid string, adju
 // Distribute bonus rewards to providers across all chains based on performance
 func (k Keeper) distributeMonthlyBonusRewards(ctx sdk.Context) {
 	details := map[string]string{}
-	total := k.TotalPoolTokens(ctx, types.ProviderRewardsDistributionPool)
+	total := k.TotalPoolTokens(ctx, types.ProviderRewardsDistributionPool, k.stakingKeeper.BondDenom(ctx))
 	totalRewarded := sdk.ZeroInt()
 	// specs emissions from the total reward pool base on stake
 	specs := k.specEmissionParts(ctx)

x/rewards/keeper/providers_test.go

@@ -124,7 +124,7 @@ func TestSpecAllocationProvidersRewards(t *testing.T) {
 
 	// now the provider should get all of the provider allocation
 	ts.AdvanceMonths(1)
-	distBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ProviderRewardsDistributionPool)
+	distBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ProviderRewardsDistributionPool, ts.BondDenom())
 	ts.AdvanceEpoch()
 
 	res, err = ts.QueryDualstakingDelegatorRewards(providerAcc.Addr.String(), providerAcc.Addr.String(), "")
@@ -172,7 +172,7 @@ func TestProvidersDiminishingRewards(t *testing.T) {
 
 	// now the provider should get all of the provider allocation
 	ts.AdvanceMonths(1)
-	distBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ProviderRewardsDistributionPool)
+	distBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ProviderRewardsDistributionPool, ts.BondDenom())
 	ts.AdvanceEpoch()
 
 	res, err = ts.QueryDualstakingDelegatorRewards(providerAcc.Addr.String(), providerAcc.Addr.String(), "")
@@ -280,7 +280,7 @@ func Test2SpecsZeroShares(t *testing.T) {
 
 	// now the provider should get all of the provider allocation
 	ts.AdvanceMonths(1)
-	distBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ProviderRewardsDistributionPool)
+	distBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ProviderRewardsDistributionPool, ts.BondDenom())
 	ts.AdvanceEpoch()
 
 	res, err = ts.QueryDualstakingDelegatorRewards(providerAcc.Addr.String(), providerAcc.Addr.String(), "")
@@ -414,7 +414,7 @@ func TestBonusRewards3Providers(t *testing.T) {
 
 	// now the provider should get all of the provider allocation
 	ts.AdvanceMonths(1)
-	distBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ProviderRewardsDistributionPool)
+	distBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ProviderRewardsDistributionPool, ts.BondDenom())
 	ts.AdvanceEpoch()
 
 	res1, err := ts.QueryDualstakingDelegatorRewards(providerAcc1.Addr.String(), "", "")
@@ -596,7 +596,7 @@ func TestBonusRewardsEquall5Providers(t *testing.T) {
 
 	// now the provider should get all of the provider allocation
 	ts.AdvanceMonths(1)
-	distBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ProviderRewardsDistributionPool)
+	distBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ProviderRewardsDistributionPool, ts.BondDenom())
 	ts.AdvanceEpoch()
 
 	for _, providerAcc := range providerAccs {
@@ -663,7 +663,7 @@ func TestBonusRewards5Providers(t *testing.T) {
 
 	// now the provider should get all of the provider allocation
 	ts.AdvanceMonths(1)
-	distBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ProviderRewardsDistributionPool)
+	distBalance := ts.Keepers.Rewards.TotalPoolTokens(ts.Ctx, types.ProviderRewardsDistributionPool, ts.BondDenom())
 	ts.AdvanceEpoch()
 
 	// distribution pool divided between all providers (5) equally (they served the same amount of CU in total)

x/rewards/keeper/rewards.go

@@ -20,7 +20,7 @@ func (k Keeper) DistributeBlockReward(ctx sdk.Context) {
 	blocksToNextTimerExpiry := k.BlocksToNextTimerExpiry(ctx)
 
 	// get validator distribution pool balance
-	distributionPoolBalance := k.TotalPoolTokens(ctx, types.ValidatorsRewardsDistributionPoolName)
+	distributionPoolBalance := k.TotalPoolTokens(ctx, types.ValidatorsRewardsDistributionPoolName, k.stakingKeeper.BondDenom(ctx))
 
 	if blocksToNextTimerExpiry == 0 {
 		utils.LavaFormatWarning("blocksToNextTimerExpiry is zero", fmt.Errorf("critical: Attempt to divide by zero"),
@@ -84,8 +84,8 @@ func (k Keeper) RefillRewardsPools(ctx sdk.Context, _ []byte, data []byte) {
 	nextMonth := utils.NextMonth(ctx.BlockTime()).UTC()
 	k.refillRewardsPoolTS.AddTimerByBlockTime(ctx, uint64(nextMonth.Unix()), []byte(types.RefillRewardsPoolTimerName), monthsLeftBytes)
 
-	valDistPoolBalance := k.TotalPoolTokens(ctx, types.ValidatorsRewardsDistributionPoolName).Int64()
-	providerDistPoolBalance := k.TotalPoolTokens(ctx, types.ProviderRewardsDistributionPool).Int64()
+	valDistPoolBalance := k.TotalPoolTokens(ctx, types.ValidatorsRewardsDistributionPoolName, k.stakingKeeper.BondDenom(ctx)).Int64()
+	providerDistPoolBalance := k.TotalPoolTokens(ctx, types.ProviderRewardsDistributionPool, k.stakingKeeper.BondDenom(ctx)).Int64()
 	nextRefillBlock := k.blocksToNextTimerExpiry(ctx, nextMonth.Unix()-ctx.BlockTime().UTC().Unix()) + ctx.BlockHeight()
 	details := map[string]string{
 		"allocation_pool_remaining_lifetime":   strconv.FormatUint(monthsLeft, 10),
@@ -101,8 +101,8 @@ func (k Keeper) RefillRewardsPools(ctx sdk.Context, _ []byte, data []byte) {
 
 func (k Keeper) refillDistributionPool(ctx sdk.Context, monthsLeft uint64, allocationPool types.Pool, distributionPool types.Pool, burnRate sdkmath.LegacyDec) {
 	// burn remaining tokens in the distribution pool
-	tokensToBurn := burnRate.MulInt(k.TotalPoolTokens(ctx, distributionPool)).TruncateInt()
-	err := k.BurnPoolTokens(ctx, distributionPool, tokensToBurn)
+	tokensToBurn := burnRate.MulInt(k.TotalPoolTokens(ctx, distributionPool, k.stakingKeeper.BondDenom(ctx))).TruncateInt()
+	err := k.BurnPoolTokens(ctx, distributionPool, tokensToBurn, k.stakingKeeper.BondDenom(ctx))
 	if err != nil {
 		utils.LavaFormatError("critical - could not burn distribution pool tokens", err,
 			utils.Attribute{Key: "distribution_pool", Value: string(distributionPool)},
@@ -111,7 +111,7 @@ func (k Keeper) refillDistributionPool(ctx sdk.Context, monthsLeft uint64, alloc
 	}
 
 	// transfer the new monthly quota (if allocation pool is expired, rewards=0)
-	allocPoolBalance := k.TotalPoolTokens(ctx, allocationPool)
+	allocPoolBalance := k.TotalPoolTokens(ctx, allocationPool, k.stakingKeeper.BondDenom(ctx))
 	if monthsLeft != 0 && !allocPoolBalance.IsZero() {
 		monthlyQuota := sdk.Coin{Denom: k.stakingKeeper.BondDenom(ctx), Amount: allocPoolBalance.QuoRaw(int64(monthsLeft))}