[ValSet-Pref] Add WithdrawDelegationRewards message

app/keepers/keepers.go

@@ -363,6 +363,7 @@ func (appKeepers *AppKeepers) InitNormalKeepers(
 		appKeepers.keys[valsetpreftypes.StoreKey],
 		appKeepers.GetSubspace(valsetpreftypes.ModuleName),
 		appKeepers.StakingKeeper,
+		appKeepers.DistrKeeper,
 	)
 
 	appKeepers.ValidatorSetPreferenceKeeper = &validatorSetPreferenceKeeper

x/valset-pref/keeper.go

@@ -13,19 +13,22 @@ import (
 )
 
 type Keeper struct {
-	storeKey      sdk.StoreKey
-	paramSpace    paramtypes.Subspace
-	stakingKeeper types.StakingInterface
+	storeKey           sdk.StoreKey
+	paramSpace         paramtypes.Subspace
+	stakingKeeper      types.StakingInterface
+	distirbutionKeeper types.DistributionKeeper
 }
 
 func NewKeeper(storeKey sdk.StoreKey,
 	paramSpace paramtypes.Subspace,
 	stakingKeeper types.StakingInterface,
+	distirbutionKeeper types.DistributionKeeper,
 ) Keeper {
 	return Keeper{
-		storeKey:      storeKey,
-		paramSpace:    paramSpace,
-		stakingKeeper: stakingKeeper,
+		storeKey:           storeKey,
+		paramSpace:         paramSpace,
+		stakingKeeper:      stakingKeeper,
+		distirbutionKeeper: distirbutionKeeper,
 	}
 }
 

x/valset-pref/keeper_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 	"github.com/osmosis-labs/osmosis/v13/app/apptesting"
 	"github.com/osmosis-labs/osmosis/v13/x/valset-pref/types"
 	"github.com/stretchr/testify/suite"
@@ -41,6 +42,64 @@ func (suite *KeeperTestSuite) PrepareDelegateToValidatorSet() []types.ValidatorP
 	return valPreferences
 }
 
+func (suite *KeeperTestSuite) GetDelegationRewards(ctx sdk.Context, valAddrStr string, delegator sdk.AccAddress) (sdk.DecCoins, stakingtypes.Validator) {
+	valAddr, err := sdk.ValAddressFromBech32(valAddrStr)
+	suite.Require().NoError(err)
+
+	validator, found := suite.App.StakingKeeper.GetValidator(ctx, valAddr)
+	suite.Require().True(found)
+
+	endingPeriod := suite.App.DistrKeeper.IncrementValidatorPeriod(ctx, validator)
+
+	delegation, found := suite.App.StakingKeeper.GetDelegation(ctx, delegator, valAddr)
+	suite.Require().True(found)
+
+	rewards := suite.App.DistrKeeper.CalculateDelegationRewards(ctx, validator, delegation, endingPeriod)
+
+	return rewards, validator
+}
+
+func (suite *KeeperTestSuite) SetupExistingValidatorDelegations(ctx sdk.Context, valAddrStr string, delegator sdk.AccAddress, delegateAmt sdk.Int) {
+	valAddr, err := sdk.ValAddressFromBech32(valAddrStr)
+	suite.Require().NoError(err)
+
+	validator, found := suite.App.StakingKeeper.GetValidator(ctx, valAddr)
+	suite.Require().True(found)
+
+	_, err = suite.App.StakingKeeper.Delegate(ctx, delegator, delegateAmt, stakingtypes.Unbonded, validator, true)
+	suite.Require().NoError(err)
+
+}
+
+func (suite *KeeperTestSuite) SetupDelegationReward(ctx sdk.Context, delegator sdk.AccAddress, preferences []types.ValidatorPreference, existingValAddrStr string, setValSetDel, setExistingdel bool) {
+	// incrementing the blockheight by 1 for reward
+	ctx = suite.Ctx.WithBlockHeight(suite.Ctx.BlockHeight() + 1)
+
+	if setValSetDel {
+		// only necessary if there are tokens delegated
+		for _, val := range preferences {
+			suite.AllocateRewards(ctx, delegator, val.ValOperAddress)
+		}
+	}
+
+	if setExistingdel {
+		suite.AllocateRewards(ctx, delegator, existingValAddrStr)
+	}
+}
+
+func (suite *KeeperTestSuite) AllocateRewards(ctx sdk.Context, delegator sdk.AccAddress, valAddrStr string) {
+	// check that there is enough reward to withdraw
+	_, validator := suite.GetDelegationRewards(ctx, valAddrStr, delegator)
+
+	// allocate some rewards
+	tokens := sdk.NewDecCoins(sdk.NewInt64DecCoin(sdk.DefaultBondDenom, 10))
+	suite.App.DistrKeeper.AllocateTokensToValidator(ctx, validator, tokens)
+
+	rewardsAfterAllocation, _ := suite.GetDelegationRewards(ctx, valAddrStr, delegator)
+	suite.Require().NotNil(rewardsAfterAllocation)
+	suite.Require().NotZero(rewardsAfterAllocation[0].Amount)
+}
+
 func TestKeeperTestSuite(t *testing.T) {
 	suite.Run(t, new(KeeperTestSuite))
 }

