feat(evmstaking): two pointer loop for validator delegations

client/app/upgrades.go

@@ -8,12 +8,15 @@ import (
  sdk "github.com/cosmos/cosmos-sdk/types"
 
  "github.com/piplabs/story/client/app/upgrades"
+ "github.com/piplabs/story/client/app/upgrades/v0_10_0"
 )
 
 var (
  // `Upgrades` defines the upgrade handlers and store loaders for the application.
  // New upgrades should be added to this slice after they are implemented.
- Upgrades = []upgrades.Upgrade{}
+ Upgrades = []upgrades.Upgrade{
+ v0_10_0.Upgrade,
+ }
  // Forks are for hard forks that breaks backward compatibility.
  Forks = []upgrades.Fork{}
 )

client/app/upgrades/v0_10_0/constants.go

@@ -0,0 +1,16 @@
+//nolint:revive,stylecheck // version underscores
+package v0_10_0
+
+import (
+ storetypes "cosmossdk.io/store/types"
+
+ "github.com/piplabs/story/client/app/upgrades"
+)
+
+const UpgradeName = "v0.10.0"
+
+var Upgrade = upgrades.Upgrade{
+ UpgradeName: UpgradeName,
+ CreateUpgradeHandler: CreateUpgradeHandler,
+ StoreUpgrades: storetypes.StoreUpgrades{},
+}

client/app/upgrades/v0_10_0/upgrades.go

@@ -0,0 +1,64 @@
+//nolint:revive,stylecheck // version underscores
+package v0_10_0
+
+import (
+ "context"
+
+ upgradetypes "cosmossdk.io/x/upgrade/types"
+
+ "github.com/cosmos/cosmos-sdk/types/module"
+
+ "github.com/piplabs/story/client/app/keepers"
+ "github.com/piplabs/story/lib/errors"
+ clog "github.com/piplabs/story/lib/log"
+)
+
+const (
+ NewMaxSweepPerBlock = 1024
+)
+
+func CreateUpgradeHandler(
+ mm *module.Manager,
+ configurator module.Configurator,
+ keepers *keepers.Keepers,
+) upgradetypes.UpgradeHandler {
+ return func(ctx context.Context, _ upgradetypes.Plan, vm module.VersionMap) (module.VersionMap, error) {
+ clog.Info(ctx, "Starting module migrations...")
+
+ vm, err := mm.RunMigrations(ctx, configurator, vm)
+ if err != nil {
+ return vm, errors.Wrap(err, "run migrations")
+ }
+
+ clog.Info(ctx, "Setting NextValidatorDelegationSweepIndex parameter...")
+ nextValIndex, err := keepers.EvmStakingKeeper.GetOldValidatorSweepIndex(ctx)
+ if err != nil {
+ return vm, errors.Wrap(err, "get old validator sweep index")
+ }
+
+ nextValDelIndex := uint64(0)
+ if err := keepers.EvmStakingKeeper.SetValidatorSweepIndex(
+ ctx,
+ nextValIndex,
+ nextValDelIndex,
+ ); err != nil {
+ return vm, errors.Wrap(err, "set evmstaking NextValidatorDelegationSweepIndex")
+ }
+
+ // Update MaxSweepPerBlock
+ clog.Info(ctx, "Updating MaxSweepPerBlock parameter...")
+ params, err := keepers.EvmStakingKeeper.GetParams(ctx)
+ if err != nil {
+ return vm, errors.Wrap(err, "get evmstaking params")
+ }
+
+ params.MaxSweepPerBlock = NewMaxSweepPerBlock
+ if err = keepers.EvmStakingKeeper.SetParams(ctx, params); err != nil {
+ return vm, errors.Wrap(err, "set evmstaking params")
+ }
+
+ clog.Info(ctx, "Upgrade v0.10.0 complete")
+
+ return vm, nil
+ }
+}

client/x/evmstaking/keeper/params.go

