feat: refund remaining rewards to the service's address when a plan ends

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
 	}

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) {

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
+}

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())
+}

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

x/rewards/types/expected_keepers.go

@@ -25,6 +25,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

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)
 	}