x/valset-pref/msg_server.go

@@ -92,5 +92,12 @@ func (server msgServer) RedelegateValidatorSet(goCtx context.Context, msg *types
 }
 
 func (server msgServer) WithdrawDelegationRewards(goCtx context.Context, msg *types.MsgWithdrawDelegationRewards) (*types.MsgWithdrawDelegationRewardsResponse, error) {
+	ctx := sdk.UnwrapSDKContext(goCtx)
+
+	err := server.keeper.WithdrawDelegationRewards(ctx, msg.Delegator)
+	if err != nil {
+		return nil, err
+	}
+
 	return &types.MsgWithdrawDelegationRewardsResponse{}, nil
 }

x/valset-pref/msg_server_test.go

@@ -470,3 +470,122 @@ func (suite *KeeperTestSuite) TestRedelegateValidatorSet() {
 		})
 	}
 }
+
+func (suite *KeeperTestSuite) TestWithdrawDelegationRewards() {
+	tests := []struct {
+		name                 string
+		delegator            sdk.AccAddress
+		coinsToDelegate      sdk.Coin
+		setValSetDelegation  bool
+		setStakingDelegation bool
+		expectPass           bool
+	}{
+		{
+			name:                 "Withdraw all rewards with existing valset delegations, and existing staking position",
+			delegator:            sdk.AccAddress([]byte("addr1---------------")),
+			coinsToDelegate:      sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(20_000_000)), // delegate 20osmo
+			setValSetDelegation:  true,
+			setStakingDelegation: true,
+			expectPass:           true,
+		},
+		{
+			name:                 "Withdraw all rewards with no existing valset delegation, but existing staking position",
+			delegator:            sdk.AccAddress([]byte("addr2---------------")),
+			coinsToDelegate:      sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(20_000_000)), // delegate 20osmo
+			setValSetDelegation:  false,
+			setStakingDelegation: true,
+			expectPass:           true,
+		},
+		{
+			name:                 "Withdraw all rewards with existing valset delegation, but no existing staking position",
+			delegator:            sdk.AccAddress([]byte("addr3---------------")),
+			coinsToDelegate:      sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(20_000_000)), // delegate 20osmo
+			setValSetDelegation:  true,
+			setStakingDelegation: false,
+			expectPass:           true,
+		},
+		{
+			name:                 "Withdraw all rewards with no existing valset delegation, no existing staking position",
+			delegator:            sdk.AccAddress([]byte("addr4---------------")),
+			coinsToDelegate:      sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(20_000_000)), // delegate 20osmo
+			setValSetDelegation:  false,
+			setStakingDelegation: false,
+			expectPass:           false,
+		},
+	}
+
+	for _, test := range tests {
+		suite.Run(test.name, func() {
+			suite.SetupTest()
+
+			suite.FundAcc(test.delegator, sdk.Coins{sdk.NewInt64Coin(sdk.DefaultBondDenom, 100_000_000)}) // 100 osmo
+
+			// setup message server
+			msgServer := valPref.NewMsgServerImpl(suite.App.ValidatorSetPreferenceKeeper)
+			c := sdk.WrapSDKContext(suite.Ctx)
+
+			// call the create validator set preference
+			preferences := suite.PrepareDelegateToValidatorSet()
+			// setup extra validator for non val-set staking position
+			valAddrs := suite.SetupMultipleValidators(1)
+			ctx := suite.Ctx
+
+			// setup test for only valset delegation
+			if test.setValSetDelegation && !test.setStakingDelegation {
+				// delegators have to set val-set before delegating tokens
+				_, err := msgServer.SetValidatorSetPreference(c, types.NewMsgSetValidatorSetPreference(test.delegator, preferences))
+				suite.Require().NoError(err)
+
+				// call the delegate to validator set preference message
+				_, err = msgServer.DelegateToValidatorSet(c, types.NewMsgDelegateToValidatorSet(test.delegator, test.coinsToDelegate))
+				suite.Require().NoError(err)
+
+				suite.SetupDelegationReward(ctx, test.delegator, preferences, "", test.setValSetDelegation, test.setStakingDelegation)
+			}
+
+			// setup test for only existing staking position
+			if test.setStakingDelegation && !test.setValSetDelegation {
+				suite.SetupExistingValidatorDelegations(suite.Ctx, valAddrs[0], test.delegator, test.coinsToDelegate.Amount)
+
+				suite.SetupDelegationReward(ctx, test.delegator, nil, valAddrs[0], test.setValSetDelegation, test.setStakingDelegation)
+			}
+
+			// setup test for existing staking position and valset delegation position
+			if test.setStakingDelegation && test.setValSetDelegation {
+				suite.SetupExistingValidatorDelegations(suite.Ctx, valAddrs[0], test.delegator, test.coinsToDelegate.Amount)
+
+				// delegators have to set val-set before delegating tokens
+				_, err := msgServer.SetValidatorSetPreference(c, types.NewMsgSetValidatorSetPreference(test.delegator, preferences))
+				suite.Require().NoError(err)
+
+				// call the delegate to validator set preference message
+				_, err = msgServer.DelegateToValidatorSet(c, types.NewMsgDelegateToValidatorSet(test.delegator, test.coinsToDelegate))
+				suite.Require().NoError(err)
+
+				suite.SetupDelegationReward(ctx, test.delegator, preferences, valAddrs[0], test.setValSetDelegation, test.setStakingDelegation)
+			}
+
+			_, err := msgServer.WithdrawDelegationRewards(c, types.NewMsgWithdrawDelegationRewards(test.delegator))
+			if test.expectPass {
+				suite.Require().NoError(err)
+
+				// the rewards for valset and exising delegations should be nil
+				if test.setValSetDelegation {
+					for _, val := range preferences {
+						rewardAfterWithdrawValSet, _ := suite.GetDelegationRewards(ctx, val.ValOperAddress, test.delegator)
+						suite.Require().Nil(rewardAfterWithdrawValSet)
+					}
+				}
+
+				if test.setStakingDelegation {
+					rewardAfterWithdrawExistingSet, _ := suite.GetDelegationRewards(ctx, valAddrs[0], test.delegator)
+					suite.Require().Nil(rewardAfterWithdrawExistingSet)
+				}
+
+			} else {
+				suite.Require().Error(err)
+
+			}
+		})
+	}
+}

