Skip to content

Commit

Permalink
feat(zetacore): add upgrade tracker (#2095)
Browse files Browse the repository at this point in the history
* feat(zetacore): add develop store upgrade tracker

* add ibc upgrade

* fix gosec annotation?

* feedback

- directly call from upgrade handlers
- add test coverage for bad state

* ibc crosschain migration

* review feedback

* fix version
  • Loading branch information
gartnera authored May 16, 2024
1 parent d9a5d44 commit 2e7dac2
Show file tree
Hide file tree
Showing 4 changed files with 285 additions and 36 deletions.
1 change: 1 addition & 0 deletions app/app.go
Original file line number Diff line number Diff line change
Expand Up @@ -850,6 +850,7 @@ func (app *App) InitChainer(ctx sdk.Context, req abci.RequestInitChain) abci.Res
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)
}

Expand Down
117 changes: 81 additions & 36 deletions app/setup_handlers.go
Original file line number Diff line number Diff line change
@@ -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"
Expand Down Expand Up @@ -53,27 +57,87 @@ func SetupHandlers(app *App) {
}
}
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)
})

Expand All @@ -82,29 +146,10 @@ func SetupHandlers(app *App) {
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
}
96 changes: 96 additions & 0 deletions app/upgrade_tracker.go
Original file line number Diff line number Diff line change
@@ -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
}
107 changes: 107 additions & 0 deletions app/upgrade_tracker_test.go
Original file line number Diff line number Diff line change
@@ -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)
}

0 comments on commit 2e7dac2

Please sign in to comment.