PLT-980: Populate validator-indexed delegation store at the v6.7 upgrade

app/upgrade_test.go

@@ -7,6 +7,7 @@ import (
 	"github.com/sei-protocol/sei-chain/app"
 	"github.com/sei-protocol/sei-chain/sei-cosmos/crypto/keys/secp256k1"
 	sdk "github.com/sei-protocol/sei-chain/sei-cosmos/types"
+	stakingtypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/staking/types"
 	"github.com/sei-protocol/sei-chain/sei-cosmos/x/upgrade/types"
 	storekeys "github.com/sei-protocol/sei-chain/sei-db/common/keys"
 	abci "github.com/sei-protocol/sei-chain/sei-tendermint/abci/types"
@@ -114,3 +115,34 @@ func TestSkipOptimisticProcessingOnUpgrade(t *testing.T) {
 		require.False(t, testWrapper.App.GetOptimisticProcessingInfo().Aborted)
 	})
 }
+
+func TestV67PopulatesDelegationByValIndex(t *testing.T) {
+	t.Setenv("UPGRADE_VERSION_LIST", "v6.7")
+	tm := time.Now().UTC()
+	valPub := secp256k1.GenPrivKey().PubKey()
+	testWrapper := app.NewTestWrapper(t, tm, valPub, false)
+	testWrapper.App.RegisterUpgradeHandlers()
+
+	ctx := testWrapper.Ctx
+	stakingKeeper := testWrapper.App.StakingKeeper
+	delAddr := sdk.AccAddress(secp256k1.GenPrivKey().PubKey().Address())
+	valAddr := sdk.ValAddress(secp256k1.GenPrivKey().PubKey().Address())
+	delegation := stakingtypes.NewDelegation(delAddr, valAddr, sdk.NewDec(1))
+
+	// Seed through the store rather than SetDelegation, so the index entry can only
+	// come from the upgrade handler.
+	store := ctx.KVStore(stakingKeeper.GetStoreKey())
+	store.Set(
+		stakingtypes.GetDelegationKey(delAddr, valAddr),
+		stakingtypes.MustMarshalDelegation(testWrapper.App.AppCodec(), delegation),
+	)
+	require.False(t, stakingKeeper.DelegationByValIndexReady(ctx))
+
+	testWrapper.App.UpgradeKeeper.ApplyUpgrade(ctx, types.Plan{
+		Name:   "v6.7",
+		Height: ctx.BlockHeight(),
+	})
+
+	require.True(t, stakingKeeper.DelegationByValIndexReady(ctx))
+	require.True(t, store.Has(stakingtypes.GetDelegationByValIndexKey(delAddr, valAddr)))
+}

app/upgrades.go

@@ -102,6 +102,9 @@ func (app *App) RegisterUpgradeHandlers() {
 				app.UpgradeKeeper.DeleteModuleVersion(ctx, capabilityModuleName)
 				app.UpgradeKeeper.DeleteModuleVersion(ctx, feegrantModuleName)
 				app.UpgradeKeeper.DeleteModuleVersion(ctx, transferModuleName)
+				if err := migrateDelegationByValIndex(ctx, app); err != nil {
+					return nil, err
+				}
 				return newVM, nil
 			}
 
@@ -111,3 +114,24 @@ func (app *App) RegisterUpgradeHandlers() {
 }
 
 const v606UpgradeHeight = 151573570
+
+// migrateDelegationByValIndex populates the validator-indexed delegation store and
+// marks it ready, so the staking precompile's validatorDelegations can answer from a
+// per-validator prefix instead of scanning every delegation.
+//
+// It runs on an infinite gas meter: the cost is a property of chain size at the
+// upgrade height, not of anything a transaction chose to spend.
+func migrateDelegationByValIndex(ctx sdk.Context, app *App) error {
+	result, err := app.StakingKeeper.MigrateDelegationByValIndex(ctx.WithGasMeter(sdk.NewInfiniteGasMeter(1, 1)))
+	if err != nil {
+		return err
+	}
+	logger.Info(
+		"populated delegation-by-validator index",
+		"total_delegations", result.TotalDelegations,
+		"index_written", result.IndexWritten,
+		"already_ready", result.AlreadyReady,
+		"elapsed", result.Elapsed.String(),
+	)
+	return nil
+}

