Skip to content

Commit

Permalink
refactor(x/staking): Migrate Validators to use Collections (#17123)
Browse files Browse the repository at this point in the history
  • Loading branch information
likhita-809 authored Aug 28, 2023
1 parent a734071 commit 70b3e85
Show file tree
Hide file tree
Showing 7 changed files with 95 additions and 30 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ Ref: https://keepachangelog.com/en/1.0.0/

### API Breaking Changes

* (x/staking) [#17123](https://github.com/cosmos/cosmos-sdk/pull/17123) Use collections for `Validators`
* (x/staking) [#17270](https://github.com/cosmos/cosmos-sdk/pull/17270) Use collections for `UnbondingDelegation`:
* remove from `types`: `GetUBDsKey`
* remove from `Keeper`: `IterateUnbondingDelegations`, `IterateDelegatorUnbondingDelegations`
Expand Down
2 changes: 2 additions & 0 deletions x/staking/keeper/keeper.go
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ type Keeper struct {
Redelegations collections.Map[collections.Triple[[]byte, []byte, []byte], types.Redelegation]
Delegations collections.Map[collections.Pair[sdk.AccAddress, sdk.ValAddress], types.Delegation]
UnbondingIndex collections.Map[uint64, []byte]
Validators collections.Map[[]byte, types.Validator]
UnbondingDelegations collections.Map[collections.Pair[[]byte, []byte], types.UnbondingDelegation]
RedelegationsByValDst collections.Map[collections.Triple[[]byte, []byte, []byte], []byte]
RedelegationsByValSrc collections.Map[collections.Triple[[]byte, []byte, []byte], []byte]
Expand Down Expand Up @@ -146,6 +147,7 @@ func NewKeeper(
),
collections.BytesValue,
),
Validators: collections.NewMap(sb, types.ValidatorsKey, "validators", sdk.LengthPrefixedBytesKey, codec.CollValue[types.Validator](cdc)), // sdk.LengthPrefixedBytesKey is needed to retain state compatibility
UnbondingDelegations: collections.NewMap(
sb, types.UnbondingDelegationKey,
"unbonding_delegation",
Expand Down
89 changes: 78 additions & 11 deletions x/staking/keeper/keeper_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import (
"github.com/cosmos/cosmos-sdk/baseapp"
"github.com/cosmos/cosmos-sdk/codec"
"github.com/cosmos/cosmos-sdk/codec/address"
codectypes "github.com/cosmos/cosmos-sdk/codec/types"
"github.com/cosmos/cosmos-sdk/runtime"
"github.com/cosmos/cosmos-sdk/testutil"
simtestutil "github.com/cosmos/cosmos-sdk/testutil/sims"
Expand Down Expand Up @@ -166,6 +167,20 @@ func getREDsFromValSrcIndexKey(valSrcAddr sdk.ValAddress) []byte {
return append(redelegationByValSrcIndexKey, addresstypes.MustLengthPrefix(valSrcAddr)...)
}

// 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 {
unbondingDelegationKey := []byte{0x32}
return append(append(unbondingDelegationKey, addresstypes.MustLengthPrefix(delAddr)...), addresstypes.MustLengthPrefix(valAddr)...)
}

// getValidatorKey creates the key for the validator with address
// VALUE: staking/Validator
func getValidatorKey(operatorAddr sdk.ValAddress) []byte {
validatorsKey := []byte{0x21}
return append(validatorsKey, addresstypes.MustLengthPrefix(operatorAddr)...)
}

func (s *KeeperTestSuite) TestSrcRedelegationsMigrationToColls() {
s.SetupTest()

Expand Down Expand Up @@ -230,17 +245,6 @@ func (s *KeeperTestSuite) TestDstRedelegationsMigrationToColls() {
s.Require().NoError(err)
}

func TestKeeperTestSuite(t *testing.T) {
suite.Run(t, new(KeeperTestSuite))
}

// 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 {
unbondingDelegationKey := []byte{0x32}
return append(append(unbondingDelegationKey, addresstypes.MustLengthPrefix(delAddr)...), addresstypes.MustLengthPrefix(valAddr)...)
}

func (s *KeeperTestSuite) TestUnbondingDelegationsMigrationToColls() {
s.SetupTest()

Expand Down Expand Up @@ -292,6 +296,69 @@ func (s *KeeperTestSuite) TestUnbondingDelegationsMigrationToColls() {
},
"d03ca412f3f6849b5148a2ca49ac2555f65f90b7fab6a289575ed337f15c0f4b",
)
s.Require().NoError(err)
}

func TestKeeperTestSuite(t *testing.T) {
suite.Run(t, new(KeeperTestSuite))
}

func (s *KeeperTestSuite) TestValidatorsMigrationToColls() {
s.SetupTest()
pkAny, err := codectypes.NewAnyWithValue(PKs[0])
s.Require().NoError(err)

_, valAddrs := createValAddrs(100)

err = testutil.DiffCollectionsMigration(
s.ctx,
s.key,
100,
func(i int64) {
val := stakingtypes.Validator{
OperatorAddress: valAddrs[i].String(),
ConsensusPubkey: pkAny,
Jailed: false,
Status: stakingtypes.Bonded,
Tokens: sdk.DefaultPowerReduction,
DelegatorShares: math.LegacyOneDec(),
Description: stakingtypes.Description{},
UnbondingHeight: int64(0),
UnbondingTime: time.Unix(0, 0).UTC(),
Commission: stakingtypes.NewCommission(math.LegacyZeroDec(), math.LegacyZeroDec(), math.LegacyZeroDec()),
MinSelfDelegation: math.ZeroInt(),
}
valBz := stakingtypes.MustMarshalValidator(s.cdc, &val)
// legacy Set method
s.ctx.KVStore(s.key).Set(getValidatorKey(valAddrs[i]), valBz)
},
"6a8737af6309d53494a601e900832ec27763adefd7fe8ff104477d8130d7405f",
)
s.Require().NoError(err)

err = testutil.DiffCollectionsMigration(
s.ctx,
s.key,
100,
func(i int64) {
val := stakingtypes.Validator{
OperatorAddress: valAddrs[i].String(),
ConsensusPubkey: pkAny,
Jailed: false,
Status: stakingtypes.Bonded,
Tokens: sdk.DefaultPowerReduction,
DelegatorShares: math.LegacyOneDec(),
Description: stakingtypes.Description{},
UnbondingHeight: int64(0),
UnbondingTime: time.Unix(0, 0).UTC(),
Commission: stakingtypes.NewCommission(math.LegacyZeroDec(), math.LegacyZeroDec(), math.LegacyZeroDec()),
MinSelfDelegation: math.ZeroInt(),
}

err := s.stakingKeeper.SetValidator(s.ctx, val)
s.Require().NoError(err)
},
"6a8737af6309d53494a601e900832ec27763adefd7fe8ff104477d8130d7405f",
)
s.Require().NoError(err)
}
21 changes: 8 additions & 13 deletions x/staking/keeper/validator.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,17 +22,14 @@ import (

// GetValidator gets a single validator
func (k Keeper) GetValidator(ctx context.Context, addr sdk.ValAddress) (validator types.Validator, err error) {
store := k.storeService.OpenKVStore(ctx)
value, err := store.Get(types.GetValidatorKey(addr))
validator, err = k.Validators.Get(ctx, addr)
if err != nil {
if errors.Is(err, collections.ErrNotFound) {
return types.Validator{}, types.ErrNoValidatorFound
}
return validator, err
}

if value == nil {
return validator, types.ErrNoValidatorFound
}

return types.UnmarshalValidator(k.cdc, value)
return validator, nil
}

func (k Keeper) mustGetValidator(ctx context.Context, addr sdk.ValAddress) types.Validator {
Expand Down Expand Up @@ -69,13 +66,11 @@ func (k Keeper) mustGetValidatorByConsAddr(ctx context.Context, consAddr sdk.Con

// SetValidator sets the main record holding validator details
func (k Keeper) SetValidator(ctx context.Context, validator types.Validator) error {
store := k.storeService.OpenKVStore(ctx)
bz := types.MustMarshalValidator(k.cdc, &validator)
str, err := k.ValidatorAddressCodec().StringToBytes(validator.GetOperator())
valBz, err := k.ValidatorAddressCodec().StringToBytes(validator.GetOperator())
if err != nil {
return err
}
return store.Set(types.GetValidatorKey(str), bz)
return k.Validators.Set(ctx, sdk.ValAddress(valBz), validator)
}

// SetValidatorByConsAddr sets a validator by conesensus address
Expand Down Expand Up @@ -233,7 +228,7 @@ func (k Keeper) RemoveValidator(ctx context.Context, address sdk.ValAddress) err

// delete the old validator record
store := k.storeService.OpenKVStore(ctx)
if err = store.Delete(types.GetValidatorKey(address)); err != nil {
if err = k.Validators.Remove(ctx, address); err != nil {
return err
}

Expand Down
6 changes: 5 additions & 1 deletion x/staking/migrations/v2/store_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ func TestStoreMigration(t *testing.T) {
{
"ValidatorsKey",
v1.GetValidatorKey(valAddr1),
types.GetValidatorKey(valAddr1),
getValidatorKey(valAddr1),
},
{
"ValidatorsByConsAddrKey",
Expand Down Expand Up @@ -141,6 +141,10 @@ func TestStoreMigration(t *testing.T) {
}
}

func getValidatorKey(operatorAddr sdk.ValAddress) []byte {
return append(types.ValidatorsKey, sdkaddress.MustLengthPrefix(operatorAddr)...)
}

func unbondingKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte {
return append(append(types.UnbondingDelegationKey, sdkaddress.MustLengthPrefix(delAddr)...), sdkaddress.MustLengthPrefix(valAddr)...)
}
4 changes: 0 additions & 4 deletions x/staking/simulation/decoder_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,15 +25,12 @@ func TestDecodeStore(t *testing.T) {
cdc := testutil.MakeTestEncodingConfig().Codec
dec := simulation.NewDecodeStore(cdc)

val, err := types.NewValidator(valAddr1.String(), delPk1, types.NewDescription("test", "test", "test", "test", "test"))
require.NoError(t, err)
oneIntBz, err := math.OneInt().Marshal()
require.NoError(t, err)

kvPairs := kv.Pairs{
Pairs: []kv.Pair{
{Key: types.LastTotalPowerKey, Value: oneIntBz},
{Key: types.GetValidatorKey(valAddr1), Value: cdc.MustMarshal(&val)},
{Key: types.LastValidatorPowerKey, Value: valAddr1.Bytes()},
{Key: []byte{0x99}, Value: []byte{0x99}},
},
Expand All @@ -44,7 +41,6 @@ func TestDecodeStore(t *testing.T) {
expectedLog string
}{
{"LastTotalPower", fmt.Sprintf("%v\n%v", math.OneInt(), math.OneInt())},
{"Validator", fmt.Sprintf("%v\n%v", val, val)},
{"LastValidatorPower/ValidatorsByConsAddr/ValidatorsByPowerIndex", fmt.Sprintf("%v\n%v", valAddr1, valAddr1)},
{"other", ""},
}
Expand Down
2 changes: 1 addition & 1 deletion x/staking/types/keys.go
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ var (
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
ValidatorsKey = collections.NewPrefix(33) // prefix for each key to a validator
ValidatorsByConsAddrKey = collections.NewPrefix(34) // prefix for each key to a validator index, by pubkey
ValidatorsByPowerIndexKey = []byte{0x23} // prefix for each key to a validator index, sorted by power

Expand Down

0 comments on commit 70b3e85

Please sign in to comment.