Skip to content

Commit

Permalink
refactor(x/staking): migrate RedelegationByValSrcIndexKey key to co…
Browse files Browse the repository at this point in the history
…llections (#17332)
  • Loading branch information
atheeshp authored Aug 22, 2023
1 parent 393dcc1 commit 222d28e
Show file tree
Hide file tree
Showing 8 changed files with 122 additions and 105 deletions.
29 changes: 12 additions & 17 deletions x/staking/keeper/delegation.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
}

Expand Down Expand Up @@ -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
}

Expand Down
22 changes: 8 additions & 14 deletions x/staking/keeper/grpc_query.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down Expand Up @@ -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) {
Expand Down
12 changes: 12 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]
RedelegationsByValSrc collections.Map[collections.Triple[[]byte, []byte, []byte], []byte]
}

// NewKeeper creates a new staking Keeper instance
Expand Down Expand Up @@ -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()
Expand Down
60 changes: 60 additions & 0 deletions x/staking/keeper/keeper_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand All @@ -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"
Expand All @@ -39,13 +41,15 @@ type KeeperTestSuite struct {
accountKeeper *stakingtestutil.MockAccountKeeper
queryClient stakingtypes.QueryClient
msgServer stakingtypes.MsgServer
key *storetypes.KVStoreKey
}

func (s *KeeperTestSuite) SetupTest() {
require := s.Require()
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()

Expand Down Expand Up @@ -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))
}
32 changes: 28 additions & 4 deletions x/staking/migrations/v2/keys.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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)...)
}
2 changes: 1 addition & 1 deletion x/staking/migrations/v2/store_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
43 changes: 1 addition & 42 deletions x/staking/types/keys.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand All @@ -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
Expand Down Expand Up @@ -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 {
Expand Down
27 changes: 0 additions & 27 deletions x/staking/types/keys_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit 222d28e

Please sign in to comment.