@@ -3,8 +3,6 @@ package keeper
 import (
  "context"
 
- "cosmossdk.io/math"
-
  sdk "github.com/cosmos/cosmos-sdk/types"
 
  "github.com/piplabs/story/client/x/evmstaking/types"
@@ -73,36 +71,60 @@ func (k Keeper) GetParams(ctx context.Context) (params types.Params, err error)
  return params, nil
 }
 
-func (k Keeper) SetNextValidatorSweepIndex(ctx context.Context, nextValIndex sdk.IntProto) error {
+func (k Keeper) SetValidatorSweepIndex(ctx context.Context, nextValIndex uint64, nextValDelIndex uint64) error {
  store := k.storeService.OpenKVStore(ctx)
- bz, err := k.cdc.Marshal(&nextValIndex)
+ bz, err := k.cdc.Marshal(&types.ValidatorSweepIndex{
+ NextValIndex: nextValIndex,
+ NextValDelIndex: nextValDelIndex,
+ })
  if err != nil {
- return errors.Wrap(err, "marshal next validator sweep index")
+ return errors.Wrap(err, "marshal validator sweep index")
  }
 
- err = store.Set(types.NextValidatorSweepIndexKey, bz)
+ err = store.Set(types.ValidatorSweepIndexKey, bz)
  if err != nil {
- return errors.Wrap(err, "set next validator sweep index")
+ return errors.Wrap(err, "set validator sweep index")
  }
 
  return nil
 }
 
-func (k Keeper) GetNextValidatorSweepIndex(ctx context.Context) (nextValIndex sdk.IntProto, err error) {
+func (k Keeper) GetValidatorSweepIndex(ctx context.Context) (nextValIndex uint64, nextValDelIndex uint64, err error) {
+ store := k.storeService.OpenKVStore(ctx)
+ bz, err := store.Get(types.ValidatorSweepIndexKey)
+ if err != nil {
+ return nextValIndex, nextValDelIndex, errors.Wrap(err, "get validator sweep index")
+ }
+
+ if bz == nil {
+ return nextValIndex, nextValDelIndex, nil
+ }
+
+ var sweepIndex types.ValidatorSweepIndex
+ err = k.cdc.Unmarshal(bz, &sweepIndex)
+ if err != nil {
+ return nextValIndex, nextValDelIndex, errors.Wrap(err, "unmarshal validator sweep index")
+ }
+
+ return sweepIndex.NextValIndex, sweepIndex.NextValDelIndex, nil
+}
+
+func (k Keeper) GetOldValidatorSweepIndex(ctx context.Context) (nextValIndex uint64, err error) {
  store := k.storeService.OpenKVStore(ctx)
- bz, err := store.Get(types.NextValidatorSweepIndexKey)
+ bz, err := store.Get(types.ValidatorSweepIndexKey)
  if err != nil {
  return nextValIndex, errors.Wrap(err, "get next validator sweep index")
  }
 
  if bz == nil {
- return sdk.IntProto{Int: math.NewInt(0)}, nil
+ return nextValIndex, nil
  }
 
- err = k.cdc.Unmarshal(bz, &nextValIndex)
+ var nextValIndexProto sdk.IntProto
+ err = k.cdc.Unmarshal(bz, &nextValIndexProto)
  if err != nil {
  return nextValIndex, errors.Wrap(err, "unmarshal next validator sweep index")
  }
 
- return nextValIndex, nil
+ return nextValIndexProto.Int.Uint64(), nil
 }

client/x/evmstaking/keeper/withdraw.go

@@ -3,8 +3,6 @@ package keeper
 import (
  "context"
 
- "cosmossdk.io/math"
-
  sdk "github.com/cosmos/cosmos-sdk/types"
  dtypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
  skeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
@@ -20,19 +18,18 @@ import (
 )
 
 func (k Keeper) ExpectedPartialWithdrawals(ctx context.Context) ([]estypes.Withdrawal, error) {
- // TODO: user more fine-grained cursor with next delegator sweep index.
- nextValSweepIndex, err := k.GetNextValidatorSweepIndex(ctx)
+ nextValIndex, nextValDelIndex, err := k.GetValidatorSweepIndex(ctx)
  if err != nil {
  return nil, err
  }
- nextValIndex := nextValSweepIndex.Int.Int64()
+
  // Get all validators first, and then do a circular sweep
  validatorSet, err := (k.stakingKeeper.(*skeeper.Keeper)).GetAllValidators(ctx)
  if err != nil {
  return nil, errors.Wrap(err, "get all validators")
  }
 
- if nextValIndex >= int64(len(validatorSet)) {
+ if nextValIndex >= uint64(len(validatorSet)) {
  // TODO: TBD
  log.Warn(
  ctx, "NextValidatorIndex exceeds the validator set size",
@@ -41,37 +38,34 @@ func (k Keeper) ExpectedPartialWithdrawals(ctx context.Context) ([]estypes.Withd
  "next_validator_index", nextValIndex,
  )
  nextValIndex = 0
+ nextValDelIndex = 0
  }
 
  // Iterate all validators from `nextValidatorIndex` to find out eligible partial withdrawals.
  var (
  swept uint32
  withdrawals []estypes.Withdrawal
  )
+
  // Get sweep limit per block.
  sweepBound, err := k.MaxSweepPerBlock(ctx)
  if err != nil {
  return nil, err
  }
+
  // Get minimal partial withdrawal amount.
  minPartialWithdrawalAmount, err := k.MinPartialWithdrawalAmount(ctx)
  if err != nil {
  return nil, err
  }
- log.Debug(
- ctx, "partial withdrawal params",
- "min_partial_withdraw_amount", minPartialWithdrawalAmount,
- "max_sweep_per_block", sweepBound,
- )
+
  // Sweep and get eligible partial withdrawals.
  for range validatorSet {
- if swept > sweepBound {
- break
- }
-
  if validatorSet[nextValIndex].IsJailed() {
  // nextValIndex should be updated, even if the validator is jailed, to progress to the sweep.
- nextValIndex = (nextValIndex + 1) % int64(len(validatorSet))
+ nextValIndex = (nextValIndex + 1) % uint64(len(validatorSet))
+ nextValDelIndex = 0
+
  continue
  }
 
@@ -82,86 +76,100 @@ func (k Keeper) ExpectedPartialWithdrawals(ctx context.Context) ([]estypes.Withd
  }
  valAddr := sdk.ValAddress(valBz)
  valAccAddr := sdk.AccAddress(valAddr)
+
  // Get validator commissions.
  valCommission, err := k.distributionKeeper.GetValidatorAccumulatedCommission(ctx, valAddr)
  if err != nil {
  return nil, err
  }
- log.Debug(
- ctx, "Get validator commission",
- "val_addr", valAddr.String(),
- "commission_amount", valCommission.Commission.String(),
- )
+
  // Get all delegators of the validator.
  delegators, err := (k.stakingKeeper.(*skeeper.Keeper)).GetValidatorDelegations(ctx, valAddr)
  if err != nil {
  return nil, errors.Wrap(err, "get validator delegations")
  }
- swept += uint32(len(delegators))
- log.Debug(
- ctx, "Get all delegators of validator",
- "val_addr", valAddr.String(),
- "delegator_amount", len(delegators),
- )
- // Get delegator rewards.
- for i := range delegators {
+
+ if nextValDelIndex >= uint64(len(delegators)) {
+ nextValIndex = (nextValIndex + 1) % uint64(len(validatorSet))
+ nextValDelIndex = 0
+
+ continue
+ }
+
+ nextDelegators := delegators[nextValDelIndex:]
+ var shouldStopPrematurely bool
+
+ // Check if the sweep should stop prematurely as the current delegator loop exceeds the sweep bound while sweeping.
+ remainingSweep := sweepBound - swept
+ if uint32(len(nextDelegators)) > remainingSweep {
+ nextDelegators = nextDelegators[:remainingSweep]
+ shouldStopPrematurely = true
+ }
+
+ // Iterate on the validator's delegator rewards in the range [nextValDelIndex, len(delegators)].
+ for i := range nextDelegators {
  // Get end current period and calculate rewards.
  endingPeriod, err := k.distributionKeeper.IncrementValidatorPeriod(ctx, validatorSet[nextValIndex])
  if err != nil {
  return nil, err
  }
- delRewards, err := k.distributionKeeper.CalculateDelegationRewards(ctx, validatorSet[nextValIndex], delegators[i], endingPeriod)
+
+ delRewards, err := k.distributionKeeper.CalculateDelegationRewards(ctx, validatorSet[nextValIndex], nextDelegators[i], endingPeriod)
  if err != nil {
  return nil, err
  }
- if delegators[i].DelegatorAddress == valAccAddr.String() {
+
+ if nextDelegators[i].DelegatorAddress == valAccAddr.String() {
  delRewards = delRewards.Add(valCommission.Commission...)
  }
  delRewardsTruncated, _ := delRewards.TruncateDecimal()
  bondDenomAmount := delRewardsTruncated.AmountOf(sdk.DefaultBondDenom).Uint64()
 
- log.Debug(
- ctx, "Calculate delegator rewards",
- "val_addr", valAddr.String(),
- "del_addr", delegators[i].DelegatorAddress,
- "rewards_amount", bondDenomAmount,
- "ending_period", endingPeriod,
- )
-
  if bondDenomAmount >= minPartialWithdrawalAmount {
- delEvmAddr, err := k.DelegatorMap.Get(ctx, delegators[i].DelegatorAddress)
+ delEvmAddr, err := k.DelegatorMap.Get(ctx, nextDelegators[i].DelegatorAddress)
  if err != nil {
  return nil, errors.Wrap(err, "map delegator pubkey to evm address")
  }
+
  withdrawals = append(withdrawals, estypes.NewWithdrawal(
  uint64(sdk.UnwrapSDKContext(ctx).BlockHeight()),
- delegators[i].DelegatorAddress,
+ nextDelegators[i].DelegatorAddress,
  valAddr.String(),
  delEvmAddr,
  bondDenomAmount,
  ))
-
- log.Debug(
- ctx, "Found an eligible partial withdrawal",
- "val_addr", valAddr.String(),
- "del_addr", delegators[i].DelegatorAddress,
- "del_evm_addr", delEvmAddr,
- "rewards_amount", bondDenomAmount,
- )
  }
+
+ nextValDelIndex++
+ }
+
+ // If the validator's delegation loop was stopped prematurely, we break from the validator sweep loop.
+ if shouldStopPrematurely {
+ break
  }
- nextValIndex = (nextValIndex + 1) % int64(len(validatorSet))
+
+ // Here, we have looped through all delegators of the validator (since we did not prematurely stop in the loop above).
+ // Thus, we signal to progress to the next validator by resetting the nextValDelIndex and circularly incrementing the nextValIndex
+ nextValIndex = (nextValIndex + 1) % uint64(len(validatorSet))
+ nextValDelIndex = 0
+
+ // Increase the total swept amount.
+ swept += uint32(len(nextDelegators))
  }
- // Update the nextValidatorSweepIndex.
- if err := k.SetNextValidatorSweepIndex(
+
+ // Update the validator sweep index.
+ if err := k.SetValidatorSweepIndex(
  ctx,
- sdk.IntProto{Int: math.NewInt(nextValIndex)},
+ nextValIndex,
+ nextValDelIndex,
  ); err != nil {
  return nil, err
  }
+
  log.Debug(
  ctx, "Finish validator sweep for partial withdrawals",
  "next_validator_index", nextValIndex,
+ "next_validator_delegation_index", nextValDelIndex,
  "partial_withdrawals", len(withdrawals),
  )