Skip to content

Commit

Permalink
refactor: migrate LastTotalPower to use collections in x/staking (#17026
Browse files Browse the repository at this point in the history
)
  • Loading branch information
likhita-809 authored Jul 17, 2023
1 parent b68e7f5 commit 1a986f0
Show file tree
Hide file tree
Showing 6 changed files with 24 additions and 38 deletions.
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,8 @@ Ref: https://keepachangelog.com/en/1.0.0/

### API Breaking Changes

* (x/staking) [#17026](https://github.com/cosmos/cosmos-sdk/pull/17026) Use collections for `LastTotalPower`:
* remove `Keeper`: `SetLastTotalPower`, `GetLastTotalPower`
* (staking) [#16959](https://github.com/cosmos/cosmos-sdk/pull/16959) Add validator and consensus address codec as staking keeper arguments.
* (types) [#16272](https://github.com/cosmos/cosmos-sdk/pull/16272) From now the `FeeGranter` in the `FeeTx` interface takes the byte type instead of string.
* (testutil) [#16899](https://github.com/cosmos/cosmos-sdk/pull/16899) The *cli testutil* `QueryBalancesExec` has been removed. Use the gRPC or REST query instead.
Expand Down
4 changes: 2 additions & 2 deletions x/staking/keeper/genesis.go
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ func (k Keeper) InitGenesis(ctx context.Context, data *types.GenesisState) (res
panic(err)
}

if err := k.SetLastTotalPower(ctx, data.LastTotalPower); err != nil {
if err := k.LastTotalPower.Set(ctx, data.LastTotalPower); err != nil {
panic(err)
}

Expand Down Expand Up @@ -238,7 +238,7 @@ func (k Keeper) ExportGenesis(ctx sdk.Context) *types.GenesisState {
panic(err)
}

totalPower, err := k.GetLastTotalPower(ctx)
totalPower, err := k.LastTotalPower.Get(ctx)
if err != nil {
panic(err)
}
Expand Down
45 changes: 14 additions & 31 deletions x/staking/keeper/keeper.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (

abci "github.com/cometbft/cometbft/abci/types"

"cosmossdk.io/collections"
addresscodec "cosmossdk.io/core/address"
storetypes "cosmossdk.io/core/store"
"cosmossdk.io/log"
Expand All @@ -32,6 +33,9 @@ type Keeper struct {
authority string
validatorAddressCodec addresscodec.Codec
consensusAddressCodec addresscodec.Codec

Schema collections.Schema
LastTotalPower collections.Item[math.Int]
}

// NewKeeper creates a new staking Keeper instance
Expand All @@ -44,6 +48,7 @@ func NewKeeper(
validatorAddressCodec addresscodec.Codec,
consensusAddressCodec addresscodec.Codec,
) *Keeper {
sb := collections.NewSchemaBuilder(storeService)
// ensure bonded and not bonded module accounts are set
if addr := ak.GetModuleAddress(types.BondedPoolName); addr == nil {
panic(fmt.Sprintf("%s module account has not been set", types.BondedPoolName))
Expand All @@ -62,7 +67,7 @@ func NewKeeper(
panic("validator and/or consensus address codec are nil")
}

return &Keeper{
k := &Keeper{
storeService: storeService,
cdc: cdc,
authKeeper: ak,
Expand All @@ -71,7 +76,15 @@ func NewKeeper(
authority: authority,
validatorAddressCodec: validatorAddressCodec,
consensusAddressCodec: consensusAddressCodec,
LastTotalPower: collections.NewItem(sb, types.LastTotalPowerKey, "last_total_power", sdk.IntValue),
}

schema, err := sb.Build()
if err != nil {
panic(err)
}
k.Schema = schema
return k
}

// Logger returns a module-specific logger.
Expand Down Expand Up @@ -100,36 +113,6 @@ func (k *Keeper) SetHooks(sh types.StakingHooks) {
k.hooks = sh
}

// GetLastTotalPower loads the last total validator power.
func (k Keeper) GetLastTotalPower(ctx context.Context) (math.Int, error) {
store := k.storeService.OpenKVStore(ctx)
bz, err := store.Get(types.LastTotalPowerKey)
if err != nil {
return math.ZeroInt(), err
}

if bz == nil {
return math.ZeroInt(), nil
}

var power math.Int
if err = power.Unmarshal(bz); err != nil {
return math.ZeroInt(), err
}

return power, nil
}

// SetLastTotalPower sets the last total validator power.
func (k Keeper) SetLastTotalPower(ctx context.Context, power math.Int) error {
store := k.storeService.OpenKVStore(ctx)
bz, err := power.Marshal()
if err != nil {
return err
}
return store.Set(types.LastTotalPowerKey, bz)
}

// GetAuthority returns the x/staking module's authority.
func (k Keeper) GetAuthority() string {
return k.authority
Expand Down
4 changes: 2 additions & 2 deletions x/staking/keeper/keeper_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -103,8 +103,8 @@ func (s *KeeperTestSuite) TestLastTotalPower() {
require := s.Require()

expTotalPower := math.NewInt(10 ^ 9)
require.NoError(keeper.SetLastTotalPower(ctx, expTotalPower))
resTotalPower, err := keeper.GetLastTotalPower(ctx)
require.NoError(keeper.LastTotalPower.Set(ctx, expTotalPower))
resTotalPower, err := keeper.LastTotalPower.Get(ctx)
require.NoError(err)
require.True(expTotalPower.Equal(resTotalPower))
}
Expand Down
2 changes: 1 addition & 1 deletion x/staking/keeper/val_state_change.go
Original file line number Diff line number Diff line change
Expand Up @@ -251,7 +251,7 @@ func (k Keeper) ApplyAndReturnValidatorSetUpdates(ctx context.Context) (updates

// set total power on lookup index if there are any updates
if len(updates) > 0 {
if err = k.SetLastTotalPower(ctx, totalPower); err != nil {
if err = k.LastTotalPower.Set(ctx, totalPower); err != nil {
return nil, err
}
}
Expand Down
5 changes: 3 additions & 2 deletions x/staking/types/keys.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"fmt"
"time"

"cosmossdk.io/collections"
"cosmossdk.io/math"

sdk "github.com/cosmos/cosmos-sdk/types"
Expand All @@ -30,8 +31,8 @@ const (
var (
// Keys for store prefixes
// Last* values are constant during a block.
LastValidatorPowerKey = []byte{0x11} // prefix for each key to a validator index, for bonded validators
LastTotalPowerKey = []byte{0x12} // prefix for the total power
LastValidatorPowerKey = []byte{0x11} // prefix for each key to a validator index, for bonded validators
LastTotalPowerKey = collections.NewPrefix(18) // prefix for the total power

ValidatorsKey = []byte{0x21} // prefix for each key to a validator
ValidatorsByConsAddrKey = []byte{0x22} // prefix for each key to a validator index, by pubkey
Expand Down

0 comments on commit 1a986f0

Please sign in to comment.