x/valset-pref/types/expected_interfaces.go

@@ -15,8 +15,17 @@ type StakingInterface interface {
 	GetDelegation(ctx sdk.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress) (delegation stakingtypes.Delegation, found bool)
 	Undelegate(ctx sdk.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress, sharesAmount sdk.Dec) (time.Time, error)
 	BeginRedelegation(ctx sdk.Context, delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress, sharesAmount sdk.Dec) (completionTime time.Time, err error)
+	GetDelegatorDelegations(ctx sdk.Context, delegator sdk.AccAddress, maxRetrieve uint16) (delegations []stakingtypes.Delegation)
 }
 
 type BankKeeper interface {
 	GetBalance(ctx sdk.Context, addr sdk.AccAddress, denom string) sdk.Coin
 }
+
+// For testing only
+type DistributionKeeper interface {
+	WithdrawDelegationRewards(ctx sdk.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress) (sdk.Coins, error)
+	IncrementValidatorPeriod(ctx sdk.Context, val stakingtypes.ValidatorI) uint64
+	CalculateDelegationRewards(ctx sdk.Context, val stakingtypes.ValidatorI, del stakingtypes.DelegationI, endingPeriod uint64) (rewards sdk.DecCoins)
+	AllocateTokensToValidator(ctx sdk.Context, val stakingtypes.ValidatorI, tokens sdk.DecCoins)
+}

