From 257d3e5e7cc1c6c6fc9a3c29bc95a3cce8e485ee Mon Sep 17 00:00:00 2001 From: atheeshp <59333759+atheeshp@users.noreply.github.com> Date: Fri, 28 Jul 2023 14:16:12 +0530 Subject: [PATCH] refactor(x/staking): migrate delegations by validator to use collections (#17155) --- CHANGELOG.md | 4 +- .../staking/keeper/validator_bench_test.go | 22 +++++----- x/staking/keeper/delegation.go | 32 +++++++-------- x/staking/keeper/grpc_query.go | 25 ++++++----- x/staking/keeper/keeper.go | 15 +++++-- x/staking/migrations/v5/keys.go | 30 ++++++++++++++ x/staking/migrations/v5/migrations_test.go | 2 +- x/staking/types/keys.go | 41 ------------------- 8 files changed, 85 insertions(+), 86 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7c072c8303..932f3bff54 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -39,7 +39,9 @@ Ref: https://keepachangelog.com/en/1.0.0/ ## [Unreleased] ### Features + * (x/bank) [#16795](https://github.com/cosmos/cosmos-sdk/pull/16852) Add `DenomMetadataByQueryString` query in bank module to support metadata query by query string. + ### Improvements * (x/group, x/gov) [#17109](https://github.com/cosmos/cosmos-sdk/pull/17109) Let proposal summary be 40x longer than metadata limit. @@ -53,7 +55,7 @@ Ref: https://keepachangelog.com/en/1.0.0/ ### API Breaking Changes -* (x/slashing) [17063](https://github.com/cosmos/cosmos-sdk/pull/17063) Use collections for `HistoricalInfo`: +* (x/staking) [17063](https://github.com/cosmos/cosmos-sdk/pull/17063) Use collections for `HistoricalInfo`: * remove `Keeper`: `GetHistoricalInfo`, `SetHistoricalInfo`, * (x/staking) [17062](https://github.com/cosmos/cosmos-sdk/pull/17062) Use collections for `ValidatorUpdates`: * remove `Keeper`: `SetValidatorUpdates`, `GetValidatorUpdates` diff --git a/tests/integration/staking/keeper/validator_bench_test.go b/tests/integration/staking/keeper/validator_bench_test.go index 06bca8ee87..70f38055f2 100644 --- a/tests/integration/staking/keeper/validator_bench_test.go +++ b/tests/integration/staking/keeper/validator_bench_test.go @@ -2,9 +2,11 @@ package keeper_test import ( "bytes" + "errors" "fmt" "testing" + "cosmossdk.io/collections" "cosmossdk.io/math" storetypes "cosmossdk.io/store/types" @@ -155,29 +157,27 @@ func updateValidatorDelegations(f *fixture, existingValAddr, newValAddr sdk.ValA store := f.sdkCtx.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) - } + rng := collections.NewPrefixedPairRange[sdk.ValAddress, sdk.AccAddress](existingValAddr) + err := k.DelegationsByValidator.Walk(f.sdkCtx, rng, func(key collections.Pair[sdk.ValAddress, sdk.AccAddress], _ []byte) (stop bool, err error) { + valAddr, delAddr := key.K1(), key.K2() bz := store.Get(types.GetDelegationKey(delAddr, valAddr)) delegation := types.MustUnmarshalDelegation(cdc, bz) // remove old operator addr from delegation if err := k.RemoveDelegation(f.sdkCtx, delegation); err != nil { - panic(err) + return true, err } delegation.ValidatorAddress = newValAddr.String() // add with new operator addr if err := k.SetDelegation(f.sdkCtx, delegation); err != nil { - panic(err) + return true, err } + return false, nil + }) + if err != nil && !errors.Is(err, collections.ErrInvalidIterator) { + panic(err) } } diff --git a/x/staking/keeper/delegation.go b/x/staking/keeper/delegation.go index 7d6210eb7c..f7761c9f91 100644 --- a/x/staking/keeper/delegation.go +++ b/x/staking/keeper/delegation.go @@ -7,6 +7,7 @@ import ( "fmt" "time" + "cosmossdk.io/collections" corestore "cosmossdk.io/core/store" errorsmod "cosmossdk.io/errors" "cosmossdk.io/math" @@ -65,32 +66,29 @@ func (k Keeper) GetAllDelegations(ctx context.Context) (delegations []types.Dele // GetValidatorDelegations returns all delegations to a specific validator. // Useful for querier. -func (k Keeper) GetValidatorDelegations(ctx context.Context, valAddr sdk.ValAddress) (delegations []types.Delegation, err error) { +func (k Keeper) GetValidatorDelegations(ctx context.Context, valAddr sdk.ValAddress) ([]types.Delegation, error) { store := k.storeService.OpenKVStore(ctx) - prefix := types.GetDelegationsByValPrefixKey(valAddr) - iterator, err := store.Iterator(prefix, storetypes.PrefixEndBytes(prefix)) - if err != nil { - return delegations, err - } - defer iterator.Close() - for ; iterator.Valid(); iterator.Next() { + var delegations []types.Delegation + rng := collections.NewPrefixedPairRange[sdk.ValAddress, sdk.AccAddress](valAddr) + err := k.DelegationsByValidator.Walk(ctx, rng, func(key collections.Pair[sdk.ValAddress, sdk.AccAddress], _ []byte) (stop bool, err error) { var delegation types.Delegation - valAddr, delAddr, err := types.ParseDelegationsByValKey(iterator.Key()) - if err != nil { - return delegations, err - } - + valAddr, delAddr := key.K1(), key.K2() bz, err := store.Get(types.GetDelegationKey(delAddr, valAddr)) if err != nil { - return delegations, err + return true, err } if err := k.cdc.Unmarshal(bz, &delegation); err != nil { - return delegations, err + return true, err } delegations = append(delegations, delegation) + + return false, nil + }) + if err != nil && !errors.Is(err, collections.ErrInvalidIterator) { + return nil, err } return delegations, nil @@ -142,7 +140,7 @@ func (k Keeper) SetDelegation(ctx context.Context, delegation types.Delegation) } // set the delegation in validator delegator index - return store.Set(types.GetDelegationsByValKey(valAddr, delegatorAddress), []byte{}) + return k.DelegationsByValidator.Set(ctx, collections.Join(sdk.ValAddress(valAddr), sdk.AccAddress(delegatorAddress)), []byte{}) } // RemoveDelegation removes a delegation @@ -168,7 +166,7 @@ func (k Keeper) RemoveDelegation(ctx context.Context, delegation types.Delegatio return err } - return store.Delete(types.GetDelegationsByValKey(valAddr, delegatorAddress)) + return k.DelegationsByValidator.Remove(ctx, collections.Join(sdk.ValAddress(valAddr), sdk.AccAddress(delegatorAddress))) } // GetUnbondingDelegations returns a given amount of all the delegator unbonding-delegations. diff --git a/x/staking/keeper/grpc_query.go b/x/staking/keeper/grpc_query.go index e80b5de859..39f628dee8 100644 --- a/x/staking/keeper/grpc_query.go +++ b/x/staking/keeper/grpc_query.go @@ -7,6 +7,7 @@ import ( "google.golang.org/grpc/codes" "google.golang.org/grpc/status" + "cosmossdk.io/collections" "cosmossdk.io/store/prefix" storetypes "cosmossdk.io/store/types" @@ -101,24 +102,26 @@ func (k Querier) ValidatorDelegations(ctx context.Context, req *types.QueryValid } store := runtime.KVStoreAdapter(k.storeService.OpenKVStore(ctx)) - delStore := prefix.NewStore(store, types.GetDelegationsByValPrefixKey(valAddr)) 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 - } + dels, pageRes, err = query.CollectionPaginate(ctx, k.DelegationsByValidator, + req.Pagination, func(key collections.Pair[sdk.ValAddress, sdk.AccAddress], _ []byte) (types.Delegation, error) { + valAddr, delAddr := key.K1(), key.K2() + var delegation types.Delegation + + bz := store.Get(types.GetDelegationKey(delAddr, valAddr)) + err = k.cdc.Unmarshal(bz, &delegation) + if err != nil { + return types.Delegation{}, err + } + + return delegation, nil + }, query.WithCollectionPaginationPairPrefix[sdk.ValAddress, sdk.AccAddress](valAddr)) - dels = append(dels, delegation) - return nil - }) if err != nil { delegations, pageResponse, err := k.getValidatorDelegationsLegacy(ctx, req) if err != nil { diff --git a/x/staking/keeper/keeper.go b/x/staking/keeper/keeper.go index 8154b1b532..64ad687d1f 100644 --- a/x/staking/keeper/keeper.go +++ b/x/staking/keeper/keeper.go @@ -32,10 +32,11 @@ type Keeper struct { validatorAddressCodec addresscodec.Codec consensusAddressCodec addresscodec.Codec - Schema collections.Schema - HistoricalInfo collections.Map[uint64, types.HistoricalInfo] - LastTotalPower collections.Item[math.Int] - ValidatorUpdates collections.Item[types.ValidatorUpdates] + Schema collections.Schema + HistoricalInfo collections.Map[uint64, types.HistoricalInfo] + LastTotalPower collections.Item[math.Int] + ValidatorUpdates collections.Item[types.ValidatorUpdates] + DelegationsByValidator collections.Map[collections.Pair[sdk.ValAddress, sdk.AccAddress], []byte] } // NewKeeper creates a new staking Keeper instance @@ -79,6 +80,12 @@ func NewKeeper( LastTotalPower: collections.NewItem(sb, types.LastTotalPowerKey, "last_total_power", sdk.IntValue), HistoricalInfo: collections.NewMap(sb, types.HistoricalInfoKey, "historical_info", collections.Uint64Key, codec.CollValue[types.HistoricalInfo](cdc)), ValidatorUpdates: collections.NewItem(sb, types.ValidatorUpdatesKey, "validator_updates", codec.CollValue[types.ValidatorUpdates](cdc)), + DelegationsByValidator: collections.NewMap( + sb, types.DelegationByValIndexKey, + "delegations_by_validator", + collections.PairKeyCodec(sdk.LengthPrefixedAddressKey(sdk.ValAddressKey), sdk.AccAddressKey), // nolint: staticcheck // sdk.LengthPrefixedAddressKey is needed to retain state compatibility + collections.BytesValue, + ), } schema, err := sb.Build() diff --git a/x/staking/migrations/v5/keys.go b/x/staking/migrations/v5/keys.go index 5c11c8735c..69e763ce61 100644 --- a/x/staking/migrations/v5/keys.go +++ b/x/staking/migrations/v5/keys.go @@ -71,3 +71,33 @@ func GetDelegationsByValPrefixKey(valAddr sdk.ValAddress) []byte { func GetDelegationsByValKey(valAddr sdk.ValAddress, delAddr sdk.AccAddress) []byte { return append(GetDelegationsByValPrefixKey(valAddr), delAddr...) } + +// ParseDelegationsByValKey parses given key and returns validator, delegator address bytes +func ParseDelegationsByValKey(bz []byte) (sdk.ValAddress, sdk.AccAddress, error) { + prefixLength := len(DelegationByValIndexKey) + if prefix := bz[:prefixLength]; !bytes.Equal(prefix, DelegationByValIndexKey) { + return nil, nil, fmt.Errorf("invalid prefix; expected: %X, got: %x", DelegationByValIndexKey, prefix) + } + + bz = bz[prefixLength:] // remove the prefix byte + if len(bz) == 0 { + return nil, nil, fmt.Errorf("no bytes left to parse: %X", bz) + } + + valAddrLen := bz[0] + bz = bz[1:] // remove the length byte of validator address. + if len(bz) == 0 { + return nil, nil, fmt.Errorf("no bytes left to parse validator address: %X", bz) + } + + val := bz[0:int(valAddrLen)] + + bz = bz[int(valAddrLen):] // remove the delegator bytes + if len(bz) == 0 { + return nil, nil, fmt.Errorf("no bytes left to parse delegator address: %X", bz) + } + + del := bz + + return val, del, nil +} diff --git a/x/staking/migrations/v5/migrations_test.go b/x/staking/migrations/v5/migrations_test.go index ef9ca6e8d7..c27180eb65 100644 --- a/x/staking/migrations/v5/migrations_test.go +++ b/x/staking/migrations/v5/migrations_test.go @@ -115,7 +115,7 @@ func getValDelegations(ctx sdk.Context, cdc codec.Codec, storeKey storetypes.Sto iterator := storetypes.KVStorePrefixIterator(store, v5.GetDelegationsByValPrefixKey(valAddr)) for ; iterator.Valid(); iterator.Next() { var delegation stakingtypes.Delegation - valAddr, delAddr, err := stakingtypes.ParseDelegationsByValKey(iterator.Key()) + valAddr, delAddr, err := v5.ParseDelegationsByValKey(iterator.Key()) if err != nil { panic(err) } diff --git a/x/staking/types/keys.go b/x/staking/types/keys.go index 8de7195ee6..067073c1bd 100644 --- a/x/staking/types/keys.go +++ b/x/staking/types/keys.go @@ -215,47 +215,6 @@ func GetDelegationKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte { return append(GetDelegationsKey(delAddr), address.MustLengthPrefix(valAddr)...) } -// GetDelegationsByValKey creates the key for delegations by validator address -// VALUE: staking/Delegation -func GetDelegationsByValKey(valAddr sdk.ValAddress, delAddr sdk.AccAddress) []byte { - return append(GetDelegationsByValPrefixKey(valAddr), delAddr...) -} - -// GetDelegationsByValPrefixKey builds a prefix key bytes with the given validator address bytes. -func GetDelegationsByValPrefixKey(valAddr sdk.ValAddress) []byte { - return append(DelegationByValIndexKey, address.MustLengthPrefix(valAddr)...) -} - -// ParseDelegationsByValKey parses given key and returns validator, delegator address bytes -func ParseDelegationsByValKey(bz []byte) (sdk.ValAddress, sdk.AccAddress, error) { - prefixLength := len(DelegationByValIndexKey) - if prefix := bz[:prefixLength]; !bytes.Equal(prefix, DelegationByValIndexKey) { - return nil, nil, fmt.Errorf("invalid prefix; expected: %X, got: %x", DelegationByValIndexKey, prefix) - } - - bz = bz[prefixLength:] // remove the prefix byte - if len(bz) == 0 { - return nil, nil, fmt.Errorf("no bytes left to parse: %X", bz) - } - - valAddrLen := bz[0] - bz = bz[1:] // remove the length byte of validator address. - if len(bz) == 0 { - return nil, nil, fmt.Errorf("no bytes left to parse validator address: %X", bz) - } - - val := bz[0:int(valAddrLen)] - - bz = bz[int(valAddrLen):] // remove the delegator bytes - if len(bz) == 0 { - return nil, nil, fmt.Errorf("no bytes left to parse delegator address: %X", bz) - } - - del := bz - - return val, del, nil -} - // GetDelegationsKey creates the prefix for a delegator for all validators func GetDelegationsKey(delAddr sdk.AccAddress) []byte { return append(DelegationKey, address.MustLengthPrefix(delAddr)...)