[SKI-23] Add v5.0.0 enable vote extensions upgrade handler (#1308)

* Enable VE upgrade handler

* gofumpt

* Update round, test

* Comment updates

protocol/app/upgrades.go

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

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

@@ -9,6 +9,8 @@ import (
 
 const (
 	UpgradeName = "v5.0.0"
+	// VEEnableHeightDelta is the number of blocks after the v5.0.0 upgrade to turn Vote Extensions on.
+	VEEnableHeightDelta = int64(4)
 )
 
 var Upgrade = upgrades.Upgrade{

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

@@ -6,13 +6,15 @@ import (
 	"math/big"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	consensustypes "github.com/cosmos/cosmos-sdk/x/consensus/types"
 
 	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"
+	consensusparamkeeper "github.com/cosmos/cosmos-sdk/x/consensus/keeper"
 	"github.com/dydxprotocol/v4-chain/protocol/lib"
 )
 
@@ -204,6 +206,37 @@ func initializeOIMFCaps(
 	}
 }
 
+func voteExtensionsUpgrade(
+	ctx sdk.Context,
+	keeper consensusparamkeeper.Keeper,
+) {
+	currentParams, err := keeper.Params(ctx, &consensustypes.QueryParamsRequest{})
+	if err != nil {
+		panic(fmt.Sprintf("failed to retrieve existing consensus params in VE upgrade handler: %s", err))
+	}
+	if currentParams.Params.Abci.VoteExtensionsEnableHeight != 0 {
+		panic(fmt.Sprintf(
+			"unable to update VE Enable Height since its current value of %d is already non-zero",
+			currentParams.Params.Abci.VoteExtensionsEnableHeight))
+	}
+	currentParams.Params.Abci.VoteExtensionsEnableHeight = ctx.BlockHeight() + VEEnableHeightDelta
+	_, err = keeper.UpdateParams(ctx, &consensustypes.MsgUpdateParams{
+		Authority: keeper.GetAuthority(),
+		Block:     currentParams.Params.Block,
+		Evidence:  currentParams.Params.Evidence,
+		Validator: currentParams.Params.Validator,
+		Abci:      currentParams.Params.Abci,
+	})
+	if err != nil {
+		panic(fmt.Sprintf("failed to update consensus params : %s", err))
+	}
+	ctx.Logger().Info(
+		"Successfully set VoteExtensionsEnableHeight",
+		"consensus_params",
+		currentParams.Params.String(),
+	)
+}
+
 // Initialize open interest for perpetuals
 func initializePerpOpenInterest(
 	ctx sdk.Context,
@@ -273,6 +306,7 @@ func CreateUpgradeHandler(
 	perpetualsKeeper perptypes.PerpetualsKeeper,
 	clobKeeper clobtypes.ClobKeeper,
 	subaccountsKeeper satypes.SubaccountsKeeper,
+	consensusParamKeeper consensusparamkeeper.Keeper,
 ) upgradetypes.UpgradeHandler {
 	return func(ctx context.Context, plan upgradetypes.Plan, vm module.VersionMap) (module.VersionMap, error) {
 		sdkCtx := lib.UnwrapSDKContext(ctx, "app/upgrades")
@@ -298,6 +332,9 @@ func CreateUpgradeHandler(
 		// Initialize liquidity tier with lower and upper OI caps.
 		initializeOIMFCaps(sdkCtx, perpetualsKeeper)
 
+		// Set vote extension enable height
+		voteExtensionsUpgrade(sdkCtx, consensusParamKeeper)
+
 		// TODO(TRA-93): Initialize `x/vault` module.
 
 		return mm.RunMigrations(ctx, configurator, vm)

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

@@ -5,6 +5,7 @@ package v_5_0_0_test
 import (
 	"testing"
 
+	consensus "github.com/cosmos/cosmos-sdk/x/consensus/types"
 	"github.com/cosmos/gogoproto/proto"
 	v_5_0_0 "github.com/dydxprotocol/v4-chain/protocol/app/upgrades/v5.0.0"
 	"github.com/dydxprotocol/v4-chain/protocol/dtypes"
@@ -45,16 +46,32 @@ func preUpgradeSetups(node *containertest.Node, t *testing.T) {
 
 func preUpgradeChecks(node *containertest.Node, t *testing.T) {
 	preUpgradeCheckPerpetualMarketType(node, t)
+	preUpgradeCheckVoteExtensions(node, t)
 	// Add test for your upgrade handler logic below
 }
 
 func postUpgradeChecks(node *containertest.Node, t *testing.T) {
 	postUpgradecheckPerpetualMarketType(node, t)
 	postUpgradeCheckLiquidityTiers(node, t)
 	postUpgradePerpetualOIs(node, t)
+	postUpgradeCheckVoteExtensions(node, t)
 	// Add test for your upgrade handler logic below
 }
 
+func preUpgradeCheckVoteExtensions(node *containertest.Node, t *testing.T) {
+	consensusParams := &consensus.QueryParamsResponse{}
+	resp, err := containertest.Query(
+		node,
+		consensus.NewQueryClient,
+		consensus.QueryClient.Params,
+		&consensus.QueryParamsRequest{},
+	)
+	require.NoError(t, err)
+	err = proto.UnmarshalText(resp.String(), consensusParams)
+	require.NoError(t, err)
+	assert.Equal(t, int64(0), consensusParams.Params.Abci.VoteExtensionsEnableHeight)
+}
+
 func placeOrders(node *containertest.Node, t *testing.T) {
 	require.NoError(t, containertest.BroadcastTx(
 		node,
@@ -120,6 +137,21 @@ func preUpgradeCheckPerpetualMarketType(node *containertest.Node, t *testing.T)
 	}
 }
 
+func postUpgradeCheckVoteExtensions(node *containertest.Node, t *testing.T) {
+	consensusParams := &consensus.QueryParamsResponse{}
+	resp, err := containertest.Query(
+		node,
+		consensus.NewQueryClient,
+		consensus.QueryClient.Params,
+		&consensus.QueryParamsRequest{},
+	)
+	require.NoError(t, err)
+	err = proto.UnmarshalText(resp.String(), consensusParams)
+	require.NoError(t, err)
+	// testnet_utils.go::UpgradeTesnet has a Plan for upgrading on height 10
+	assert.Equal(t, int64(10)+v_5_0_0.VEEnableHeightDelta, consensusParams.Params.Abci.VoteExtensionsEnableHeight)
+}
+
 func postUpgradecheckPerpetualMarketType(node *containertest.Node, t *testing.T) {
 	perpetualsList := &perpetuals.QueryAllPerpetualsResponse{}
 	resp, err := containertest.Query(