x/valset-pref/validator_set.go

@@ -2,6 +2,7 @@ package keeper
 
 import (
 	"fmt"
+	"math"
 	"sort"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -207,6 +208,84 @@ func (k Keeper) PreformRedelegation(ctx sdk.Context, delegator sdk.AccAddress, e
 	return nil
 }
 
+// WithdrawDelegationRewards withdraws all the delegation rewards from the validator in the val-set.
+// Delegation reward is collected by the validator and in doing so, they can charge commission to the delegators.
+// Rewards are calculated per period, and is updated each time validator delegation changes. For ex: when a delegator
+// receives new delgation the rewards can be calculated by taking (total rewards before new delegation - the total current rewards).
+func (k Keeper) WithdrawDelegationRewards(ctx sdk.Context, delegatorAddr string) error {
+	delegator, err := sdk.AccAddressFromBech32(delegatorAddr)
+	if err != nil {
+		return err
+	}
+
+	// check if there is existing staking position that's not val-set
+	delegations := k.stakingKeeper.GetDelegatorDelegations(ctx, delegator, math.MaxUint16)
+
+	// get the existing validator set preference
+	existingSet, found := k.GetValidatorSetPreference(ctx, delegatorAddr)
+	if !found && len(delegations) == 0 {
+		return fmt.Errorf("user %s doesn't have validator set or existing delegations", delegatorAddr)
+	}
+
+	// there is existing staking position, but it's not valset
+	if !found && len(delegations) != 0 {
+		err := k.withdrawExistingStakingPosition(ctx, delegator, delegations)
+		if err != nil {
+			return err
+		}
+		return nil
+	}
+
+	// there is no existing staking position, but there is val-set delegation
+	if found && len(delegations) == 0 {
+		err := k.withdrawExistingValSetStakingPosition(ctx, delegator, existingSet.Preferences)
+		if err != nil {
+			return err
+		}
+		return nil
+	}
+
+	// there is staking position delegation, as well as val-set delegation
+	err = k.withdrawExistingStakingPosition(ctx, delegator, delegations)
+	if err != nil {
+		return err
+	}
+
+	err = k.withdrawExistingValSetStakingPosition(ctx, delegator, existingSet.Preferences)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// withdrawExistingStakingPosition takes the existing staking delegator delegations and withdraws the rewards.
+func (k Keeper) withdrawExistingStakingPosition(ctx sdk.Context, delegator sdk.AccAddress, delegations []stakingtypes.Delegation) error {
+	for _, dels := range delegations {
+		_, err := k.distirbutionKeeper.WithdrawDelegationRewards(ctx, delegator, dels.GetValidatorAddr())
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// withdrawExistingValSetStakingPosition takes the existing valset delegator delegations and withdraws the rewards.
+func (k Keeper) withdrawExistingValSetStakingPosition(ctx sdk.Context, delegator sdk.AccAddress, delegations []types.ValidatorPreference) error {
+	for _, dels := range delegations {
+		valAddr, err := sdk.ValAddressFromBech32(dels.ValOperAddress)
+		if err != nil {
+			return fmt.Errorf("validator address not formatted")
+		}
+
+		_, err = k.distirbutionKeeper.WithdrawDelegationRewards(ctx, delegator, valAddr)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // GetValAddrAndVal checks if the validator address is valid and the validator provided exists on chain.
 func (k Keeper) getValAddrAndVal(ctx sdk.Context, valOperAddress string) (sdk.ValAddress, stakingtypes.Validator, error) {
 	valAddr, err := sdk.ValAddressFromBech32(valOperAddress)