From 0f1335e166f6f1851e6374ba2652761f533eafb2 Mon Sep 17 00:00:00 2001 From: codchen Date: Fri, 10 Feb 2023 14:08:06 +0800 Subject: [PATCH] Delete contract-related states unpon unregister --- x/dex/keeper/common.go | 12 +++++++++++ x/dex/keeper/long_book.go | 4 ++++ x/dex/keeper/match_result.go | 8 +++++++ .../msg_server_unregister_contract.go | 7 +++++++ .../msg_server_unregister_contract_test.go | 21 +++++++++++++++++++ x/dex/keeper/order_placement.go | 6 ++++++ x/dex/keeper/pair.go | 4 ++++ x/dex/keeper/price.go | 4 ++++ x/dex/keeper/short_book.go | 4 ++++ x/dex/keeper/trigger_book.go | 4 ++++ x/dex/types/keys.go | 18 +++++++++++----- 11 files changed, 87 insertions(+), 5 deletions(-) create mode 100644 x/dex/keeper/common.go diff --git a/x/dex/keeper/common.go b/x/dex/keeper/common.go new file mode 100644 index 0000000000..3668362b7a --- /dev/null +++ b/x/dex/keeper/common.go @@ -0,0 +1,12 @@ +package keeper + +import sdk "github.com/cosmos/cosmos-sdk/types" + +func (k Keeper) removeAllForPrefix(ctx sdk.Context, prefix []byte) { + store := ctx.KVStore(k.storeKey) + iter := sdk.KVStorePrefixIterator(store, prefix) + defer iter.Close() + for ; iter.Valid(); iter.Next() { + store.Delete(iter.Key()) + } +} diff --git a/x/dex/keeper/long_book.go b/x/dex/keeper/long_book.go index 971248faa4..7c1cbeef35 100644 --- a/x/dex/keeper/long_book.go +++ b/x/dex/keeper/long_book.go @@ -104,6 +104,10 @@ func (k Keeper) GetLongAllocationForOrderID(ctx sdk.Context, contractAddr string return nil, false } +func (k Keeper) RemoveAllLongBooksForContract(ctx sdk.Context, contractAddr string) { + k.removeAllForPrefix(ctx, types.OrderBookContractPrefix(true, contractAddr)) +} + func GetKeyForLongBook(longBook types.LongBook) []byte { return GetKeyForPrice(longBook.Entry.Price) } diff --git a/x/dex/keeper/match_result.go b/x/dex/keeper/match_result.go index ca15ef12eb..09424196cf 100644 --- a/x/dex/keeper/match_result.go +++ b/x/dex/keeper/match_result.go @@ -35,3 +35,11 @@ func (k Keeper) GetMatchResultState(ctx sdk.Context, contractAddr string) (*type } return &result, true } + +func (k Keeper) DeleteMatchResultState(ctx sdk.Context, contractAddr string) { + store := prefix.NewStore( + ctx.KVStore(k.storeKey), + types.MatchResultPrefix(contractAddr), + ) + store.Delete([]byte(MatchResultKey)) +} diff --git a/x/dex/keeper/msgserver/msg_server_unregister_contract.go b/x/dex/keeper/msgserver/msg_server_unregister_contract.go index 5636a0ef77..cc268a9e75 100644 --- a/x/dex/keeper/msgserver/msg_server_unregister_contract.go +++ b/x/dex/keeper/msgserver/msg_server_unregister_contract.go @@ -31,5 +31,12 @@ func (k msgServer) UnregisterContract(goCtx context.Context, msg *types.MsgUnreg return nil, err } k.DeleteContract(ctx, msg.ContractAddr) + k.RemoveAllLongBooksForContract(ctx, msg.ContractAddr) + k.RemoveAllShortBooksForContract(ctx, msg.ContractAddr) + k.RemoveAllPricesForContract(ctx, msg.ContractAddr) + k.DeleteMatchResultState(ctx, msg.ContractAddr) + k.DeleteNextOrderID(ctx, msg.ContractAddr) + k.DeleteAllRegisteredPairsForContract(ctx, msg.ContractAddr) + k.RemoveAllTriggeredOrders(ctx, msg.ContractAddr) return &types.MsgUnregisterContractResponse{}, nil } diff --git a/x/dex/keeper/msgserver/msg_server_unregister_contract_test.go b/x/dex/keeper/msgserver/msg_server_unregister_contract_test.go index 28733bdb28..b97d1b6401 100644 --- a/x/dex/keeper/msgserver/msg_server_unregister_contract_test.go +++ b/x/dex/keeper/msgserver/msg_server_unregister_contract_test.go @@ -67,6 +67,25 @@ func TestUnregisterContractSetSiblings(t *testing.T) { require.Equal(t, int64(8900000), balance.Amount.Int64()) handler := dex.NewHandler(keeper) + tickSize := sdk.OneDec() + _, err = handler(ctx, &types.MsgRegisterPairs{ + Creator: testAccount.String(), + Batchcontractpair: []types.BatchContractPair{ + { + ContractAddr: contractAddr.String(), + Pairs: []*types.Pair{ + { + PriceDenom: "usei", + AssetDenom: "uatom", + PriceTicksize: &tickSize, + QuantityTicksize: &tickSize, + }, + }, + }, + }, + }) + require.NoError(t, err) + _, err = handler(ctx, &types.MsgUnregisterContract{ Creator: testAccount.String(), ContractAddr: contractAddr.String(), @@ -76,4 +95,6 @@ func TestUnregisterContractSetSiblings(t *testing.T) { require.Error(t, err) balance = keeper.BankKeeper.GetBalance(ctx, testAccount, "usei") require.Equal(t, int64(9900000), balance.Amount.Int64()) + pairs := keeper.GetAllRegisteredPairs(ctx, contractAddr.String()) + require.Empty(t, pairs) } diff --git a/x/dex/keeper/order_placement.go b/x/dex/keeper/order_placement.go index 3dab2f3880..b6d5fcec60 100644 --- a/x/dex/keeper/order_placement.go +++ b/x/dex/keeper/order_placement.go @@ -25,3 +25,9 @@ func (k Keeper) SetNextOrderID(ctx sdk.Context, contractAddr string, nextID uint binary.BigEndian.PutUint64(bz, nextID) store.Set(byteKey, bz) } + +func (k Keeper) DeleteNextOrderID(ctx sdk.Context, contractAddr string) { + store := prefix.NewStore(ctx.KVStore(k.storeKey), types.NextOrderIDPrefix(contractAddr)) + byteKey := types.KeyPrefix(types.NextOrderIDKey) + store.Delete(byteKey) +} diff --git a/x/dex/keeper/pair.go b/x/dex/keeper/pair.go index 9a2ee0f2ba..914131cc92 100644 --- a/x/dex/keeper/pair.go +++ b/x/dex/keeper/pair.go @@ -50,3 +50,7 @@ func (k Keeper) GetAllRegisteredPairs(ctx sdk.Context, contractAddr string) []ty return list } + +func (k Keeper) DeleteAllRegisteredPairsForContract(ctx sdk.Context, contractAddr string) { + k.removeAllForPrefix(ctx, types.RegisteredPairPrefix(contractAddr)) +} diff --git a/x/dex/keeper/price.go b/x/dex/keeper/price.go index 721d37c9de..2e2bc80cdc 100644 --- a/x/dex/keeper/price.go +++ b/x/dex/keeper/price.go @@ -59,6 +59,10 @@ func (k Keeper) GetAllPrices(ctx sdk.Context, contractAddr string, pair types.Pa return } +func (k Keeper) RemoveAllPricesForContract(ctx sdk.Context, contractAddr string) { + k.removeAllForPrefix(ctx, types.PriceContractPrefix(contractAddr)) +} + func GetKeyForTs(ts uint64) []byte { tsKey := make([]byte, 8) binary.BigEndian.PutUint64(tsKey, ts) diff --git a/x/dex/keeper/short_book.go b/x/dex/keeper/short_book.go index 9a449b1ef8..409190f655 100644 --- a/x/dex/keeper/short_book.go +++ b/x/dex/keeper/short_book.go @@ -89,6 +89,10 @@ func (k Keeper) GetShortAllocationForOrderID(ctx sdk.Context, contractAddr strin return nil, false } +func (k Keeper) RemoveAllShortBooksForContract(ctx sdk.Context, contractAddr string) { + k.removeAllForPrefix(ctx, types.OrderBookContractPrefix(false, contractAddr)) +} + func GetKeyForShortBook(shortBook types.ShortBook) []byte { return GetKeyForPrice(shortBook.Entry.Price) } diff --git a/x/dex/keeper/trigger_book.go b/x/dex/keeper/trigger_book.go index f54b81adfc..1d5e584583 100644 --- a/x/dex/keeper/trigger_book.go +++ b/x/dex/keeper/trigger_book.go @@ -60,6 +60,10 @@ func (k Keeper) GetAllTriggeredOrders(ctx sdk.Context, contractAddr string) (lis return } +func (k Keeper) RemoveAllTriggeredOrders(ctx sdk.Context, contractAddr string) { + k.removeAllForPrefix(ctx, types.ContractKeyPrefix(types.TriggerBookKey, contractAddr)) +} + func GetKeyForOrderID(orderID uint64) []byte { key := make([]byte, 8) binary.LittleEndian.PutUint64(key, orderID) diff --git a/x/dex/types/keys.go b/x/dex/types/keys.go index 2571431a96..1a2ad9356d 100644 --- a/x/dex/types/keys.go +++ b/x/dex/types/keys.go @@ -45,16 +45,20 @@ func PairPrefix(priceDenom string, assetDenom string) []byte { } func OrderBookPrefix(long bool, contractAddr string, priceDenom string, assetDenom string) []byte { + return append( + OrderBookContractPrefix(long, contractAddr), + PairPrefix(priceDenom, assetDenom)..., + ) +} + +func OrderBookContractPrefix(long bool, contractAddr string) []byte { var prefix []byte if long { prefix = KeyPrefix(LongBookKey) } else { prefix = KeyPrefix(ShortBookKey) } - return append( - append(prefix, AddressKeyPrefix(contractAddr)...), - PairPrefix(priceDenom, assetDenom)..., - ) + return append(prefix, AddressKeyPrefix(contractAddr)...) } func TriggerOrderBookPrefix(contractAddr string, priceDenom string, assetDenom string) []byte { @@ -73,11 +77,15 @@ func TwapPrefix(contractAddr string) []byte { // `Price` constant + contract + price denom + asset denom func PricePrefix(contractAddr string, priceDenom string, assetDenom string) []byte { return append( - append(KeyPrefix(PriceKey), AddressKeyPrefix(contractAddr)...), + PriceContractPrefix(contractAddr), PairPrefix(priceDenom, assetDenom)..., ) } +func PriceContractPrefix(contractAddr string) []byte { + return append(KeyPrefix(PriceKey), AddressKeyPrefix(contractAddr)...) +} + func SettlementEntryPrefix(contractAddr string, priceDenom string, assetDenom string) []byte { return append( append(KeyPrefix(SettlementEntryKey), AddressKeyPrefix(contractAddr)...),