feat: add --unsafe-remove-modules flag to Rollback cmd (#546)

* add --unsafe-remove-modules flag to Rollback cmd

* adds public DeleteKVStore method on root multistore
* rollback command accepts list of store keys names to forcibly delete

this is useful for rolling back an upgrade that adds modules.
rollbacks are performed by loading & committing the previous version.
without this new functionality, the rollback will fail because no store
version will exist for modules added during the upgrade.

to properly rollback the state, pass in a list of the added module names
and they will be completely removed before the rollback of pre-existing
modules takes place:
```
chain rollback --unsafe-remove-modules mynewmodule,othernewmodule
```

* add DeleteVersionsFrom to Tree

like DeleteVersionsTo, but deletes the given version & upwards

* write deletion of current & future versions

* rename DeleteKVStore -> DeleteLatestVersion

* remove unnecessary deleteKVStore call

the DeleteVersionsFrom on the unwrapped KVStore should be sufficient

server/rollback.go

@@ -8,11 +8,13 @@ import (
 
 	"github.com/cosmos/cosmos-sdk/client/flags"
 	"github.com/cosmos/cosmos-sdk/server/types"
+	"github.com/cosmos/cosmos-sdk/store/rootmulti"
 )
 
 // NewRollbackCmd creates a command to rollback tendermint and multistore state by one height.
 func NewRollbackCmd(appCreator types.AppCreator, defaultNodeHome string) *cobra.Command {
 	var removeBlock bool
+	moduleKeysToDelete := make([]string, 0)
 
 	cmd := &cobra.Command{
 		Use:   "rollback",
@@ -39,6 +41,17 @@ application.
 			if err != nil {
 				return fmt.Errorf("failed to rollback tendermint state: %w", err)
 			}
+
+			// for rolling back upgrades that add modules, we must first forcibly delete those.
+			// otherwise, the rollback will panic because no version of new modules will exist.
+			store := app.CommitMultiStore().(*rootmulti.Store)
+			for _, key := range moduleKeysToDelete {
+				fmt.Printf("deleting latest version of KVStore with key %s\n", key)
+				if err := store.DeleteLatestVersion(key); err != nil {
+					return err
+				}
+			}
+
 			// rollback the multistore
 
 			if err := app.CommitMultiStore().RollbackToVersion(height); err != nil {
@@ -52,5 +65,6 @@ application.
 
 	cmd.Flags().String(flags.FlagHome, defaultNodeHome, "The application home directory")
 	cmd.Flags().BoolVar(&removeBlock, "hard", false, "remove last block as well as state")
+	cmd.Flags().StringSliceVar(&moduleKeysToDelete, "unsafe-remove-modules", []string{}, "force delete KV stores with the provided prefix. useful for rolling back an upgrade that adds a module")
 	return cmd
 }

store/iavl/store.go

@@ -229,6 +229,11 @@ func (st *Store) DeleteVersionsTo(version int64) error {
 	return st.tree.DeleteVersionsTo(version)
 }
 
+// DeleteVersionsFrom deletes from the given version upwards form the MutableTree.
+func (st *Store) DeleteVersionsFrom(version int64) error {
+	return st.tree.DeleteVersionsFrom(version)
+}
+
 // 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) {

store/iavl/tree.go

@@ -30,6 +30,7 @@ type (
 		WorkingHash() []byte
 		VersionExists(version int64) bool
 		DeleteVersionsTo(version int64) error
+		DeleteVersionsFrom(version int64) error
 		GetVersioned(key []byte, version int64) ([]byte, error)
 		GetImmutable(version int64) (*iavl.ImmutableTree, error)
 		SetInitialVersion(version uint64)
@@ -67,6 +68,10 @@ func (it *immutableTree) DeleteVersionsTo(_ int64) error {
 	panic("cannot call 'DeleteVersionsTo' on an immutable IAVL tree")
 }
 
+func (it *immutableTree) DeleteVersionsFrom(_ int64) error {
+	panic("cannot call 'DeleteVersionsFrom' on an immutable IAVL tree")
+}
+
 func (it *immutableTree) SetInitialVersion(_ uint64) {
 	panic("cannot call 'SetInitialVersion' on an immutable IAVL tree")
 }

store/rootmulti/store.go

@@ -992,6 +992,68 @@ func (rs *Store) buildCommitInfo(version int64) *types.CommitInfo {
 	}
 }
 
+// DeleteLatestVersion finds a store with the given key name and deletes its latest version.
+// The store is deregistered from the rootmulti store. Calls to buildCommitInfo will not include it.
+// For stores with IAVL types, the deletion is written to the disk.
+// This is a destructive operation to be used with caution.
+// The reason it was added was to allow for rollbacks of upgrades that add modules.
+// Stores that do not exist in the version prior to upgrade can be forcibly deleted
+// before calling Rollback()
+func (rs *Store) DeleteLatestVersion(keyName string) error {
+	ver := GetLatestVersion(rs.db)
+	if ver == 0 {
+		return fmt.Errorf("unable to delete KVStore with key name %s, latest version is 0", keyName)
+	}
+
+	// Find the store key with the provided name
+	var key types.StoreKey = nil
+	for k := range rs.storesParams {
+		if k.Name() == keyName {
+			key = k
+			break
+		}
+	}
+	if key == nil {
+		return fmt.Errorf("no store found with key name %s", keyName)
+	}
+
+	// Get the KVStore for that key
+	cInfo, err := rs.GetCommitInfo(ver)
+	if err != nil {
+		return err
+	}
+	infos := make(map[string]types.StoreInfo)
+	for _, storeInfo := range cInfo.StoreInfos {
+		infos[storeInfo.Name] = storeInfo
+	}
+	commitID := rs.getCommitID(infos, key.Name())
+	store, err := rs.loadCommitStoreFromParams(key, commitID, rs.storesParams[key])
+	if err != nil {
+		return errors.Wrap(err, "failed to load store")
+	}
+
+	rs.logger.Debug("deleting KVStore", "key", key.Name(), "latest version", ver)
+
+	// for IAVL stores, commit the deletion of the latest version to disk.
+	if store.GetStoreType() == types.StoreTypeIAVL {
+		// unwrap the caching layer
+		store = rs.GetCommitKVStore(key)
+		if err := store.(*iavl.Store).DeleteVersionsFrom(ver); err != nil {
+			return errors.Wrapf(err, "failed to delete versions %d onwards of %s store", ver, key.Name())
+		}
+	}
+
+	// deregister store from the rootmulti store
+	// Any future buildCommitInfo will no longer include the store.
+	if _, ok := rs.stores[key]; ok {
+		delete(rs.stores, key)
+		delete(rs.storesParams, key)
+		delete(rs.keysByName, key.Name())
+	}
+
+	return nil
+}
+
 // RollbackToVersion delete the versions after `target` and update the latest version.
 func (rs *Store) RollbackToVersion(target int64) error {
 	if target <= 0 {
@@ -1003,9 +1065,10 @@ func (rs *Store) RollbackToVersion(target int64) error {
 			// 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)
+			rs.logger.Debug("loading version %d for store with key %s (%s)\n", target, key.String(), key.Name())
 			_, err := store.(*iavl.Store).LoadVersionForOverwriting(target)
 			if err != nil {
-				return err
+				return errors.Wrapf(err, "failed loading version %d for store with key name '%s'", target, key.Name())
 			}
 		}
 	}