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

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 delegationByValIndexActive(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 delegationByValIndexActive(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,59 @@
+package keeper
+
+import (
+	"time"
+
+	"golang.org/x/mod/semver"
+
+	sdk "github.com/sei-protocol/sei-chain/sei-cosmos/types"
+	"github.com/sei-protocol/sei-chain/sei-cosmos/x/staking/types"
+)
+
+// DelegationByValIndexUpgrade is the upgrade at which SetDelegation and
+// RemoveDelegation begin maintaining the validator-indexed delegation store.
+const DelegationByValIndexUpgrade = "v6.7"
+
+func delegationByValIndexActive(ctx sdk.Context) bool {
+	return semver.Compare(ctx.ClosestUpgradeName(), DelegationByValIndexUpgrade) >= 0
+}
+
+// BackfillDelegationByValIndexResult reports the outcome of a delegation index backfill.
+type BackfillDelegationByValIndexResult struct {
+	TotalDelegations int
+	IndexWritten     int
+	AlreadyIndexed   int
+	DryRun           bool
+	Elapsed          time.Duration
+}
+
+// BackfillDelegationByValIndex writes validator-indexed delegation keys for existing
+// delegations. When dryRun is true, delegations are counted but no store writes occur.
+func (k Keeper) BackfillDelegationByValIndex(ctx sdk.Context, dryRun bool) BackfillDelegationByValIndexResult {
+	start := time.Now()
+	result := BackfillDelegationByValIndexResult{DryRun: dryRun}
+
+	store := ctx.KVStore(k.storeKey)
+	iterator := sdk.KVStorePrefixIterator(store, types.DelegationKey)
+	defer func() { _ = iterator.Close() }()
+
+	for ; iterator.Valid(); iterator.Next() {
+		delegation := types.MustUnmarshalDelegation(k.cdc, iterator.Value())
+		result.TotalDelegations++
+
+		delegatorAddress := sdk.MustAccAddressFromBech32(delegation.DelegatorAddress)
+		valAddr := delegation.GetValidatorAddr()
+		indexKey := types.GetDelegationByValIndexKey(delegatorAddress, valAddr)
+		if store.Has(indexKey) {
+			result.AlreadyIndexed++
+			continue
+		}
+
+		if !dryRun {
+			store.Set(indexKey, []byte{})
+		}
+		result.IndexWritten++
+	}
+
+	result.Elapsed = time.Since(start)
+	return result
+}

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

@@ -0,0 +1,98 @@
+package keeper_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	sdk "github.com/sei-protocol/sei-chain/sei-cosmos/types"
+	"github.com/sei-protocol/sei-chain/sei-cosmos/x/staking/types"
+)
+
+func TestDelegationByValIndexDualWrite(t *testing.T) {
+	_, app, ctx := createTestInput(t)
+	ctx = ctx.WithClosestUpgradeName("v6.7")
+
+	addrDels, valAddrs := generateAddresses(app, ctx, 1)
+	delegation := types.NewDelegation(addrDels[0], valAddrs[0], sdk.NewDec(1))
+
+	store := ctx.KVStore(app.StakingKeeper.GetStoreKey())
+	indexKey := types.GetDelegationByValIndexKey(addrDels[0], valAddrs[0])
+
+	require.False(t, store.Has(indexKey))
+
+	app.StakingKeeper.SetDelegation(ctx, delegation)
+	require.True(t, store.Has(indexKey))
+	require.NotNil(t, store.Get(types.GetDelegationKey(addrDels[0], valAddrs[0])))
+
+	app.StakingKeeper.RemoveDelegation(ctx, delegation)
+	require.False(t, store.Has(indexKey))
+	require.Nil(t, store.Get(types.GetDelegationKey(addrDels[0], valAddrs[0])))
+}
+
+func TestDelegationByValIndexPreUpgradeNoDualWrite(t *testing.T) {
+	_, app, ctx := createTestInput(t)
+	ctx = ctx.WithClosestUpgradeName("v6.6")
+
+	addrDels, valAddrs := generateAddresses(app, ctx, 1)
+	delegation := types.NewDelegation(addrDels[0], valAddrs[0], sdk.NewDec(1))
+
+	store := ctx.KVStore(app.StakingKeeper.GetStoreKey())
+	indexKey := types.GetDelegationByValIndexKey(addrDels[0], valAddrs[0])
+
+	app.StakingKeeper.SetDelegation(ctx, delegation)
+	require.False(t, store.Has(indexKey))
+	require.NotNil(t, store.Get(types.GetDelegationKey(addrDels[0], valAddrs[0])))
+
+	app.StakingKeeper.RemoveDelegation(ctx, delegation)
+	require.False(t, store.Has(indexKey))
+	require.Nil(t, store.Get(types.GetDelegationKey(addrDels[0], valAddrs[0])))
+}
+
+func TestBackfillDelegationByValIndex(t *testing.T) {
+	_, app, ctx := createTestInput(t)
+
+	addrDels, valAddrs := generateAddresses(app, ctx, 2)
+	delegations := []types.Delegation{
+		types.NewDelegation(addrDels[0], valAddrs[0], sdk.NewDec(1)),
+		types.NewDelegation(addrDels[0], valAddrs[1], sdk.NewDec(2)),
+		types.NewDelegation(addrDels[1], valAddrs[0], sdk.NewDec(3)),
+	}
+
+	store := ctx.KVStore(app.StakingKeeper.GetStoreKey())
+	for _, delegation := range delegations {
+		delegatorAddress := sdk.MustAccAddressFromBech32(delegation.DelegatorAddress)
+		valAddr := delegation.GetValidatorAddr()
+		store.Set(
+			types.GetDelegationKey(delegatorAddress, valAddr),
+			types.MustMarshalDelegation(app.AppCodec(), delegation),
+		)
+	}
+
+	dryRunResult := app.StakingKeeper.BackfillDelegationByValIndex(ctx, true)
+	require.Equal(t, 3, dryRunResult.TotalDelegations)
+	require.Equal(t, 3, dryRunResult.IndexWritten)
+	require.Equal(t, 0, dryRunResult.AlreadyIndexed)
+	require.True(t, dryRunResult.DryRun)
+	for _, delegation := range delegations {
+		delegatorAddress := sdk.MustAccAddressFromBech32(delegation.DelegatorAddress)
+		valAddr := delegation.GetValidatorAddr()
+		require.False(t, store.Has(types.GetDelegationByValIndexKey(delegatorAddress, valAddr)))
+	}
+
+	writeResult := app.StakingKeeper.BackfillDelegationByValIndex(ctx, false)
+	require.Equal(t, 3, writeResult.TotalDelegations)
+	require.Equal(t, 3, writeResult.IndexWritten)
+	require.Equal(t, 0, writeResult.AlreadyIndexed)
+	require.False(t, writeResult.DryRun)
+	for _, delegation := range delegations {
+		delegatorAddress := sdk.MustAccAddressFromBech32(delegation.DelegatorAddress)
+		valAddr := delegation.GetValidatorAddr()
+		require.True(t, store.Has(types.GetDelegationByValIndexKey(delegatorAddress, valAddr)))
+	}
+
+	repeatResult := app.StakingKeeper.BackfillDelegationByValIndex(ctx, false)
+	require.Equal(t, 3, repeatResult.TotalDelegations)
+	require.Equal(t, 0, repeatResult.IndexWritten)
+	require.Equal(t, 3, repeatResult.AlreadyIndexed)
+}

sei-cosmos/x/staking/simulation/decoder.go

@@ -41,6 +41,8 @@ func NewDecodeStore(cdc codec.Codec) func(kvA, kvB kv.Pair) string {
 			cdc.MustUnmarshal(kvB.Value, &delegationB)
 
 			return fmt.Sprintf("%v\n%v", delegationA, delegationB)
+		case bytes.Equal(kvA.Key[:1], types.DelegationByValIndexKey):
+			return fmt.Sprintf("%X\n%X", kvA.Key, kvB.Key)
 		case bytes.Equal(kvA.Key[:1], types.UnbondingDelegationKey),
 			bytes.Equal(kvA.Key[:1], types.UnbondingDelegationByValIndexKey):
 			var ubdA, ubdB types.UnbondingDelegation

sei-cosmos/x/staking/types/keys.go

@@ -39,6 +39,7 @@ var (
 	DelegationKey                    = []byte{0x31} // key for a delegation
 	UnbondingDelegationKey           = []byte{0x32} // key for an unbonding-delegation
 	UnbondingDelegationByValIndexKey = []byte{0x33} // prefix for each key for an unbonding-delegation, by validator operator
+	DelegationByValIndexKey          = []byte{0x37} // prefix for each key for a delegation, by validator operator
 	RedelegationKey                  = []byte{0x34} // key for a redelegation
 	RedelegationByValSrcIndexKey     = []byte{0x35} // prefix for each key for an redelegation, by source validator operator
 	RedelegationByValDstIndexKey     = []byte{0x36} // prefix for each key for an redelegation, by destination validator operator
@@ -198,6 +199,31 @@ func GetDelegationsKey(delAddr sdk.AccAddress) []byte {
 	return append(DelegationKey, address.MustLengthPrefix(delAddr)...)
 }
 
+// GetDelegationByValIndexKey creates the index-key for a delegation, stored by validator-index.
+// VALUE: none (key rearrangement used)
+func GetDelegationByValIndexKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte {
+	return append(GetDelegationsByValIndexKey(valAddr), address.MustLengthPrefix(delAddr)...)
+}
+
+// GetDelegationKeyFromValIndexKey rearranges the ValIndexKey to get the DelegationKey.
+func GetDelegationKeyFromValIndexKey(indexKey []byte) []byte {
+	kv.AssertKeyAtLeastLength(indexKey, 2)
+	addrs := indexKey[1:] // remove prefix bytes
+
+	valAddrLen := addrs[0]
+	kv.AssertKeyAtLeastLength(addrs, 2+int(valAddrLen))
+	valAddr := addrs[1 : 1+valAddrLen]
+	kv.AssertKeyAtLeastLength(addrs, 3+int(valAddrLen))
+	delAddr := addrs[valAddrLen+2:]
+
+	return GetDelegationKey(delAddr, valAddr)
+}
+
+// GetDelegationsByValIndexKey creates the prefix keyspace for delegation indexes by validator.
+func GetDelegationsByValIndexKey(valAddr sdk.ValAddress) []byte {
+	return append(DelegationByValIndexKey, address.MustLengthPrefix(valAddr)...)
+}
+
 // 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 {

sei-cosmos/x/staking/types/keys_test.go

@@ -49,6 +49,16 @@ func TestGetValidatorPowerRank(t *testing.T) {
 	}
 }
 
+func TestGetDelegationKeyFromValIndexKey(t *testing.T) {
+	delAddr := sdk.AccAddress(keysAddr1)
+	valAddr := sdk.ValAddress(keysAddr2)
+
+	delegationKey := types.GetDelegationKey(delAddr, valAddr)
+	indexKey := types.GetDelegationByValIndexKey(delAddr, valAddr)
+
+	require.Equal(t, delegationKey, types.GetDelegationKeyFromValIndexKey(indexKey))
+}
+
 func TestGetREDByValDstIndexKey(t *testing.T) {
 	tests := []struct {
 		delAddr    sdk.AccAddress

tools/cmd.go

@@ -3,6 +3,7 @@ package tools
 import (
 	"github.com/spf13/cobra"
 
+	stakingcmd "github.com/sei-protocol/sei-chain/tools/staking/cmd"
 	scanner "github.com/sei-protocol/sei-chain/tools/tx-scanner/cmd"
 )
 
@@ -12,5 +13,6 @@ func ToolCmd() *cobra.Command {
 		Short: "A set of useful tools for sei chain",
 	}
 	toolsCmd.AddCommand(scanner.ScanCmd())
+	toolsCmd.AddCommand(stakingcmd.StakingCmd())
 	return toolsCmd
 }

tools/staking/cmd/backfill.go

@@ -0,0 +1,85 @@
+package cmd
+
+import (
+	"fmt"
+	"os"
+
+	tmproto "github.com/sei-protocol/sei-chain/sei-tendermint/proto/tendermint/types"
+	"github.com/spf13/cobra"
+
+	"github.com/sei-protocol/sei-chain/app"
+	stakingtool "github.com/sei-protocol/sei-chain/tools/staking"
+)
+
+// BackfillDelegationIndexCmd backfills validator-indexed delegation keys for benchmarking.
+func BackfillDelegationIndexCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "backfill-delegation-index",
+		Short: "Backfill validator-indexed delegation keys (dev/benchmark)",
+		Long: `Iterate all delegations and write validator-indexed secondary index keys.
+
+By default this runs in dry-run mode and does not modify state. Pass --write to
+persist index keys. When using --write, stop the node and operate on a copy of
+the data directory. --write cannot be combined with --height.`,
+		RunE: runBackfillDelegationIndex,
+	}
+
+	stakingtool.BindServerFlags(cmd, app.DefaultNodeHome)
+	cmd.Flags().Int64("height", 0, "Load state at this height (0 = latest committed)")
+	cmd.Flags().Bool("write", false, "Write index keys and commit state (requires stopped node)")
+	return cmd
+}
+
+func runBackfillDelegationIndex(cmd *cobra.Command, _ []string) error {
+	height, err := cmd.Flags().GetInt64("height")
+	if err != nil {
+		return err
+	}
+	write, err := cmd.Flags().GetBool("write")
+	if err != nil {
+		return err
+	}
+	dryRun := !write
+
+	if write && height > 0 {
+		return fmt.Errorf("--write cannot be used with --height")
+	}
+
+	if write {
+		fmt.Fprintln(os.Stderr, "WARNING: --write mutates application state. Stop the node and use a data copy.")
+	}
+
+	seiApp, err := stakingtool.LoadApp(cmd, height)
+	if err != nil {
+		return err
+	}
+	defer seiApp.Close()
+
+	blockHeight := seiApp.LastBlockHeight()
+	if blockHeight == 0 {
+		blockHeight = 1
+	}
+
+	ctx := seiApp.NewUncachedContext(false, tmproto.Header{Height: blockHeight})
+	result := seiApp.StakingKeeper.BackfillDelegationByValIndex(ctx, dryRun)
+
+	fmt.Printf("dry_run=%t height=%d\n", result.DryRun, blockHeight)
+	fmt.Printf("total_delegations=%d index_written=%d already_indexed=%d elapsed=%s\n",
+		result.TotalDelegations,
+		result.IndexWritten,
+		result.AlreadyIndexed,
+		result.Elapsed,
+	)
+	if result.TotalDelegations > 0 && result.Elapsed > 0 {
+		perSec := float64(result.TotalDelegations) / result.Elapsed.Seconds()
+		fmt.Printf("delegations_per_second=%.2f\n", perSec)
+	}
+
+	if !write {
+		return nil
+	}
+
+	commitID := seiApp.CommitMultiStore().Commit(true)
+	fmt.Printf("committed_version=%d\n", commitID.Version)
+	return nil
+}

tools/staking/cmd/staking.go

@@ -0,0 +1,15 @@
+package cmd
+
+import (
+	"github.com/spf13/cobra"
+)
+
+// StakingCmd returns offline staking maintenance commands.
+func StakingCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "staking",
+		Short: "Offline staking maintenance tools",
+	}
+	cmd.AddCommand(BackfillDelegationIndexCmd())
+	return cmd
+}