feat(zetacore): add upgrade tracker

app/app.go

@@ -850,6 +850,7 @@
 	if err := tmjson.Unmarshal(req.AppStateBytes, &genesisState); err != nil {
 		panic(err)
 	}
+	app.UpgradeKeeper.SetModuleVersionMap(ctx, app.mm.GetVersionMap())
 	return app.mm.InitGenesis(ctx, app.appCodec, genesisState)
 }
 

app/setup_handlers.go

@@ -1,6 +1,10 @@
 package app
 
 import (
+	"os"
+
+	"golang.org/x/exp/slices"
+
 	"github.com/cosmos/cosmos-sdk/baseapp"
 	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -53,27 +57,87 @@
 		}
 	}
 	baseAppLegacySS := app.ParamsKeeper.Subspace(baseapp.Paramspace).WithKeyTable(paramstypes.ConsensusParamsKeyTable())
+	needsForcedMigration := []string{
+		authtypes.ModuleName,
+		banktypes.ModuleName,
+		stakingtypes.ModuleName,
+		distrtypes.ModuleName,
+		slashingtypes.ModuleName,
+		govtypes.ModuleName,
+		crisistypes.ModuleName,
+		emissionstypes.ModuleName,
+	}
+	allUpgrades := upgradeTracker{
+		upgrades: []upgradeTrackerItem{
+			{
+				index: 1714664193,
+				storeUpgrade: &storetypes.StoreUpgrades{
+					Added: []string{consensustypes.ModuleName, crisistypes.ModuleName},
+				},
+				upgradeHandler: func(ctx sdk.Context, vm module.VersionMap) (module.VersionMap, error) {
+					// Migrate Tendermint consensus parameters from x/params module to a dedicated x/consensus module
+					// https://docs.cosmos.network/main/build/migrations/upgrading#xconsensus
+					baseapp.MigrateParams(ctx, baseAppLegacySS, &app.ConsensusParamsKeeper)
+
+					// empty version map happens when upgrading from old versions which did not correctly call
+					// app.UpgradeKeeper.SetModuleVersionMap(ctx, app.mm.GetVersionMap()) in InitChainer.
+					// we must populate the version map if we detect this scenario
+					//
+					// this will only happen on the first upgrade. mainnet and testnet will not require this condition.
+					if len(vm) == 0 {
+						for m, mb := range app.mm.Modules {
+							if module, ok := mb.(module.HasConsensusVersion); ok {
+								if slices.Contains(needsForcedMigration, m) {
+									vm[m] = module.ConsensusVersion() - 1
+								} else {
+									vm[m] = module.ConsensusVersion()
+								}
+							}
+						}
+					}
+					return vm, nil
+				},
+			},
+			{
+				index: 1715624665,
+				storeUpgrade: &storetypes.StoreUpgrades{
+					Added: []string{capabilitytypes.ModuleName, ibcexported.ModuleName, ibctransfertypes.ModuleName},
+				},
+			},
+			{
+				index: 1715707436,
+				storeUpgrade: &storetypes.StoreUpgrades{
+					Added: []string{ibccrosschaintypes.ModuleName},
+				},
+			},
+		},
+		stateFileDir: DefaultNodeHome,
+	}
+
+	var upgradeHandlerFns []upgradeHandlerFn
+	var storeUpgrades *storetypes.StoreUpgrades
+	var err error
+	_, useIncrementalTracker := os.LookupEnv("ZETACORED_USE_INCREMENTAL_UPGRADE_TRACKER")
+	if useIncrementalTracker {
+		upgradeHandlerFns, storeUpgrades, err = allUpgrades.getIncrementalUpgrades()
+		if err != nil {
+			panic(err)
+		}
+	} else {
+		upgradeHandlerFns, storeUpgrades = allUpgrades.mergeAllUpgrades()
+	}
+
 	app.UpgradeKeeper.SetUpgradeHandler(constant.Version, func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) (module.VersionMap, error) {
 		app.Logger().Info("Running upgrade handler for " + constant.Version)
-		// Migrate Tendermint consensus parameters from x/params module to a dedicated x/consensus module.
-		baseapp.MigrateParams(ctx, baseAppLegacySS, &app.ConsensusParamsKeeper)
-		// Updated version map to the latest consensus versions from each module
-		for m, mb := range app.mm.Modules {
-			if module, ok := mb.(module.HasConsensusVersion); ok {
-				vm[m] = module.ConsensusVersion()
+
+		var err error
+		for _, upgradeHandler := range upgradeHandlerFns {
+			vm, err = upgradeHandler(ctx, vm)
+			if err != nil {
+				return vm, err
 			}
 		}
 
-		VersionMigrator{v: vm}.TriggerMigration(authtypes.ModuleName)
-		VersionMigrator{v: vm}.TriggerMigration(banktypes.ModuleName)
-		VersionMigrator{v: vm}.TriggerMigration(stakingtypes.ModuleName)
-		VersionMigrator{v: vm}.TriggerMigration(distrtypes.ModuleName)
-		VersionMigrator{v: vm}.TriggerMigration(slashingtypes.ModuleName)
-		VersionMigrator{v: vm}.TriggerMigration(govtypes.ModuleName)
-		VersionMigrator{v: vm}.TriggerMigration(crisistypes.ModuleName)
-
-		VersionMigrator{v: vm}.TriggerMigration(emissionstypes.ModuleName)
-
 		return app.mm.RunMigrations(ctx, app.configurator, vm)
 	})
 
@@ -82,29 +146,10 @@
 		panic(err)
 	}
 	if upgradeInfo.Name == constant.Version && !app.UpgradeKeeper.IsSkipHeight(upgradeInfo.Height) {
-		storeUpgrades := storetypes.StoreUpgrades{
-			Added: []string{
-				consensustypes.ModuleName,
-				crisistypes.ModuleName,
-				capabilitytypes.ModuleName,
-				ibcexported.ModuleName,
-				ibctransfertypes.ModuleName,
-				ibccrosschaintypes.ModuleName,
-			},
-		}
 		// Use upgrade store loader for the initial loading of all stores when app starts,
 		// it checks if version == upgradeHeight and applies store upgrades before loading the stores,
 		// so that new stores start with the correct version (the current height of chain),
 		// instead the default which is the latest version that store last committed i.e 0 for new stores.
-		app.SetStoreLoader(types.UpgradeStoreLoader(upgradeInfo.Height, &storeUpgrades))
+		app.SetStoreLoader(types.UpgradeStoreLoader(upgradeInfo.Height, storeUpgrades))
 	}
 }
