diff --git a/x/staking/keeper/delegation.go b/x/staking/keeper/delegation.go index e6eae66f4866..2da5bae3a799 100644 --- a/x/staking/keeper/delegation.go +++ b/x/staking/keeper/delegation.go @@ -501,25 +501,20 @@ func (k Keeper) GetRedelegations(ctx context.Context, delegator sdk.AccAddress, // GetRedelegationsFromSrcValidator returns all redelegations from a particular // validator. func (k Keeper) GetRedelegationsFromSrcValidator(ctx context.Context, valAddr sdk.ValAddress) (reds []types.Redelegation, err error) { - store := k.storeService.OpenKVStore(ctx) - prefix := types.GetREDsFromValSrcIndexKey(valAddr) - iterator, err := store.Iterator(prefix, storetypes.PrefixEndBytes(prefix)) - if err != nil { - return nil, err - } - defer iterator.Close() + rng := collections.NewPrefixedTripleRange[[]byte, []byte, []byte](valAddr) + err = k.RedelegationsByValSrc.Walk(ctx, rng, func(key collections.Triple[[]byte, []byte, []byte], value []byte) (stop bool, err error) { + valSrcAddr, delAddr, valDstAddr := key.K1(), key.K2(), key.K3() - for ; iterator.Valid(); iterator.Next() { - key := types.GetREDKeyFromValSrcIndexKey(iterator.Key()) - value, err := store.Get(key) - if err != nil { - return nil, err - } - red, err := types.UnmarshalRED(k.cdc, value) + red, err := k.Redelegations.Get(ctx, collections.Join3(delAddr, valSrcAddr, valDstAddr)) if err != nil { - return nil, err + return true, err } reds = append(reds, red) + + return false, nil + }) + if err != nil { + return nil, err } return reds, nil @@ -576,7 +571,7 @@ func (k Keeper) SetRedelegation(ctx context.Context, red types.Redelegation) err return err } - if err = store.Set(types.GetREDByValSrcIndexKey(delegatorAddress, valSrcAddr, valDestAddr), []byte{}); err != nil { + if err = k.RedelegationsByValSrc.Set(ctx, collections.Join3(valSrcAddr, delegatorAddress, valDestAddr), []byte{}); err != nil { return err } @@ -664,7 +659,7 @@ func (k Keeper) RemoveRedelegation(ctx context.Context, red types.Redelegation) return err } - if err = store.Delete(types.GetREDByValSrcIndexKey(delegatorAddress, valSrcAddr, valDestAddr)); err != nil { + if err = k.RedelegationsByValSrc.Remove(ctx, collections.Join3(valSrcAddr, delegatorAddress, valDestAddr)); err != nil { return err } diff --git a/x/staking/keeper/grpc_query.go b/x/staking/keeper/grpc_query.go index 3803c4ffd690..961221998d6b 100644 --- a/x/staking/keeper/grpc_query.go +++ b/x/staking/keeper/grpc_query.go @@ -410,7 +410,7 @@ func (k Querier) Redelegations(ctx context.Context, req *types.QueryRedelegation case req.DelegatorAddr != "" && req.SrcValidatorAddr != "" && req.DstValidatorAddr != "": redels, err = queryRedelegation(ctx, k, req) case req.DelegatorAddr == "" && req.SrcValidatorAddr != "" && req.DstValidatorAddr == "": - redels, pageRes, err = queryRedelegationsFromSrcValidator(store, k, req) + redels, pageRes, err = queryRedelegationsFromSrcValidator(ctx, store, k, req) default: redels, pageRes, err = queryAllRedelegations(ctx, store, k, req) } @@ -517,26 +517,20 @@ func queryRedelegation(ctx context.Context, k Querier, req *types.QueryRedelegat return redels, nil } -func queryRedelegationsFromSrcValidator(store storetypes.KVStore, k Querier, req *types.QueryRedelegationsRequest) (redels types.Redelegations, res *query.PageResponse, err error) { +func queryRedelegationsFromSrcValidator(ctx context.Context, store storetypes.KVStore, k Querier, req *types.QueryRedelegationsRequest) (types.Redelegations, *query.PageResponse, error) { valAddr, err := k.validatorAddressCodec.StringToBytes(req.SrcValidatorAddr) if err != nil { return nil, nil, err } - srcValPrefix := types.GetREDsFromValSrcIndexKey(valAddr) - redStore := prefix.NewStore(store, srcValPrefix) - res, err = query.Paginate(redStore, req.Pagination, func(key, value []byte) error { - storeKey := types.GetREDKeyFromValSrcIndexKey(append(srcValPrefix, key...)) - storeValue := store.Get(storeKey) - red, err := types.UnmarshalRED(k.cdc, storeValue) + return query.CollectionPaginate(ctx, k.RedelegationsByValSrc, req.Pagination, func(key collections.Triple[[]byte, []byte, []byte], val []byte) (types.Redelegation, error) { + valSrcAddr, delAddr, valDstAddr := key.K1(), key.K2(), key.K3() + red, err := k.Keeper.Redelegations.Get(ctx, collections.Join3(delAddr, valSrcAddr, valDstAddr)) if err != nil { - return err + return types.Redelegation{}, err } - redels = append(redels, red) - return nil - }) - - return redels, res, err + return red, nil + }, query.WithCollectionPaginationTriplePrefix[[]byte, []byte, []byte](valAddr)) } func queryAllRedelegations(ctx context.Context, store storetypes.KVStore, k Querier, req *types.QueryRedelegationsRequest) (redels types.Redelegations, res *query.PageResponse, err error) { diff --git a/x/staking/keeper/keeper.go b/x/staking/keeper/keeper.go index a2d8f6ff3985..81c00360712a 100644 --- a/x/staking/keeper/keeper.go +++ b/x/staking/keeper/keeper.go @@ -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] + RedelegationsByValSrc collections.Map[collections.Triple[[]byte, []byte, []byte], []byte] } // NewKeeper creates a new staking Keeper instance @@ -120,6 +121,17 @@ func NewKeeper( codec.CollValue[types.Redelegation](cdc), ), UnbondingIndex: collections.NewMap(sb, types.UnbondingIndexKey, "unbonding_index", collections.Uint64Key, collections.BytesValue), + // key format is: 53 | lengthPrefixedBytes(DstValAddr) | lengthPrefixedBytes(AccAddr) | lengthPrefixedBytes(SrcValAddr) + RedelegationsByValSrc: collections.NewMap( + sb, types.RedelegationByValSrcIndexKey, + "redelegations_by_val_src", + collections.TripleKeyCodec( + collections.BytesKey, + collections.BytesKey, + sdk.LengthPrefixedBytesKey, // sdk.LengthPrefixedBytesKey is needed to retain state compatibility + ), + collections.BytesValue, + ), } schema, err := sb.Build() diff --git a/x/staking/keeper/keeper_test.go b/x/staking/keeper/keeper_test.go index 46aabdcabb4c..f12a55fc04e4 100644 --- a/x/staking/keeper/keeper_test.go +++ b/x/staking/keeper/keeper_test.go @@ -8,6 +8,7 @@ import ( "github.com/golang/mock/gomock" "github.com/stretchr/testify/suite" + "cosmossdk.io/collections" "cosmossdk.io/math" storetypes "cosmossdk.io/store/types" @@ -17,6 +18,7 @@ import ( "github.com/cosmos/cosmos-sdk/testutil" simtestutil "github.com/cosmos/cosmos-sdk/testutil/sims" sdk "github.com/cosmos/cosmos-sdk/types" + addresstypes "github.com/cosmos/cosmos-sdk/types/address" moduletestutil "github.com/cosmos/cosmos-sdk/types/module/testutil" authtypes "github.com/cosmos/cosmos-sdk/x/auth/types" stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper" @@ -39,6 +41,7 @@ type KeeperTestSuite struct { accountKeeper *stakingtestutil.MockAccountKeeper queryClient stakingtypes.QueryClient msgServer stakingtypes.MsgServer + key *storetypes.KVStoreKey } func (s *KeeperTestSuite) SetupTest() { @@ -46,6 +49,7 @@ func (s *KeeperTestSuite) SetupTest() { key := storetypes.NewKVStoreKey(stakingtypes.StoreKey) storeService := runtime.NewKVStoreService(key) testCtx := testutil.DefaultContextWithDB(s.T(), key, storetypes.NewTransientStoreKey("transient_test")) + s.key = key ctx := testCtx.Ctx.WithBlockHeader(cmtproto.Header{Time: cmttime.Now()}) encCfg := moduletestutil.MakeTestEncodingConfig() @@ -109,6 +113,62 @@ func (s *KeeperTestSuite) TestLastTotalPower() { require.True(expTotalPower.Equal(resTotalPower)) } +// GetREDByValSrcIndexKey creates the index-key for a redelegation, stored by source-validator-index +// VALUE: none (key rearrangement used) +func getREDByValSrcIndexKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress) []byte { + REDSFromValsSrcKey := getREDsFromValSrcIndexKey(valSrcAddr) + offset := len(REDSFromValsSrcKey) + + // key is of the form REDSFromValsSrcKey || delAddrLen (1 byte) || delAddr || valDstAddrLen (1 byte) || valDstAddr + key := make([]byte, offset+2+len(delAddr)+len(valDstAddr)) + copy(key[0:offset], REDSFromValsSrcKey) + key[offset] = byte(len(delAddr)) + copy(key[offset+1:offset+1+len(delAddr)], delAddr.Bytes()) + key[offset+1+len(delAddr)] = byte(len(valDstAddr)) + copy(key[offset+2+len(delAddr):], valDstAddr.Bytes()) + + return key +} + +// GetREDsFromValSrcIndexKey returns a key prefix for indexing a redelegation to +// a source validator. +func getREDsFromValSrcIndexKey(valSrcAddr sdk.ValAddress) []byte { + redelegationByValSrcIndexKey := []byte{0x35} + return append(redelegationByValSrcIndexKey, addresstypes.MustLengthPrefix(valSrcAddr)...) +} + +func (s *KeeperTestSuite) TestRedelegationsMigrationToColls() { + s.SetupTest() + + addrs, valAddrs := createValAddrs(101) + + err := testutil.DiffCollectionsMigration( + s.ctx, + s.key, + 100, + func(i int64) { + // legacy method to set in the state + s.ctx.KVStore(s.key).Set(getREDByValSrcIndexKey(addrs[i], valAddrs[i], valAddrs[i+1]), []byte{}) + }, + "cb7b7086b1e03add24f85f894531fb36b3b9746f2e661e1640ec528a4f23a3d9", + ) + s.Require().NoError(err) + + err = testutil.DiffCollectionsMigration( + s.ctx, + s.key, + 100, + func(i int64) { + // using collections + err := s.stakingKeeper.RedelegationsByValSrc.Set(s.ctx, collections.Join3(valAddrs[i].Bytes(), addrs[i].Bytes(), valAddrs[i+1].Bytes()), []byte{}) + s.Require().NoError(err) + }, + "cb7b7086b1e03add24f85f894531fb36b3b9746f2e661e1640ec528a4f23a3d9", + ) + + s.Require().NoError(err) +} + func TestKeeperTestSuite(t *testing.T) { suite.Run(t, new(KeeperTestSuite)) } diff --git a/x/staking/migrations/v2/keys.go b/x/staking/migrations/v2/keys.go index bc2d6d923c39..b3da66ccb547 100644 --- a/x/staking/migrations/v2/keys.go +++ b/x/staking/migrations/v2/keys.go @@ -13,10 +13,11 @@ const ( ) var ( - ValidatorsByConsAddrKey = []byte{0x22} // prefix for validators by consensus address - RedelegationKey = []byte{0x34} // key for a redelegation - DelegationKey = []byte{0x31} // prefix for the delegation - HistoricalInfoKey = []byte{0x50} // prefix for the historical info + ValidatorsByConsAddrKey = []byte{0x22} // prefix for validators by consensus address + DelegationKey = []byte{0x31} // prefix for the delegation + RedelegationKey = []byte{0x34} // key for a redelegation + RedelegationByValSrcIndexKey = []byte{0x35} // prefix for each key for an redelegation, by source validator operator + HistoricalInfoKey = []byte{0x50} // prefix for the historical info ) // GetHistoricalInfoKey returns a key prefix for indexing HistoricalInfo objects. @@ -61,3 +62,26 @@ func GetREDKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress) [] func GetREDsKey(delAddr sdk.AccAddress) []byte { return append(RedelegationKey, address.MustLengthPrefix(delAddr)...) } + +// GetREDByValSrcIndexKey creates the index-key for a redelegation, stored by source-validator-index +// VALUE: none (key rearrangement used) +func GetREDByValSrcIndexKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress) []byte { + REDSFromValsSrcKey := GetREDsFromValSrcIndexKey(valSrcAddr) + offset := len(REDSFromValsSrcKey) + + // key is of the form REDSFromValsSrcKey || delAddrLen (1 byte) || delAddr || valDstAddrLen (1 byte) || valDstAddr + key := make([]byte, offset+2+len(delAddr)+len(valDstAddr)) + copy(key[0:offset], REDSFromValsSrcKey) + key[offset] = byte(len(delAddr)) + copy(key[offset+1:offset+1+len(delAddr)], delAddr.Bytes()) + key[offset+1+len(delAddr)] = byte(len(valDstAddr)) + copy(key[offset+2+len(delAddr):], valDstAddr.Bytes()) + + return key +} + +// GetREDsFromValSrcIndexKey returns a key prefix for indexing a redelegation to +// a source validator. +func GetREDsFromValSrcIndexKey(valSrcAddr sdk.ValAddress) []byte { + return append(RedelegationByValSrcIndexKey, address.MustLengthPrefix(valSrcAddr)...) +} diff --git a/x/staking/migrations/v2/store_test.go b/x/staking/migrations/v2/store_test.go index 7e8750877d0d..605f561d30bd 100644 --- a/x/staking/migrations/v2/store_test.go +++ b/x/staking/migrations/v2/store_test.go @@ -90,7 +90,7 @@ func TestStoreMigration(t *testing.T) { { "RedelegationByValSrcIndexKey", v1.GetREDByValSrcIndexKey(addr4, valAddr1, valAddr2), - types.GetREDByValSrcIndexKey(addr4, valAddr1, valAddr2), + v2.GetREDByValSrcIndexKey(addr4, valAddr1, valAddr2), }, { "RedelegationByValDstIndexKey", diff --git a/x/staking/types/keys.go b/x/staking/types/keys.go index 961f9b33f3da..3106e56d126a 100644 --- a/x/staking/types/keys.go +++ b/x/staking/types/keys.go @@ -43,7 +43,7 @@ var ( UnbondingDelegationKey = []byte{0x32} // key for an unbonding-delegation UnbondingDelegationByValIndexKey = []byte{0x33} // prefix for each key for an unbonding-delegation, by validator operator RedelegationKey = collections.NewPrefix(52) // key for a redelegation - RedelegationByValSrcIndexKey = []byte{0x35} // prefix for each key for an redelegation, by source validator operator + RedelegationByValSrcIndexKey = collections.NewPrefix(53) // prefix for each key for an redelegation, by source validator operator RedelegationByValDstIndexKey = []byte{0x36} // prefix for each key for an redelegation, by destination validator operator UnbondingIDKey = collections.NewPrefix(55) // key for the counter for the incrementing id for UnbondingOperations @@ -247,23 +247,6 @@ func GetREDKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress) [] return key } -// GetREDByValSrcIndexKey creates the index-key for a redelegation, stored by source-validator-index -// VALUE: none (key rearrangement used) -func GetREDByValSrcIndexKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress) []byte { - REDSFromValsSrcKey := GetREDsFromValSrcIndexKey(valSrcAddr) - offset := len(REDSFromValsSrcKey) - - // key is of the form REDSFromValsSrcKey || delAddrLen (1 byte) || delAddr || valDstAddrLen (1 byte) || valDstAddr - key := make([]byte, offset+2+len(delAddr)+len(valDstAddr)) - copy(key[0:offset], REDSFromValsSrcKey) - key[offset] = byte(len(delAddr)) - copy(key[offset+1:offset+1+len(delAddr)], delAddr.Bytes()) - key[offset+1+len(delAddr)] = byte(len(valDstAddr)) - copy(key[offset+2+len(delAddr):], valDstAddr.Bytes()) - - return key -} - // GetREDByValDstIndexKey creates the index-key for a redelegation, stored by destination-validator-index // VALUE: none (key rearrangement used) func GetREDByValDstIndexKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress) []byte { @@ -281,24 +264,6 @@ func GetREDByValDstIndexKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.V return key } -// GetREDKeyFromValSrcIndexKey rearranges the ValSrcIndexKey to get the REDKey -func GetREDKeyFromValSrcIndexKey(indexKey []byte) []byte { - // note that first byte is prefix byte, which we remove - kv.AssertKeyAtLeastLength(indexKey, 2) - addrs := indexKey[1:] - - valSrcAddrLen := addrs[0] - kv.AssertKeyAtLeastLength(addrs, int(valSrcAddrLen)+2) - valSrcAddr := addrs[1 : valSrcAddrLen+1] - delAddrLen := addrs[valSrcAddrLen+1] - kv.AssertKeyAtLeastLength(addrs, int(valSrcAddrLen)+int(delAddrLen)+2) - delAddr := addrs[valSrcAddrLen+2 : valSrcAddrLen+2+delAddrLen] - kv.AssertKeyAtLeastLength(addrs, int(valSrcAddrLen)+int(delAddrLen)+4) - valDstAddr := addrs[valSrcAddrLen+delAddrLen+3:] - - return GetREDKey(delAddr, valSrcAddr, valDstAddr) -} - // GetREDKeyFromValDstIndexKey rearranges the ValDstIndexKey to get the REDKey func GetREDKeyFromValDstIndexKey(indexKey []byte) []byte { // note that first byte is prefix byte, which we remove @@ -330,12 +295,6 @@ func GetREDsKey(delAddr sdk.AccAddress) []byte { return append(RedelegationKey, address.MustLengthPrefix(delAddr)...) } -// GetREDsFromValSrcIndexKey returns a key prefix for indexing a redelegation to -// a source validator. -func GetREDsFromValSrcIndexKey(valSrcAddr sdk.ValAddress) []byte { - return append(RedelegationByValSrcIndexKey, address.MustLengthPrefix(valSrcAddr)...) -} - // GetREDsToValDstIndexKey returns a key prefix for indexing a redelegation to a // destination (target) validator. func GetREDsToValDstIndexKey(valDstAddr sdk.ValAddress) []byte { diff --git a/x/staking/types/keys_test.go b/x/staking/types/keys_test.go index 681292452ded..172a4cd94e7c 100644 --- a/x/staking/types/keys_test.go +++ b/x/staking/types/keys_test.go @@ -79,33 +79,6 @@ func TestGetREDByValDstIndexKey(t *testing.T) { } } -func TestGetREDByValSrcIndexKey(t *testing.T) { - tests := []struct { - delAddr sdk.AccAddress - valSrcAddr sdk.ValAddress - valDstAddr sdk.ValAddress - wantHex string - }{ - { - sdk.AccAddress(keysAddr1), sdk.ValAddress(keysAddr1), sdk.ValAddress(keysAddr1), - "351463d771218209d8bd03c482f69dfba57310f086091463d771218209d8bd03c482f69dfba57310f086091463d771218209d8bd03c482f69dfba57310f08609", - }, - { - sdk.AccAddress(keysAddr1), sdk.ValAddress(keysAddr2), sdk.ValAddress(keysAddr3), - "35145ef3b5f25c54946d4a89fc0d09d2f126614540f21463d771218209d8bd03c482f69dfba57310f08609143ab62f0d93849be495e21e3e9013a517038f45bd", - }, - { - sdk.AccAddress(keysAddr2), sdk.ValAddress(keysAddr1), sdk.ValAddress(keysAddr3), - "351463d771218209d8bd03c482f69dfba57310f08609145ef3b5f25c54946d4a89fc0d09d2f126614540f2143ab62f0d93849be495e21e3e9013a517038f45bd", - }, - } - for i, tt := range tests { - got := hex.EncodeToString(types.GetREDByValSrcIndexKey(tt.delAddr, tt.valSrcAddr, tt.valDstAddr)) - - require.Equal(t, tt.wantHex, got, "Keys did not match on test case %d", i) - } -} - func TestGetValidatorQueueKey(t *testing.T) { ts := time.Now() height := int64(1024)