refactor(x/staking): Migrate ValidatorQueue to use Collections

CHANGELOG.md

@@ -59,6 +59,9 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### API Breaking Changes
 
+* (x/staking) [#17562](https://github.com/cosmos/cosmos-sdk/pull/17562) Use collections for `ValidatorQueue`
+    * remove from `types`: `GetValidatorQueueKey`, `ParseValidatorQueueKey`
+    * remove from `Keeper`: `ValidatorQueueIterator`
 * (x/staking) [#17498](https://github.com/cosmos/cosmos-sdk/pull/17498) Use collections for `LastValidatorPower`:
     * remove from `types`: `GetLastValidatorPowerKey`
     * remove from `Keeper`: `LastValidatorsIterator`, `IterateLastValidators`

x/staking/keeper/keeper.go

@@ -68,6 +68,8 @@ type Keeper struct {
 	RedelegationsByValSrc collections.Map[collections.Triple[[]byte, []byte, []byte], []byte]
 	// UnbondingDelegationByValIndex key: valAddr+delAddr | value: none used (index key for UnbondingDelegations stored by validator index)
 	UnbondingDelegationByValIndex collections.Map[collections.Pair[[]byte, []byte], []byte]
+	// ValidatorQueue key: len(timestamp bytes)+timestamp+height | value: ValAddresses
+	ValidatorQueue collections.Map[collections.Triple[uint64, time.Time, uint64], types.ValAddresses]
 	// LastValidatorPower key: valAddr | value: power(gogotypes.Int64Value())
 	LastValidatorPower collections.Map[[]byte, gogotypes.Int64Value]
 }
@@ -186,7 +188,20 @@ func NewKeeper(
 				collections.BytesKey,
 				sdk.LengthPrefixedBytesKey, // sdk.LengthPrefixedBytesKey is needed to retain state compatibility
 			),
-			codec.CollValue[types.UnbondingDelegation](cdc)),
+			codec.CollValue[types.UnbondingDelegation](cdc),
+		),
+		// key format is: 67 | length(timestamp Bytes) | timestamp | height
+		// Note: We use 3 keys here because we prefixed time bytes with its length previously and to retain state compatibility we remain to use the same
+		ValidatorQueue: collections.NewMap(
+			sb, types.ValidatorQueueKey,
+			"validator_queue",
+			collections.TripleKeyCodec(
+				collections.Uint64Key,
+				sdk.TimeKey,
+				collections.Uint64Key,
+			),
+			codec.CollValue[types.ValAddresses](cdc),
+		),
 	}
 
 	schema, err := sb.Build()

x/staking/keeper/keeper_test.go

@@ -202,6 +202,33 @@ func getLastValidatorPowerKey(operator sdk.ValAddress) []byte {
 	return append(lastValidatorPowerKey, addresstypes.MustLengthPrefix(operator)...)
 }
 
