From 3d35dc2dd9843d848ea1d83515d6f90c76df1034 Mon Sep 17 00:00:00 2001 From: Hanjun Kim Date: Thu, 5 Sep 2024 22:57:22 +0900 Subject: [PATCH] feat: refund remaining rewards to the service's address when a plan ends (#101) ## Description Closes: MILK-82 --- ### Author Checklist *All items are required. Please add a note to the item if the item is not applicable and please add links to any relevant follow up issues.* I have... - [ ] included the correct [type prefix](https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.json) in the PR title - [ ] added `!` to the type prefix if API or client breaking change - [ ] targeted the correct branch (see [PR Targeting](https://github.com/milkyway-labs/milkyway/blob/master/CONTRIBUTING.md#pr-targeting)) - [ ] provided a link to the relevant issue or specification - [ ] followed the guidelines for [building modules](https://docs.cosmos.network/v0.44/building-modules/intro.html) - [ ] included the necessary unit and integration [tests](https://github.com/milkyway-labs/milkyway/blob/master/CONTRIBUTING.md#testing) - [ ] added a changelog entry to `CHANGELOG.md` - [ ] included comments for [documenting Go code](https://blog.golang.org/godoc) - [ ] updated the relevant documentation or specification - [ ] reviewed "Files changed" and left comments if necessary - [ ] confirmed all CI checks have passed ### Reviewers Checklist *All items are required. Please add a note if the item is not applicable and please add your handle next to the items reviewed if you only reviewed selected items.* I have... - [ ] confirmed the correct [type prefix](https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.json) in the PR title - [ ] confirmed `!` in the type prefix if API or client breaking change - [ ] confirmed all author checklist items have been addressed - [ ] reviewed state machine logic - [ ] reviewed API design and naming - [ ] reviewed documentation is accurate - [ ] reviewed tests and test coverage - [ ] manually tested (if applicable) --- x/rewards/keeper/abci.go | 10 +++- x/rewards/keeper/allocation.go | 2 +- x/rewards/keeper/rewards_plan.go | 70 ++++++++++++++++++++++++++- x/rewards/keeper/rewards_plan_test.go | 44 +++++++++++++++++ x/rewards/types/events.go | 14 +++--- x/rewards/types/expected_keepers.go | 1 + x/rewards/types/models.go | 5 +- 7 files changed, 134 insertions(+), 12 deletions(-) diff --git a/x/rewards/keeper/abci.go b/x/rewards/keeper/abci.go index 9f2c84b6..f1e55ab7 100644 --- a/x/rewards/keeper/abci.go +++ b/x/rewards/keeper/abci.go @@ -4,9 +4,15 @@ import ( "context" ) -// BeginBlocker allocates restaking rewards for the previous block. +// BeginBlocker is called every block and is used to terminate ended rewards +// plans and allocate restaking rewards for the previous block. func (k *Keeper) BeginBlocker(ctx context.Context) error { - err := k.AllocateRewards(ctx) + err := k.TerminateEndedRewardsPlans(ctx) + if err != nil { + return err + } + + err = k.AllocateRewards(ctx) if err != nil { return err } diff --git a/x/rewards/keeper/allocation.go b/x/rewards/keeper/allocation.go index 197b180f..784d3194 100644 --- a/x/rewards/keeper/allocation.go +++ b/x/rewards/keeper/allocation.go @@ -129,7 +129,7 @@ func (k *Keeper) AllocateRewardsByPlan( rewards = sdk.NewDecCoinsFromCoins(rewardsTruncated...) // Check if the rewards pool has enough coins to allocate rewards. - planRewardsPoolAddr := plan.MustGetRewardsPoolAddress() + planRewardsPoolAddr := plan.MustGetRewardsPoolAddress(k.accountKeeper.AddressCodec()) balances := k.bankKeeper.GetAllBalances(ctx, planRewardsPoolAddr) sdkCtx := sdk.UnwrapSDKContext(ctx) if !balances.IsAllGTE(rewardsTruncated) { diff --git a/x/rewards/keeper/rewards_plan.go b/x/rewards/keeper/rewards_plan.go index 61c44a27..dd93e894 100644 --- a/x/rewards/keeper/rewards_plan.go +++ b/x/rewards/keeper/rewards_plan.go @@ -2,6 +2,7 @@ package keeper import ( "context" + "fmt" "time" "cosmossdk.io/errors" @@ -75,7 +76,7 @@ func (k *Keeper) CreateRewardsPlan( // types.UsersDistributionTypeBasic only which doesn't need a validation. // Create a rewards pool account if it doesn't exist - k.createAccountIfNotExists(ctx, plan.MustGetRewardsPoolAddress()) + k.createAccountIfNotExists(ctx, plan.MustGetRewardsPoolAddress(k.accountKeeper.AddressCodec())) // Store the rewards plan err = k.RewardsPlans.Set(ctx, planID, plan) @@ -112,3 +113,70 @@ func (k *Keeper) validateDistributionDelegationTargets(ctx context.Context, dist func (k *Keeper) GetRewardsPlan(ctx context.Context, planID uint64) (types.RewardsPlan, error) { return k.RewardsPlans.Get(ctx, planID) } + +// terminateRewardsPlan removes a rewards plan and transfers the remaining +// rewards in the plan's rewards pool to the service's address. +func (k *Keeper) terminateRewardsPlan(ctx context.Context, plan types.RewardsPlan) error { + sdkCtx := sdk.UnwrapSDKContext(ctx) + + // Transfer remaining rewards in the plan's rewards pool to the service's + // address. + rewardsPoolAddr := plan.MustGetRewardsPoolAddress(k.accountKeeper.AddressCodec()) + remaining := k.bankKeeper.GetAllBalances(ctx, rewardsPoolAddr) + if remaining.IsAllPositive() { + // Get the service's address. + service, found := k.servicesKeeper.GetService(sdkCtx, plan.ServiceID) + if !found { + return servicestypes.ErrServiceNotFound + } + serviceAddr, err := k.accountKeeper.AddressCodec().StringToBytes(service.Address) + if err != nil { + return err + } + + // Transfer all the remaining rewards to the service's address. + err = k.bankKeeper.SendCoins(ctx, rewardsPoolAddr, serviceAddr, remaining) + if err != nil { + return err + } + } + + // Remove the plan. + err := k.RewardsPlans.Remove(ctx, plan.ID) + if err != nil { + return err + } + + sdkCtx.EventManager().EmitEvents(sdk.Events{ + sdk.NewEvent( + types.EventTypeTerminateRewardsPlan, + sdk.NewAttribute(types.AttributeKeyRewardsPlanID, fmt.Sprint(plan.ID)), + sdk.NewAttribute(types.AttributeKeyRemainingRewards, remaining.String()), + ), + }) + + return nil +} + +// TerminateEndedRewardsPlans terminates all rewards plans that have ended. +func (k *Keeper) TerminateEndedRewardsPlans(ctx context.Context) error { + sdkCtx := sdk.UnwrapSDKContext(ctx) + // Get the current block time + blockTime := sdkCtx.BlockTime() + + // Iterate over all rewards plans + err := k.RewardsPlans.Walk(ctx, nil, func(planID uint64, plan types.RewardsPlan) (stop bool, err error) { + // If the plan has already ended, terminate it + if !blockTime.Before(plan.EndTime) { + err = k.terminateRewardsPlan(ctx, plan) + if err != nil { + return false, err + } + } + return false, nil + }) + if err != nil { + return err + } + return nil +} diff --git a/x/rewards/keeper/rewards_plan_test.go b/x/rewards/keeper/rewards_plan_test.go index 1025b0a6..428aae8d 100644 --- a/x/rewards/keeper/rewards_plan_test.go +++ b/x/rewards/keeper/rewards_plan_test.go @@ -3,6 +3,8 @@ package keeper_test import ( "time" + "cosmossdk.io/collections" + "github.com/milkyway-labs/milkyway/app/testutil" "github.com/milkyway-labs/milkyway/utils" rewardskeeper "github.com/milkyway-labs/milkyway/x/rewards/keeper" @@ -66,3 +68,45 @@ func (suite *KeeperTestSuite) TestCreateRewardsPlan_PoolOrOperatorNotFound() { )) suite.Require().EqualError(err, "cannot get delegation target 2: operator not found: not found") } + +func (suite *KeeperTestSuite) TestTerminateEndedRewardsPlans() { + // Cache the context to avoid errors + ctx, _ := suite.Ctx.CacheContext() + + service, _ := suite.setupSampleServiceAndOperator(ctx) + + // Create an active rewards plan. + plan := suite.CreateBasicRewardsPlan( + ctx, + service.ID, + utils.MustParseCoins("100_000000service"), + time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC), + time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC), + utils.MustParseCoins("10000_000000service"), + ) + + rewardsPoolAddr := plan.MustGetRewardsPoolAddress(suite.App.AccountKeeper.AddressCodec()) + remaining := suite.App.BankKeeper.GetAllBalances(ctx, rewardsPoolAddr) + suite.Require().Equal("10000000000service", remaining.String()) + + // Change the block time so that the plan becomes no more active. + ctx = ctx.WithBlockTime(time.Date(2025, 1, 2, 0, 0, 0, 0, time.UTC)) + + // Terminate the ended rewards plans + err := suite.keeper.TerminateEndedRewardsPlans(ctx) + suite.Require().NoError(err) + + // The plan is removed. + _, err = suite.keeper.GetRewardsPlan(ctx, plan.ID) + suite.Require().ErrorIs(err, collections.ErrNotFound) + + // All remaining rewards are transferred to the service's address. + remaining = suite.App.BankKeeper.GetAllBalances(ctx, rewardsPoolAddr) + suite.Require().True(remaining.IsZero()) + + // Check the service's address balances. + serviceAddr, err := suite.App.AccountKeeper.AddressCodec().StringToBytes(service.Address) + suite.Require().NoError(err) + serviceBalances := suite.App.BankKeeper.GetAllBalances(ctx, serviceAddr) + suite.Require().Equal("10000000000service", serviceBalances.String()) +} diff --git a/x/rewards/types/events.go b/x/rewards/types/events.go index e245429f..7fcea42c 100644 --- a/x/rewards/types/events.go +++ b/x/rewards/types/events.go @@ -1,17 +1,19 @@ package types const ( - EventTypeCreateRewardsPlan = "create_rewards_plan" - EventTypeSetWithdrawAddress = "set_withdraw_address" - EventTypeRewards = "rewards" - EventTypeCommission = "commission" - EventTypeWithdrawRewards = "withdraw_rewards" - EventTypeWithdrawCommission = "withdraw_commission" + EventTypeCreateRewardsPlan = "create_rewards_plan" + EventTypeSetWithdrawAddress = "set_withdraw_address" + EventTypeRewards = "rewards" + EventTypeCommission = "commission" + EventTypeWithdrawRewards = "withdraw_rewards" + EventTypeWithdrawCommission = "withdraw_commission" + EventTypeTerminateRewardsPlan = "terminate_rewards_plan" AttributeKeyRewardsPlanID = "rewards_plan_id" AttributeKeyWithdrawAddress = "withdraw_address" AttributeKeyDelegationType = "delegation_type" AttributeKeyDelegationTargetID = "delegation_target_id" + AttributeKeyRemainingRewards = "remaining_rewards" // AttributeKeyAmountPerPool represents the amount of rewards per pool (per denom). // See https://github.com/initia-labs/initia/blob/v0.2.10/x/distribution/types/events.go#L3-L6 diff --git a/x/rewards/types/expected_keepers.go b/x/rewards/types/expected_keepers.go index 10471272..ebf66d37 100644 --- a/x/rewards/types/expected_keepers.go +++ b/x/rewards/types/expected_keepers.go @@ -27,6 +27,7 @@ type AccountKeeper interface { type BankKeeper interface { GetAllBalances(ctx context.Context, addr sdk.AccAddress) sdk.Coins + SendCoins(ctx context.Context, fromAddr, toAddr sdk.AccAddress, amt sdk.Coins) error SendCoinsFromModuleToAccount(ctx context.Context, moduleName string, addr sdk.AccAddress, amt sdk.Coins) error SendCoinsFromAccountToModule(ctx context.Context, addr sdk.AccAddress, moduleName string, amt sdk.Coins) error BlockedAddr(addr sdk.AccAddress) bool diff --git a/x/rewards/types/models.go b/x/rewards/types/models.go index 7b1a8fbe..4d9f0231 100644 --- a/x/rewards/types/models.go +++ b/x/rewards/types/models.go @@ -4,6 +4,7 @@ import ( "fmt" "time" + coreaddress "cosmossdk.io/core/address" codectypes "github.com/cosmos/cosmos-sdk/codec/types" sdk "github.com/cosmos/cosmos-sdk/types" "github.com/cosmos/cosmos-sdk/types/address" @@ -63,8 +64,8 @@ func (plan RewardsPlan) IsActiveAt(t time.Time) bool { } // MustGetRewardsPoolAddress returns the rewards pool address. -func (plan RewardsPlan) MustGetRewardsPoolAddress() sdk.AccAddress { - addr, err := sdk.AccAddressFromBech32(plan.RewardsPool) +func (plan RewardsPlan) MustGetRewardsPoolAddress(addressCodec coreaddress.Codec) sdk.AccAddress { + addr, err := addressCodec.StringToBytes(plan.RewardsPool) if err != nil { panic(err) }