Add in-place store migrations

simapp/app.go

@@ -198,6 +198,9 @@ type SimApp struct {
 
 	// simulation manager
 	sm *module.SimulationManager
+
+	// the configurator
+	configurator module.Configurator
 }
 
 func init() {
@@ -393,7 +396,8 @@ func NewSimApp(
 
 	app.mm.RegisterInvariants(&app.CrisisKeeper)
 	app.mm.RegisterRoutes(app.Router(), app.QueryRouter(), encodingConfig.Amino)
-	app.mm.RegisterServices(module.NewConfigurator(app.MsgServiceRouter(), app.GRPCQueryRouter()))
+	app.configurator = module.NewConfigurator(app.MsgServiceRouter(), app.GRPCQueryRouter(), app.keys)
+	app.mm.RegisterServices(app.configurator)
 
 	// add test gRPC service for testing gRPC queries in isolation
 	testdata.RegisterQueryServer(app.GRPCQueryRouter(), testdata.QueryImpl{})
@@ -598,6 +602,29 @@ func (app *SimApp) RegisterTendermintService(clientCtx client.Context) {
 	tmservice.RegisterTendermintService(app.BaseApp.GRPCQueryRouter(), clientCtx, app.interfaceRegistry)
 }
 
+// RunMigrations performs in-place store migrations for all modules. This
+// function is not called automatically, it is meant to be called from an
+// x/upgrade UpgradeHandler.
+//
+// `migrationsMap` is a map of moduleName to fromVersion (unit64), where
+// fromVersion denotes the version from which we should migrate the module, the
+// target version being the module's latest ConsensusVersion.
+//
+// Example:
+//   cfg := module.NewConfigurator(...)
+//   app.UpgradeKeeper.SetUpgradeHandler("store-migration", func(ctx sdk.Context, plan upgradetypes.Plan) {
+//       err := app.RunMigrations(ctx, map[string]unint64{
+//           "bank": 1,     // Migrate x/bank from v1 to current x/bank's ConsensusVersion
+//           "staking": 8,  // Migrate x/staking from v8 to current x/staking's ConsensusVersion
+//      })
+//      if err != nil {
+//           panic(err)
+//      }
+//   })
+func (app *SimApp) RunMigrations(ctx sdk.Context, migrationsMap map[string]uint64) error {
+	return app.mm.RunMigrations(ctx, app.configurator, migrationsMap)
+}
+
 // RegisterSwaggerAPI registers swagger route with API Server
 func RegisterSwaggerAPI(ctx client.Context, rtr *mux.Router) {
 	statikFS, err := fs.New()

simapp/app_test.go

@@ -6,10 +6,13 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	abci "github.com/tendermint/tendermint/abci/types"
 	"github.com/tendermint/tendermint/libs/log"
+	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
 	dbm "github.com/tendermint/tm-db"
 
-	abci "github.com/tendermint/tendermint/abci/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
 )
 
 func TestSimAppExportAndBlockedAddrs(t *testing.T) {
@@ -45,3 +48,82 @@ func TestGetMaccPerms(t *testing.T) {
 	dup := GetMaccPerms()
 	require.Equal(t, maccPerms, dup, "duplicated module account permissions differed from actual module account permissions")
 }
+
+func TestRunMigrations(t *testing.T) {
+	db := dbm.NewMemDB()
+	app := NewSimApp(log.NewTMLogger(log.NewSyncWriter(os.Stdout)), db, nil, true, map[int64]bool{}, DefaultNodeHome, 0, MakeTestEncodingConfig(), EmptyAppOptions{})
+
+	// Create a new configurator for the purpose of this test.
+	app.configurator = module.NewConfigurator(app.MsgServiceRouter(), app.GRPCQueryRouter(), app.keys)
+
+	testCases := []struct {
+		name         string
+		moduleName   string
+		expRegErr    bool // errors while registering migration
+		expRegErrMsg string
+		expRunErr    bool // errors while running migration
+		expRunErrMsg string
+		expCalled    int
+	}{
+		{
+			"cannot register migration for non-existant module",
+			"foo",
+			true, "store key for module foo not found: not found", false, "", 0,
+		},
+		{
+			"throws error on RunMigrations if no migration registered for bank",
+			"",
+			false, "", true, "no migration found for module bank from version 0 to version 1: not found", 0,
+		},
+		{
+			"can register and run migration handler for x/bank",
+			"bank",
+			false, "", false, "", 1,
+		},
+		{
+			"cannot register migration handler for same module & fromVersion",
+			"bank",
+			true, "another migration for module bank and version 0 already exists: internal logic error", false, "", 0,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			var err error
+
+			// Since it's very hard to test actual in-place store migrations in
+			// tests (due to the difficulty of maintaing multiple versions of a
+			// module), we're just testing here that the migration logic is
+			// called.
+			called := 0
+
+			if tc.moduleName != "" {
+				err = app.configurator.RegisterMigration(tc.moduleName, 0, func(sdk.KVStore) error {
+					called++
+
+					return nil
+				})
+			}
+
+			if tc.expRegErr {
+				require.Error(t, err)
+				require.Equal(t, tc.expRegErrMsg, err.Error())
+
+				return
+			}
+			require.NoError(t, err)
+
+			err = app.RunMigrations(
+				app.NewContext(true, tmproto.Header{Height: app.LastBlockHeight()}),
+				map[string]uint64{"bank": 0},
+			)
+			if tc.expRunErr {
+				require.Error(t, err)
+				require.Equal(t, tc.expRunErrMsg, err.Error())
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, tc.expCalled, called)
+			}
+		})
+	}
+}

testutil/context.go

@@ -0,0 +1,25 @@
+package testutil
+
+import (
+	"github.com/tendermint/tendermint/libs/log"
+	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
+	dbm "github.com/tendermint/tm-db"
+
+	"github.com/cosmos/cosmos-sdk/store"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+// DefaultContext creates a sdk.Context that can be used in tests.
+func DefaultContext(key sdk.StoreKey, tkey sdk.StoreKey) sdk.Context {
+	db := dbm.NewMemDB()
+	cms := store.NewCommitMultiStore(db)
+	cms.MountStoreWithDB(key, sdk.StoreTypeIAVL, db)
+	cms.MountStoreWithDB(tkey, sdk.StoreTypeTransient, db)
+	err := cms.LoadLatestVersion()
+	if err != nil {
+		panic(err)
+	}
+	ctx := sdk.NewContext(cms, tmproto.Header{}, false, log.NewNopLogger())
+
+	return ctx
+}

types/context_test.go

@@ -8,13 +8,11 @@ import (
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/suite"
 	abci "github.com/tendermint/tendermint/abci/types"
-	"github.com/tendermint/tendermint/libs/log"
 	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
-	dbm "github.com/tendermint/tm-db"
 
 	"github.com/cosmos/cosmos-sdk/crypto/keys/secp256k1"
-	"github.com/cosmos/cosmos-sdk/store"
 	"github.com/cosmos/cosmos-sdk/tests/mocks"
+	"github.com/cosmos/cosmos-sdk/testutil"
 	"github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -26,23 +24,14 @@ func TestContextTestSuite(t *testing.T) {
 	suite.Run(t, new(contextTestSuite))
 }
 
-func (s *contextTestSuite) defaultContext(key types.StoreKey) types.Context {
-	db := dbm.NewMemDB()
-	cms := store.NewCommitMultiStore(db)
-	cms.MountStoreWithDB(key, types.StoreTypeIAVL, db)
-	s.Require().NoError(cms.LoadLatestVersion())
-	ctx := types.NewContext(cms, tmproto.Header{}, false, log.NewNopLogger())
-	return ctx
-}
-
 func (s *contextTestSuite) TestCacheContext() {
 	key := types.NewKVStoreKey(s.T().Name() + "_TestCacheContext")
 	k1 := []byte("hello")
 	v1 := []byte("world")
 	k2 := []byte("key")
 	v2 := []byte("value")
 
-	ctx := s.defaultContext(key)
+	ctx := testutil.DefaultContext(key, types.NewTransientStoreKey("transient_"+s.T().Name()))
 	store := ctx.KVStore(key)
 	store.Set(k1, v1)
 	s.Require().Equal(v1, store.Get(k1))
@@ -64,7 +53,7 @@ func (s *contextTestSuite) TestCacheContext() {
 
 func (s *contextTestSuite) TestLogContext() {
 	key := types.NewKVStoreKey(s.T().Name())
-	ctx := s.defaultContext(key)
+	ctx := testutil.DefaultContext(key, types.NewTransientStoreKey("transient_"+s.T().Name()))
 	ctrl := gomock.NewController(s.T())
 	s.T().Cleanup(ctrl.Finish)
 

types/module/configurator.go

@@ -1,6 +1,10 @@
 package module
 
-import "github.com/gogo/protobuf/grpc"
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
+	"github.com/gogo/protobuf/grpc"
+)
 
 // Configurator provides the hooks to allow modules to configure and register
 // their services in the RegisterServices method. It is designed to eventually
@@ -15,16 +19,31 @@ type Configurator interface {
 	// QueryServer returns a grpc.Server instance which allows registering services
 	// that will be exposed as gRPC services as well as ABCI query handlers.
 	QueryServer() grpc.Server
+
+	// RegisterMigration registers an in-place store migration for a module. The
+	// handler is a migration script to perform in-place migrations from version
+	// `fromVersion` to version `fromVersion+1`.
+	RegisterMigration(moduleName string, fromVersion uint64, handler func(store sdk.KVStore) error) error
 }
 
 type configurator struct {
 	msgServer   grpc.Server
 	queryServer grpc.Server
+
+	// storeKeys is used to access module stores inside in-place store migrations.
+	storeKeys map[string]*sdk.KVStoreKey
+	// migrations is a map of moduleName -> fromVersion -> migration script handler
+	migrations map[string]map[uint64]func(store sdk.KVStore) error
 }
 
 // NewConfigurator returns a new Configurator instance
-func NewConfigurator(msgServer grpc.Server, queryServer grpc.Server) Configurator {
-	return configurator{msgServer: msgServer, queryServer: queryServer}
+func NewConfigurator(msgServer grpc.Server, queryServer grpc.Server, storeKeys map[string]*sdk.KVStoreKey) Configurator {
+	return configurator{
+		msgServer:   msgServer,
+		queryServer: queryServer,
+		storeKeys:   storeKeys,
+		migrations:  map[string]map[uint64]func(store sdk.KVStore) error{},
+	}
 }
 
 var _ Configurator = configurator{}
@@ -38,3 +57,23 @@ func (c configurator) MsgServer() grpc.Server {
 func (c configurator) QueryServer() grpc.Server {
 	return c.queryServer
 }
+
+// RegisterMigration implements the Configurator.RegisterMigration method
+func (c configurator) RegisterMigration(moduleName string, fromVersion uint64, handler func(store sdk.KVStore) error) error {
+	if c.migrations[moduleName][fromVersion] != nil {
+		return sdkerrors.Wrapf(sdkerrors.ErrLogic, "another migration for module %s and version %d already exists", moduleName, fromVersion)
+	}
+
+	_, found := c.storeKeys[moduleName]
+	if !found {
+		return sdkerrors.Wrapf(sdkerrors.ErrNotFound, "store key for module %s not found", moduleName)
+	}
+
+	if c.migrations[moduleName] == nil {
+		c.migrations[moduleName] = map[uint64]func(store sdk.KVStore) error{}
+	}
+
+	c.migrations[moduleName][fromVersion] = handler
+
+	return nil
+}

types/module/module.go

@@ -40,6 +40,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/codec"
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 )
 
 //__________________________________________________________________________________________
@@ -174,6 +175,9 @@ type AppModule interface {
 	// RegisterServices allows a module to register services
 	RegisterServices(Configurator)
 
+	// ConsensusVersion tracks state-breaking versions of the module
+	ConsensusVersion() uint64
+
 	// ABCI
 	BeginBlock(sdk.Context, abci.RequestBeginBlock)
 	EndBlock(sdk.Context, abci.RequestEndBlock) []abci.ValidatorUpdate
@@ -208,6 +212,9 @@ func (gam GenesisOnlyAppModule) LegacyQuerierHandler(*codec.LegacyAmino) sdk.Que
 // RegisterServices registers all services.
 func (gam GenesisOnlyAppModule) RegisterServices(Configurator) {}
 
+// ConsensusVersion tracks state-breaking versions of the module.
+func (gam GenesisOnlyAppModule) ConsensusVersion() uint64 { return 0 }
+
 // BeginBlock returns an empty module begin-block
 func (gam GenesisOnlyAppModule) BeginBlock(ctx sdk.Context, req abci.RequestBeginBlock) {}
 
@@ -328,6 +335,23 @@ func (m *Manager) ExportGenesis(ctx sdk.Context, cdc codec.JSONMarshaler) map[st
 	return genesisData
 }
 
+// RunMigrations performs in-place store migrations for all modules.
+func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, migrationsMap map[string]uint64) error {
+	c, ok := cfg.(configurator)
+	if !ok {
+		return sdkerrors.Wrapf(sdkerrors.ErrInvalidType, "expected %T, got %T", configurator{}, cfg)
+	}
+
+	for moduleName, module := range m.Modules {
+		err := runMigrations(c, ctx, moduleName, migrationsMap[moduleName], module.ConsensusVersion())
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // BeginBlock performs begin block functionality for all modules. It creates a
 // child context with an event manager to aggregate events emitted from all
 // modules.
@@ -369,3 +393,35 @@ func (m *Manager) EndBlock(ctx sdk.Context, req abci.RequestEndBlock) abci.Respo
 		Events:           ctx.EventManager().ABCIEvents(),
 	}
 }
+
+// runMigrations runs all in-place store migrations for one given module from a
+// version to another version.
+func runMigrations(cfg configurator, ctx sdk.Context, moduleName string, fromVersion, toVersion uint64) error {
+	// No-op if toVersion is the initial version.
+	// Some modules don't have a store key (e.g. vesting), in this case, their
+	// ConsensusVersion will always stay at 0, and running migrations on
+	// those modules will be skipped on this line.
+	if toVersion == 0 {
+		return nil
+	}
+
+	storeKey, found := cfg.storeKeys[moduleName]
+	if !found {
+		return sdkerrors.Wrapf(sdkerrors.ErrNotFound, "store key for module %s not found", moduleName)
+	}
+
+	// Run in-place migrations for the module sequentially until toVersion.
+	for i := fromVersion; i < toVersion; i++ {
+		migrateFn, found := cfg.migrations[moduleName][i]
+		if !found {
+			return sdkerrors.Wrapf(sdkerrors.ErrNotFound, "no migration found for module %s from version %d to version %d", moduleName, i, i+1)
+		}
+
+		err := migrateFn(ctx.KVStore(storeKey))
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}

types/module/module_test.go

@@ -179,7 +179,7 @@ func TestManager_RegisterQueryServices(t *testing.T) {
 
 	msgRouter := mocks.NewMockServer(mockCtrl)
 	queryRouter := mocks.NewMockServer(mockCtrl)
-	cfg := module.NewConfigurator(msgRouter, queryRouter)
+	cfg := module.NewConfigurator(msgRouter, queryRouter, nil)
 	mockAppModule1.EXPECT().RegisterServices(cfg).Times(1)
 	mockAppModule2.EXPECT().RegisterServices(cfg).Times(1)