decommission vaults at the beginning of a block

protocol/x/vault/abci.go

@@ -5,6 +5,13 @@ import (
 	"github.com/dydxprotocol/v4-chain/protocol/x/vault/keeper"
 )
 
+func BeginBlocker(
+	ctx sdk.Context,
+	keeper *keeper.Keeper,
+) {
+	keeper.DecommissionVaults(ctx)
+}
+
 func EndBlocker(
 	ctx sdk.Context,
 	keeper *keeper.Keeper,

protocol/x/vault/keeper/vault.go

@@ -0,0 +1,83 @@
+package keeper
+
+import (
+	"math/big"
+
+	"cosmossdk.io/store/prefix"
+	storetypes "cosmossdk.io/store/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/dydxprotocol/v4-chain/protocol/lib/log"
+	satypes "github.com/dydxprotocol/v4-chain/protocol/x/subaccounts/types"
+	"github.com/dydxprotocol/v4-chain/protocol/x/vault/types"
+)
+
+// GetVaultEquity returns the equity of a vault (in quote quantums).
+func (k Keeper) GetVaultEquity(
+	ctx sdk.Context,
+	vaultId types.VaultId,
+) (*big.Int, error) {
+	netCollateral, _, _, err := k.subaccountsKeeper.GetNetCollateralAndMarginRequirements(
+		ctx,
+		satypes.Update{
+			SubaccountId: *vaultId.ToSubaccountId(),
+		},
+	)
+	if err != nil {
+		return nil, err
+	}
+	return netCollateral, nil
+}
+
+// DecommissionVaults decommissions all vaults with positive shares and non-positive equity.
+func (k Keeper) DecommissionVaults(
+	ctx sdk.Context,
+) {
+	// Iterate through all vaults.
+	totalSharesIterator := k.getTotalSharesIterator(ctx)
+	defer totalSharesIterator.Close()
+	for ; totalSharesIterator.Valid(); totalSharesIterator.Next() {
+		var totalShares types.NumShares
+		k.cdc.MustUnmarshal(totalSharesIterator.Value(), &totalShares)
+
+		// Skip if TotalShares is non-positive.
+		totalSharesRat, err := totalShares.ToBigRat()
+		if err != nil || totalSharesRat.Sign() <= 0 {
+			continue
+		}
+
+		// Get vault equity.
+		vaultId, err := types.GetVaultIdFromStateKey(totalSharesIterator.Key())
+		if err != nil {
+			log.ErrorLogWithError(ctx, "Failed to get vault ID from state key", err)
+			continue
+		}
+		equity, err := k.GetVaultEquity(ctx, *vaultId)
+		if err != nil {
+			log.ErrorLogWithError(ctx, "Failed to get vault equity", err)
+			continue
+		}
+
+		// Decommission vault if equity is non-positive.
+		if equity.Sign() <= 0 {
+			k.DecommissionVault(ctx, *vaultId)
+		}
+	}
+}
+
+// DecommissionVault decommissions a vault by deleting its tota shares and owner shares.
+func (k Keeper) DecommissionVault(
+	ctx sdk.Context,
+	vaultId types.VaultId,
+) {
+	// Delete TotalShares of the vault.
+	totalSharesStore := prefix.NewStore(ctx.KVStore(k.storeKey), []byte(types.TotalSharesKeyPrefix))
+	totalSharesStore.Delete(vaultId.ToStateKey())
+
+	// Delete all OwnerShares of the vault.
+	ownerSharesStore := k.getVaultOwnerSharesStore(ctx, vaultId)
+	ownerSharesIterator := storetypes.KVStorePrefixIterator(ownerSharesStore, []byte{})
+	defer ownerSharesIterator.Close()
+	for ; ownerSharesIterator.Valid(); ownerSharesIterator.Next() {
+		ownerSharesStore.Delete(ownerSharesIterator.Key())
+	}
+}

protocol/x/vault/keeper/vault_test.go

@@ -0,0 +1,135 @@
+package keeper_test
+
+import (
+	"math/big"
+	"testing"
+
+	"github.com/cometbft/cometbft/types"
+	"github.com/dydxprotocol/v4-chain/protocol/dtypes"
+	testapp "github.com/dydxprotocol/v4-chain/protocol/testutil/app"
+	"github.com/dydxprotocol/v4-chain/protocol/testutil/constants"
+	assettypes "github.com/dydxprotocol/v4-chain/protocol/x/assets/types"
+	satypes "github.com/dydxprotocol/v4-chain/protocol/x/subaccounts/types"
+	vaulttypes "github.com/dydxprotocol/v4-chain/protocol/x/vault/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDecomissionVaults(t *testing.T) {
+	vault0 := constants.Vault_Clob_0
+	vault1 := constants.Vault_Clob_1
+	// Initialize vault 0 with positive equity.
+	tApp := testapp.NewTestAppBuilder(t).WithGenesisDocFn(func() (genesis types.GenesisDoc) {
+		genesis = testapp.DefaultGenesis()
+		testapp.UpdateGenesisDocWithAppStateForModule(
+			&genesis,
+			func(genesisState *satypes.GenesisState) {
+				genesisState.Subaccounts = []satypes.Subaccount{
+					{
+						Id: vault0.ToSubaccountId(),
+						AssetPositions: []*satypes.AssetPosition{
+							{
+								AssetId:  assettypes.AssetUsdc.Id,
+								Quantums: dtypes.NewInt(1),
+							},
+						},
+					},
+				}
+			},
+		)
+		return genesis
+	}).Build()
+	ctx := tApp.InitChain()
+	k := tApp.App.VaultKeeper
+
+	// Set total shares and owner shares for both vaults.
+	shares := vaulttypes.BigRatToNumShares(
+		big.NewRat(7, 1),
+	)
+	err := k.SetTotalShares(
+		ctx,
+		vault0,
+		shares,
+	)
+	require.NoError(t, err)
+	err = k.SetOwnerShares(
+		ctx,
+		vault0,
+		constants.Alice_Num0.Owner,
+		shares,
+	)
+	require.NoError(t, err)
+	err = k.SetTotalShares(
+		ctx,
+		vault1,
+		shares,
+	)
+	require.NoError(t, err)
+	err = k.SetOwnerShares(
+		ctx,
+		vault1,
+		constants.Bob_Num0.Owner,
+		shares,
+	)
+	require.NoError(t, err)
+
+	// Decomission all vaults.
+	k.DecommissionVaults(ctx)
+
+	// Check that total shares and owner shares are not deleted for vault 0.
+	got, exists := k.GetTotalShares(ctx, vault0)
+	require.Equal(t, true, exists)
+	require.Equal(t, shares, got)
+	got, exists = k.GetOwnerShares(ctx, vault0, constants.Alice_Num0.Owner)
+	require.Equal(t, true, exists)
+	require.Equal(t, shares, got)
+	// Check that total shares and owner shares are deleted for vault 1.
+	_, exists = k.GetTotalShares(ctx, vault1)
+	require.Equal(t, false, exists)
+	_, exists = k.GetOwnerShares(ctx, vault1, constants.Bob_Num0.Owner)
+	require.Equal(t, false, exists)
+}
+
+func TestDecomissionVault(t *testing.T) {
+	tApp := testapp.NewTestAppBuilder(t).Build()
+	ctx := tApp.InitChain()
+	k := tApp.App.VaultKeeper
+
+	// Decomission a non-existent vault.
+	k.DecommissionVault(ctx, constants.Vault_Clob_0)
+
+	// Set total shares and owner shares for two owners of a vault.
+	shares := vaulttypes.BigRatToNumShares(
+		big.NewRat(7, 1),
+	)
+	err := k.SetTotalShares(
+		ctx,
+		constants.Vault_Clob_0,
+		shares,
+	)
+	require.NoError(t, err)
+	err = k.SetOwnerShares(
+		ctx,
+		constants.Vault_Clob_0,
+		constants.Alice_Num0.Owner,
+		shares,
+	)
+	require.NoError(t, err)
+	err = k.SetOwnerShares(
+		ctx,
+		constants.Vault_Clob_0,
+		constants.Bob_Num0.Owner,
+		shares,
+	)
+	require.NoError(t, err)
+
+	// Decomission above vault.
+	k.DecommissionVault(ctx, constants.Vault_Clob_0)
+
+	// Check that total shares and owner shares are deleted.
+	_, exists := k.GetTotalShares(ctx, constants.Vault_Clob_0)
+	require.Equal(t, false, exists)
+	_, exists = k.GetOwnerShares(ctx, constants.Vault_Clob_0, constants.Alice_Num0.Owner)
+	require.Equal(t, false, exists)
+	_, exists = k.GetOwnerShares(ctx, constants.Vault_Clob_0, constants.Bob_Num0.Owner)
+	require.Equal(t, false, exists)
+}

protocol/x/vault/module.go

@@ -28,6 +28,7 @@ var (
 	_ module.HasGenesisBasics = AppModuleBasic{}
 
 	_ appmodule.AppModule        = AppModule{}
+	_ appmodule.HasBeginBlocker  = AppModule{}
 	_ appmodule.HasEndBlocker    = AppModule{}
 	_ module.HasConsensusVersion = AppModule{}
 	_ module.HasGenesis          = AppModule{}
@@ -147,6 +148,16 @@ func (am AppModule) ExportGenesis(ctx sdk.Context, cdc codec.JSONCodec) json.Raw
 // be set to 1.
 func (AppModule) ConsensusVersion() uint64 { return 1 }
 
+// BeginBlock executes all ABCI BeginBlock logic respective to the vault module.
+func (am AppModule) BeginBlock(ctx context.Context) error {
+	defer telemetry.ModuleMeasureSince(am.Name(), time.Now(), telemetry.MetricKeyBeginBlocker)
+	BeginBlocker(
+		lib.UnwrapSDKContext(ctx, types.ModuleName),
+		&am.keeper,
+	)
+	return nil
+}
+
 // EndBlock executes all ABCI EndBlock logic respective to the vault module.
 func (am AppModule) EndBlock(ctx context.Context) error {
 	defer telemetry.ModuleMeasureSince(am.Name(), time.Now(), telemetry.MetricKeyEndBlocker)