fix: rollback command don't actually delete multistore versions

baseapp/baseapp.go

@@ -283,11 +283,8 @@ func DefaultStoreLoader(ms sdk.CommitMultiStore) error {
 
 // CommitMultiStore returns the root multi-store.
 // App constructor can use this to access the `cms`.
-// UNSAFE: only safe to use during app initialization.
+// UNSAFE: must not be used during the abci life cycle.
 func (app *BaseApp) CommitMultiStore() sdk.CommitMultiStore {
-	if app.sealed {
-		panic("cannot call CommitMultiStore() after baseapp is sealed")
-	}
 	return app.cms
 }
 

server/mock/store.go

@@ -144,6 +144,10 @@ func (ms multiStore) Restore(
 	panic("not implemented")
 }
 
+func (ms multiStore) RollbackToVersion(version int64) error {
+	panic("not implemented")
+}
+
 var _ sdk.KVStore = kvStore{}
 
 type kvStore struct {

server/rollback.go

@@ -4,13 +4,13 @@ import (
 	"fmt"
 
 	"github.com/cosmos/cosmos-sdk/client/flags"
-	"github.com/cosmos/cosmos-sdk/store/rootmulti"
+	"github.com/cosmos/cosmos-sdk/server/types"
 	"github.com/spf13/cobra"
 	tmcmd "github.com/tendermint/tendermint/cmd/tendermint/commands"
 )
 
 // NewRollbackCmd creates a command to rollback tendermint and multistore state by one height.
-func NewRollbackCmd(defaultNodeHome string) *cobra.Command {
+func NewRollbackCmd(appCreator types.AppCreator, defaultNodeHome string) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "rollback",
 		Short: "rollback cosmos-sdk and tendermint state by one height",
@@ -30,14 +30,17 @@ application.
 			if err != nil {
 				return err
 			}
+			app := appCreator(ctx.Logger, db, nil, ctx.Viper)
 			// rollback tendermint state
 			height, hash, err := tmcmd.RollbackState(ctx.Config)
 			if err != nil {
 				return fmt.Errorf("failed to rollback tendermint state: %w", err)
 			}
 			// rollback the multistore
-			cms := rootmulti.NewStore(db, ctx.Logger)
-			cms.RollbackToVersion(height)
+
+			if err := app.CommitMultiStore().RollbackToVersion(height); err != nil {
+				return fmt.Errorf("failed to rollback to version: %w", err)
+			}
 
 			fmt.Printf("Rolled back state to height %d and hash %X", height, hash)
 			return nil

server/types/app.go

@@ -15,6 +15,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/client"
 	"github.com/cosmos/cosmos-sdk/server/api"
 	"github.com/cosmos/cosmos-sdk/server/config"
+	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
 // ServerStartTime defines the time duration that the server need to stay running after startup
@@ -51,6 +52,9 @@ type (
 
 		// RegisterTendermintService registers the gRPC Query service for tendermint queries.
 		RegisterTendermintService(clientCtx client.Context)
+
+		// Return the multistore instance
+		CommitMultiStore() sdk.CommitMultiStore
 	}
 
 	// AppCreator is a function that allows us to lazily initialize an

server/util.go

@@ -290,7 +290,7 @@ func AddCommands(rootCmd *cobra.Command, defaultNodeHome string, appCreator type
 		tendermintCmd,
 		ExportCmd(appExport, defaultNodeHome),
 		version.NewVersionCommand(),
-		NewRollbackCmd(defaultNodeHome),
+		NewRollbackCmd(appCreator, defaultNodeHome),
 	)
 }
 

store/iavl/store.go

@@ -234,6 +234,12 @@ func (st *Store) DeleteVersions(versions ...int64) error {
 	return st.tree.DeleteVersions(versions...)
 }
 
+// LoadVersionForOverwriting attempts to load a tree at a previously committed
+// version, or the latest version below it. Any versions greater than targetVersion will be deleted.
+func (st *Store) LoadVersionForOverwriting(targetVersion int64) (int64, error) {
+	return st.tree.LoadVersionForOverwriting(targetVersion)
+}
+
 // Implements types.KVStore.
 func (st *Store) Iterator(start, end []byte) types.Iterator {
 	iterator, err := st.tree.Iterator(start, end, true)

store/iavl/tree.go

@@ -34,6 +34,7 @@ type (
 		SetInitialVersion(version uint64)
 		Iterator(start, end []byte, ascending bool) (types.Iterator, error)
 		AvailableVersions() []int
+		LoadVersionForOverwriting(targetVersion int64) (int64, error)
 	}
 
 	// immutableTree is a simple wrapper around a reference to an iavl.ImmutableTree
@@ -99,3 +100,7 @@ func (it *immutableTree) GetImmutable(version int64) (*iavl.ImmutableTree, error
 func (it *immutableTree) AvailableVersions() []int {
 	return []int{}
 }
+
+func (it *immutableTree) LoadVersionForOverwriting(targetVersion int64) (int64, error) {
+	panic("cannot call 'LoadVersionForOverwriting' on an immutable IAVL tree")
+}

store/rootmulti/rollback_test.go

@@ -0,0 +1,69 @@
+package rootmulti_test
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/cosmos/cosmos-sdk/simapp"
+	simtestutil "github.com/cosmos/cosmos-sdk/testutil/sims"
+	"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"
+)
+
+func TestRollback(t *testing.T) {
+	db := dbm.NewMemDB()
+	options := simapp.SetupOptions{
+		Logger:  log.NewNopLogger(),
+		DB:      db,
+		AppOpts: simtestutil.NewAppOptionsWithFlagHome(simapp.DefaultNodeHome),
+	}
+	app := simapp.NewSimappWithCustomOptions(t, false, options)
+	app.Commit()
+	ver0 := app.LastBlockHeight()
+	// commit 10 blocks
+	for i := int64(1); i <= 10; i++ {
+		header := tmproto.Header{
+			Height:  ver0 + i,
+			AppHash: app.LastCommitID().Hash,
+		}
+		app.BeginBlock(abci.RequestBeginBlock{Header: header})
+		ctx := app.NewContext(false, header)
+		store := ctx.KVStore(app.GetKey("bank"))
+		store.Set([]byte("key"), []byte(fmt.Sprintf("value%d", i)))
+		app.Commit()
+	}
+
+	require.Equal(t, ver0+10, app.LastBlockHeight())
+	store := app.NewContext(true, tmproto.Header{}).KVStore(app.GetKey("bank"))
+	require.Equal(t, []byte("value10"), store.Get([]byte("key")))
+
+	// rollback 5 blocks
+	target := ver0 + 5
+	require.NoError(t, app.CommitMultiStore().RollbackToVersion(target))
+	require.Equal(t, target, app.LastBlockHeight())
+
+	// recreate app to have clean check state
+	app = simapp.NewSimApp(options.Logger, options.DB, nil, true, simtestutil.NewAppOptionsWithFlagHome(simapp.DefaultNodeHome))
+	store = app.NewContext(true, tmproto.Header{}).KVStore(app.GetKey("bank"))
+	require.Equal(t, []byte("value5"), store.Get([]byte("key")))
+
+	// commit another 5 blocks with different values
+	for i := int64(6); i <= 10; i++ {
+		header := tmproto.Header{
+			Height:  ver0 + i,
+			AppHash: app.LastCommitID().Hash,
+		}
+		app.BeginBlock(abci.RequestBeginBlock{Header: header})
+		ctx := app.NewContext(false, header)
+		store := ctx.KVStore(app.GetKey("bank"))
+		store.Set([]byte("key"), []byte(fmt.Sprintf("VALUE%d", i)))
+		app.Commit()
+	}
+
+	require.Equal(t, ver0+10, app.LastBlockHeight())
+	store = app.NewContext(true, tmproto.Header{}).KVStore(app.GetKey("bank"))
+	require.Equal(t, []byte("VALUE10"), store.Get([]byte("key")))
+}

store/rootmulti/store.go

@@ -933,29 +933,26 @@ func (rs *Store) buildCommitInfo(version int64) *types.CommitInfo {
 }
 
 // RollbackToVersion delete the versions after `target` and update the latest version.
-func (rs *Store) RollbackToVersion(target int64) int64 {
-	if target < 0 {
-		panic("Negative rollback target")
-	}
-	current := getLatestVersion(rs.db)
-	if target >= current {
-		return current
-	}
-	for ; current > target; current-- {
-		rs.pruningManager.HandleHeight(current)
-	}
-	if err := rs.pruneStores(); err != nil {
-		panic(err)
+func (rs *Store) RollbackToVersion(target int64) error {
+	if target <= 0 {
+		return fmt.Errorf("invalid rollback height target: %d", target)
 	}
 
-	// update latest height
-	bz, err := gogotypes.StdInt64Marshal(current)
-	if err != nil {
-		panic(err)
+	for key, store := range rs.stores {
+		if store.GetStoreType() == types.StoreTypeIAVL {
+			// If the store is wrapped with an inter-block cache, we must first unwrap
+			// it to get the underlying IAVL store.
+			store = rs.GetCommitKVStore(key)
+			_, err := store.(*iavl.Store).LoadVersionForOverwriting(target)
+			if err != nil {
+				return err
+			}
+		}
 	}
 
-	rs.db.Set([]byte(latestVersionKey), bz)
-	return current
+	rs.flushMetadata(rs.db, target, rs.buildCommitInfo(target))
+
+	return rs.LoadLatestVersion()
 }
 
 func (rs *Store) flushMetadata(db dbm.DB, version int64, cInfo *types.CommitInfo) {

store/types/store.go

@@ -188,6 +188,9 @@ type CommitMultiStore interface {
 
 	// SetIAVLCacheSize sets the cache size of the IAVL tree.
 	SetIAVLCacheSize(size int)
+
+	// RollbackToVersion rollback the db to specific version(height).
+	RollbackToVersion(version int64) error
 }
 
 //---------subsp-------------------------------