[TRA-621] Sweep bank funds of megavault into subaccount. (#2293)

protocol/app/app.go

@@ -1178,6 +1178,8 @@ func New(
 	app.VaultKeeper = *vaultmodulekeeper.NewKeeper(
 		appCodec,
 		keys[vaultmoduletypes.StoreKey],
+		app.AssetsKeeper,
+		app.BankKeeper,
 		app.ClobKeeper,
 		app.DelayMsgKeeper,
 		app.PerpetualsKeeper,

protocol/mocks/Makefile

@@ -63,3 +63,4 @@ mock-gen:
 	@go run github.com/vektra/mockery/v2 --name=OracleClient --dir=$(GOPATH)/pkg/mod/github.com/skip-mev/slinky@$(SLINKY_VERSION)/service/clients/oracle --recursive --output=./mocks
 	@go run github.com/vektra/mockery/v2 --name=ExtendVoteHandler --dir=$(GOPATH)/pkg/mod/github.com/dydxprotocol/cosmos-sdk@$(COSMOS_VERSION)/types --recursive --output=./mocks
 	@go run github.com/vektra/mockery/v2 --name=UpdateMarketPriceTxDecoder --dir=./app/process --recursive --output=./mocks
+	@go run github.com/vektra/mockery/v2 --name=AssetsKeeper --dir=./x//types --recursive --output=./mocks

protocol/testutil/keeper/vault.go

@@ -62,6 +62,8 @@ func createVaultKeeper(
 	k := keeper.NewKeeper(
 		cdc,
 		storeKey,
+		&mocks.AssetsKeeper{},
+		&mocks.BankKeeper{},
 		&mocks.ClobKeeper{},
 		&mocks.DelayMsgKeeper{},
 		&mocks.PerpetualsKeeper{},

protocol/x/assets/types/types.go

@@ -1 +1,32 @@
 package types
+
+import (
+	"math/big"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+type AssetsKeeper interface {
+	ConvertAssetToCoin(ctx sdk.Context, assetId uint32, quantums *big.Int) (*big.Int, sdk.Coin, error)
+	CreateAsset(
+		ctx sdk.Context,
+		assetId uint32,
+		symbol string,
+		denom string,
+		denomExponent int32,
+		hasMarket bool,
+		marketId uint32,
+		atomicResolution int32,
+	) (
+		Asset,
+		error,
+	)
+
+	GetAsset(ctx sdk.Context, id uint32) (Asset, bool)
+
+	GetAllAssets(ctx sdk.Context) []Asset
+
+	IsPositionUpdatable(ctx sdk.Context, id uint32) (bool, error)
+
+	ModifyAsset(ctx sdk.Context, id uint32, hasMarket bool, marketId uint32) (Asset, error)
+}

protocol/x/vault/abci.go

@@ -36,6 +36,7 @@ func EndBlocker(
 	// halting the chain.
 	if err := abci.RunCached(ctx, func(ctx sdk.Context) error {
 		keeper.RefreshAllVaultOrders(ctx)
+		keeper.SweepMainVaultBankBalance(ctx)
 		return nil
 	}); err != nil {
 		log.ErrorLog(

protocol/x/vault/keeper/keeper.go

@@ -16,6 +16,8 @@ type (
 	Keeper struct {
 		cdc                 codec.BinaryCodec
 		storeKey            storetypes.StoreKey
+		assetsKeeper        types.AssetsKeeper
+		bankKeeper          types.BankKeeper
 		clobKeeper          types.ClobKeeper
 		delayMsgKeeper      types.DelayMsgKeeper
 		perpetualsKeeper    types.PerpetualsKeeper
@@ -30,6 +32,8 @@ type (
 func NewKeeper(
 	cdc codec.BinaryCodec,
 	storeKey storetypes.StoreKey,
+	assetsKeeper types.AssetsKeeper,
+	bankKeeper types.BankKeeper,
 	clobKeeper types.ClobKeeper,
 	delayMsgKeeper types.DelayMsgKeeper,
 	perpetualsKeeper types.PerpetualsKeeper,
@@ -42,6 +46,8 @@ func NewKeeper(
 	return &Keeper{
 		cdc:                 cdc,
 		storeKey:            storeKey,
+		assetsKeeper:        assetsKeeper,
+		bankKeeper:          bankKeeper,
 		clobKeeper:          clobKeeper,
 		delayMsgKeeper:      delayMsgKeeper,
 		perpetualsKeeper:    perpetualsKeeper,

protocol/x/vault/keeper/sweep_funds.go

@@ -0,0 +1,48 @@
+package keeper
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/dydxprotocol/v4-chain/protocol/lib/log"
+	assettypes "github.com/dydxprotocol/v4-chain/protocol/x/assets/types"
+	"github.com/dydxprotocol/v4-chain/protocol/x/vault/types"
+)
+
+// SweepMainVaultBankBalances deposits any usdc balance from the Megavault main vault bank balance
+// into the Megavault main vault subaccount balance.
+func (k Keeper) SweepMainVaultBankBalance(
+	ctx sdk.Context,
+) {
+	usdcAsset, exists := k.assetsKeeper.GetAsset(ctx, assettypes.AssetUsdc.Id)
+	if !exists {
+		log.ErrorLog(
+			ctx,
+			"SweepMainVaultBankBalance: Usdc asset not found in state",
+		)
+		return
+	}
+	mainVaultBalance := k.bankKeeper.GetBalance(
+		ctx,
+		types.MegavaultMainAddress,
+		usdcAsset.Denom,
+	)
+	// No funds to sweep
+	if mainVaultBalance.Amount.BigInt().Sign() <= 0 {
+		return
+	}
+
+	err := k.subaccountsKeeper.DepositFundsFromAccountToSubaccount(
+		ctx,
+		types.MegavaultMainAddress,
+		types.MegavaultMainSubaccount,
+		usdcAsset.Id,
+		mainVaultBalance.Amount.BigInt(),
+	)
+	if err != nil {
+		log.ErrorLogWithError(
+			ctx,
+			"SweepMainVaultBankBalance: Failed to sweep funds from main vault bank balance to subaccount",
+			err,
+		)
+		return
+	}
+}