Skip to content

Commit

Permalink
refactor(x/staking): migrate redelegation key to use collections (#17315
Browse files Browse the repository at this point in the history 
)
  • Loading branch information
@atheeshp
atheeshp committed Aug 17, 2023
1 parent cd02384 commit 09db699
Show file tree
Hide file tree
Showing 12 changed files with 113 additions and 137 deletions.
4 changes: 3 additions & 1 deletion CHANGELOG.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,10 +71,12 @@ Ref: https://keepachangelog.com/en/1.0.0/

### API Breaking Changes

* (x/staking) [#17315](https://github.com/cosmos/cosmos-sdk/pull/17044) Use collections for `RedelegationKey`:
* remove from `keeper`: `GetRedelegation`
* (types) `module.BeginBlockAppModule` has been replaced by Core API `appmodule.HasBeginBlocker`.
* (types) `module.EndBlockAppModule` has been replaced by Core API `appmodule.HasEndBlocker` or `module.HasABCIEndBlock` when needing validator updates.
* (types) [#17358](https://github.com/cosmos/cosmos-sdk/pull/17358) Remove deprecated `sdk.Handler`, use `baseapp.MsgServiceHandler` instead.
* (x/slashing) [17044](https://github.com/cosmos/cosmos-sdk/pull/17044) Use collections for `AddrPubkeyRelation`:
* (x/slashing) [#17044](https://github.com/cosmos/cosmos-sdk/pull/17044) Use collections for `AddrPubkeyRelation`:
* remove from `types`: `AddrPubkeyRelationKey`
* remove from `Keeper`: `AddPubkey`
* (x/staking) [#17260](https://github.com/cosmos/cosmos-sdk/pull/17260) Use collections for `DelegationKey`:
Expand Down
2 changes: 1 addition & 1 deletion tests/integration/staking/keeper/grpc_query_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -820,7 +820,7 @@ func TestGRPCQueryRedelegations(t *testing.T) {
assert.NilError(t, err)
applyValidatorSetUpdates(t, ctx, f.stakingKeeper, -1)

redel, found := f.stakingKeeper.GetRedelegation(ctx, addrAcc1, val1bz, val2bz)
redel, found := f.stakingKeeper.Redelegations.Get(ctx, collections.Join3(addrAcc1.Bytes(), valAddrs[0].Bytes(), valAddrs[1].Bytes()))
assert.Assert(t, found)

var req *types.QueryRedelegationsRequest
Expand Down
12 changes: 6 additions & 6 deletions tests/integration/staking/keeper/slash_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,7 @@ func TestSlashRedelegation(t *testing.T) {
slashAmount, err = f.stakingKeeper.SlashRedelegation(f.sdkCtx, validator, rd, 0, fraction)
assert.NilError(t, err)
assert.Assert(t, slashAmount.Equal(math.NewInt(5)))
rd, found = f.stakingKeeper.GetRedelegation(f.sdkCtx, addrDels[0], addrVals[0], addrVals[1])
rd, found = f.stakingKeeper.Redelegations.Get(f.sdkCtx, collections.Join3(addrDels[0].Bytes(), addrVals[0].Bytes(), addrVals[1].Bytes()))
assert.Assert(t, found)
assert.Assert(t, len(rd.Entries) == 1)

Expand Down Expand Up @@ -438,7 +438,7 @@ func TestSlashWithRedelegation(t *testing.T) {
oldBonded = f.bankKeeper.GetBalance(f.sdkCtx, bondedPool.GetAddress(), bondDenom).Amount

// read updating redelegation
rd, found = f.stakingKeeper.GetRedelegation(f.sdkCtx, addrDels[0], addrVals[0], addrVals[1])
rd, found = f.stakingKeeper.Redelegations.Get(f.sdkCtx, collections.Join3(addrDels[0].Bytes(), addrVals[0].Bytes(), addrVals[1].Bytes()))
assert.Assert(t, found)
assert.Assert(t, len(rd.Entries) == 1)
// read updated validator
Expand Down Expand Up @@ -475,7 +475,7 @@ func TestSlashWithRedelegation(t *testing.T) {
oldBonded = f.bankKeeper.GetBalance(f.sdkCtx, bondedPool.GetAddress(), bondDenom).Amount

// read updating redelegation
rd, found = f.stakingKeeper.GetRedelegation(f.sdkCtx, addrDels[0], addrVals[0], addrVals[1])
rd, found = f.stakingKeeper.Redelegations.Get(f.sdkCtx, collections.Join3(addrDels[0].Bytes(), addrVals[0].Bytes(), addrVals[1].Bytes()))
assert.Assert(t, found)
assert.Assert(t, len(rd.Entries) == 1)
// read updated validator
Expand Down Expand Up @@ -506,7 +506,7 @@ func TestSlashWithRedelegation(t *testing.T) {
oldBonded = f.bankKeeper.GetBalance(f.sdkCtx, bondedPool.GetAddress(), bondDenom).Amount

// read updating redelegation
rd, found = f.stakingKeeper.GetRedelegation(f.sdkCtx, addrDels[0], addrVals[0], addrVals[1])
rd, found = f.stakingKeeper.Redelegations.Get(f.sdkCtx, collections.Join3(addrDels[0].Bytes(), addrVals[0].Bytes(), addrVals[1].Bytes()))
assert.Assert(t, found)
assert.Assert(t, len(rd.Entries) == 1)
// apply TM updates
Expand Down Expand Up @@ -536,7 +536,7 @@ func TestSlashWithRedelegation(t *testing.T) {
assert.Assert(math.IntEq(t, oldNotBonded, notBondedPoolBalance))

// read updating redelegation
rd, found = f.stakingKeeper.GetRedelegation(f.sdkCtx, addrDels[0], addrVals[0], addrVals[1])
rd, found = f.stakingKeeper.Redelegations.Get(f.sdkCtx, collections.Join3(addrDels[0].Bytes(), addrVals[0].Bytes(), addrVals[1].Bytes()))
assert.Assert(t, found)
assert.Assert(t, len(rd.Entries) == 1)
// read updated validator
Expand Down Expand Up @@ -607,7 +607,7 @@ func TestSlashBoth(t *testing.T) {
assert.Assert(math.IntEq(t, oldNotBonded.Sub(burnedNotBondedAmount), notBondedPoolBalance))

// read updating redelegation
rdA, found = f.stakingKeeper.GetRedelegation(f.sdkCtx, addrDels[0], addrVals[0], addrVals[1])
rdA, found = f.stakingKeeper.Redelegations.Get(f.sdkCtx, collections.Join3(addrDels[0].Bytes(), addrVals[0].Bytes(), addrVals[1].Bytes()))
assert.Assert(t, found)
assert.Assert(t, len(rdA.Entries) == 1)
// read updated validator
Expand Down
103 changes: 30 additions & 73 deletions x/staking/keeper/delegation.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -291,27 +291,6 @@ func (k Keeper) IterateDelegatorDelegations(ctx context.Context, delegator sdk.A
return nil
}

// IterateDelegatorRedelegations iterates through one delegator's redelegations.
func (k Keeper) IterateDelegatorRedelegations(ctx context.Context, delegator sdk.AccAddress, cb func(red types.Redelegation) (stop bool)) error {
store := k.storeService.OpenKVStore(ctx)
delegatorPrefixKey := types.GetREDsKey(delegator)
iterator, err := store.Iterator(delegatorPrefixKey, storetypes.PrefixEndBytes(delegatorPrefixKey))
if err != nil {
return err
}

for ; iterator.Valid(); iterator.Next() {
red, err := types.UnmarshalRED(k.cdc, iterator.Value())
if err != nil {
return err
}
if cb(red) {
break
}
}
return nil
}

// HasMaxUnbondingDelegationEntries checks if unbonding delegation has maximum number of entries.
func (k Keeper) HasMaxUnbondingDelegationEntries(ctx context.Context, delegatorAddr sdk.AccAddress, validatorAddr sdk.ValAddress) (bool, error) {
ubd, err := k.GetUnbondingDelegation(ctx, delegatorAddr, validatorAddr)
Expand Down Expand Up @@ -500,41 +479,23 @@ func (k Keeper) DequeueAllMatureUBDQueue(ctx context.Context, currTime time.Time
func (k Keeper) GetRedelegations(ctx context.Context, delegator sdk.AccAddress, maxRetrieve uint16) (redelegations []types.Redelegation, err error) {
redelegations = make([]types.Redelegation, maxRetrieve)

store := k.storeService.OpenKVStore(ctx)
delegatorPrefixKey := types.GetREDsKey(delegator)
iterator, err := store.Iterator(delegatorPrefixKey, storetypes.PrefixEndBytes(delegatorPrefixKey))
if err != nil {
return nil, err
}

i := 0
for ; iterator.Valid() && i < int(maxRetrieve); iterator.Next() {
redelegation, err := types.UnmarshalRED(k.cdc, iterator.Value())
if err != nil {
return nil, err
rng := collections.NewPrefixedTripleRange[[]byte, []byte, []byte](delegator)
err = k.Redelegations.Walk(ctx, rng, func(key collections.Triple[[]byte, []byte, []byte], redelegation types.Redelegation) (stop bool, err error) {
if i >= int(maxRetrieve) {
return true, nil
}

redelegations[i] = redelegation
i++
}

return redelegations[:i], nil // trim if the array length < maxRetrieve
}

// GetRedelegation returns a redelegation.
func (k Keeper) GetRedelegation(ctx context.Context, delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress) (red types.Redelegation, err error) {
store := k.storeService.OpenKVStore(ctx)
key := types.GetREDKey(delAddr, valSrcAddr, valDstAddr)
return false, nil
})

value, err := store.Get(key)
if err != nil {
return red, err
}

if value == nil {
return red, types.ErrNoRedelegation
return nil, err
}

return types.UnmarshalRED(k.cdc, value)
return redelegations[:i], nil // trim if the array length < maxRetrieve
}

// GetRedelegationsFromSrcValidator returns all redelegations from a particular
Expand Down Expand Up @@ -578,9 +539,9 @@ func (k Keeper) HasReceivingRedelegation(ctx context.Context, delAddr sdk.AccAdd

// HasMaxRedelegationEntries checks if the redelegation entries reached maximum limit.
func (k Keeper) HasMaxRedelegationEntries(ctx context.Context, delegatorAddr sdk.AccAddress, validatorSrcAddr, validatorDstAddr sdk.ValAddress) (bool, error) {
red, err := k.GetRedelegation(ctx, delegatorAddr, validatorSrcAddr, validatorDstAddr)
red, err := k.Redelegations.Get(ctx, collections.Join3(delegatorAddr.Bytes(), validatorSrcAddr.Bytes(), validatorDstAddr.Bytes()))
if err != nil {
if err == types.ErrNoRedelegation {
if errors.Is(err, collections.ErrNotFound) {
return false, nil
}

Expand All @@ -602,7 +563,6 @@ func (k Keeper) SetRedelegation(ctx context.Context, red types.Redelegation) err
}

store := k.storeService.OpenKVStore(ctx)
bz := types.MustMarshalRED(k.cdc, red)
valSrcAddr, err := k.validatorAddressCodec.StringToBytes(red.ValidatorSrcAddress)
if err != nil {
return err
Expand All @@ -611,8 +571,8 @@ func (k Keeper) SetRedelegation(ctx context.Context, red types.Redelegation) err
if err != nil {
return err
}
key := types.GetREDKey(delegatorAddress, valSrcAddr, valDestAddr)
if err = store.Set(key, bz); err != nil {

if err = k.Redelegations.Set(ctx, collections.Join3(delegatorAddress, valSrcAddr, valDestAddr), red); err != nil {
return err
}

Expand All @@ -636,10 +596,10 @@ func (k Keeper) SetRedelegationEntry(ctx context.Context,
return types.Redelegation{}, err
}

red, err := k.GetRedelegation(ctx, delegatorAddr, validatorSrcAddr, validatorDstAddr)
red, err := k.Redelegations.Get(ctx, collections.Join3(delegatorAddr.Bytes(), validatorSrcAddr.Bytes(), validatorDstAddr.Bytes()))
if err == nil {
red.AddEntry(creationHeight, minTime, balance, sharesDst, id)
} else if errors.Is(err, types.ErrNoRedelegation) {
} else if errors.Is(err, collections.ErrNotFound) {
red = types.NewRedelegation(delegatorAddr, validatorSrcAddr,
validatorDstAddr, creationHeight, minTime, balance, sharesDst, id, k.validatorAddressCodec, k.authKeeper.AddressCodec())
} else {
Expand All @@ -665,22 +625,19 @@ func (k Keeper) SetRedelegationEntry(ctx context.Context,

// IterateRedelegations iterates through all redelegations.
func (k Keeper) IterateRedelegations(ctx context.Context, fn func(index int64, red types.Redelegation) (stop bool)) error {
store := k.storeService.OpenKVStore(ctx)
iterator, err := store.Iterator(types.RedelegationKey, storetypes.PrefixEndBytes(types.RedelegationKey))
if err != nil {
return err
}
defer iterator.Close()
var i int64
err := k.Redelegations.Walk(ctx, nil,
func(key collections.Triple[[]byte, []byte, []byte], red types.Redelegation) (bool, error) {
if stop := fn(i, red); stop {
return true, nil
}
i++

for i := int64(0); iterator.Valid(); iterator.Next() {
red, err := types.UnmarshalRED(k.cdc, iterator.Value())
if err != nil {
return err
}
if stop := fn(i, red); stop {
break
}
i++
return false, nil
},
)
if err != nil && !errors.Is(err, collections.ErrInvalidIterator) {
return err
}

return nil
Expand All @@ -702,8 +659,8 @@ func (k Keeper) RemoveRedelegation(ctx context.Context, red types.Redelegation)
if err != nil {
return err
}
redKey := types.GetREDKey(delegatorAddress, valSrcAddr, valDestAddr)
if err = store.Delete(redKey); err != nil {

if err = k.Redelegations.Remove(ctx, collections.Join3(delegatorAddress, valSrcAddr, valDestAddr)); err != nil {
return err
}

Expand Down Expand Up @@ -1250,7 +1207,7 @@ func (k Keeper) BeginRedelegation(
func (k Keeper) CompleteRedelegation(
ctx context.Context, delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress,
) (sdk.Coins, error) {
red, err := k.GetRedelegation(ctx, delAddr, valSrcAddr, valDstAddr)
red, err := k.Redelegations.Get(ctx, collections.Join3(delAddr.Bytes(), valSrcAddr.Bytes(), valDstAddr.Bytes()))
if err != nil {
return nil, err
}
Expand Down
18 changes: 9 additions & 9 deletions x/staking/keeper/delegation_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -682,7 +682,7 @@ func (s *KeeperTestSuite) TestGetRedelegationsFromSrcValidator() {
// set and retrieve a record
err := keeper.SetRedelegation(ctx, rd)
require.NoError(err)
resBond, err := keeper.GetRedelegation(ctx, addrDels[0], addrVals[0], addrVals[1])
resBond, err := keeper.Redelegations.Get(ctx, collections.Join3(addrDels[0].Bytes(), addrVals[0].Bytes(), addrVals[1].Bytes()))
require.NoError(err)

// get the redelegations one time
Expand Down Expand Up @@ -717,7 +717,7 @@ func (s *KeeperTestSuite) TestRedelegation() {
// set and retrieve a record
err = keeper.SetRedelegation(ctx, rd)
require.NoError(err)
resRed, err := keeper.GetRedelegation(ctx, addrDels[0], addrVals[0], addrVals[1])
resRed, err := keeper.Redelegations.Get(ctx, collections.Join3(addrDels[0].Bytes(), addrVals[0].Bytes(), addrVals[1].Bytes()))
require.NoError(err)

redelegations, err := keeper.GetRedelegationsFromSrcValidator(ctx, addrVals[0])
Expand All @@ -735,7 +735,7 @@ func (s *KeeperTestSuite) TestRedelegation() {
require.Equal(1, len(redelegations))
require.Equal(redelegations[0], resRed)

// check if has the redelegation
// check if it has the redelegation
has, err = keeper.HasReceivingRedelegation(ctx, addrDels[0], addrVals[1])
require.NoError(err)
require.True(has)
Expand All @@ -745,7 +745,7 @@ func (s *KeeperTestSuite) TestRedelegation() {
err = keeper.SetRedelegation(ctx, rd)
require.NoError(err)

resRed, err = keeper.GetRedelegation(ctx, addrDels[0], addrVals[0], addrVals[1])
resRed, err = keeper.Redelegations.Get(ctx, collections.Join3(addrDels[0].Bytes(), addrVals[0].Bytes(), addrVals[1].Bytes()))
require.NoError(err)
require.Equal(rd, resRed)

Expand All @@ -762,8 +762,8 @@ func (s *KeeperTestSuite) TestRedelegation() {
// delete a record
err = keeper.RemoveRedelegation(ctx, rd)
require.NoError(err)
_, err = keeper.GetRedelegation(ctx, addrDels[0], addrVals[0], addrVals[1])
require.ErrorIs(err, stakingtypes.ErrNoRedelegation)
_, err = keeper.Redelegations.Get(ctx, collections.Join3(addrDels[0].Bytes(), addrVals[0].Bytes(), addrVals[1].Bytes()))
require.ErrorIs(err, collections.ErrNotFound)

redelegations, err = keeper.GetRedelegations(ctx, addrDels[0], 5)
require.NoError(err)
Expand Down Expand Up @@ -976,7 +976,7 @@ func (s *KeeperTestSuite) TestRedelegateFromUnbondingValidator() {
require.NoError(err)

// retrieve the unbonding delegation
ubd, err := keeper.GetRedelegation(ctx, addrDels[1], addrVals[0], addrVals[1])
ubd, err := keeper.Redelegations.Get(ctx, collections.Join3(addrDels[1].Bytes(), addrVals[0].Bytes(), addrVals[1].Bytes()))
require.NoError(err)
require.Len(ubd.Entries, 1)
require.Equal(blockHeight, ubd.Entries[0].CreationHeight)
Expand Down Expand Up @@ -1050,8 +1050,8 @@ func (s *KeeperTestSuite) TestRedelegateFromUnbondedValidator() {
require.NoError(err)

// no red should have been found
red, err := keeper.GetRedelegation(ctx, addrDels[0], addrVals[0], addrVals[1])
require.ErrorIs(err, stakingtypes.ErrNoRedelegation, "%v", red)
red, err := keeper.Redelegations.Get(ctx, collections.Join3(addrDels[0].Bytes(), addrVals[0].Bytes(), addrVals[1].Bytes()))
require.ErrorIs(err, collections.ErrNotFound, "%v", red)
}

func (s *KeeperTestSuite) TestUnbondingDelegationAddEntry() {
Expand Down
21 changes: 9 additions & 12 deletions x/staking/keeper/grpc_query.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -412,7 +412,7 @@ func (k Querier) Redelegations(ctx context.Context, req *types.QueryRedelegation
case req.DelegatorAddr == "" && req.SrcValidatorAddr != "" && req.DstValidatorAddr == "":
redels, pageRes, err = queryRedelegationsFromSrcValidator(store, k, req)
default:
redels, pageRes, err = queryAllRedelegations(store, k, req)
redels, pageRes, err = queryAllRedelegations(ctx, store, k, req)
}
if err != nil {
return nil, status.Error(codes.Internal, err.Error())
Expand Down Expand Up @@ -505,7 +505,7 @@ func queryRedelegation(ctx context.Context, k Querier, req *types.QueryRedelegat
return nil, err
}

redel, err := k.GetRedelegation(ctx, delAddr, srcValAddr, dstValAddr)
redel, err := k.Keeper.Redelegations.Get(ctx, collections.Join3(delAddr, srcValAddr, dstValAddr))
if err != nil {
return nil, status.Errorf(
codes.NotFound,
Expand Down Expand Up @@ -539,21 +539,18 @@ func queryRedelegationsFromSrcValidator(store storetypes.KVStore, k Querier, req
return redels, res, err
}

func queryAllRedelegations(store storetypes.KVStore, k Querier, req *types.QueryRedelegationsRequest) (redels types.Redelegations, res *query.PageResponse, err error) {
func queryAllRedelegations(ctx context.Context, store storetypes.KVStore, k Querier, req *types.QueryRedelegationsRequest) (redels types.Redelegations, res *query.PageResponse, err error) {
delAddr, err := k.authKeeper.AddressCodec().StringToBytes(req.DelegatorAddr)
if err != nil {
return nil, nil, err
}

redStore := prefix.NewStore(store, types.GetREDsKey(delAddr))
res, err = query.Paginate(redStore, req.Pagination, func(key, value []byte) error {
redelegation, err := types.UnmarshalRED(k.cdc, value)
if err != nil {
return err
}
redels = append(redels, redelegation)
return nil
})
redels, res, err = query.CollectionPaginate(ctx, k.Keeper.Redelegations, req.Pagination, func(_ collections.Triple[[]byte, []byte, []byte], red types.Redelegation) (types.Redelegation, error) {
return red, nil
}, query.WithCollectionPaginationTriplePrefix[[]byte, []byte, []byte](delAddr))
if err != nil {
return nil, nil, err
}

return redels, res, err
}
Expand Down
13 changes: 12 additions & 1 deletion x/staking/keeper/keeper.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ type Keeper struct {
UnbondingID collections.Sequence
ValidatorByConsensusAddress collections.Map[sdk.ConsAddress, sdk.ValAddress]
UnbondingType collections.Map[uint64, uint64]
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]
}
Expand Down Expand Up @@ -107,7 +108,17 @@ func NewKeeper(
sdk.LengthPrefixedAddressKey(sdk.ConsAddressKey), // nolint: staticcheck // sdk.LengthPrefixedAddressKey is needed to retain state compatibility
collcodec.KeyToValueCodec(sdk.ValAddressKey),
),
UnbondingType: collections.NewMap(sb, types.UnbondingTypeKey, "unbonding_type", collections.Uint64Key, collections.Uint64Value),
UnbondingType: collections.NewMap(sb, types.UnbondingTypeKey, "unbonding_type", collections.Uint64Key, collections.Uint64Value),
Redelegations: collections.NewMap(
sb, types.RedelegationKey,
"redelegations",
collections.TripleKeyCodec(
collections.BytesKey,
collections.BytesKey,
sdk.LengthPrefixedBytesKey, // sdk.LengthPrefixedBytesKey is needed to retain state compatibility
),
codec.CollValue[types.Redelegation](cdc),
),
UnbondingIndex: collections.NewMap(sb, types.UnbondingIndexKey, "unbonding_index", collections.Uint64Key, collections.BytesValue),
}

Expand Down
Loading

0 comments on commit 09db699

Please sign in to comment.