refactor: add delegations by validator index (#15731)

Co-authored-by: Marko <marbar3778@yahoo.com>
Co-authored-by: Aleksandr Bezobchuk <alexanderbez@users.noreply.github.com>

CHANGELOG.md

@@ -110,6 +110,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (x/slashing) [#15580](https://github.com/cosmos/cosmos-sdk/pull/15580) The validator slashing window now stores "chunked" bitmap entries for each validator's signing window instead of a single boolean entry per signing window index.
 * (x/feegrant) [#14294](https://github.com/cosmos/cosmos-sdk/pull/14294) Moved the logic of rejecting duplicate grant from `msg_server` to `keeper` method.
 * (x/staking) [#14590](https://github.com/cosmos/cosmos-sdk/pull/14590) `MsgUndelegateResponse` now includes undelegated amount. `x/staking` module's `keeper.Undelegate` now returns 3 values (completionTime,undelegateAmount,error)  instead of 2.
+* (x/staking) (#15731) (https://github.com/cosmos/cosmos-sdk/pull/15731) Introducing a new index to retrieve the delegations by validator efficiently.
 
 ### API Breaking Changes
 

tests/integration/staking/keeper/determinstic_test.go

@@ -327,7 +327,7 @@ func TestGRPCValidatorDelegations(t *testing.T) {
 		ValidatorAddr: validator.OperatorAddress,
 	}
 
-	testdata.DeterministicIterations(f.ctx, t, req, f.queryClient.ValidatorDelegations, 11985, false)
+	testdata.DeterministicIterations(f.ctx, t, req, f.queryClient.ValidatorDelegations, 14475, false)
 }
 
 func TestGRPCValidatorUnbondingDelegations(t *testing.T) {

tests/integration/staking/keeper/validator_bench_test.go

@@ -1,6 +1,15 @@
 package keeper_test
 
-import "testing"
+import (
+	"fmt"
+	"testing"
+
+	"cosmossdk.io/simapp"
+	storetypes "cosmossdk.io/store/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	banktestutil "github.com/cosmos/cosmos-sdk/x/bank/testutil"
+	"github.com/cosmos/cosmos-sdk/x/staking/types"
+)
 
 func BenchmarkGetValidator(b *testing.B) {
 	// 900 is the max number we are allowed to use in order to avoid simtestutil.CreateTestPubKeys
@@ -27,3 +36,117 @@ func BenchmarkGetValidator(b *testing.B) {
 		}
 	}
 }
+
+func BenchmarkGetValidatorDelegations(b *testing.B) {
+	var totalPower int64
+	powersNumber := 10
+
+	powers := make([]int64, powersNumber)
+	for i := range powers {
+		powers[i] = int64(i)
+		totalPower += int64(i)
+	}
+
+	app, ctx, _, valAddrs, vals := initValidators(b, totalPower, len(powers), powers)
+	for _, validator := range vals {
+		app.StakingKeeper.SetValidator(ctx, validator)
+	}
+
+	delegationsNum := 1000
+	for _, val := range valAddrs {
+		for i := 0; i < delegationsNum; i++ {
+			delegator := sdk.AccAddress(fmt.Sprintf("address%d", i))
+			banktestutil.FundAccount(app.BankKeeper, ctx, delegator,
+				sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(int64(i)))))
+			NewDel := types.NewDelegation(delegator, val, sdk.NewDec(int64(i)))
+			app.StakingKeeper.SetDelegation(ctx, NewDel)
+		}
+	}
+
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		updateValidatorDelegations(ctx, app, valAddrs[0], sdk.ValAddress("val"))
+	}
+}
+
+func BenchmarkGetValidatorDelegationsLegacy(b *testing.B) {
+	var totalPower int64
+	powersNumber := 10
+
+	powers := make([]int64, powersNumber)
+	for i := range powers {
+		powers[i] = int64(i)
+		totalPower += int64(i)
+	}
+
+	app, ctx, _, valAddrs, vals := initValidators(b, totalPower, len(powers), powers)
+
+	for _, validator := range vals {
+		app.StakingKeeper.SetValidator(ctx, validator)
+	}
+
+	delegationsNum := 1000
+	for _, val := range valAddrs {
+		for i := 0; i < delegationsNum; i++ {
+			delegator := sdk.AccAddress(fmt.Sprintf("address%d", i))
+			banktestutil.FundAccount(app.BankKeeper, ctx, delegator,
+				sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(int64(i)))))
+			NewDel := types.NewDelegation(delegator, val, sdk.NewDec(int64(i)))
+			app.StakingKeeper.SetDelegation(ctx, NewDel)
+		}
+	}
+
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		updateValidatorDelegationsLegacy(ctx, app, valAddrs[0], sdk.ValAddress("val"))
+	}
+}
+
+func updateValidatorDelegationsLegacy(ctx sdk.Context, app *simapp.SimApp, existingValAddr, newValAddr sdk.ValAddress) {
+	storeKey := app.GetKey(types.StoreKey)
+	cdc, k := app.AppCodec(), app.StakingKeeper
+
+	store := ctx.KVStore(storeKey)
+
+	iterator := storetypes.KVStorePrefixIterator(store, types.DelegationKey)
+	defer iterator.Close()
+
+	for ; iterator.Valid(); iterator.Next() {
+		delegation := types.MustUnmarshalDelegation(cdc, iterator.Value())
+		if delegation.GetValidatorAddr().Equals(existingValAddr) {
+			k.RemoveDelegation(ctx, delegation)
+			delegation.ValidatorAddress = newValAddr.String()
+			k.SetDelegation(ctx, delegation)
+		}
+	}
+}
+
+func updateValidatorDelegations(ctx sdk.Context, app *simapp.SimApp, existingValAddr, newValAddr sdk.ValAddress) {
+	storeKey := app.GetKey(types.StoreKey)
+	cdc, k := app.AppCodec(), app.StakingKeeper
+
+	store := ctx.KVStore(storeKey)
+
+	itr := storetypes.KVStorePrefixIterator(store, types.GetDelegationsByValPrefixKey(existingValAddr))
+	defer itr.Close()
+
+	for ; itr.Valid(); itr.Next() {
+		key := itr.Key()
+		valAddr, delAddr, err := types.ParseDelegationsByValKey(key)
+		if err != nil {
+			panic(err)
+		}
+
+		bz := store.Get(types.GetDelegationKey(delAddr, valAddr))
+		delegation := types.MustUnmarshalDelegation(cdc, bz)
+
+		// remove old operator addr from delegation
+		if err := k.RemoveDelegation(ctx, delegation); err != nil {
+			panic(err)
+		}
+
+		delegation.ValidatorAddress = newValAddr.String()
+		// add with new operator addr
+		k.SetDelegation(ctx, delegation)
+	}
+}

