backport v5.2 upgrade handler (backport #1946) (#1981)

Co-authored-by: Tian Qin <[email protected]>

protocol/app/upgrades/v5.2.0/constants.go

@@ -0,0 +1,13 @@
+package v_5_2_0
+
+import (
+	"github.com/dydxprotocol/v4-chain/protocol/app/upgrades"
+)
+
+const (
+	UpgradeName = "v5.2.0"
+)
+
+var Upgrade = upgrades.Upgrade{
+	UpgradeName: UpgradeName,
+}

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

@@ -0,0 +1,61 @@
+package v_5_2_0
+
+import (
+	"context"
+	"fmt"
+
+	upgradetypes "cosmossdk.io/x/upgrade/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
+	"github.com/dydxprotocol/v4-chain/protocol/lib"
+	clobtypes "github.com/dydxprotocol/v4-chain/protocol/x/clob/types"
+	vaultkeeper "github.com/dydxprotocol/v4-chain/protocol/x/vault/keeper"
+)
+
+func CreateUpgradeHandler(
+	mm *module.Manager,
+	configurator module.Configurator,
+	clobKeeper clobtypes.ClobKeeper,
+	vaultKeeper vaultkeeper.Keeper,
+) upgradetypes.UpgradeHandler {
+	return func(ctx context.Context, plan upgradetypes.Plan, vm module.VersionMap) (module.VersionMap, error) {
+		sdkCtx := lib.UnwrapSDKContext(ctx, "app/upgrades")
+		sdkCtx.Logger().Info(fmt.Sprintf("Running %s Upgrade...", UpgradeName))
+
+		// Migrate x/clob.
+		clobKeeper.UnsafeMigrateOrderExpirationState(sdkCtx)
+
+		// Migrate x/vault.
+		setVaultOrderExpirationSecondsToSixty(sdkCtx, vaultKeeper)
+		addAllVaultsToVaultAddressStore(sdkCtx, vaultKeeper)
+
+		return mm.RunMigrations(ctx, configurator, vm)
+	}
+}
+
+// addAllVaultsToVaultAddressStore adds all existing vaults to Vault Address store.
+func addAllVaultsToVaultAddressStore(
+	ctx sdk.Context,
+	vaultKeeper vaultkeeper.Keeper,
+) {
+	allVaults := vaultKeeper.GetAllVaults(ctx)
+	for _, vault := range allVaults {
+		vaultKeeper.AddVaultToAddressStore(ctx, *vault.VaultId)
+	}
+}
+
+// setVaultOrderExpirationSecondsToSixty sets vault module param `OrderExpirationSeconds` to 60.
+func setVaultOrderExpirationSecondsToSixty(
+	ctx sdk.Context,
+	vaultKeeper vaultkeeper.Keeper,
+) {
+	params := vaultKeeper.GetParams(ctx)
+	params.OrderExpirationSeconds = 60
+	err := vaultKeeper.SetParams(
+		ctx,
+		params,
+	)
+	if err != nil {
+		panic(err)
+	}
+}

protocol/x/clob/types/clob_keeper.go

@@ -145,4 +145,6 @@ type ClobKeeper interface {
 	MigratePruneableOrders(ctx sdk.Context)
 	GetAllStatefulOrders(ctx sdk.Context) []Order
 	ResetAllDeliveredOrderIds(ctx sdk.Context)
+	// Migrate order expiration state (for upgrading to 5.2 only)
+	UnsafeMigrateOrderExpirationState(ctx sdk.Context)
 }