precompiles/staking/staking.go

@@ -1002,7 +1002,7 @@ func (p PrecompileExecutor) validatorDelegations(ctx sdk.Context, method *abi.Me
 		},
 	}
 
-	response, err := p.stakingQuerier.ValidatorDelegations(sdk.WrapSDKContext(ctx), request)
+	response, err := p.stakingQuerier.ValidatorDelegationsIndexed(sdk.WrapSDKContext(ctx), request)
 	if err != nil {
 		return nil, 0, err
 	}

precompiles/staking/staking_test.go

@@ -297,6 +297,10 @@ func (tq *TestStakingQuerier) ValidatorDelegations(c context.Context, _ *staking
 	return tq.ValidatorDelegationsResponse, tq.Err
 }
 
+func (tq *TestStakingQuerier) ValidatorDelegationsIndexed(c context.Context, _ *stakingtypes.QueryValidatorDelegationsRequest) (*stakingtypes.QueryValidatorDelegationsResponse, error) {
+	return tq.ValidatorDelegationsResponse, tq.Err
+}
+
 func (tq *TestStakingQuerier) ValidatorUnbondingDelegations(c context.Context, _ *stakingtypes.QueryValidatorUnbondingDelegationsRequest) (*stakingtypes.QueryValidatorUnbondingDelegationsResponse, error) {
 	return tq.ValidatorUnbondingDelegationsResponse, tq.Err
 }

precompiles/utils/expected_keepers.go

