[TRA-167] v5.0.0 state migration negative tnc subaccount store. (#1226)

protocol/app/upgrades.go

@@ -32,6 +32,7 @@ func (app *App) setupUpgradeHandlers() {
 			app.configurator,
 			app.PerpetualsKeeper,
 			app.ClobKeeper,
+			app.SubaccountsKeeper,
 		),
 	)
 }

protocol/app/upgrades/v5.0.0/upgrade.go

@@ -8,6 +8,7 @@ import (
 
 	clobtypes "github.com/dydxprotocol/v4-chain/protocol/x/clob/types"
 	perptypes "github.com/dydxprotocol/v4-chain/protocol/x/perpetuals/types"
+	satypes "github.com/dydxprotocol/v4-chain/protocol/x/subaccounts/types"
 
 	upgradetypes "cosmossdk.io/x/upgrade/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
@@ -98,6 +99,61 @@ func blockRateLimitConfigUpdate(
 	)
 }
 
+func negativeTncSubaccountSeenAtBlockUpgrade(
+	ctx sdk.Context,
+	perpetualsKeeper perptypes.PerpetualsKeeper,
+	subaccountsKeeper satypes.SubaccountsKeeper,
+) {
+	// Get block height stored by v4.x.x.
+	blockHeight, exists := subaccountsKeeper.LegacyGetNegativeTncSubaccountSeenAtBlock(ctx)
+	ctx.Logger().Info(
+		fmt.Sprintf(
+			"Retrieved block height from store for negative tnc subaccount seen at block: %d, exists: %t\n",
+			blockHeight,
+			exists,
+		),
+	)
+	// If no block height was stored in the legacy store, no migration needed.
+	if !exists {
+		return
+	}
+
+	// If there are no perpetuals, then no new state needs to be stored, as there can be no
+	// negative tnc subaccounts w/o perpetuals.
+	perpetuals := perpetualsKeeper.GetAllPerpetuals(ctx)
+	ctx.Logger().Info(
+		fmt.Sprintf(
+			"Retrieved all perpetuals for negative tnc subaccount migration, # of perpetuals is %d\n",
+			len(perpetuals),
+		),
+	)
+	if len(perpetuals) == 0 {
+		return
+	}
+
+	ctx.Logger().Info(
+		fmt.Sprintf(
+			"Migrating negative tnc subaccount seen store, storing block height %d for perpetual %d\n",
+			perpetuals[0].Params.Id,
+			blockHeight,
+		),
+	)
+	// Migrate the value from the legacy store to the new store.
+	if err := subaccountsKeeper.SetNegativeTncSubaccountSeenAtBlock(
+		ctx,
+		perpetuals[0].Params.Id, // must be a cross-margined perpetual due to `perpetualsUpgrade`.
+		blockHeight,
+	); err != nil {
+		panic(fmt.Sprintf("failed to set negative tnc subaccount seen at block with value %d: %s", blockHeight, err))
+	}
+	ctx.Logger().Info(
+		fmt.Sprintf(
+			"Successfully migrated negative tnc subaccount seen at block with block height %d\n",
+			blockHeight,
+		),
+	)
+}
+
 // Initialize soft and upper caps for OIMF
 func initializeOIMFCaps(
 	ctx sdk.Context,
@@ -152,6 +208,7 @@ func CreateUpgradeHandler(
 	configurator module.Configurator,
 	perpetualsKeeper perptypes.PerpetualsKeeper,
 	clobKeeper clobtypes.ClobKeeper,
+	subaccountsKeeper satypes.SubaccountsKeeper,
 ) upgradetypes.UpgradeHandler {
 	return func(ctx context.Context, plan upgradetypes.Plan, vm module.VersionMap) (module.VersionMap, error) {
 		sdkCtx := lib.UnwrapSDKContext(ctx, "app/upgrades")
@@ -167,6 +224,10 @@ func CreateUpgradeHandler(
 		// Set block rate limit configuration
 		blockRateLimitConfigUpdate(sdkCtx, clobKeeper)
 
+		// Migrate state from legacy store for negative tnc subaccount seen to new store for
+		// negative tnc subaccount seen.
+		// Note, must be done after the upgrade to perpetuals to cross market type.
+		negativeTncSubaccountSeenAtBlockUpgrade(sdkCtx, perpetualsKeeper, subaccountsKeeper)
 		// Initialize liquidity tier with lower and upper OI caps.
 		initializeOIMFCaps(sdkCtx, perpetualsKeeper)
 

protocol/x/subaccounts/keeper/negative_tnc_subaccount.go

@@ -176,3 +176,25 @@ func (k Keeper) getLastBlockNegativeSubaccountSeen(
 	}
 	return lastBlockNegativeSubaccountSeen, negativeSubaccountExists, nil
 }
+
+// LegacyGetNegativeTncSubaccountSeenAtBlock gets the last block height a negative TNC subaccount was
+// seen in state and a boolean for whether it exists in state.
+// Deprecated: This is the legacy implementation and meant to be used for the v5.0.0 state migration.
+// Use `GetNegativeTncSubaccountSeenAtBlock` instead.
+func (k Keeper) LegacyGetNegativeTncSubaccountSeenAtBlock(
+	ctx sdk.Context,
+) (uint32, bool) {
+	store := ctx.KVStore(k.storeKey)
+	b := store.Get(
+		// Key used in v4.0.0.
+		[]byte("NegSA:"),
+	)
+	blockHeight := gogotypes.UInt32Value{Value: 0}
+	exists := false
+	if b != nil {
+		k.cdc.MustUnmarshal(b, &blockHeight)
+		exists = true
+	}
+
+	return blockHeight.Value, exists
+}

protocol/x/subaccounts/types/types.go

@@ -63,4 +63,14 @@ type SubaccountsKeeper interface {
 		ctx sdk.Context,
 		id SubaccountId,
 	) (val Subaccount)
+	LegacyGetNegativeTncSubaccountSeenAtBlock(ctx sdk.Context) (uint32, bool)
+	GetNegativeTncSubaccountSeenAtBlock(
+		ctx sdk.Context,
+		perpetualId uint32,
+	) (uint32, bool, error)
+	SetNegativeTncSubaccountSeenAtBlock(
+		ctx sdk.Context,
+		perpetualId uint32,
+		blockHeight uint32,
+	) error
 }