PLT-980: Index delegations by validator at the v6.8 upgrade

app/testdata/upgrade_v68_offline_source_test.go

@@ -22,6 +22,7 @@ import (
 	vestingtypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/auth/vesting/types"
 	banktypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/bank/types"
 	govtypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/gov/types"
+	stakingtypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/staking/types"
 	upgradetypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/upgrade/types"
 	abci "github.com/sei-protocol/sei-chain/sei-tendermint/abci/types"
 	tmproto "github.com/sei-protocol/sei-chain/sei-tendermint/proto/tendermint/types"
@@ -60,6 +61,7 @@ func TestV68OfflineUpgradeSource(t *testing.T) {
 	seedV68IBCProposal(t, testApp, ctx, &retained)
 	seedV68UpgradedIBCState(t, testApp, ctx, &retained)
 	seedV68Voucher(t, testApp, ctx, &retained)
+	seedV68OfflineDelegations(testApp, ctx)
 	requireV68OfflineBalancesLocked(t, testApp, ctx, retained)
 	upgradeHeight := ctx.BlockHeight() + 2
 	require.NoError(t, testApp.UpgradeKeeper.ScheduleUpgrade(ctx, upgradetypes.Plan{
@@ -204,6 +206,20 @@ func seedV68Voucher(t *testing.T, testApp *App, ctx sdk.Context, retained *offli
 	retained.VoucherSupply = testApp.BankKeeper.GetSupply(ctx, voucher.Denom).Amount.String()
 }
 
+// seedV68OfflineDelegations stores delegations to two validators, which v6.7
+// writes without any delegation-by-validator index entry.
+func seedV68OfflineDelegations(testApp *App, ctx sdk.Context) {
+	validators := []sdk.ValAddress{
+		sdk.ValAddress(secp256k1.GenPrivKey().PubKey().Address()),
+		sdk.ValAddress(secp256k1.GenPrivKey().PubKey().Address()),
+	}
+	for i := 0; i < 3; i++ {
+		delegator := sdk.AccAddress(secp256k1.GenPrivKey().PubKey().Address())
+		testApp.StakingKeeper.SetDelegation(ctx,
+			stakingtypes.NewDelegation(delegator, validators[i%len(validators)], sdk.NewDec(int64(i+1))))
+	}
+}
+
 // seedV68OfflineVestingAccounts writes one account of every vesting type
 // through the v6.7 keepers, each keyed and holding a balance its schedule
 // locks. The delayed vesting account is recorded as the sender of a

app/upgrade_v68_offline_target_test.go

@@ -23,6 +23,7 @@ import (
 	distrtypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/distribution/types"
 	govtypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/gov/types"
 	paramstypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/params/types"
+	stakingtypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/staking/types"
 	upgradetypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/upgrade/types"
 	"github.com/sei-protocol/sei-chain/sei-db/state_db/sc/memiavl"
 	abci "github.com/sei-protocol/sei-chain/sei-tendermint/abci/types"
@@ -137,7 +138,8 @@ func applyV68OfflineUpgradeCrashReplay(t *testing.T, root string, artifact offli
 // accounts are rewritten, the retired IBC proposal reads back as a text
 // proposal, the upgraded IBC client record, the denom allowlists and the
 // tokenfactory parameters are pruned, the IBC voucher balance and supply are
-// intact and spendable, and nothing else in the retained stores changed.
+// intact and spendable, every delegation is indexed by validator, and nothing
+// else in the retained stores changed.
 func requireV68OfflineMigrated(t *testing.T, testApp *App, artifact offlineUpgradeArtifact) {
 	t.Helper()
 	requireV68OfflineAppliedName(t, testApp, artifact)
@@ -149,6 +151,7 @@ func requireV68OfflineMigrated(t *testing.T, testApp *App, artifact offlineUpgra
 	requireV68OfflineDenomAllowListsPruned(t, testApp, artifact.Stores[banktypes.StoreKey])
 	requireV68OfflineTokenfactoryParamsPruned(t, testApp, artifact.Stores[paramstypes.StoreKey])
 	requireV68OfflineVoucher(t, testApp, artifact.Retained)
+	requireV68OfflineDelegationIndex(t, testApp)
 	requireOfflineUpgradeRetainedStoresExcept(t, testApp, artifact.Stores, v68OfflineTouchedKey(t, artifact))
 }
 
@@ -281,6 +284,31 @@ func requireV68OfflineStoresDeleted(t *testing.T, testApp *App) {
 	}
 }
 
+// requireV68OfflineDelegationIndex requires the delegation-by-validator index
+// to be ready and to hold exactly one entry per stored delegation.
+func requireV68OfflineDelegationIndex(t *testing.T, testApp *App) {
+	t.Helper()
+	ctx := offlineUpgradeReadContext(testApp, testApp.LastBlockHeight())
+	require.True(t, testApp.StakingKeeper.DelegationByValIndexReady(ctx))
+	store := ctx.KVStore(testApp.GetKey(stakingtypes.StoreKey))
+
+	delegations := map[string]struct{}{}
+	iterator := sdk.KVStorePrefixIterator(store, stakingtypes.DelegationKey)
+	for ; iterator.Valid(); iterator.Next() {
+		delegations[string(iterator.Key())] = struct{}{}
+	}
+	require.NoError(t, iterator.Close())
+	require.NotEmpty(t, delegations, "the source phase stored no delegations")
+
+	indexed := map[string]struct{}{}
+	iterator = sdk.KVStorePrefixIterator(store, stakingtypes.DelegationByValIndexKey)
+	for ; iterator.Valid(); iterator.Next() {
+		indexed[string(stakingtypes.GetDelegationKeyFromValIndexKey(iterator.Key()))] = struct{}{}
+	}
+	require.NoError(t, iterator.Close())
+	require.Equal(t, delegations, indexed, "the index does not hold exactly one entry per delegation")
+}
+
 func requireV68OfflineProposalRewritten(t *testing.T, testApp *App, retained offlineUpgradeRetainedState) {
 	t.Helper()
 	ctx := offlineUpgradeReadContext(testApp, testApp.LastBlockHeight())

app/upgrade_v68_test.go

@@ -23,6 +23,9 @@ import (
 	banktypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/bank/types"
 	govtypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/gov/types"
 	paramstypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/params/types"
+	"github.com/sei-protocol/sei-chain/sei-cosmos/x/staking"
+	stakingkeeper "github.com/sei-protocol/sei-chain/sei-cosmos/x/staking/keeper"
+	stakingtypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/staking/types"
 	upgradetypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/upgrade/types"
 	abci "github.com/sei-protocol/sei-chain/sei-tendermint/abci/types"
 	"github.com/sei-protocol/sei-chain/testutil/processblock"
@@ -41,7 +44,9 @@ import (
 // base accounts they embed. These tests cover the rewrite of every vesting
 // account type, spending a balance a schedule had locked, the messages the
 // removed modules served, the allowlist and parameter pruning, genesis export
-// over the migrated store, and the handler itself.
+// over the migrated store, and the handler itself. The handler also starts the
+// staking delegation-by-validator index backfill, which the staking EndBlocker
+// completes.
 
 // v68RemovedStoreModules are the module version map entries the v6.8 handler
 // deletes along with their stores.
@@ -331,6 +336,87 @@ func TestV68ApplyUpgradeTwice(t *testing.T) {
 		"second ApplyUpgrade is not a no-op: ApplyUpgrade increments protocol version on every call")
 }
 
+// unreadyV68DelegationIndex deletes the delegation-by-validator index and its
+// markers, leaving the staking store as v6.7 left it.
+func unreadyV68DelegationIndex(t *testing.T, a *processblock.App) {
+	t.Helper()
+	store := a.Ctx().KVStore(a.GetKey(stakingtypes.StoreKey))
+	iterator := sdk.KVStorePrefixIterator(store, stakingtypes.DelegationByValIndexKey)
+	var keys [][]byte
+	for ; iterator.Valid(); iterator.Next() {
+		keys = append(keys, iterator.Key())
+	}
+	require.NoError(t, iterator.Close())
+	for _, key := range keys {
+		store.Delete(key)
+	}
+	store.Delete(stakingtypes.DelegationByValIndexReadyKey)
+	store.Delete(stakingtypes.DelegationByValIndexCursorKey)
+}
+
+// v68DelegationKeys returns the keys of every stored delegation.
+func v68DelegationKeys(t *testing.T, store sdk.KVStore) map[string]struct{} {
+	t.Helper()
+	return v68PrefixKeys(t, store, stakingtypes.DelegationKey, func(key []byte) []byte { return key })
+}
+
+// v68IndexedDelegationKeys returns the delegation key every index entry points at.
+func v68IndexedDelegationKeys(t *testing.T, store sdk.KVStore) map[string]struct{} {
+	t.Helper()
+	return v68PrefixKeys(t, store, stakingtypes.DelegationByValIndexKey, stakingtypes.GetDelegationKeyFromValIndexKey)
+}
+
+func v68PrefixKeys(t *testing.T, store sdk.KVStore, prefix []byte, resolve func([]byte) []byte) map[string]struct{} {
+	t.Helper()
+	iterator := sdk.KVStorePrefixIterator(store, prefix)
+	defer func() { require.NoError(t, iterator.Close()) }()
+	keys := map[string]struct{}{}
+	for ; iterator.Valid(); iterator.Next() {
+		keys[string(resolve(iterator.Key()))] = struct{}{}
+	}
+	return keys
+}
+
+func TestV68PopulatesDelegationByValIndex(t *testing.T) {
+	a := newV68Chain(t)
+	unreadyV68DelegationIndex(t, a)
+	ctx := a.Ctx()
+	store := ctx.KVStore(a.GetKey(stakingtypes.StoreKey))
+
+	validators := a.StakingKeeper.GetAllValidators(ctx)
+	require.NotEmpty(t, validators)
+	valAddr := validators[0].GetOperator()
+	delAddr := a.NewAccount()
+	a.StakingKeeper.SetDelegation(ctx, stakingtypes.NewDelegation(delAddr, valAddr, sdk.NewDec(68)))
+	require.False(t, store.Has(stakingtypes.GetDelegationByValIndexKey(delAddr, valAddr)),
+		"a delegation written before v6.8 is indexed")
+
+	applyV68(t, a)
+
+	require.False(t, a.StakingKeeper.DelegationByValIndexReady(ctx),
+		"v6.8 populated the index inside the upgrade handler")
+	require.True(t, store.Has(stakingtypes.DelegationByValIndexCursorKey),
+		"v6.8 did not start the index backfill")
+	staking.EndBlocker(ctx, a.StakingKeeper)
+
+	require.True(t, a.StakingKeeper.DelegationByValIndexReady(ctx))
+	require.False(t, store.Has(stakingtypes.DelegationByValIndexCursorKey))
+	delegations := v68DelegationKeys(t, store)
+	require.Contains(t, delegations, string(stakingtypes.GetDelegationKey(delAddr, valAddr)))
+	require.Equal(t, delegations, v68IndexedDelegationKeys(t, store),
+		"the index does not hold exactly one entry per delegation")
+
+	querier := stakingkeeper.Querier{Keeper: a.StakingKeeper}
+	res, err := querier.ValidatorDelegations(sdk.WrapSDKContext(ctx),
+		&stakingtypes.QueryValidatorDelegationsRequest{ValidatorAddr: valAddr.String()})
+	require.NoError(t, err)
+	var delegators []string
+	for _, response := range res.DelegationResponses {
+		delegators = append(delegators, response.Delegation.DelegatorAddress)
+	}
+	require.Contains(t, delegators, delAddr.String())
+}
+
 // TestV68RejectsOracleTxsWithoutCharging pins that retired oracle transactions
 // are refused before fees are charged.
 func TestV68RejectsOracleTxsWithoutCharging(t *testing.T) {

app/upgrades.go

@@ -169,6 +169,7 @@ func (app *App) RegisterUpgradeHandlers() {
 				app.pruneUpgradedIBCState(ctx)
 				app.pruneDenomAllowLists(ctx)
 				app.pruneTokenfactoryParams(ctx)
+				app.StakingKeeper.StartDelegationByValIndexBackfill(ctx)
 				return newVM, nil
 			}
 

evmrpc/tests/regression_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/sei-protocol/sei-chain/evmrpc"
 	sdk "github.com/sei-protocol/sei-chain/sei-cosmos/types"
 	govtypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/gov/types"
+	stakingtypes "github.com/sei-protocol/sei-chain/sei-cosmos/x/staking/types"
 	tmproto "github.com/sei-protocol/sei-chain/sei-tendermint/proto/tendermint/types"
 	"github.com/stretchr/testify/require"
 )
@@ -387,7 +388,7 @@ func testTx(t *testing.T, txHash string, version string, expectedGasUsed string,
 		blockHeight := mockStatesFromTxJson(ctx, txHash, a, mc)
 		ctx = setLegacySstoreIfNeeded(ctx, a, version)
 		ctx = withCapturedConsensusParams(ctx, mc, blockHeight)
-		removeFutureGovernanceActivation(ctx, a, version)
+		removeFutureActivations(ctx, a, version)
 		return ctx.WithBlockHeight(blockHeight)
 	})
 	s.Run(
@@ -428,7 +429,7 @@ func testBlock(
 			)
 			ctx = setLegacySstoreIfNeeded(ctx, a, version)
 			ctx = withCapturedConsensusParams(ctx, mc, blockHeight)
-			removeFutureGovernanceActivation(ctx, a, version)
+			removeFutureActivations(ctx, a, version)
 			return ctx.WithBlockHeight(blockHeight)
 		},
 	)
@@ -460,9 +461,12 @@ func setLegacySstoreIfNeeded(ctx sdk.Context, a *app.App, version string) sdk.Co
 	return ctx
 }
 
-func removeFutureGovernanceActivation(ctx sdk.Context, a *app.App, version string) {
+// removeFutureActivations clears the markers genesis sets for features that mainnet
+// only turns on at a later upgrade than version.
+func removeFutureActivations(ctx sdk.Context, a *app.App, version string) {
 	if semver.Compare(version, "v6.8") < 0 {
 		ctx.KVStore(a.GetKey(govtypes.StoreKey)).Delete(govtypes.IncrementalTallyEnabledKey)
+		ctx.KVStore(a.GetKey(stakingtypes.StoreKey)).Delete(stakingtypes.DelegationByValIndexReadyKey)
 	}
 }
 

sei-cosmos/x/staking/abci.go

@@ -19,12 +19,28 @@ func BeginBlocker(ctx sdk.Context, k keeper.Keeper) {
 	k.TrackHistoricalInfo(ctx)
 }
 
-// Called every block, update validator set
+// backfillDelegationByValIndex indexes the next batch of delegations while the
+// delegation-by-validator index is backfilled. It runs after every transaction in the
+// block, so no transaction iterates the staking store while the batch's writes are
+// pending in it. A delegation that cannot be decoded halts the chain, as it would at
+// any other read.
+func backfillDelegationByValIndex(ctx sdk.Context, k keeper.Keeper) {
+	// The batch is block-level work, so it is not charged to any gas meter.
+	ctx = ctx.WithGasMeter(sdk.NewInfiniteGasMeter(1, 1))
+	if _, err := k.BackfillDelegationByValIndex(ctx, keeper.DelegationByValIndexBackfillBatchSize); err != nil {
+		panic(err)
+	}
+}
+
+// Called every block, update validator set. It also advances the
+// delegation-by-validator index backfill while one is in progress.
 func EndBlocker(ctx sdk.Context, k keeper.Keeper) []abci.ValidatorUpdate {
 	endBlockerStart := time.Now()
 	defer func() {
 		stakingMetrics.endBlockerDuration.Record(ctx.Context(), time.Since(endBlockerStart).Seconds())
 	}()
 
-	return k.BlockValidatorUpdates(ctx)
+	updates := k.BlockValidatorUpdates(ctx)
+	backfillDelegationByValIndex(ctx, k)
+	return updates
 }

sei-cosmos/x/staking/genesis.go

@@ -76,6 +76,10 @@ func InitGenesis(
 		}
 	}
 
+	if _, err := keeper.MigrateDelegationByValIndex(ctx); err != nil {
+		panic(err)
+	}
+
 	for _, ubd := range data.UnbondingDelegations {
 		keeper.SetUnbondingDelegation(ctx, ubd)
 

sei-cosmos/x/staking/genesis_test.go

@@ -97,6 +97,36 @@ func TestInitGenesis(t *testing.T) {
 	require.Equal(t, abcivals, vals)
 }
 
+func TestInitGenesisIndexesDelegationsByValidator(t *testing.T) {
+	app, ctx, addrs := bootstrapGenesisTest(t, 2)
+	store := ctx.KVStore(app.StakingKeeper.GetStoreKey())
+	store.Delete(types.DelegationByValIndexReadyKey)
+
+	valTokens := app.StakingKeeper.TokensFromConsensusPower(ctx, 1)
+	params := app.StakingKeeper.GetParams(ctx)
+	pk, err := codectypes.NewAnyWithValue(PKs[0])
+	require.NoError(t, err)
+	valAddr := sdk.ValAddress(addrs[0])
+	validator := types.Validator{
+		OperatorAddress: valAddr.String(),
+		ConsensusPubkey: pk,
+		Status:          types.Bonded,
+		Tokens:          valTokens,
+		DelegatorShares: valTokens.ToDec(),
+		Description:     types.NewDescription("hoop", "", "", "", ""),
+	}
+	validators := append(app.StakingKeeper.GetAllValidators(ctx), validator)
+	require.NoError(t, apptesting.FundModuleAccount(app.BankKeeper, ctx, types.BondedPoolName,
+		sdk.NewCoins(sdk.NewCoin(params.BondDenom, valTokens.MulRaw(int64(len(validators)))))))
+	delegation := types.NewDelegation(addrs[1], valAddr, valTokens.ToDec())
+
+	staking.InitGenesis(ctx, app.StakingKeeper, app.AccountKeeper, app.BankKeeper,
+		types.NewGenesisState(params, validators, []types.Delegation{delegation}))
+
+	require.True(t, app.StakingKeeper.DelegationByValIndexReady(ctx))
+	require.True(t, store.Has(types.GetDelegationByValIndexKey(addrs[1], valAddr)))
+}
+
 func TestInitGenesis_PoolsBalanceMismatch(t *testing.T) {
 	app := seiapp.Setup(t, false, false, false)
 	ctx := app.NewContext(false, tmproto.Header{})

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.delegationByValIndexMaintained(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.delegationByValIndexMaintained(ctx) {
+		store.Delete(types.GetDelegationByValIndexKey(delegatorAddress, valAddr))
+	}
 }
 
 // GetUnbondingDelegations returns a given amount of all the delegator unbonding-delegations.