Automatically unregister failed contracts in Dex EndBlock (#739)

* Auto unregister contract if it fails execution

* rebase

* fix test

* Remove refund from failed contract unregisteration

---------

Co-authored-by: Cyson <[email protected]>
Co-authored-by: Xingchen Liao <[email protected]>

Author: 3 people

x/dex/keeper/contract.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/cosmos/cosmos-sdk/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	appparams "github.com/sei-protocol/sei-chain/app/params"
 	"github.com/sei-protocol/sei-chain/x/dex/types"
 )
 
@@ -113,3 +114,25 @@ func (k Keeper) GetRentsForContracts(ctx sdk.Context, contractAddrs []string) ma
 	}
 	return res
 }
+
+// Unregistrate and refund the creator
+func (k Keeper) DoUnregisterContractWithRefund(ctx sdk.Context, contract types.ContractInfoV2) error {
+	k.DoUnregisterContract(ctx, contract)
+	creatorAddr, _ := sdk.AccAddressFromBech32(contract.Creator)
+	if err := k.BankKeeper.SendCoins(ctx, k.AccountKeeper.GetModuleAddress(types.ModuleName), creatorAddr, sdk.NewCoins(sdk.NewCoin(appparams.BaseCoinUnit, sdk.NewInt(int64(contract.RentBalance))))); err != nil {
+		return err
+	}
+	return nil
+}
+
+// Contract unregistration will remove all orderbook data stored for the contract
+func (k Keeper) DoUnregisterContract(ctx sdk.Context, contract types.ContractInfoV2) {
+	k.DeleteContract(ctx, contract.ContractAddr)
+	k.RemoveAllLongBooksForContract(ctx, contract.ContractAddr)
+	k.RemoveAllShortBooksForContract(ctx, contract.ContractAddr)
+	k.RemoveAllPricesForContract(ctx, contract.ContractAddr)
+	k.DeleteMatchResultState(ctx, contract.ContractAddr)
+	k.DeleteNextOrderID(ctx, contract.ContractAddr)
+	k.DeleteAllRegisteredPairsForContract(ctx, contract.ContractAddr)
+	k.RemoveAllTriggeredOrders(ctx, contract.ContractAddr)
+}

x/dex/keeper/msgserver/msg_server_unregister_contract.go

@@ -6,7 +6,6 @@ import (
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
-	appparams "github.com/sei-protocol/sei-chain/app/params"
 	"github.com/sei-protocol/sei-chain/x/dex/types"
 )
 
@@ -25,19 +24,9 @@ func (k msgServer) UnregisterContract(goCtx context.Context, msg *types.MsgUnreg
 	if contract.Creator != msg.Creator {
 		return nil, sdkerrors.ErrUnauthorized
 	}
-	// refund remaining rent to the creator
-	creatorAddr, _ := sdk.AccAddressFromBech32(contract.Creator)
-	if err := k.BankKeeper.SendCoins(ctx, k.AccountKeeper.GetModuleAddress(types.ModuleName), creatorAddr, sdk.NewCoins(sdk.NewCoin(appparams.BaseCoinUnit, sdk.NewInt(int64(contract.RentBalance))))); err != nil {
+	if err := k.DoUnregisterContractWithRefund(ctx, contract); err != nil {
 		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)
 
 	ctx.EventManager().EmitEvent(sdk.NewEvent(
 		types.EventTypeUnregisterContract,

x/dex/module.go

@@ -273,6 +273,7 @@ func (am AppModule) EndBlock(ctx sdk.Context, _ abci.RequestEndBlock) (ret []abc
 	defer dexutils.GetMemState(ctx.Context()).Clear(ctx)
 
 	validContractsInfo := am.getAllContractInfo(ctx)
+	validContractsInfoAtBeginning := validContractsInfo
 	// Each iteration is atomic. If an iteration finishes without any error, it will return,
 	// otherwise it will rollback any state change, filter out contracts that cause the error,
 	// and proceed to the next iteration. The loop is guaranteed to finish since
@@ -294,6 +295,15 @@ func (am AppModule) EndBlock(ctx sdk.Context, _ abci.RequestEndBlock) (ret []abc
 			break
 		}
 	}
+	validContractAddrs := map[string]struct{}{}
+	for _, c := range validContractsInfo {
+		validContractAddrs[c.ContractAddr] = struct{}{}
+	}
+	for _, c := range validContractsInfoAtBeginning {
+		if _, ok := validContractAddrs[c.ContractAddr]; !ok {
+			am.keeper.DoUnregisterContract(ctx, c)
+		}
+	}
 
 	return []abci.ValidatorUpdate{}
 }

x/dex/module_test.go

@@ -458,6 +458,9 @@ func TestEndBlockPartialRollback(t *testing.T) {
 	// No state change should've been persisted for bad contract
 	matchResult, _ := dexkeeper.GetMatchResultState(ctx, keepertest.TestContract)
 	require.Equal(t, &types.MatchResult{}, matchResult)
+	// bad contract should be unregistered
+	_, err = dexkeeper.GetContract(ctx, keepertest.TestContract)
+	require.Equal(t, types.ErrContractNotExists, err)
 	// state change should've been persisted for good contract
 	matchResult, _ = dexkeeper.GetMatchResultState(ctx, contractAddr.String())
 	require.Equal(t, 1, len(matchResult.Orders))
@@ -615,14 +618,17 @@ func TestEndBlockRollbackWithRentCharge(t *testing.T) {
 	)
 
 	ctx = ctx.WithBlockHeight(1)
+	creatorBalanceBefore := bankkeeper.GetBalance(ctx, testAccount, "usei")
 	testApp.EndBlocker(ctx, abci.RequestEndBlock{})
 	// no state change should've been persisted for good contract because it should've run out of gas
 	matchResult, _ := dexkeeper.GetMatchResultState(ctx, contractAddr.String())
 	require.Equal(t, 0, len(matchResult.Orders))
-	// rent should still be charged even if the contract failed
-	contract, err := dexkeeper.GetContract(ctx, contractAddr.String())
-	require.Nil(t, err)
-	require.Zero(t, contract.RentBalance)
+	// rent should still be charged even if the contract failed, so no rent should be sent to the creator after
+	// auto unregister
+	_, err = dexkeeper.GetContract(ctx, contractAddr.String())
+	require.NotNil(t, err) // auto-unregistered
+	creatorBalanceAfter := bankkeeper.GetBalance(ctx, testAccount, "usei")
+	require.Equal(t, creatorBalanceBefore, creatorBalanceAfter)
 }
 
 func TestEndBlockContractWithoutPair(t *testing.T) {