-
-type VersionMigrator struct {
-	v module.VersionMap
-}
-
-func (v VersionMigrator) TriggerMigration(moduleName string) module.VersionMap {
-	v.v[moduleName] = v.v[moduleName] - 1
-	return v.v
-}

app/upgrade_tracker.go

@@ -0,0 +1,96 @@
+package app
+
+import (
+	"fmt"
+	"os"
+	"path"
+	"strconv"
+
+	storetypes "github.com/cosmos/cosmos-sdk/store/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
+)
+
+const incrementalUpgradeTrackerStateFile = "incrementalupgradetracker"
+
+type upgradeHandlerFn func(ctx sdk.Context, vm module.VersionMap) (module.VersionMap, error)
+
+type upgradeTrackerItem struct {
+	// Monotonically increasing index to order and track migrations. Typically the current unix epoch timestamp.
+	index int64
+	// Function that will run during the SetUpgradeHandler callback. The VersionMap must always be returned.
+	upgradeHandler upgradeHandlerFn
+	// StoreUpgrades that will be provided to UpgradeStoreLoader
+	storeUpgrade *storetypes.StoreUpgrades
+}
+
+// upgradeTracker allows us to track needed upgrades/migrations across both release and develop builds
+type upgradeTracker struct {
+	upgrades []upgradeTrackerItem
+	// directory the incremental state file is stored
+	stateFileDir string
+}
+
+// getIncrementalUpgrades gets all upgrades that have not been been applied. This is typically
+// used for developnet upgrades since we need to run migrations as the are committed rather than
+// all at once during a release
+func (t upgradeTracker) getIncrementalUpgrades() ([]upgradeHandlerFn, *storetypes.StoreUpgrades, error) {
+	neededUpgrades := &storetypes.StoreUpgrades{}
+	neededUpgradeHandlers := []upgradeHandlerFn{}
+	stateFilePath := path.Join(t.stateFileDir, incrementalUpgradeTrackerStateFile)
+
+	currentIndex := int64(0)
+	stateFileContents, err := os.ReadFile(stateFilePath) // #nosec G304 -- stateFilePath is not user controllable
+	if err == nil {
+		currentIndex, err = strconv.ParseInt(string(stateFileContents), 10, 64)
+		if err != nil {
+			return nil, nil, fmt.Errorf("unable to decode upgrade tracker: %w", err)
+		}
+	} else {
+		fmt.Printf("unable to load upgrade tracker: %v\n", err)
+	}
+
+	maxIndex := currentIndex
+	for _, item := range t.upgrades {
+		index := item.index
+		upgrade := item.storeUpgrade
+		upgradeHandler := item.upgradeHandler
+		if index <= currentIndex {
+			continue
+		}
+		if upgradeHandler != nil {
+			neededUpgradeHandlers = append(neededUpgradeHandlers, upgradeHandler)
+		}
+		if upgrade != nil {
+			neededUpgrades.Added = append(neededUpgrades.Added, upgrade.Added...)
+			neededUpgrades.Deleted = append(neededUpgrades.Deleted, upgrade.Deleted...)
+			neededUpgrades.Renamed = append(neededUpgrades.Renamed, upgrade.Renamed...)
+		}
+		maxIndex = index
+	}
+	err = os.WriteFile(stateFilePath, []byte(strconv.FormatInt(maxIndex, 10)), 0o600)
+	if err != nil {
+		return nil, nil, fmt.Errorf("unable to write upgrade state file: %w", err)
+	}
+	return neededUpgradeHandlers, neededUpgrades, nil
+}
+
+// mergeAllUpgrades unconditionally merges all upgrades. Typically used to gather the
+// migrations used during a release upgrade.
+func (t upgradeTracker) mergeAllUpgrades() ([]upgradeHandlerFn, *storetypes.StoreUpgrades) {
+	upgrades := &storetypes.StoreUpgrades{}
+	upgradeHandlers := []upgradeHandlerFn{}
+	for _, item := range t.upgrades {
+		upgrade := item.storeUpgrade
+		versionModifier := item.upgradeHandler
+		if versionModifier != nil {
+			upgradeHandlers = append(upgradeHandlers, versionModifier)
+		}
+		if upgrade != nil {
+			upgrades.Added = append(upgrades.Added, upgrade.Added...)
+			upgrades.Deleted = append(upgrades.Deleted, upgrade.Deleted...)
+			upgrades.Renamed = append(upgrades.Renamed, upgrade.Renamed...)
+		}
+	}
+	return upgradeHandlers, upgrades
+}