+// getValidatorQueueKey returns the prefix key used for getting a set of unbonding
+// validators whose unbonding completion occurs at the given time and height.
+func getValidatorQueueKey(timestamp time.Time, height int64) []byte {
+	validatorQueueKey := []byte{0x43}
+
+	heightBz := sdk.Uint64ToBigEndian(uint64(height))
+	timeBz := sdk.FormatTimeBytes(timestamp)
+	timeBzL := len(timeBz)
+	prefixL := len(validatorQueueKey)
+
+	bz := make([]byte, prefixL+8+timeBzL+8)
+
+	// copy the prefix
+	copy(bz[:prefixL], validatorQueueKey)
+
+	// copy the encoded time bytes length
+	copy(bz[prefixL:prefixL+8], sdk.Uint64ToBigEndian(uint64(timeBzL)))
+
+	// copy the encoded time bytes
+	copy(bz[prefixL+8:prefixL+8+timeBzL], timeBz)
+
+	// copy the encoded height
+	copy(bz[prefixL+8+timeBzL:], heightBz)
+
+	return bz
+}
+
 func (s *KeeperTestSuite) TestLastTotalPowerMigrationToColls() {
 	s.SetupTest()
 
@@ -445,6 +472,44 @@ func (s *KeeperTestSuite) TestValidatorsMigrationToColls() {
 	s.Require().NoError(err)
 }
 
+func (s *KeeperTestSuite) TestValidatorQueueMigrationToColls() {
+	s.SetupTest()
+	_, valAddrs := createValAddrs(100)
+	endTime := time.Unix(0, 0).UTC()
+	endHeight := int64(10)
+	err := testutil.DiffCollectionsMigration(
+		s.ctx,
+		s.key,
+		100,
+		func(i int64) {
+			var addrs []string
+			addrs = append(addrs, valAddrs[i].String())
+			bz, err := s.cdc.Marshal(&stakingtypes.ValAddresses{Addresses: addrs})
+			s.Require().NoError(err)
+
+			// legacy Set method
+			s.ctx.KVStore(s.key).Set(getValidatorQueueKey(endTime, endHeight), bz)
+		},
+		"8cf5fb4def683e83ea4cc4f14d8b2abc4c6af66709ad8af391dc749e63ef7524",
+	)
+	s.Require().NoError(err)
+
+	err = testutil.DiffCollectionsMigration(
+		s.ctx,
+		s.key,
+		100,
+		func(i int64) {
+			var addrs []string
+			addrs = append(addrs, valAddrs[i].String())
+
+			err := s.stakingKeeper.SetUnbondingValidatorsQueue(s.ctx, endTime, endHeight, addrs)
+			s.Require().NoError(err)
+		},
+		"8cf5fb4def683e83ea4cc4f14d8b2abc4c6af66709ad8af391dc749e63ef7524",
+	)
+	s.Require().NoError(err)
+}
+
 func TestKeeperTestSuite(t *testing.T) {
 	suite.Run(t, new(KeeperTestSuite))
 }

x/staking/keeper/validator.go

@@ -403,34 +403,25 @@ func (k Keeper) GetLastValidators(ctx context.Context) (validators []types.Valid
 // GetUnbondingValidators returns a slice of mature validator addresses that
 // complete their unbonding at a given time and height.
 func (k Keeper) GetUnbondingValidators(ctx context.Context, endTime time.Time, endHeight int64) ([]string, error) {
-	store := k.storeService.OpenKVStore(ctx)
-
-	bz, err := store.Get(types.GetValidatorQueueKey(endTime, endHeight))
+	timeSize := sdk.TimeKey.Size(endTime)
+	valAddrs, err := k.ValidatorQueue.Get(ctx, collections.Join3(uint64(timeSize), endTime, uint64(endHeight)))
 	if err != nil {
-		return nil, err
-	}
-
-	if bz == nil {
+		if !errors.Is(err, collections.ErrNotFound) {
+			return nil, err
+		}
 		return []string{}, nil
 	}
 
-	addrs := types.ValAddresses{}
-	if err = k.cdc.Unmarshal(bz, &addrs); err != nil {
-		return nil, err
-	}
-
-	return addrs.Addresses, nil
+	return valAddrs.Addresses, nil
 }
 
 // SetUnbondingValidatorsQueue sets a given slice of validator addresses into
 // the unbonding validator queue by a given height and time.
 func (k Keeper) SetUnbondingValidatorsQueue(ctx context.Context, endTime time.Time, endHeight int64, addrs []string) error {
-	store := k.storeService.OpenKVStore(ctx)
-	bz, err := k.cdc.Marshal(&types.ValAddresses{Addresses: addrs})
-	if err != nil {
-		return err
-	}
-	return store.Set(types.GetValidatorQueueKey(endTime, endHeight), bz)
+	valAddrs := types.ValAddresses{Addresses: addrs}
+	timeBz := sdk.FormatTimeBytes(endTime)
+	timeBzL := len(timeBz)
+	return k.ValidatorQueue.Set(ctx, collections.Join3(uint64(timeBzL), endTime, uint64(endHeight)), valAddrs)
 }
 
 // InsertUnbondingValidatorQueue inserts a given unbonding validator address into
@@ -447,8 +438,9 @@ func (k Keeper) InsertUnbondingValidatorQueue(ctx context.Context, val types.Val
 // DeleteValidatorQueueTimeSlice deletes all entries in the queue indexed by a
 // given height and time.
 func (k Keeper) DeleteValidatorQueueTimeSlice(ctx context.Context, endTime time.Time, endHeight int64) error {
-	store := k.storeService.OpenKVStore(ctx)
-	return store.Delete(types.GetValidatorQueueKey(endTime, endHeight))
+	timeBz := sdk.FormatTimeBytes(endTime)
+	timeBzL := len(timeBz)
+	return k.ValidatorQueue.Remove(ctx, collections.Join3(uint64(timeBzL), endTime, uint64(endHeight)))
 }
 
 // DeleteValidatorQueue removes a validator by address from the unbonding queue
@@ -485,92 +477,85 @@ func (k Keeper) DeleteValidatorQueue(ctx context.Context, val types.Validator) e
 	return k.SetUnbondingValidatorsQueue(ctx, val.UnbondingTime, val.UnbondingHeight, newAddrs)
 }
 
-// ValidatorQueueIterator returns an interator ranging over validators that are
-// unbonding whose unbonding completion occurs at the given height and time.
-func (k Keeper) ValidatorQueueIterator(ctx context.Context, endTime time.Time, endHeight int64) (corestore.Iterator, error) {
-	store := k.storeService.OpenKVStore(ctx)
-	return store.Iterator(types.ValidatorQueueKey, storetypes.InclusiveEndBytes(types.GetValidatorQueueKey(endTime, endHeight)))
-}
-
 // UnbondAllMatureValidators unbonds all the mature unbonding validators that
 // have finished their unbonding period.
 func (k Keeper) UnbondAllMatureValidators(ctx context.Context) error {
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
 	blockTime := sdkCtx.BlockTime()
-	blockHeight := sdkCtx.BlockHeight()
-
-	// unbondingValIterator will contains all validator addresses indexed under
-	// the ValidatorQueueKey prefix. Note, the entire index key is composed as
-	// ValidatorQueueKey | timeBzLen (8-byte big endian) | timeBz | heightBz (8-byte big endian),
-	// so it may be possible that certain validator addresses that are iterated
-	// over are not ready to unbond, so an explicit check is required.
-	unbondingValIterator, err := k.ValidatorQueueIterator(ctx, blockTime, blockHeight)
-	if err != nil {
-		return err
+	blockHeight := uint64(sdkCtx.BlockHeight())
+
+	rng := new(collections.Range[collections.Triple[uint64, time.Time, uint64]]).
+		EndInclusive(collections.Join3(uint64(29), blockTime, blockHeight))
+
+	return k.ValidatorQueue.Walk(ctx, rng, func(key collections.Triple[uint64, time.Time, uint64], value types.ValAddresses) (stop bool, err error) {
+		return false, k.unbondMatureValidators(ctx, blockHeight, blockTime, key, value)
+	})
+}
+
+func (k Keeper) unbondMatureValidators(
+	ctx context.Context,
+	blockHeight uint64,
+	blockTime time.Time,
+	key collections.Triple[uint64, time.Time, uint64],
+	addrs types.ValAddresses,
+) error {
+	keyTime, keyHeight := key.K2(), key.K3()
+
+	// All addresses for the given key have the same unbonding height and time.
+	// We only unbond if the height and time are less than the current height
+	// and time.
+	if keyHeight > blockHeight || keyTime.After(blockTime) {
+		return nil
 	}
-	defer unbondingValIterator.Close()
 
-	for ; unbondingValIterator.Valid(); unbondingValIterator.Next() {
-		key := unbondingValIterator.Key()
-		keyTime, keyHeight, err := types.ParseValidatorQueueKey(key)
+	// finalize unbonding
+	for _, valAddr := range addrs.Addresses {
+		addr, err := k.validatorAddressCodec.StringToBytes(valAddr)
 		if err != nil {
-			return fmt.Errorf("failed to parse unbonding key: %w", err)
+			return err
+		}
+		val, err := k.GetValidator(ctx, addr)
+		if err != nil {
+			return errorsmod.Wrap(err, "validator in the unbonding queue was not found")
+		}
+
+		if !val.IsUnbonding() {
+			return fmt.Errorf("unexpected validator in unbonding queue; status was not unbonding")
+		}
+
+		// if the ref count is not zero, early exit.
+		if val.UnbondingOnHoldRefCount != 0 {
+			return nil
 		}
 
-		// All addresses for the given key have the same unbonding height and time.
-		// We only unbond if the height and time are less than the current height
-		// and time.
-		if keyHeight <= blockHeight && (keyTime.Before(blockTime) || keyTime.Equal(blockTime)) {
-			addrs := types.ValAddresses{}
-			if err = k.cdc.Unmarshal(unbondingValIterator.Value(), &addrs); err != nil {
+		// otherwise do proper unbonding
+		for _, id := range val.UnbondingIds {
+			if err = k.DeleteUnbondingIndex(ctx, id); err != nil {
 				return err
 			}
+		}
+
+		val, err = k.UnbondingToUnbonded(ctx, val)
+		if err != nil {
+			return err
+		}
 
-			for _, valAddr := range addrs.Addresses {
-				addr, err := k.validatorAddressCodec.StringToBytes(valAddr)
-				if err != nil {
-					return err
-				}
-				val, err := k.GetValidator(ctx, addr)
-				if err != nil {
-					return errorsmod.Wrap(err, "validator in the unbonding queue was not found")
-				}
-
-				if !val.IsUnbonding() {
-					return fmt.Errorf("unexpected validator in unbonding queue; status was not unbonding")
-				}
-
-				if val.UnbondingOnHoldRefCount == 0 {
-					for _, id := range val.UnbondingIds {
-						if err = k.DeleteUnbondingIndex(ctx, id); err != nil {
-							return err
-						}
-					}
-
-					val, err = k.UnbondingToUnbonded(ctx, val)
-					if err != nil {
-						return err
-					}
-
-					if val.GetDelegatorShares().IsZero() {
-						str, err := k.validatorAddressCodec.StringToBytes(val.GetOperator())
-						if err != nil {
-							return err
-						}
-						if err = k.RemoveValidator(ctx, str); err != nil {
-							return err
-						}
-					} else {
-						// remove unbonding ids
-						val.UnbondingIds = []uint64{}
-					}
-
-					// remove validator from queue
-					if err = k.DeleteValidatorQueue(ctx, val); err != nil {
-						return err
-					}
-				}
+		if val.GetDelegatorShares().IsZero() {
+			str, err := k.validatorAddressCodec.StringToBytes(val.GetOperator())
+			if err != nil {
+				return err
 			}
+			if err = k.RemoveValidator(ctx, str); err != nil {
+				return err
+			}
+		} else {
+			// remove unbonding ids
+			val.UnbondingIds = []uint64{}
+		}
+
+		// remove validator from queue
+		if err = k.DeleteValidatorQueue(ctx, val); err != nil {
+			return err
 		}
 	}
 	return nil

x/staking/migrations/v2/store_test.go

@@ -111,7 +111,7 @@ func TestStoreMigration(t *testing.T) {
 		{
 			"ValidatorQueueKey",
 			v1.GetValidatorQueueKey(now, 4),
-			types.GetValidatorQueueKey(now, 4),
+			getValidatorQueueKey(now, 4),
 		},
 		{
 			"HistoricalInfoKey",
@@ -161,3 +161,26 @@ func getValidatorKey(operatorAddr sdk.ValAddress) []byte {
 func unbondingKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte {
 	return append(append(types.UnbondingDelegationKey, sdkaddress.MustLengthPrefix(delAddr)...), sdkaddress.MustLengthPrefix(valAddr)...)
 }
+
+func getValidatorQueueKey(timestamp time.Time, height int64) []byte {
+	heightBz := sdk.Uint64ToBigEndian(uint64(height))
+	timeBz := sdk.FormatTimeBytes(timestamp)
+	timeBzL := len(timeBz)
+	prefixL := len(types.ValidatorQueueKey)
+
+	bz := make([]byte, prefixL+8+timeBzL+8)
+
+	// copy the prefix
+	copy(bz[:prefixL], types.ValidatorQueueKey)
+
+	// copy the encoded time bytes length
+	copy(bz[prefixL:prefixL+8], sdk.Uint64ToBigEndian(uint64(timeBzL)))
+
+	// copy the encoded time bytes
+	copy(bz[prefixL+8:prefixL+8+timeBzL], timeBz)
+
+	// copy the encoded height
+	copy(bz[prefixL+8+timeBzL:], heightBz)
+
+	return bz
+}