x/staking/keeper/delegation.go

@@ -60,14 +60,22 @@ func (k Keeper) GetAllDelegations(ctx sdk.Context) (delegations []types.Delegati
 func (k Keeper) GetValidatorDelegations(ctx sdk.Context, valAddr sdk.ValAddress) (delegations []types.Delegation) {
 	store := ctx.KVStore(k.storeKey)
 
-	iterator := storetypes.KVStorePrefixIterator(store, types.DelegationKey)
+	iterator := storetypes.KVStorePrefixIterator(store, types.GetDelegationsByValPrefixKey(valAddr))
 	defer iterator.Close()
 
 	for ; iterator.Valid(); iterator.Next() {
-		delegation := types.MustUnmarshalDelegation(k.cdc, iterator.Value())
-		if delegation.GetValidatorAddr().Equals(valAddr) {
-			delegations = append(delegations, delegation)
+		var delegation types.Delegation
+		valAddr, delAddr, err := types.ParseDelegationsByValKey(iterator.Key())
+		if err != nil {
+			panic(err)
+		}
+
+		bz := store.Get(types.GetDelegationKey(delAddr, valAddr))
+		if err := k.cdc.Unmarshal(bz, &delegation); err != nil {
+			panic(err)
 		}
+
+		delegations = append(delegations, delegation)
 	}
 
 	return delegations
@@ -103,6 +111,9 @@ func (k Keeper) SetDelegation(ctx sdk.Context, delegation types.Delegation) {
 	store := ctx.KVStore(k.storeKey)
 	b := types.MustMarshalDelegation(k.cdc, delegation)
 	store.Set(types.GetDelegationKey(delegatorAddress, delegation.GetValidatorAddr()), b)
+
+	// set the delegation in validator delegator index
+	store.Set(types.GetDelegationsByValKey(delegation.GetValidatorAddr(), delegatorAddress), []byte{})
 }
 
 // RemoveDelegation removes a delegation
@@ -119,6 +130,8 @@ func (k Keeper) RemoveDelegation(ctx sdk.Context, delegation types.Delegation) e
 
 	store := ctx.KVStore(k.storeKey)
 	store.Delete(types.GetDelegationKey(delegatorAddress, delegation.GetValidatorAddr()))
+	store.Delete(types.GetDelegationsByValKey(delegation.GetValidatorAddr(), delegatorAddress))
+
 	return nil
 }
 

x/staking/keeper/delegation_test.go

@@ -144,6 +144,83 @@ func (s *KeeperTestSuite) TestDelegation() {
 	require.Equal(0, len(resBonds))
 }
 
+func (s *KeeperTestSuite) TestDelegationsByValIndex() {
+	ctx, keeper := s.ctx, s.stakingKeeper
+	require := s.Require()
+
+	addrDels, valAddrs := createValAddrs(3)
+
+	for _, addr := range addrDels {
+		s.accountKeeper.EXPECT().StringToBytes(addr.String()).Return(addr, nil).AnyTimes()
+		s.accountKeeper.EXPECT().BytesToString(addr).Return(addr.String(), nil).AnyTimes()
+		s.bankKeeper.EXPECT().DelegateCoinsFromAccountToModule(gomock.Any(), addr, gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	}
+
+	// construct the validators
+	amts := []math.Int{sdk.NewInt(9), sdk.NewInt(8), sdk.NewInt(7)}
+	var validators [3]stakingtypes.Validator
+	for i, amt := range amts {
+		validators[i] = testutil.NewValidator(s.T(), valAddrs[i], PKs[i])
+		validators[i], _ = validators[i].AddTokensFromDel(amt)
+
+		validators[i] = stakingkeeper.TestingUpdateValidator(keeper, ctx, validators[i], true)
+	}
+
+	// delegate 2 tokens
+	//
+	// total delegations after delegating: del1 -> 2stake
+	_, err := s.msgServer.Delegate(ctx, stakingtypes.NewMsgDelegate(addrDels[0], valAddrs[0], sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(2))))
+	require.NoError(err)
+
+	dels := s.stakingKeeper.GetValidatorDelegations(ctx, valAddrs[0])
+	require.Len(dels, 1)
+
+	// delegate 4 tokens
+	//
+	// total delegations after delegating: del1 -> 2stake, del2 -> 4stake
+	_, err = s.msgServer.Delegate(ctx, stakingtypes.NewMsgDelegate(addrDels[1], valAddrs[0], sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(4))))
+	require.NoError(err)
+
+	dels = s.stakingKeeper.GetValidatorDelegations(ctx, valAddrs[0])
+	require.Len(dels, 2)
+
+	// undelegate 1 token from del1
+	//
+	// total delegations after undelegating: del1 -> 1stake, del2 -> 4stake
+	_, err = s.msgServer.Undelegate(ctx, stakingtypes.NewMsgUndelegate(addrDels[0], valAddrs[0], sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(1))))
+	require.NoError(err)
+
+	dels = s.stakingKeeper.GetValidatorDelegations(ctx, valAddrs[0])
+	require.Len(dels, 2)
+
+	// undelegate 1 token from del1
+	//
+	// total delegations after undelegating: del2 -> 4stake
+	_, err = s.msgServer.Undelegate(ctx, stakingtypes.NewMsgUndelegate(addrDels[0], valAddrs[0], sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(1))))
+	require.NoError(err)
+
+	dels = s.stakingKeeper.GetValidatorDelegations(ctx, valAddrs[0])
+	require.Len(dels, 1)
+
+	// undelegate 2 tokens from del2
+	//
+	// total delegations after undelegating: del2 -> 2stake
+	_, err = s.msgServer.Undelegate(ctx, stakingtypes.NewMsgUndelegate(addrDels[1], valAddrs[0], sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(2))))
+	require.NoError(err)
+
+	dels = s.stakingKeeper.GetValidatorDelegations(ctx, valAddrs[0])
+	require.Len(dels, 1)
+
+	// undelegate 2 tokens from del2
+	//
+	// total delegations after undelegating: []
+	_, err = s.msgServer.Undelegate(ctx, stakingtypes.NewMsgUndelegate(addrDels[1], valAddrs[0], sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(2))))
+	require.NoError(err)
+
+	dels = s.stakingKeeper.GetValidatorDelegations(ctx, valAddrs[0])
+	require.Len(dels, 0)
+}
+
 // tests Get/Set/Remove UnbondingDelegation
 func (s *KeeperTestSuite) TestUnbondingDelegation() {
 	ctx, keeper := s.ctx, s.stakingKeeper

x/staking/keeper/grpc_query.go

@@ -92,31 +92,44 @@ func (k Querier) ValidatorDelegations(c context.Context, req *types.QueryValidat
 	if req.ValidatorAddr == "" {
 		return nil, status.Error(codes.InvalidArgument, "validator address cannot be empty")
 	}
+
+	valAddr, err := sdk.ValAddressFromBech32(req.ValidatorAddr)
+	if err != nil {
+		return nil, err
+	}
 	ctx := sdk.UnwrapSDKContext(c)
 
 	store := ctx.KVStore(k.storeKey)
-	valStore := prefix.NewStore(store, types.DelegationKey)
-	delegations, pageRes, err := query.GenericFilteredPaginate(k.cdc, valStore, req.Pagination, func(key []byte, delegation *types.Delegation) (*types.Delegation, error) {
-		valAddr, err := sdk.ValAddressFromBech32(req.ValidatorAddr)
-		if err != nil {
-			return nil, err
-		}
+	delStore := prefix.NewStore(store, types.GetDelegationsByValPrefixKey(valAddr))
 
-		if !delegation.GetValidatorAddr().Equals(valAddr) {
-			return nil, nil
+	var (
+		dels    types.Delegations
+		pageRes *query.PageResponse
+	)
+	pageRes, err = query.Paginate(delStore, req.Pagination, func(delAddr, value []byte) error {
+		bz := store.Get(types.GetDelegationKey(delAddr, valAddr))
+
+		var delegation types.Delegation
+		err = k.cdc.Unmarshal(bz, &delegation)
+		if err != nil {
+			return err
 		}
 
-		return delegation, nil
-	}, func() *types.Delegation {
-		return &types.Delegation{}
+		dels = append(dels, delegation)
+		return nil
 	})
 	if err != nil {
-		return nil, status.Error(codes.Internal, err.Error())
-	}
+		delegations, pageResponse, err := k.getValidatorDelegationsLegacy(ctx, req)
+		if err != nil {
+			return nil, status.Error(codes.Internal, err.Error())
+		}
+
+		dels = types.Delegations{}
+		for _, d := range delegations {
+			dels = append(dels, *d)
+		}
 
-	dels := types.Delegations{}
-	for _, d := range delegations {
-		dels = append(dels, *d)
+		pageRes = pageResponse
 	}
 
 	delResponses, err := DelegationsToDelegationResponses(ctx, k.Keeper, dels)
@@ -129,6 +142,26 @@ func (k Querier) ValidatorDelegations(c context.Context, req *types.QueryValidat
 	}, nil
 }
 
+func (k Querier) getValidatorDelegationsLegacy(ctx sdk.Context, req *types.QueryValidatorDelegationsRequest) ([]*types.Delegation, *query.PageResponse, error) {
+	store := ctx.KVStore(k.storeKey)
+
+	valStore := prefix.NewStore(store, types.DelegationKey)
+	return query.GenericFilteredPaginate(k.cdc, valStore, req.Pagination, func(key []byte, delegation *types.Delegation) (*types.Delegation, error) {
+		valAddr, err := sdk.ValAddressFromBech32(req.ValidatorAddr)
+		if err != nil {
+			return nil, err
+		}
+
+		if !delegation.GetValidatorAddr().Equals(valAddr) {
+			return nil, nil
+		}
+
+		return delegation, nil
+	}, func() *types.Delegation {
+		return &types.Delegation{}
+	})
+}
+
 // ValidatorUnbondingDelegations queries unbonding delegations of a validator
 func (k Querier) ValidatorUnbondingDelegations(c context.Context, req *types.QueryValidatorUnbondingDelegationsRequest) (*types.QueryValidatorUnbondingDelegationsResponse, error) {
 	if req == nil {

x/staking/keeper/migrations.go

@@ -40,5 +40,5 @@ func (m Migrator) Migrate3to4(ctx sdk.Context) error {
 
 // Migrate4to5 migrates x/staking state from consensus version 4 to 5.
 func (m Migrator) Migrate4to5(ctx sdk.Context) error {
-	return v5.MigrateStore(ctx, m.keeper.storeKey)
+	return v5.MigrateStore(ctx, m.keeper.storeKey, m.keeper.cdc)
 }