Fix: Fix valset pref unbonding bugs

CHANGELOG.md

@@ -88,7 +88,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * [#5890](https://github.com/osmosis-labs/osmosis/pull/5890) feat: CreateCLPool & LinkCFMMtoCL pool into one gov-prop
 * [#5959](https://github.com/osmosis-labs/osmosis/pull/5959) allow testing with different chain-id's in E2E testing
 * [#5964](https://github.com/osmosis-labs/osmosis/pull/5964) fix e2e test concurrency bugs
-* [#5948](https://github.com/osmosis-labs/osmosis/pull/5948) Parameterizing Pool Type Information in Protorev
+* [#5967](https://github.com/osmosis-labs/osmosis/pull/5967) fix ValSet undelegate API out of sync with existing staking
 * [#6001](https://github.com/osmosis-labs/osmosis/pull/6001) feat: improve set-env CLI cmd
 * [#6012](https://github.com/osmosis-labs/osmosis/pull/6012) chore: add autocomplete to makefile
 

x/valset-pref/README.md

@@ -25,15 +25,54 @@ How does this module work?
 - If the delegator has not set a validator-set preference list then the validator set, then it defaults to their current validator set.
 - If a user has no preference list and has not staked, then these messages / queries return errors.
 
-## Calculations
-
-Staking Calculation 
-
-- The user provides an amount to delegate and our `MsgDelegateToValidatorSet` divides the amount based on validator weight distribution.
-  For example: Stake 100osmo with validator-set {ValA -> 0.5, ValB -> 0.3, ValC -> 0.2}
-  our delegate logic will attempt to delegate (100 * 0.5) 50osmo for ValA , (100 * 0.3) 30osmo from ValB and (100 * 0.2) 20osmo from ValC.
-
-UnStaking Calculation 
+## Logic
+
+### Delegate to validator set
+
+This is pretty straight-forward, theres not really any edge cases here.
+
+The user provides an amount to delegate and our `MsgDelegateToValidatorSet` divides the amount based on validator weight distribution.
+
+For example: Stake 100osmo with validator-set {ValA -> 0.5, ValB -> 0.3, ValC -> 0.2}
+our delegate logic will attempt to delegate (100 * 0.5) 50osmo for ValA , (100 * 0.3) 30osmo from ValB and (100 * 0.2) 20osmo from ValC.
+
+### Undelegate from validator set 
+
+We can imagine describing undelegate from validator set in two cases:
+- Users existing delegation distribution matches their validator-set preference distribution.
+- Users existing delegation distribution does not match their validator-set preference distribution.
+
+Algorithm for undelegation; 
+unbond as true to valset ratios as possible. Undelegation should be possible.
+Idea of what we should be doing for Undelegate(valset, amount):
+1. Calculate the amount to undelegate from each validator under full valset usage
+2. If all amounts are less than amount staked to validator, undelegate all
+3. If any amount is greater than amount staked to validator (S,V),
+   fully unstake S from that validator. Recursively proceed with undelegating the
+   remaining amount from the remaining validators.
+   Undelegate(valset - V, amount - S)
+
+The above algorithm would take O(V^2) worst case, so we instead do something better
+to be O(V). 
+
+1. Calculate the amount to undelegate from each validator under full valset usage
+2. For each validator, compute V.ratio = undelegate_amount / amount_staked_to_val
+3. Sort validators by V_ratio descending If V_ratio <= 1, undelegate taret amount from each validator. (happy path)
+4. If V_ratio <= 1, undelegate target amount from each validator. (happy path)  
+5. Set target_ratio = 1, amount_remaining_to_unbond = amount
+6. While greatest V_ratio > target_ratio:
+   - Fully undelegate validator with greatest V_ratio. (Amount S)
+   - remove validator from list
+   - recalculate target_ratio = target_ratio * (1 - removed_V.target_percent) 
+     - this works, because if you recalculated target ratio scaled to 1
+		   every val's ratio would just differ by the constant factor of 1 / (1 - removed_V.target_percent)
+       doing that would take O(V), hence we change the target.
+7. Normal undelegate the remainder.
+
+
+#### Case 1: Users existing delegation distribution matches their validator-set preference distribution
+
+{Old docs below}
 
 - The user provides an amount to undelegate and our `MsgUnDelegateToValidatorSet` divides the amount based on validator weight distribution.
 - Here, the user can either undelegate the entire amount or partial amount 

x/valset-pref/msg_server.go

@@ -26,7 +26,7 @@ var _ types.MsgServer = msgServer{}
 func (server msgServer) SetValidatorSetPreference(goCtx context.Context, msg *types.MsgSetValidatorSetPreference) (*types.MsgSetValidatorSetPreferenceResponse, error) {
 	ctx := sdk.UnwrapSDKContext(goCtx)
 
-	preferences, err := server.keeper.SetValidatorSetPreference(ctx, msg.Delegator, msg.Preferences)
+	preferences, err := server.keeper.DeriveValidatorSetPreference(ctx, msg.Delegator, msg.Preferences)
 	if err != nil {
 		return nil, err
 	}
@@ -75,7 +75,7 @@ func (server msgServer) RedelegateValidatorSet(goCtx context.Context, msg *types
 	}
 
 	// Message 1: override the validator set preference set entry
-	newPreferences, err := server.keeper.SetValidatorSetPreference(ctx, msg.Delegator, msg.Preferences)
+	newPreferences, err := server.keeper.DeriveValidatorSetPreference(ctx, msg.Delegator, msg.Preferences)
 	if err != nil {
 		return nil, err
 	}

x/valset-pref/msg_server_test.go

@@ -367,7 +367,7 @@ func (s *KeeperTestSuite) TestUnDelegateFromValidatorSet() {
 			delegator:      sdk.AccAddress([]byte("addr6---------------")),
 			coinToStake:    sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(100_000_000)),
 			coinToUnStake:  sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(87_461_351)),
-			expectedShares: []sdk.Dec{sdk.NewDec(2_507_730), sdk.NewDec(4_137_755), sdk.NewDec(1_504_638), sdk.NewDec(4_388_526)}, // validatorDelegatedShares - (weight * coinToUnstake), for ex: 20_000_000 - (0.2 * 87_461_351)
+			expectedShares: []sdk.Dec{sdk.NewDec(2_507_730), sdk.NewDec(4_137_753), sdk.NewDec(1_504_638), sdk.NewDec(4_388_528)}, // validatorDelegatedShares - (weight * coinToUnstake), for ex: 20_000_000 - (0.2 * 87_461_351)
 			setValSet:      true,
 			expectPass:     true,
 		},
@@ -376,7 +376,7 @@ func (s *KeeperTestSuite) TestUnDelegateFromValidatorSet() {
 			delegator:      sdk.AccAddress([]byte("addr7---------------")),
 			coinToStake:    sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(10_000_000)),
 			coinToUnStake:  sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(1234)),
-			expectedShares: []sdk.Dec{sdk.NewDec(1_999_754), sdk.NewDec(3_299_593), sdk.NewDec(1_199_852), sdk.NewDec(3_499_567)}, // validatorDelegatedShares - (weight * coinToUnstake),
+			expectedShares: []sdk.Dec{sdk.NewDec(1_999_752), sdk.NewDec(3_299_593), sdk.NewDec(1_199_852), sdk.NewDec(3_499_569)}, // validatorDelegatedShares - (weight * coinToUnstake),
 			setValSet:      true,
 			expectPass:     true,
 		},

x/valset-pref/validator_set.go

@@ -21,6 +21,14 @@ type valSet struct {
 	Amount  sdk.Dec
 }
 
+type ValRatio struct {
+	ValAddr       sdk.ValAddress
+	Weight        sdk.Dec
+	DelegatedAmt  sdk.Int
+	UndelegateAmt sdk.Int
+	VRatio        sdk.Dec
+}
+
 // SetValidatorSetPreferences sets a new valset position for a delegator in modules state.
 func (k Keeper) SetValidatorSetPreferences(ctx sdk.Context, delegator string, validators types.ValidatorSetPreferences) {
 	store := ctx.KVStore(k.storeKey)
@@ -44,9 +52,11 @@ func (k Keeper) GetValidatorSetPreference(ctx sdk.Context, delegator string) (ty
 	return valsetPref, true
 }
 
-// SetValidatorSetPreference creates or updates delegators validator set.
+// DeriveValidatorSetPreference derives given validator set.
+// It validates the list and formats the inputs such as rounding.
 // Errors when the given preference is the same as the existing preference in state.
-func (k Keeper) SetValidatorSetPreference(ctx sdk.Context, delegator string, preferences []types.ValidatorPreference) (types.ValidatorSetPreferences, error) {
+// NOTE: this function does not add valset to the state
+func (k Keeper) DeriveValidatorSetPreference(ctx sdk.Context, delegator string, preferences []types.ValidatorPreference) (types.ValidatorSetPreferences, error) {
 	existingValSet, found := k.GetValidatorSetPreference(ctx, delegator)
 	if found {
 		// check if the new preferences is the same as the existing preferences
@@ -99,10 +109,9 @@ func (k Keeper) DelegateToValidatorSet(ctx sdk.Context, delegatorAddr string, co
 			tokenAmt = coin.Amount.Sub(totalDelAmt).ToDec().TruncateInt()
 		} else {
 			// tokenAmt takes the amount to delegate, calculated by {val_distribution_weight * tokenAmt}
-			tokenAmt = val.Weight.Mul(coin.Amount.ToDec()).TruncateInt()
+			tokenAmt = val.Weight.MulInt(coin.Amount).TruncateInt()
 			totalDelAmt = totalDelAmt.Add(tokenAmt)
 		}
-
 		// TODO: What happens here if validator unbonding
 		// Delegate the unbonded tokens
 		_, err = k.stakingKeeper.Delegate(ctx, delegator, tokenAmt, stakingtypes.Unbonded, validator, true)
@@ -119,72 +128,133 @@ func (k Keeper) DelegateToValidatorSet(ctx sdk.Context, delegatorAddr string, co
 // For ex: userA has staked 10tokens with weight {Val->0.5, ValB->0.3, ValC->0.2}
 // undelegate 6osmo with validator-set {ValA -> 0.5, ValB -> 0.3, ValC -> 0.2}
 // our undelegate logic would attempt to undelegate 3osmo from A, 1.8osmo from B, 1.2osmo from C
-// nolint: staticcheck
+// Rounding logic ensures we do not undelegate more than the user has staked with the validator set.
+// NOTE: check readme.md for more verbose description of the algorithm.
 func (k Keeper) UndelegateFromValidatorSet(ctx sdk.Context, delegatorAddr string, coin sdk.Coin) error {
-	// get the existingValSet if it exists, if not check existingStakingPosition and return it
 	existingSet, err := k.GetDelegationPreferences(ctx, delegatorAddr)
 	if err != nil {
 		return fmt.Errorf("user %s doesn't have validator set", delegatorAddr)
 	}
 
-	delegator, err := sdk.AccAddressFromBech32(delegatorAddr)
-	if err != nil {
-		return err
-	}
-
-	// the total amount the user wants to undelegate
-	tokenAmt := sdk.NewDec(coin.Amount.Int64())
-
-	err = k.CheckUndelegateTotalAmount(tokenAmt, existingSet.Preferences)
-	if err != nil {
-		return err
-	}
-
-	// totalDelAmt is the amount that keeps running track of the amount of tokens undelegated
-	totalUnDelAmt := sdk.NewInt(0)
-	amountToUnDelegate := sdk.NewInt(0)
+	delegator := sdk.MustAccAddressFromBech32(delegatorAddr)
+	// total amount the user wants to undelegate
+	tokenAmtToUndelegate := coin.Amount.ToDec()
+	// total amount user has delegated
+	totalDelegatedAmt := sdk.ZeroDec()
 
-	for i, val := range existingSet.Preferences {
+	// Step 1-2, compute the total amount delegated and the amount to undelegate for each validator
+	// under valset-ratios.
+	var valSetRatio []ValRatio
+	validators := map[string]stakingtypes.Validator{}
+	for _, val := range existingSet.Preferences {
+		amountToUnDelegate := val.Weight.Mul(tokenAmtToUndelegate).TruncateInt()
 		valAddr, validator, err := k.getValAddrAndVal(ctx, val.ValOperAddress)
 		if err != nil {
 			return err
 		}
+		validators[valAddr.String()] = validator
 
-		// in the last valset iteration we dont calculate it from shares using decimals and trucation,
-		// we use whats remaining to get more accurate value
-		if len(existingSet.Preferences)-1 == i {
-			amountToUnDelegate = coin.Amount.Sub(totalUnDelAmt).ToDec().TruncateInt()
-		} else {
-			// Calculate the amount to undelegate based on the existing weights
-			amountToUnDelegate = val.Weight.Mul(tokenAmt).TruncateInt()
-			totalUnDelAmt = totalUnDelAmt.Add(amountToUnDelegate)
+		delegation, found := k.stakingKeeper.GetDelegation(ctx, delegator, valAddr)
+		if !found {
+			return fmt.Errorf("No delegation found for thie delegator %s to this validator %s\n", delegator, valAddr)
 		}
 
 		sharesAmt, err := validator.SharesFromTokens(amountToUnDelegate)
 		if err != nil {
 			return err
 		}
 
-		_, err = k.stakingKeeper.Undelegate(ctx, delegator, valAddr, sharesAmt) // this has to be shares amount
+		valShares := sharesAmt.Quo(delegation.Shares)
+		tokenAmt := validator.TokensFromShares(delegation.Shares)
+		totalDelegatedAmt = totalDelegatedAmt.Add(tokenAmt)
+
+		valSetRatio = append(valSetRatio, ValRatio{
+			ValAddr:       valAddr,
+			UndelegateAmt: amountToUnDelegate,
+			VRatio:        valShares,
+			DelegatedAmt:  delegation.Shares.TruncateInt(),
+			Weight:        val.Weight,
+		})
+	}
+
+	if tokenAmtToUndelegate.GT(totalDelegatedAmt) {
+		return fmt.Errorf("Total tokenAmountToUndelegate more than delegated amount have %s got %s\n", totalDelegatedAmt, tokenAmtToUndelegate)
+	}
+
+	// Step 3
+	sort.Slice(valSetRatio, func(i, j int) bool {
+		return valSetRatio[i].VRatio.GT(valSetRatio[j].VRatio)
+	})
+
+	// Step 4
+	totalUnDelAmt, amountToUnDelegate := sdk.NewInt(0), sdk.NewInt(0)
+	if valSetRatio[0].VRatio.LTE(sdk.OneDec()) {
+		fmt.Println("All validators can undelegate normally, falling to happy path exit early")
+		for index, val := range valSetRatio {
+			validator := validators[val.ValAddr.String()]
+
+			// in the last valset iteration we dont calculate it from shares using decimals and trucation,
+			// we use whats remaining to get more accurate value
+			if len(existingSet.Preferences)-1 == index {
+				amountToUnDelegate = coin.Amount.Sub(totalUnDelAmt).ToDec().TruncateInt()
+			} else {
+				// Calculate the amount to undelegate based on the existing weights
+				amountToUnDelegate = val.UndelegateAmt
+				totalUnDelAmt = totalUnDelAmt.Add(amountToUnDelegate)
+			}
+
+			sharesAmt, err := validator.SharesFromTokens(amountToUnDelegate)
+			if err != nil {
+				return err
+			}
+
+			_, err = k.stakingKeeper.Undelegate(ctx, delegator, val.ValAddr, sharesAmt) // this has to be shares amount
+			if err != nil {
+				return err
+			}
+		}
+
+		return nil
+	}
+
+	// Step 5
+	// `targetRatio`: This is a threshold value that is used to decide how to unbond tokens from validators.
+	// It starts as 1 and is recalculated each time a validator is fully unbonded and removed from the unbonding process.
+	// By reducing the target ratio using the ratio of the removed validator, we adjust the proportions we are aiming for with the remaining validators.
+	targetRatio := sdk.OneDec()
+	amountRemaining := coin.Amount
+
+	// Step 6
+	for len(valSetRatio) > 0 && valSetRatio[0].VRatio.GT(targetRatio) {
+		fmt.Printf("Undelegating fully from validator %s %s\n", valSetRatio[0].ValAddr.String(), valSetRatio[0].DelegatedAmt)
+		_, err = k.stakingKeeper.Undelegate(ctx, delegator, valSetRatio[0].ValAddr, valSetRatio[0].DelegatedAmt.ToDec()) // this has to be shares amount
 		if err != nil {
 			return err
 		}
+		amountRemaining = amountRemaining.Sub(valSetRatio[0].DelegatedAmt)
+		targetRatio = targetRatio.Mul(sdk.OneDec().Sub(valSetRatio[0].Weight))
+		valSetRatio = valSetRatio[1:]
 	}
-	return nil
-}
 
-// CheckUndelegateTotalAmount checks if the tokenAmount equals the total amount calculated from valset weights.
-func (k Keeper) CheckUndelegateTotalAmount(tokenAmt sdk.Dec, existingSet []types.ValidatorPreference) error {
-	totalAmountFromWeights := sdk.NewDec(0)
-	for _, val := range existingSet {
-		totalAmountFromWeights = totalAmountFromWeights.Add(val.Weight.Mul(tokenAmt))
-	}
+	// Step 7
+	fmt.Println("Undelegating the remaining amount.")
+	for _, val := range valSetRatio {
+		// TestUndelegateFromValSetErrorCase1 specifically checks this logic
+		fmt.Printf("Undelegate %s from validator %s\n", val.UndelegateAmt, val.ValAddr.String())
+		_, validator, err := k.getValAddrAndVal(ctx, val.ValAddr.String())
+		if err != nil {
+			return err
+		}
 
-	totalAmountFromWeights = totalAmountFromWeights.RoundInt().ToDec()
-	tokenAmt = tokenAmt.RoundInt().ToDec()
+		sharesAmt, err := validator.SharesFromTokens(val.UndelegateAmt)
+		if err != nil {
+			return err
+		}
 
-	if !totalAmountFromWeights.Equal(tokenAmt) {
-		return fmt.Errorf("The undelegate total do not add up with the amount calculated from weights expected %s got %s", tokenAmt, totalAmountFromWeights)
+		_, err = k.stakingKeeper.Undelegate(ctx, delegator, val.ValAddr, sharesAmt) // this has to be shares amount
+		if err != nil {
+			return err
+		}
 	}
 
 	return nil