@@ -183,6 +183,7 @@ type StakingQuerier interface {
 	Validators(c context.Context, req *stakingtypes.QueryValidatorsRequest) (*stakingtypes.QueryValidatorsResponse, error)
 	Validator(c context.Context, req *stakingtypes.QueryValidatorRequest) (*stakingtypes.QueryValidatorResponse, error)
 	ValidatorDelegations(c context.Context, req *stakingtypes.QueryValidatorDelegationsRequest) (*stakingtypes.QueryValidatorDelegationsResponse, error)
+	ValidatorDelegationsIndexed(c context.Context, req *stakingtypes.QueryValidatorDelegationsRequest) (*stakingtypes.QueryValidatorDelegationsResponse, error)
 	ValidatorUnbondingDelegations(c context.Context, req *stakingtypes.QueryValidatorUnbondingDelegationsRequest) (*stakingtypes.QueryValidatorUnbondingDelegationsResponse, error)
 	UnbondingDelegation(c context.Context, req *stakingtypes.QueryUnbondingDelegationRequest) (*stakingtypes.QueryUnbondingDelegationResponse, error)
 	DelegatorDelegations(c context.Context, req *stakingtypes.QueryDelegatorDelegationsRequest) (*stakingtypes.QueryDelegatorDelegationsResponse, error)

sei-cosmos/x/staking/keeper/delegation.go

@@ -98,19 +98,27 @@ func (k Keeper) GetDelegatorDelegations(ctx sdk.Context, delegator sdk.AccAddres
 // SetDelegation sets a delegation.
 func (k Keeper) SetDelegation(ctx sdk.Context, delegation types.Delegation) {
 	delegatorAddress := sdk.MustAccAddressFromBech32(delegation.DelegatorAddress)
+	valAddr := delegation.GetValidatorAddr()
 
 	store := ctx.KVStore(k.storeKey)
 	b := types.MustMarshalDelegation(k.cdc, delegation)
-	store.Set(types.GetDelegationKey(delegatorAddress, delegation.GetValidatorAddr()), b)
+	store.Set(types.GetDelegationKey(delegatorAddress, valAddr), b)
+	if k.DelegationByValIndexReady(ctx) {
+		store.Set(types.GetDelegationByValIndexKey(delegatorAddress, valAddr), []byte{}) // index, store empty bytes
+	}
 }
 
 // RemoveDelegation removes a delegation.
 func (k Keeper) RemoveDelegation(ctx sdk.Context, delegation types.Delegation) {
 	delegatorAddress := sdk.MustAccAddressFromBech32(delegation.DelegatorAddress)
+	valAddr := delegation.GetValidatorAddr()
 
-	k.BeforeDelegationRemoved(ctx, delegatorAddress, delegation.GetValidatorAddr())
+	k.BeforeDelegationRemoved(ctx, delegatorAddress, valAddr)
 	store := ctx.KVStore(k.storeKey)
-	store.Delete(types.GetDelegationKey(delegatorAddress, delegation.GetValidatorAddr()))
+	store.Delete(types.GetDelegationKey(delegatorAddress, valAddr))
+	if k.DelegationByValIndexReady(ctx) {
+		store.Delete(types.GetDelegationByValIndexKey(delegatorAddress, valAddr))
+	}
 }
 
 // GetUnbondingDelegations returns a given amount of all the delegator unbonding-delegations.

sei-cosmos/x/staking/keeper/delegation_index.go

@@ -0,0 +1,67 @@
+package keeper
+
+import (
+	"fmt"
+	"time"
+
+	sdk "github.com/sei-protocol/sei-chain/sei-cosmos/types"
+	"github.com/sei-protocol/sei-chain/sei-cosmos/x/staking/types"
+)
+
+// MigrateDelegationByValIndexResult reports the outcome of populating the
+// validator-indexed delegation store.
+type MigrateDelegationByValIndexResult struct {
+	TotalDelegations int
+	IndexWritten     int
+	AlreadyReady     bool
+	Elapsed          time.Duration
+}
+
+// DelegationByValIndexReady reports whether the validator-indexed delegation store
+// is populated at the version this context reads.
+//
+// The marker is versioned state written by MigrateDelegationByValIndex, so a context
+// reading a height before that migration observes it absent. That makes the answer
+// correct for historical queries and re-traced blocks without the caller supplying
+// an upgrade name or height.
+func (k Keeper) DelegationByValIndexReady(ctx sdk.Context) bool {
+	return ctx.KVStore(k.storeKey).Has(types.DelegationByValIndexReadyKey)
+}
+
+// MigrateDelegationByValIndex writes a validator-indexed key for every existing
+// delegation and then marks the index ready. It is a no-op once the marker is set.
+func (k Keeper) MigrateDelegationByValIndex(ctx sdk.Context) (MigrateDelegationByValIndexResult, error) {
+	start := time.Now()
+	store := ctx.KVStore(k.storeKey)
+
+	if store.Has(types.DelegationByValIndexReadyKey) {
+		return MigrateDelegationByValIndexResult{AlreadyReady: true, Elapsed: time.Since(start)}, nil
+	}
+
+	result := MigrateDelegationByValIndexResult{}
+	iterator := sdk.KVStorePrefixIterator(store, types.DelegationKey)
+	defer func() { _ = iterator.Close() }()
+
+	for ; iterator.Valid(); iterator.Next() {
+		delegation, err := types.UnmarshalDelegation(k.cdc, iterator.Value())
+		if err != nil {
+			return result, fmt.Errorf("unmarshal delegation at key %X: %w", iterator.Key(), err)
+		}
+		delAddr, err := sdk.AccAddressFromBech32(delegation.DelegatorAddress)
+		if err != nil {
+			return result, fmt.Errorf("parse delegator address %q: %w", delegation.DelegatorAddress, err)
+		}
+
+		result.TotalDelegations++
+		indexKey := types.GetDelegationByValIndexKey(delAddr, delegation.GetValidatorAddr())
+		if store.Has(indexKey) {
+			continue
+		}
+		store.Set(indexKey, []byte{})
+		result.IndexWritten++
+	}
+
+	store.Set(types.DelegationByValIndexReadyKey, []byte{})
+	result.Elapsed = time.Since(start)
+	return result, nil
+}