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

@@ -53,6 +53,7 @@ type Keeper struct {
 	RedelegationsByValDst         collections.Map[collections.Triple[[]byte, []byte, []byte], []byte]
 	RedelegationsByValSrc         collections.Map[collections.Triple[[]byte, []byte, []byte], []byte]
 	UnbondingDelegationByValIndex collections.Map[collections.Pair[[]byte, []byte], []byte]
+	ValidatorQueue                collections.Map[collections.Triple[uint64, time.Time, uint64], types.ValAddresses]
 	LastValidatorPower            collections.Map[[]byte, gogotypes.Int64Value]
 }
 
@@ -170,7 +171,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,26 @@ 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))
+	timeBz := sdk.FormatTimeBytes(endTime)
+	timeBzL := len(timeBz)
+	valAddrs, err := k.ValidatorQueue.Get(ctx, collections.Join3(uint64(timeBzL), 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 +439,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,13 +478,6 @@ 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 {
@@ -504,25 +490,28 @@ func (k Keeper) UnbondAllMatureValidators(ctx context.Context) error {
 	// 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)
+
+	timeBz := sdk.FormatTimeBytes(blockTime)
+	timeBzL := len(timeBz)
+	unbondingValIterator, err := k.ValidatorQueue.Iterate(ctx, (&collections.Range[collections.Triple[uint64, time.Time, uint64]]{}).EndInclusive(collections.Join3(uint64(timeBzL), blockTime, uint64(blockHeight))))
 	if err != nil {
 		return err
 	}
 	defer unbondingValIterator.Close()
 
 	for ; unbondingValIterator.Valid(); unbondingValIterator.Next() {
-		key := unbondingValIterator.Key()
-		keyTime, keyHeight, err := types.ParseValidatorQueueKey(key)
+		key, err := unbondingValIterator.Key()
 		if err != nil {
-			return fmt.Errorf("failed to parse unbonding key: %w", err)
+			return err
 		}
+		_, keyTime, keyHeight := key.K1(), 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.Before(blockTime) || keyTime.Equal(blockTime)) {
-			addrs := types.ValAddresses{}
-			if err = k.cdc.Unmarshal(unbondingValIterator.Value(), &addrs); err != nil {
+		if keyHeight <= uint64(blockHeight) && (keyTime.Before(blockTime) || keyTime.Equal(blockTime)) {
+			addrs, err := unbondingValIterator.Value()
+			if err != nil {
 				return err
 			}
 

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

x/staking/types/keys.go

@@ -1,9 +1,7 @@
 package types
 
 import (
-	"bytes"
 	"encoding/binary"
-	"fmt"
 	"time"
 
 	"cosmossdk.io/collections"
@@ -53,7 +51,7 @@ var (
 
 	UnbondingQueueKey    = collections.NewPrefix(65) // prefix for the timestamps in unbonding queue
 	RedelegationQueueKey = []byte{0x42}              // prefix for the timestamps in redelegations queue
-	ValidatorQueueKey    = []byte{0x43}              // prefix for the timestamps in validator queue
+	ValidatorQueueKey    = collections.NewPrefix(67) // prefix for the timestamps in validator queue
 
 	HistoricalInfoKey   = collections.NewPrefix(80) // prefix for the historical info
 	ValidatorUpdatesKey = collections.NewPrefix(97) // prefix for the end block validator updates key
@@ -142,50 +140,6 @@ func ParseValidatorPowerRankKey(key []byte) (operAddr []byte) {
 	return operAddr
 }
 
-// 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 {
-	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
-}
-
-// ParseValidatorQueueKey returns the encoded time and height from a key created
-// from GetValidatorQueueKey.
-func ParseValidatorQueueKey(bz []byte) (time.Time, int64, error) {
-	prefixL := len(ValidatorQueueKey)
-	if prefix := bz[:prefixL]; !bytes.Equal(prefix, ValidatorQueueKey) {
-		return time.Time{}, 0, fmt.Errorf("invalid prefix; expected: %X, got: %X", ValidatorQueueKey, prefix)
-	}
-
-	timeBzL := sdk.BigEndianToUint64(bz[prefixL : prefixL+8])
-	ts, err := sdk.ParseTimeBytes(bz[prefixL+8 : prefixL+8+int(timeBzL)])
-	if err != nil {
-		return time.Time{}, 0, err
-	}
-
-	height := sdk.BigEndianToUint64(bz[prefixL+8+int(timeBzL):])
-
-	return ts, int64(height), nil
-}
-
 // GetUBDKey creates the key for an unbonding delegation by delegator and validator addr
 // VALUE: staking/UnbondingDelegation
 func GetUBDKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte {

x/staking/types/keys_test.go

@@ -1,11 +1,9 @@
 package types_test
 
 import (
-	"bytes"
 	"encoding/hex"
 	"math/big"
 	"testing"
-	"time"
 
 	"github.com/stretchr/testify/require"
 
@@ -47,29 +45,3 @@ func TestGetValidatorPowerRank(t *testing.T) {
 		require.Equal(t, tt.wantHex, got, "Keys did not match on test case %d", i)
 	}
 }
-
-func TestGetValidatorQueueKey(t *testing.T) {
-	ts := time.Now()
-	height := int64(1024)
-
-	bz := types.GetValidatorQueueKey(ts, height)
-	rTs, rHeight, err := types.ParseValidatorQueueKey(bz)
-	require.NoError(t, err)
-	require.Equal(t, ts.UTC(), rTs.UTC())
-	require.Equal(t, rHeight, height)
-}
-
-func TestTestGetValidatorQueueKeyOrder(t *testing.T) {
-	ts := time.Now().UTC()
-	height := int64(1000)
-
-	endKey := types.GetValidatorQueueKey(ts, height)
-
-	keyA := types.GetValidatorQueueKey(ts.Add(-10*time.Minute), height-10)
-	keyB := types.GetValidatorQueueKey(ts.Add(-5*time.Minute), height+50)
-	keyC := types.GetValidatorQueueKey(ts.Add(10*time.Minute), height+100)
-
-	require.Equal(t, -1, bytes.Compare(keyA, endKey)) // keyA <= endKey
-	require.Equal(t, -1, bytes.Compare(keyB, endKey)) // keyB <= endKey
-	require.Equal(t, 1, bytes.Compare(keyC, endKey))  // keyB >= endKey
-}