app/upgrade_tracker_test.go

@@ -0,0 +1,107 @@
+package app
+
+import (
+	"os"
+	"path"
+	"testing"
+
+	storetypes "github.com/cosmos/cosmos-sdk/store/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
+	"github.com/stretchr/testify/require"
+	authoritytypes "github.com/zeta-chain/zetacore/x/authority/types"
+	lightclienttypes "github.com/zeta-chain/zetacore/x/lightclient/types"
+)
+
+func TestUpgradeTracker(t *testing.T) {
+	r := require.New(t)
+
+	tmpdir, err := os.MkdirTemp("", "storeupgradetracker-*")
+	r.NoError(err)
+	defer os.RemoveAll(tmpdir)
+
+	allUpgrades := upgradeTracker{
+		upgrades: []upgradeTrackerItem{
+			{
+				index: 1000,
+				storeUpgrade: &storetypes.StoreUpgrades{
+					Added: []string{authoritytypes.ModuleName},
+				},
+			},
+			{
+				index: 2000,
+				storeUpgrade: &storetypes.StoreUpgrades{
+					Added: []string{lightclienttypes.ModuleName},
+				},
+				upgradeHandler: func(ctx sdk.Context, vm module.VersionMap) (module.VersionMap, error) {
+					return vm, nil
+				},
+			},
+			{
+				index: 3000,
+				upgradeHandler: func(ctx sdk.Context, vm module.VersionMap) (module.VersionMap, error) {
+					return vm, nil
+				},
+			},
+		},
+		stateFileDir: tmpdir,
+	}
+
+	upgradeHandlers, storeUpgrades := allUpgrades.mergeAllUpgrades()
+	r.Len(storeUpgrades.Added, 2)
+	r.Len(storeUpgrades.Renamed, 0)
+	r.Len(storeUpgrades.Deleted, 0)
+	r.Len(upgradeHandlers, 2)
+
+	// should return all migrations on first call
+	upgradeHandlers, storeUpgrades, err = allUpgrades.getIncrementalUpgrades()
+	r.NoError(err)
+	r.Len(storeUpgrades.Added, 2)
+	r.Len(storeUpgrades.Renamed, 0)
+	r.Len(storeUpgrades.Deleted, 0)
+	r.Len(upgradeHandlers, 2)
+
+	// should return no upgrades on second call
+	upgradeHandlers, storeUpgrades, err = allUpgrades.getIncrementalUpgrades()
+	r.NoError(err)
+	r.Len(storeUpgrades.Added, 0)
+	r.Len(storeUpgrades.Renamed, 0)
+	r.Len(storeUpgrades.Deleted, 0)
+	r.Len(upgradeHandlers, 0)
+
+	// now add a upgrade and ensure that it gets run without running
+	// the other upgrades
+	allUpgrades.upgrades = append(allUpgrades.upgrades, upgradeTrackerItem{
+		index: 4000,
+		storeUpgrade: &storetypes.StoreUpgrades{
+			Deleted: []string{"example"},
+		},
+	})
+
+	upgradeHandlers, storeUpgrades, err = allUpgrades.getIncrementalUpgrades()
+	r.NoError(err)
+	r.Len(storeUpgrades.Added, 0)
+	r.Len(storeUpgrades.Renamed, 0)
+	r.Len(storeUpgrades.Deleted, 1)
+	r.Len(upgradeHandlers, 0)
+}
+
+func TestUpgradeTrackerBadState(t *testing.T) {
+	r := require.New(t)
+
+	tmpdir, err := os.MkdirTemp("", "storeupgradetracker-*")
+	r.NoError(err)
+	defer os.RemoveAll(tmpdir)
+
+	stateFilePath := path.Join(tmpdir, incrementalUpgradeTrackerStateFile)
+
+	err = os.WriteFile(stateFilePath, []byte("badstate"), 0o600)
+	r.NoError(err)
+
+	allUpgrades := upgradeTracker{
+		upgrades:     []upgradeTrackerItem{},
+		stateFileDir: tmpdir,
+	}
+	_, _, err = allUpgrades.getIncrementalUpgrades()
+	r.Error(err)
+}