perf(staking): optimize endblock by reducing bech32 conversions (#24354)

Co-authored-by: Dev Ojha <[email protected]>
Co-authored-by: Alex | Interchain Labs <[email protected]>

Author: 3 people

CHANGELOG.md

@@ -62,6 +62,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### Improvements
 
+* (x/staking) [#24354](https://github.com/cosmos/cosmos-sdk/pull/24354) Optimize validator endblock by reducing bech32 conversions, resulting in significant performance improvement
 * (client/keys) [#18950](https://github.com/cosmos/cosmos-sdk/pull/18950) Improve `<appd> keys add`, `<appd> keys import` and `<appd> keys rename` by checking name validation.
 * (client/keys) [#18703](https://github.com/cosmos/cosmos-sdk/pull/18703) Improve `<appd> keys add` and `<appd> keys show` by checking whether there are duplicate keys in the multisig case.
 * (client/keys) [#18745](https://github.com/cosmos/cosmos-sdk/pull/18745) Improve `<appd> keys export` and `<appd> keys mnemonic` by adding --yes option to skip interactive confirmation.

x/staking/keeper/val_state_change.go

@@ -188,17 +188,13 @@ func (k Keeper) ApplyAndReturnValidatorSetUpdates(ctx context.Context) (updates
 			panic("unexpected validator status")
 		}
 
+		valAddrStr := string(valAddr)
 		// fetch the old power bytes
-		valAddrStr, err := k.validatorAddressCodec.BytesToString(valAddr)
-		if err != nil {
-			return nil, err
-		}
-		oldPowerBytes, found := last[valAddrStr]
+		oldPower, found := last[valAddrStr]
 		newPower := validator.ConsensusPower(powerReduction)
-		newPowerBytes := k.cdc.MustMarshal(&gogotypes.Int64Value{Value: newPower})
 
 		// update the validator set if power has changed
-		if !found || !bytes.Equal(oldPowerBytes, newPowerBytes) {
+		if !found || oldPower != newPower {
 			updates = append(updates, validator.ABCIValidatorUpdate(powerReduction))
 
 			if err = k.SetLastValidatorPower(ctx, valAddr, newPower); err != nil {
@@ -209,7 +205,7 @@ func (k Keeper) ApplyAndReturnValidatorSetUpdates(ctx context.Context) (updates
 		delete(last, valAddrStr)
 		count++
 
-		totalPower = totalPower.Add(math.NewInt(newPower))
+		totalPower = totalPower.AddRaw(newPower)
 	}
 
 	noLongerBonded, err := sortNoLongerBonded(last, k.validatorAddressCodec)
@@ -454,9 +450,9 @@ func (k Keeper) completeUnbondingValidator(ctx context.Context, validator types.
 	return validator, nil
 }
 
-// map of operator bech32-addresses to serialized power
-// We use bech32 strings here, because we can't have slices as keys: map[[]byte][]byte
-type validatorsByAddr map[string][]byte
+// map of operator addresses to power
+// We use (non bech32) strings here, because we can't have slices as keys: map[[]byte][]byte
+type validatorsByAddr map[string]int64
 
 // get the last validator set
 func (k Keeper) getLastValidatorsByAddr(ctx context.Context) (validatorsByAddr, error) {
@@ -468,17 +464,12 @@ func (k Keeper) getLastValidatorsByAddr(ctx context.Context) (validatorsByAddr,
 	}
 	defer iterator.Close()
 
+	var intVal gogotypes.Int64Value
 	for ; iterator.Valid(); iterator.Next() {
 		// extract the validator address from the key (prefix is 1-byte, addrLen is 1-byte)
-		valAddr := types.AddressFromLastValidatorPowerKey(iterator.Key())
-		valAddrStr, err := k.validatorAddressCodec.BytesToString(valAddr)
-		if err != nil {
-			return nil, err
-		}
-
-		powerBytes := iterator.Value()
-		last[valAddrStr] = make([]byte, len(powerBytes))
-		copy(last[valAddrStr], powerBytes)
+		valAddrStr := string(types.AddressFromLastValidatorPowerKey(iterator.Key()))
+		k.cdc.MustUnmarshal(iterator.Value(), &intVal)
+		last[valAddrStr] = intVal.GetValue()
 	}
 
 	return last, nil
@@ -492,10 +483,7 @@ func sortNoLongerBonded(last validatorsByAddr, ac address.Codec) ([][]byte, erro
 	index := 0
 
 	for valAddrStr := range last {
-		valAddrBytes, err := ac.StringToBytes(valAddrStr)
-		if err != nil {
-			return nil, err
-		}
+		valAddrBytes := []byte(valAddrStr)
 		noLongerBonded[index] = valAddrBytes
 		index++
 	}

x/staking/keeper_bench_test.go

@@ -0,0 +1,140 @@
+package staking_test
+
+import (
+	"math/rand"
+	"testing"
+
+	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
+	cmttime "github.com/cometbft/cometbft/types/time"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
+
+	sdkmath "cosmossdk.io/math"
+	storetypes "cosmossdk.io/store/types"
+
+	"github.com/cosmos/cosmos-sdk/baseapp"
+	"github.com/cosmos/cosmos-sdk/codec/address"
+	"github.com/cosmos/cosmos-sdk/runtime"
+	"github.com/cosmos/cosmos-sdk/testutil"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	moduletestutil "github.com/cosmos/cosmos-sdk/types/module/testutil"
+	simtypes "github.com/cosmos/cosmos-sdk/types/simulation"
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	govtypes "github.com/cosmos/cosmos-sdk/x/gov/types"
+	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
+	stakingtestutil "github.com/cosmos/cosmos-sdk/x/staking/testutil"
+	"github.com/cosmos/cosmos-sdk/x/staking/types"
+)
+
+func BenchmarkApplyAndReturnValidatorSetUpdates(b *testing.B) {
+	// goal of this benchmark to measure the performance changes in ApplyAndReturnValidatorSetUpdates
+	// for dropping the bench32 conversion and different index types.
+	// therefore the validator power, max or state is not modified to focus on comparing the valset
+	// for an update only.
+	const validatorCount = 150
+	testEnv := newTestEnvironment(b)
+	keeper, ctx := testEnv.stakingKeeper, testEnv.ctx
+	r := rand.New(rand.NewSource(int64(1)))
+	vals, valAddrs := setupState(b, r, validatorCount)
+
+	params, err := keeper.GetParams(ctx)
+	require.NoError(b, err)
+	params.MaxValidators = uint32(validatorCount)
+	require.NoError(b, keeper.SetParams(ctx, params))
+
+	b.Logf("vals count: %d", validatorCount)
+	for i, validator := range vals {
+		require.NoError(b, keeper.SetValidator(ctx, validator))
+		require.NoError(b, keeper.SetValidatorByConsAddr(ctx, validator))
+		require.NoError(b, keeper.SetValidatorByPowerIndex(ctx, validator))
+		require.NoError(b, keeper.SetLastValidatorPower(ctx, valAddrs[i], validator.ConsensusPower(sdk.DefaultPowerReduction)))
+	}
+	ctx, _ = testEnv.ctx.CacheContext()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, _ = keeper.ApplyAndReturnValidatorSetUpdates(ctx)
+	}
+}
+
+type KeeperTestEnvironment struct {
+	ctx           sdk.Context
+	stakingKeeper *stakingkeeper.Keeper
+	bankKeeper    *stakingtestutil.MockBankKeeper
+	accountKeeper *stakingtestutil.MockAccountKeeper
+	queryClient   types.QueryClient
+	msgServer     types.MsgServer
+}
+
+func newTestEnvironment(tb testing.TB) *KeeperTestEnvironment {
+	tb.Helper()
+	key := storetypes.NewKVStoreKey(types.StoreKey)
+	storeService := runtime.NewKVStoreService(key)
+	testCtx := testutil.DefaultContextWithDB(tb, key, storetypes.NewTransientStoreKey("transient_test"))
+	ctx := testCtx.Ctx.WithBlockHeader(cmtproto.Header{Time: cmttime.Now()})
+	encCfg := moduletestutil.MakeTestEncodingConfig()
+
+	ctrl := gomock.NewController(tb)
+	accountKeeper := stakingtestutil.NewMockAccountKeeper(ctrl)
+	accountKeeper.EXPECT().GetModuleAddress(types.BondedPoolName).
+		Return(authtypes.NewEmptyModuleAccount(types.BondedPoolName).GetAddress())
+	accountKeeper.EXPECT().GetModuleAddress(types.NotBondedPoolName).
+		Return(authtypes.NewEmptyModuleAccount(types.NotBondedPoolName).GetAddress())
+	accountKeeper.EXPECT().AddressCodec().Return(address.NewBech32Codec("cosmos")).AnyTimes()
+
+	bankKeeper := stakingtestutil.NewMockBankKeeper(ctrl)
+
+	keeper := stakingkeeper.NewKeeper(
+		encCfg.Codec,
+		storeService,
+		accountKeeper,
+		bankKeeper,
+		authtypes.NewModuleAddress(govtypes.ModuleName).String(),
+		address.NewBech32Codec("cosmosvaloper"),
+		address.NewBech32Codec("cosmosvalcons"),
+	)
+	require.NoError(tb, keeper.SetParams(ctx, types.DefaultParams()))
+
+	testEnv := &KeeperTestEnvironment{
+		ctx:           ctx,
+		stakingKeeper: keeper,
+		bankKeeper:    bankKeeper,
+		accountKeeper: accountKeeper,
+	}
+	types.RegisterInterfaces(encCfg.InterfaceRegistry)
+	queryHelper := baseapp.NewQueryServerTestHelper(ctx, encCfg.InterfaceRegistry)
+	types.RegisterQueryServer(queryHelper, stakingkeeper.Querier{Keeper: keeper})
+	testEnv.queryClient = types.NewQueryClient(queryHelper)
+	testEnv.msgServer = stakingkeeper.NewMsgServerImpl(keeper)
+	return testEnv
+}
+
+func setupState(b *testing.B, r *rand.Rand, numBonded int) ([]types.Validator, []sdk.ValAddress) {
+	b.Helper()
+	accs := simtypes.RandomAccounts(r, numBonded)
+	initialStake := sdkmath.NewInt(r.Int63n(1000) + 10)
+
+	validators := make([]types.Validator, numBonded)
+	valAddrs := make([]sdk.ValAddress, numBonded)
+
+	for i := 0; i < numBonded; i++ {
+		valAddr := sdk.ValAddress(accs[i].Address)
+		valAddrs[i] = valAddr
+
+		maxCommission := sdkmath.LegacyNewDecWithPrec(int64(simtypes.RandIntBetween(r, 1, 100)), 2)
+		commission := types.NewCommission(
+			simtypes.RandomDecAmount(r, maxCommission),
+			maxCommission,
+			simtypes.RandomDecAmount(r, maxCommission),
+		)
+
+		validator, err := types.NewValidator(valAddr.String(), accs[i].ConsKey.PubKey(), types.Description{})
+		require.NoError(b, err)
+		startStake := sdkmath.NewInt(r.Int63n(1000) + initialStake.Int64())
+		validator.Tokens = startStake.Mul(sdk.DefaultPowerReduction)
+		validator.DelegatorShares = sdkmath.LegacyNewDecFromInt(initialStake)
+		validator.Commission = commission
+		validator.Status = types.Bonded
+		validators[i] = validator
+	}
+	return validators, valAddrs
+}