refactor(x/staking): migrate RedelegationQueue to collections

CHANGELOG.md

@@ -59,6 +59,9 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### API Breaking Changes
 
+* (x/staking) [#17486](https://github.com/cosmos/cosmos-sdk/pull/17486) Use collections for `RedelegationQueueKey`:
+    * remove from `types`: `GetRedelegationTimeKey`
+    * remove from `Keeper`: `RedelegationQueueIterator`
 * (x/staking) [#17481](https://github.com/cosmos/cosmos-sdk/pull/17481) Use collections for `UnbondingQueue`:
     * remove from `Keeper`: `UBDQueueIterator`
     * remove from `types`: `GetUnbondingDelegationTimeKey`

x/staking/keeper/delegation.go

@@ -610,33 +610,21 @@ func (k Keeper) RemoveRedelegation(ctx context.Context, red types.Redelegation)
 // timeslice is a slice of DVVTriplets corresponding to redelegations that
 // expire at a certain time.
 func (k Keeper) GetRedelegationQueueTimeSlice(ctx context.Context, timestamp time.Time) (dvvTriplets []types.DVVTriplet, err error) {
-	store := k.storeService.OpenKVStore(ctx)
-	bz, err := store.Get(types.GetRedelegationTimeKey(timestamp))
+	triplets, err := k.RedelegationQueue.Get(ctx, timestamp)
 	if err != nil {
-		return nil, err
-	}
-
-	if bz == nil {
+		if !errors.Is(err, collections.ErrNotFound) {
+			return nil, err
+		}
 		return []types.DVVTriplet{}, nil
 	}
 
-	triplets := types.DVVTriplets{}
-	err = k.cdc.Unmarshal(bz, &triplets)
-	if err != nil {
-		return nil, err
-	}
-
 	return triplets.Triplets, nil
 }
 
 // SetRedelegationQueueTimeSlice sets a specific redelegation queue timeslice.
 func (k Keeper) SetRedelegationQueueTimeSlice(ctx context.Context, timestamp time.Time, keys []types.DVVTriplet) error {
-	store := k.storeService.OpenKVStore(ctx)
-	bz, err := k.cdc.Marshal(&types.DVVTriplets{Triplets: keys})
-	if err != nil {
-		return err
-	}
-	return store.Set(types.GetRedelegationTimeKey(timestamp), bz)
+	triplets := types.DVVTriplets{Triplets: keys}
+	return k.RedelegationQueue.Set(ctx, timestamp, triplets)
 }
 
 // InsertRedelegationQueue insert an redelegation delegation to the appropriate
@@ -660,38 +648,32 @@ func (k Keeper) InsertRedelegationQueue(ctx context.Context, red types.Redelegat
 	return k.SetRedelegationQueueTimeSlice(ctx, completionTime, timeSlice)
 }
 
-// RedelegationQueueIterator returns all the redelegation queue timeslices from
-// time 0 until endTime.
-func (k Keeper) RedelegationQueueIterator(ctx context.Context, endTime time.Time) (storetypes.Iterator, error) {
-	store := k.storeService.OpenKVStore(ctx)
-	return store.Iterator(types.RedelegationQueueKey, storetypes.InclusiveEndBytes(types.GetRedelegationTimeKey(endTime)))
-}
-
 // DequeueAllMatureRedelegationQueue returns a concatenated list of all the
 // timeslices inclusively previous to currTime, and deletes the timeslices from
 // the queue.
 func (k Keeper) DequeueAllMatureRedelegationQueue(ctx context.Context, currTime time.Time) (matureRedelegations []types.DVVTriplet, err error) {
-	store := k.storeService.OpenKVStore(ctx)
-
 	// gets an iterator for all timeslices from time 0 until the current Blockheader time
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
-	redelegationTimesliceIterator, err := k.RedelegationQueueIterator(ctx, sdkCtx.HeaderInfo().Time)
+	redelegationTimesliceIterator, err := k.RedelegationQueue.Iterate(ctx, (&collections.Range[time.Time]{}).EndInclusive(sdkCtx.HeaderInfo().Time))
 	if err != nil {
 		return nil, err
 	}
 	defer redelegationTimesliceIterator.Close()
 
 	for ; redelegationTimesliceIterator.Valid(); redelegationTimesliceIterator.Next() {
-		timeslice := types.DVVTriplets{}
-		value := redelegationTimesliceIterator.Value()
-		if err = k.cdc.Unmarshal(value, &timeslice); err != nil {
-			return nil, err
+		timeslice, err := redelegationTimesliceIterator.Value()
+		if err != nil {
+			return matureRedelegations, err
 		}
 
 		matureRedelegations = append(matureRedelegations, timeslice.Triplets...)
 
-		if err = store.Delete(redelegationTimesliceIterator.Key()); err != nil {
-			return nil, err
+		key, err := redelegationTimesliceIterator.Key()
+		if err != nil {
+			return matureRedelegations, err
+		}
+		if err = k.RedelegationQueue.Remove(ctx, key); err != nil {
+			return matureRedelegations, err
 		}
 	}
 

x/staking/keeper/keeper.go

@@ -45,6 +45,7 @@ type Keeper struct {
 	Redelegations               collections.Map[collections.Triple[[]byte, []byte, []byte], types.Redelegation]
 	Delegations                 collections.Map[collections.Pair[sdk.AccAddress, sdk.ValAddress], types.Delegation]
 	UnbondingIndex              collections.Map[uint64, []byte]
+	RedelegationQueue           collections.Map[time.Time, types.DVVTriplets]
 	UnbondingQueue              collections.Map[time.Time, types.DVPairs]
 	Validators                  collections.Map[[]byte, types.Validator]
 	UnbondingDelegations        collections.Map[collections.Pair[[]byte, []byte], types.UnbondingDelegation]
@@ -150,7 +151,8 @@ func NewKeeper(
 			),
 			collections.BytesValue,
 		),
-		Validators: collections.NewMap(sb, types.ValidatorsKey, "validators", sdk.LengthPrefixedBytesKey, codec.CollValue[types.Validator](cdc)), // sdk.LengthPrefixedBytesKey is needed to retain state compatibility
+		RedelegationQueue: collections.NewMap(sb, types.RedelegationQueueKey, "redelegation_queue", sdk.TimeKey, codec.CollValue[types.DVVTriplets](cdc)),
+		Validators:        collections.NewMap(sb, types.ValidatorsKey, "validators", sdk.LengthPrefixedBytesKey, codec.CollValue[types.Validator](cdc)), // sdk.LengthPrefixedBytesKey is needed to retain state compatibility
 		UnbondingDelegations: collections.NewMap(
 			sb, types.UnbondingDelegationKey,
 			"unbonding_delegation",

x/staking/keeper/keeper_test.go

@@ -167,6 +167,14 @@ func getREDsFromValSrcIndexKey(valSrcAddr sdk.ValAddress) []byte {
 	return append(redelegationByValSrcIndexKey, addresstypes.MustLengthPrefix(valSrcAddr)...)
 }
 
+// getRedelegationTimeKey returns a key prefix for indexing an unbonding
+// redelegation based on a completion time.
+func getRedelegationTimeKey(timestamp time.Time) []byte {
+	bz := sdk.FormatTimeBytes(timestamp)
+	redelegationQueueKey := []byte{0x42}
+	return append(redelegationQueueKey, bz...)
+}
+
 // 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 {
@@ -396,6 +404,35 @@ func (s *KeeperTestSuite) TestValidatorsMigrationToColls() {
 	s.Require().NoError(err)
 }
 
+func (s *KeeperTestSuite) TestRedelegationQueueMigrationToColls() {
+	s.SetupTest()
+
+	err := testutil.DiffCollectionsMigration(
+		s.ctx,
+		s.key,
+		100,
+		func(i int64) {
+			date := time.Date(2023, 8, 21, 14, 33, 1, 0, &time.Location{})
+			s.ctx.KVStore(s.key).Set(getRedelegationTimeKey(date), []byte{})
+		},
+		"035e246f9d0bf0aa3dfeb88acf1665684168256d8afb742ae065872d6334f6d6",
+	)
+	s.Require().NoError(err)
+
+	err = testutil.DiffCollectionsMigration(
+		s.ctx,
+		s.key,
+		100,
+		func(i int64) {
+			date := time.Date(2023, 8, 21, 14, 33, 1, 0, &time.Location{})
+			err := s.stakingKeeper.SetRedelegationQueueTimeSlice(s.ctx, date, nil)
+			s.Require().NoError(err)
+		},
+		"035e246f9d0bf0aa3dfeb88acf1665684168256d8afb742ae065872d6334f6d6",
+	)
+	s.Require().NoError(err)
+}
+
 func TestKeeperTestSuite(t *testing.T) {
 	suite.Run(t, new(KeeperTestSuite))
 }

x/staking/migrations/v2/store_test.go

@@ -106,7 +106,7 @@ func TestStoreMigration(t *testing.T) {
 		{
 			"RedelegationQueueKey",
 			v1.GetRedelegationTimeKey(now),
-			types.GetRedelegationTimeKey(now),
+			getRedelegationTimeKey(now),
 		},
 		{
 			"ValidatorQueueKey",
@@ -141,6 +141,11 @@ func TestStoreMigration(t *testing.T) {
 	}
 }
 
+func getRedelegationTimeKey(timestamp time.Time) []byte {
+	bz := sdk.FormatTimeBytes(timestamp)
+	return append(types.RedelegationQueueKey, bz...)
+}
+
 func getUnbondingDelegationTimeKey(timestamp time.Time) []byte {
 	bz := sdk.FormatTimeBytes(timestamp)
 	return append(types.UnbondingQueueKey, bz...)

x/staking/types/keys.go

@@ -51,7 +51,7 @@ var (
 	UnbondingTypeKey  = collections.NewPrefix(57) // prefix for an index containing the type of unbonding operations
 
 	UnbondingQueueKey    = collections.NewPrefix(65) // prefix for the timestamps in unbonding queue
-	RedelegationQueueKey = []byte{0x42}              // prefix for the timestamps in redelegations queue
+	RedelegationQueueKey = collections.NewPrefix(66) // prefix for the timestamps in redelegations queue
 	ValidatorQueueKey    = []byte{0x43}              // prefix for the timestamps in validator queue
 
 	HistoricalInfoKey   = collections.NewPrefix(80) // prefix for the historical info
@@ -236,13 +236,6 @@ func GetREDKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress) []
 	return key
 }
 
-// GetRedelegationTimeKey returns a key prefix for indexing an unbonding
-// redelegation based on a completion time.
-func GetRedelegationTimeKey(timestamp time.Time) []byte {
-	bz := sdk.FormatTimeBytes(timestamp)
-	return append(RedelegationQueueKey, bz...)
-}
-
 // GetREDsKey returns a key prefix for indexing a redelegation from a delegator
 // address.
 func GetREDsKey(delAddr sdk.AccAddress) []byte {