[TRA-115] decommission vaults at the beginning of a block (#1264)

* decommission vaults at the beginning of a block

* address comments

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.DecommissionNonPositiveEquityVaults(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) DecommissionNonPositiveEquityVaults(
+	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 total 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,223 @@
+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 TestDecommissionNonPositiveEquityVaults(t *testing.T) {
+	tests := map[string]struct {
+		/* --- Setup --- */
+		// Vault IDs.
+		vaultIds []vaulttypes.VaultId
+		// Total shares of above vaults.
+		totalShares []*big.Rat
+		// Equities of above vaults.
+		equities []*big.Int
+
+		/* --- Expectations --- */
+		// Whether the vaults are decommissioned.
+		decommissioned []bool
+	}{
+		"Decommission no vault": {
+			vaultIds: []vaulttypes.VaultId{
+				constants.Vault_Clob_0,
+				constants.Vault_Clob_1,
+			},
+			totalShares: []*big.Rat{
+				big.NewRat(7, 1),
+				big.NewRat(7, 1),
+			},
+			equities: []*big.Int{
+				big.NewInt(1),
+				big.NewInt(1),
+			},
+			decommissioned: []bool{
+				false,
+				false,
+			},
+		},
+		"Decommission one vault": {
+			vaultIds: []vaulttypes.VaultId{
+				constants.Vault_Clob_0,
+				constants.Vault_Clob_1,
+			},
+			totalShares: []*big.Rat{
+				big.NewRat(7, 1),
+				big.NewRat(7, 1),
+			},
+			equities: []*big.Int{
+				big.NewInt(1),
+				big.NewInt(0),
+			},
+			decommissioned: []bool{
+				false,
+				true, // this vault should be decommissioned.
+			},
+		},
+		"Decommission two vaults": {
+			vaultIds: []vaulttypes.VaultId{
+				constants.Vault_Clob_0,
+				constants.Vault_Clob_1,
+			},
+			totalShares: []*big.Rat{
+				big.NewRat(7, 1),
+				big.NewRat(7, 1),
+			},
+			equities: []*big.Int{
+				big.NewInt(0),
+				big.NewInt(-1),
+			},
+			decommissioned: []bool{
+				true,
+				true,
+			},
+		},
+	}
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			// Initialize vaults with their equities.
+			tApp := testapp.NewTestAppBuilder(t).WithGenesisDocFn(func() (genesis types.GenesisDoc) {
+				genesis = testapp.DefaultGenesis()
+				testapp.UpdateGenesisDocWithAppStateForModule(
+					&genesis,
+					func(genesisState *satypes.GenesisState) {
+						subaccounts := []satypes.Subaccount{}
+						for i, vaultId := range tc.vaultIds {
+							if tc.equities[i].Sign() != 0 {
+								subaccounts = append(
+									subaccounts,
+									satypes.Subaccount{
+										Id: vaultId.ToSubaccountId(),
+										AssetPositions: []*satypes.AssetPosition{
+											{
+												AssetId:  assettypes.AssetUsdc.Id,
+												Quantums: dtypes.NewIntFromBigInt(tc.equities[i]),
+											},
+										},
+									},
+								)
+							}
+						}
+						genesisState.Subaccounts = subaccounts
+					},
+				)
+				return genesis
+			}).Build()
+			ctx := tApp.InitChain()
+			k := tApp.App.VaultKeeper
+
+			// Set total shares and owner shares for all vaults.
+			testOwner := constants.Alice_Num0.Owner
+			for i, vaultId := range tc.vaultIds {
+				err := k.SetTotalShares(
+					ctx,
+					vaultId,
+					vaulttypes.BigRatToNumShares(tc.totalShares[i]),
+				)
+				require.NoError(t, err)
+				err = k.SetOwnerShares(
+					ctx,
+					vaultId,
+					testOwner,
+					vaulttypes.BigRatToNumShares(big.NewRat(7, 1)),
+				)
+				require.NoError(t, err)
+			}
+
+			// Decommission all vaults.
+			k.DecommissionNonPositiveEquityVaults(ctx)
+
+			// Check that total shares and owner shares are deleted for decommissioned
+			// vaults and not deleted for non-decommissioned vaults.
+			for i, decommissioned := range tc.decommissioned {
+				_, exists := k.GetTotalShares(ctx, tc.vaultIds[i])
+				require.Equal(t, !decommissioned, exists)
+				_, exists = k.GetOwnerShares(ctx, tc.vaultIds[i], testOwner)
+				require.Equal(t, !decommissioned, exists)
+			}
+		})
+	}
+}
+
+func TestDecommissionVault(t *testing.T) {
+	tests := map[string]struct {
+		/* --- Setup --- */
+		// Vault ID.
+		vaultId vaulttypes.VaultId
+		// Whether total shares exists.
+		totalSharesExists bool
+		// Owners.
+		owners []string
+	}{
+		"Total shares doesn't exist, no owners": {
+			vaultId: constants.Vault_Clob_0,
+		},
+		"Total shares exists, no owners": {
+			vaultId:           constants.Vault_Clob_0,
+			totalSharesExists: true,
+		},
+		"Total shares exists, one owner": {
+			vaultId:           constants.Vault_Clob_1,
+			totalSharesExists: true,
+			owners:            []string{constants.Alice_Num0.Owner},
+		},
+		"Total shares exists, two owners": {
+			vaultId:           constants.Vault_Clob_1,
+			totalSharesExists: true,
+			owners: []string{
+				constants.Alice_Num0.Owner,
+				constants.Bob_Num0.Owner,
+			},
+		},
+	}
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			tApp := testapp.NewTestAppBuilder(t).Build()
+			ctx := tApp.InitChain()
+			k := tApp.App.VaultKeeper
+
+			shares := vaulttypes.BigRatToNumShares(
+				big.NewRat(7, 1),
+			)
+
+			if tc.totalSharesExists {
+				err := k.SetTotalShares(
+					ctx,
+					tc.vaultId,
+					shares,
+				)
+				require.NoError(t, err)
+			}
+			for _, owner := range tc.owners {
+				err := k.SetOwnerShares(
+					ctx,
+					tc.vaultId,
+					owner,
+					shares,
+				)
+				require.NoError(t, err)
+			}
+
+			// Decommission vault.
+			k.DecommissionVault(ctx, tc.vaultId)
+
+			// Check that total shares and owner shares are deleted.
+			_, exists := k.GetTotalShares(ctx, tc.vaultId)
+			require.Equal(t, false, exists)
+			for _, owner := range tc.owners {
+				_, exists = k.GetOwnerShares(ctx, tc.vaultId, 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)