feat: staking end-block queue optimization

CHANGELOG.md

@@ -76,6 +76,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (enterprise/poa) [#25838](https://github.com/cosmos/cosmos-sdk/pull/25838) Add the `poa` module under the `enterprise` directory.
 * (grpc) [#25850](https://github.com/cosmos/cosmos-sdk/pull/25850) Add `GetBlockResults` and `GetLatestBlockResults` gRPC endpoints to expose CometBFT block results including `finalize_block_events`.
 * (events) [#25877](https://github.com/cosmos/cosmos-sdk/pull/25877) Add `OverrideEvents` to `EventManagerI`.
+* (staking) [#26023](https://github.com/cosmos/cosmos-sdk/pull/26023) Optimize staking end-block queue through using pending queue slots instead of iterators. 
 
 ### Improvements
 

tests/integration/staking/keeper/common_test.go

@@ -3,6 +3,7 @@ package keeper_test
 import (
 	"math/big"
 	"testing"
+	"time"
 
 	cmtprototypes "github.com/cometbft/cometbft/proto/tendermint/types"
 	"gotest.tools/v3/assert"
@@ -151,7 +152,8 @@ func initFixture(tb testing.TB) *fixture {
 		types.ModuleName:     stakingModule,
 	})
 
-	sdkCtx := sdk.UnwrapSDKContext(integrationApp.Context())
+	initialTime := time.Date(1970, 1, 1, 0, 0, 0, 0, time.UTC)
+	sdkCtx := sdk.UnwrapSDKContext(integrationApp.Context()).WithBlockTime(initialTime)
 
 	// Register MsgServer and QueryServer
 	types.RegisterMsgServer(integrationApp.MsgServiceRouter(), stakingkeeper.NewMsgServerImpl(stakingKeeper))

tests/integration/staking/keeper/slash_test.go

@@ -267,7 +267,8 @@ func TestSlashWithUnbondingDelegation(t *testing.T) {
 	// set an unbonding delegation with expiration timestamp beyond which the
 	// unbonding delegation shouldn't be slashed
 	ubdTokens := f.stakingKeeper.TokensFromConsensusPower(f.sdkCtx, 4)
-	ubd := types.NewUnbondingDelegation(addrDels[0], addrVals[0], 11, time.Unix(0, 0), ubdTokens, 0, address.NewBech32Codec("cosmosvaloper"), address.NewBech32Codec("cosmos"))
+	completionTime := f.sdkCtx.BlockTime().Add(time.Second)
+	ubd := types.NewUnbondingDelegation(addrDels[0], addrVals[0], 11, completionTime, ubdTokens, 0, address.NewBech32Codec("cosmosvaloper"), address.NewBech32Codec("cosmos"))
 	assert.NilError(t, f.stakingKeeper.SetUnbondingDelegation(f.sdkCtx, ubd))
 
 	// slash validator for the first time
@@ -397,7 +398,8 @@ func TestSlashWithRedelegation(t *testing.T) {
 
 	// set a redelegation
 	rdTokens := f.stakingKeeper.TokensFromConsensusPower(f.sdkCtx, 6)
-	rd := types.NewRedelegation(addrDels[0], addrVals[0], addrVals[1], 11, time.Unix(0, 0), rdTokens, math.LegacyNewDecFromInt(rdTokens), 0, address.NewBech32Codec("cosmosvaloper"), address.NewBech32Codec("cosmos"))
+	rdCompletionTime := f.sdkCtx.BlockTime().Add(time.Second)
+	rd := types.NewRedelegation(addrDels[0], addrVals[0], addrVals[1], 11, rdCompletionTime, rdTokens, math.LegacyNewDecFromInt(rdTokens), 0, address.NewBech32Codec("cosmosvaloper"), address.NewBech32Codec("cosmos"))
 	assert.NilError(t, f.stakingKeeper.SetRedelegation(f.sdkCtx, rd))
 
 	// set the associated delegation
@@ -555,7 +557,8 @@ func TestSlashBoth(t *testing.T) {
 	// set a redelegation with expiration timestamp beyond which the
 	// redelegation shouldn't be slashed
 	rdATokens := f.stakingKeeper.TokensFromConsensusPower(f.sdkCtx, 6)
-	rdA := types.NewRedelegation(addrDels[0], addrVals[0], addrVals[1], 11, time.Unix(0, 0), rdATokens, math.LegacyNewDecFromInt(rdATokens), 0, address.NewBech32Codec("cosmosvaloper"), address.NewBech32Codec("cosmos"))
+	rdCompletionTime := f.sdkCtx.BlockTime().Add(time.Second)
+	rdA := types.NewRedelegation(addrDels[0], addrVals[0], addrVals[1], 11, rdCompletionTime, rdATokens, math.LegacyNewDecFromInt(rdATokens), 0, address.NewBech32Codec("cosmosvaloper"), address.NewBech32Codec("cosmos"))
 	assert.NilError(t, f.stakingKeeper.SetRedelegation(f.sdkCtx, rdA))
 
 	// set the associated delegation
@@ -565,8 +568,9 @@ func TestSlashBoth(t *testing.T) {
 	// set an unbonding delegation with expiration timestamp (beyond which the
 	// unbonding delegation shouldn't be slashed)
 	ubdATokens := f.stakingKeeper.TokensFromConsensusPower(f.sdkCtx, 4)
+	ubdCompletionTime := f.sdkCtx.BlockTime().Add(time.Second)
 	ubdA := types.NewUnbondingDelegation(addrDels[0], addrVals[0], 11,
-		time.Unix(0, 0), ubdATokens, 0, address.NewBech32Codec("cosmosvaloper"), address.NewBech32Codec("cosmos"))
+		ubdCompletionTime, ubdATokens, 0, address.NewBech32Codec("cosmosvaloper"), address.NewBech32Codec("cosmos"))
 	assert.NilError(t, f.stakingKeeper.SetUnbondingDelegation(f.sdkCtx, ubdA))
 
 	bondedCoins := sdk.NewCoins(sdk.NewCoin(bondDenom, rdATokens.MulRaw(2)))

x/slashing/keeper/slash_redelegation_test.go

@@ -37,8 +37,10 @@ func TestSlashRedelegation(t *testing.T) {
 	), &stakingKeeper, &bankKeeper, &slashKeeper, &distrKeeper)
 	require.NoError(t, err)
 
+	initialTime := time.Date(1970, 1, 1, 0, 0, 0, 0, time.UTC)
+	ctx := app.NewContext(false).WithBlockTime(initialTime)
+
 	// get sdk context, staking msg server and bond denom
-	ctx := app.NewContext(false)
 	stakingMsgServer := stakingkeeper.NewMsgServerImpl(stakingKeeper)
 	bondDenom, err := stakingKeeper.BondDenom(ctx)
 	require.NoError(t, err)

x/staking/keeper/delegation.go

@@ -7,7 +7,6 @@ import (
 	"fmt"
 	"time"
 
-	corestore "cosmossdk.io/core/store"
 	errorsmod "cosmossdk.io/errors"
 	"cosmossdk.io/math"
 	storetypes "cosmossdk.io/store/types"
@@ -490,7 +489,10 @@ func (k Keeper) SetUBDQueueTimeSlice(ctx context.Context, timestamp time.Time, k
 	if err != nil {
 		return err
 	}
-	return store.Set(types.GetUnbondingDelegationTimeKey(timestamp), bz)
+	if err = store.Set(types.GetUnbondingDelegationTimeKey(timestamp), bz); err != nil {
+		return err
+	}
+	return k.AddUBDQueuePendingSlot(ctx, timestamp)
 }
 
 // InsertUBDQueue inserts an unbonding delegation to the appropriate timeslice
@@ -514,41 +516,52 @@ func (k Keeper) InsertUBDQueue(ctx context.Context, ubd types.UnbondingDelegatio
 	return k.SetUBDQueueTimeSlice(ctx, completionTime, timeSlice)
 }
 
-// UBDQueueIterator returns all the unbonding queue timeslices from time 0 until endTime.
-func (k Keeper) UBDQueueIterator(ctx context.Context, endTime time.Time) (corestore.Iterator, error) {
-	store := k.storeService.OpenKVStore(ctx)
-	return store.Iterator(types.UnbondingQueueKey,
-		storetypes.InclusiveEndBytes(types.GetUnbondingDelegationTimeKey(endTime)))
-}
-
 // DequeueAllMatureUBDQueue returns a concatenated list of all the timeslices inclusively previous to
-// currTime, and deletes the timeslices from the queue.
+// currTime, and deletes the timeslices from the queue. Uses the pending-slot index (populated by
+// Migrate5to6); slots are read once and written once (batch update).
+// Read phase collects mature timeslices; write phase deletes keys and updates pending slots so that
+// on any error no queue keys are deleted and state remains consistent.
 func (k Keeper) DequeueAllMatureUBDQueue(ctx context.Context, currTime time.Time) (matureUnbonds []types.DVPair, err error) {
 	store := k.storeService.OpenKVStore(ctx)
 
-	// gets an iterator for all timeslices from time 0 until the current Blockheader time
-	unbondingTimesliceIterator, err := k.UBDQueueIterator(ctx, currTime)
+	slots, err := k.GetUBDQueuePendingSlots(ctx)
 	if err != nil {
-		return matureUnbonds, err
+		return nil, err
+	}
+	if len(slots) == 0 {
+		return matureUnbonds, nil
 	}
-	defer unbondingTimesliceIterator.Close()
 
-	for ; unbondingTimesliceIterator.Valid(); unbondingTimesliceIterator.Next() {
-		timeslice := types.DVPairs{}
-		value := unbondingTimesliceIterator.Value()
-		if err = k.cdc.Unmarshal(value, &timeslice); err != nil {
-			return matureUnbonds, err
+	var remaining []time.Time
+	var keysToDelete [][]byte
+	for _, t := range slots {
+		if t.After(currTime) {
+			remaining = append(remaining, t)
+			continue
+		}
+		queueKey := types.GetUnbondingDelegationTimeKey(t)
+		bz, err := store.Get(queueKey)
+		if err != nil {
+			return nil, err
+		}
+		if bz == nil {
+			continue // already deleted, omit from remaining
 		}
 
+		timeslice := types.DVPairs{}
+		if err = k.cdc.Unmarshal(bz, &timeslice); err != nil {
+			return nil, err
+		}
 		matureUnbonds = append(matureUnbonds, timeslice.Pairs...)
+		keysToDelete = append(keysToDelete, queueKey)
+	}
 
-		if err = store.Delete(unbondingTimesliceIterator.Key()); err != nil {
+	for _, key := range keysToDelete {
+		if err = store.Delete(key); err != nil {
 			return matureUnbonds, err
 		}
-
 	}
-
-	return matureUnbonds, nil
+	return matureUnbonds, k.SetUBDQueuePendingSlots(ctx, remaining)
 }
 
 // GetRedelegations returns a given amount of all the delegator redelegations.
@@ -802,7 +815,10 @@ func (k Keeper) SetRedelegationQueueTimeSlice(ctx context.Context, timestamp tim
 	if err != nil {
 		return err
 	}
-	return store.Set(types.GetRedelegationTimeKey(timestamp), bz)
+	if err = store.Set(types.GetRedelegationTimeKey(timestamp), bz); err != nil {
+		return err
+	}
+	return k.AddRedelegationQueuePendingSlot(ctx, timestamp)
 }
 
 // InsertRedelegationQueue insert an redelegation delegation to the appropriate
@@ -826,42 +842,53 @@ 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.
+// the queue. Uses the pending-slot index (populated by Migrate5to6); slots are
+// read once and written once (batch update).
+// Read phase collects mature timeslices; write phase deletes keys and updates pending slots so that
+// on any error no queue keys are deleted and state remains consistent.
 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)
+	slots, err := k.GetRedelegationQueuePendingSlots(ctx)
 	if err != nil {
 		return nil, err
 	}
-	defer redelegationTimesliceIterator.Close()
+	if len(slots) == 0 {
+		return matureRedelegations, nil
+	}
 
-	for ; redelegationTimesliceIterator.Valid(); redelegationTimesliceIterator.Next() {
-		timeslice := types.DVVTriplets{}
-		value := redelegationTimesliceIterator.Value()
-		if err = k.cdc.Unmarshal(value, &timeslice); err != nil {
+	var remaining []time.Time
+	var keysToDelete [][]byte
+	for _, t := range slots {
+		if t.After(currTime) {
+			remaining = append(remaining, t)
+			continue
+		}
+		queueKey := types.GetRedelegationTimeKey(t)
+		bz, err := store.Get(queueKey)
+		if err != nil {
 			return nil, err
 		}
+		if bz == nil {
+			continue
+		}
 
-		matureRedelegations = append(matureRedelegations, timeslice.Triplets...)
-
-		if err = store.Delete(redelegationTimesliceIterator.Key()); err != nil {
+		timeslice := types.DVVTriplets{}
+		if err = k.cdc.Unmarshal(bz, &timeslice); err != nil {
 			return nil, err
 		}
+		matureRedelegations = append(matureRedelegations, timeslice.Triplets...)
+		keysToDelete = append(keysToDelete, queueKey)
 	}
 
-	return matureRedelegations, nil
+	for _, key := range keysToDelete {
+		if err = store.Delete(key); err != nil {
+			return matureRedelegations, err
+		}
+	}
+	return matureRedelegations, k.SetRedelegationQueuePendingSlots(ctx, remaining)
 }
 
 // Delegate performs a delegation, set/update everything necessary within the store.

x/staking/keeper/migrations.go

@@ -8,6 +8,7 @@ import (
 	v3 "github.com/cosmos/cosmos-sdk/x/staking/migrations/v3"
 	v4 "github.com/cosmos/cosmos-sdk/x/staking/migrations/v4"
 	v5 "github.com/cosmos/cosmos-sdk/x/staking/migrations/v5"
+	v6 "github.com/cosmos/cosmos-sdk/x/staking/migrations/v6"
 )
 
 // Migrator is a struct for handling in-place store migrations.
@@ -47,3 +48,9 @@ func (m Migrator) Migrate4to5(ctx sdk.Context) error {
 	store := runtime.KVStoreAdapter(m.keeper.storeService.OpenKVStore(ctx))
 	return v5.MigrateStore(ctx, store, m.keeper.cdc)
 }
+
+// Migrate5to6 migrates x/staking state from consensus version 5 to 6.
+func (m Migrator) Migrate5to6(ctx sdk.Context) error {
+	store := m.keeper.storeService.OpenKVStore(ctx)
+	return v6.MigrateStore(ctx, store, m.keeper)
+}