feat(store/v2): implement the feature to upgrade the store keys

core/store/upgrade.go

@@ -41,6 +41,19 @@ func (s *StoreUpgrades) IsDeleted(key string) bool {
 	return false
 }
 
+// IsRenamed returns true if the given key should be renamed
+func (s *StoreUpgrades) IsRenamed(key string) bool {
+	if s == nil {
+		return false
+	}
+	for _, re := range s.Renamed {
+		if re.OldKey == key {
+			return true
+		}
+	}
+	return false
+}
+
 // RenamedFrom returns the oldKey if it was renamed
 // Returns "" if it was not renamed
 func (s *StoreUpgrades) RenamedFrom(key string) string {

store/v2/commitment/iavl/tree.go

@@ -14,23 +14,31 @@ import (
 )
 
 var (
-	_ commitment.Tree      = (*IavlTree)(nil)
-	_ store.PausablePruner = (*IavlTree)(nil)
+	_ commitment.Tree          = (*IavlTree)(nil)
+	_ commitment.KVStoreGetter = (*IavlTree)(nil)
+	_ store.PausablePruner     = (*IavlTree)(nil)
 )
 
 // IavlTree is a wrapper around iavl.MutableTree.
 type IavlTree struct {
 	tree *iavl.MutableTree
+	db   corestore.KVStoreWithBatch
 }
 
 // NewIavlTree creates a new IavlTree instance.
 func NewIavlTree(db corestore.KVStoreWithBatch, logger log.Logger, cfg *Config) *IavlTree {
 	tree := iavl.NewMutableTree(dbm.NewWrapper(db), cfg.CacheSize, cfg.SkipFastStorageUpgrade, logger, iavl.AsyncPruningOption(true))
 	return &IavlTree{
 		tree: tree,
+		db:   db,
 	}
 }
 
+// GetKVStoreWithBatch implements the KVStoreGetter interface.
+func (t *IavlTree) GetKVStoreWithBatch() corestore.KVStoreWithBatch {
+	return t.db
+}
+
 // Remove removes the given key from the tree.
 func (t *IavlTree) Remove(key []byte) error {
 	_, _, err := t.tree.Remove(key)

store/v2/commitment/store.go

@@ -6,6 +6,7 @@
 	"fmt"
 	"io"
 	"math"
+	"sort"
 
 	protoio "github.com/cosmos/gogoproto/io"
 
@@ -23,12 +24,15 @@
 const (
 	commitInfoKeyFmt = "c/%d" // c/<version>
 	latestVersionKey = "c/latest"
+
+	batchFlushThreshold = 1 << 16 // 64KB
 )
 
 var (
 	_ store.Committer             = (*CommitStore)(nil)
 	_ snapshots.CommitSnapshotter = (*CommitStore)(nil)
 	_ store.PausablePruner        = (*CommitStore)(nil)
+	_ store.KVStoreGetter         = (*CommitStore)(nil)
 )
 
 // CommitStore is a wrapper around multiple Tree objects mapped by a unique store
@@ -126,7 +130,7 @@
 	}
 }
 
-func (c *CommitStore) LoadVersion(targetVersion uint64) error {
+func (c *CommitStore) LoadVersion(targetVersion uint64, upgrades *corestore.StoreUpgrades) error {
 	// Rollback the metadata to the target version.
 	latestVersion, err := c.GetLatestVersion()
 	if err != nil {
@@ -146,20 +150,150 @@
 		}
 	}
 
-	for _, tree := range c.multiTrees {
-		if err := tree.LoadVersion(targetVersion); err != nil {
+	// If the target version is greater than the latest version, it is the snapshot
+	// restore case, we should create a new commit info for the target version.
+	if targetVersion > latestVersion {
+		cInfo := c.WorkingCommitInfo(targetVersion)
+
+		for _, tree := range c.multiTrees {
+			if err := tree.LoadVersion(targetVersion); err != nil {
+				return err
+			}
+		}
+
+		return c.flushCommitInfo(targetVersion, cInfo)
+	}
+
+	storesKeys := make([]string, 0, len(c.multiTrees))
+	for storeKey := range c.multiTrees {
+		storesKeys = append(storesKeys, storeKey)
+	}
+
+	if upgrades != nil {
+		// deterministic iteration order for upgrades
+		// (as the underlying store may change and
+		// upgrades make store changes where the execution order may matter)
+		sort.Slice(storesKeys, func(i, j int) bool {
+			return storesKeys[i] < storesKeys[j]
+		})
+	}
+
+	commitInfo, err := c.GetCommitInfo(targetVersion)
+	if err != nil {
+		return err
+	}
+
+	for _, storeKey := range storesKeys {
+		// If it has been renamed, continue to the next store key.
+		if upgrades.IsRenamed(storeKey) {
+			continue
+		}
+
+		tree := c.multiTrees[storeKey]
+
+		removeTree := func(storeKey string) error {
+			if oldTree, ok := c.multiTrees[storeKey]; ok {
+				if err := oldTree.Close(); err != nil {
+					return err
+				}
+				delete(c.multiTrees, storeKey)
+			}
+			return nil
+		}
+
+		// If it has been deleted, remove the tree.
+		if upgrades.IsDeleted(storeKey) {
+			if err := removeTree(storeKey); err != nil {
+				return err
+			}
+			continue
+		}
+
+		// If it has been added, set the initial version.
+		if upgrades.IsAdded(storeKey) {
+			if err := tree.SetInitialVersion(targetVersion + 1); err != nil {
+				return err
+			}
+		}
+
+		commitID := commitInfo.GetStoreCommitID([]byte(storeKey))
+		// If it has been renamed, migrate the data.
+		if oldKey := upgrades.RenamedFrom(storeKey); oldKey != "" {
+			if err := c.migrateKVStore(oldKey, storeKey); err != nil {
+				return err
+			}
+			if err := removeTree(oldKey); err != nil {
+				return err
+			}
+			commitID = commitInfo.GetStoreCommitID([]byte(oldKey))
+		}
+
+		if err := tree.LoadVersion(commitID.Version); err != nil {
 			return err
 		}
 	}
 
-	// If the target version is greater than the latest version, it is the snapshot
-	// restore case, we should create a new commit info for the target version.
-	var cInfo *proof.CommitInfo
-	if targetVersion > latestVersion {
-		cInfo = c.WorkingCommitInfo(targetVersion)
+	return nil
+}
+
+// migrateKVStore migrates the data from the old key to the new key.
+// It should be synchronous to ensure that the data is migrated before the
+// old store is pruned.
+func (c *CommitStore) migrateKVStore(oldKey, newKey string) error {
+	var (
+		oldKVStore corestore.KVStoreWithBatch
+		newKVStore corestore.KVStoreWithBatch
+	)
+
+	if oldTree, ok := c.multiTrees[oldKey]; ok {
+		if getter, ok := oldTree.(KVStoreGetter); ok {
+			oldKVStore = getter.GetKVStoreWithBatch()
+		}
+	}
+
+	if newTree, ok := c.multiTrees[newKey]; ok {
+		if getter, ok := newTree.(KVStoreGetter); ok {
+			newKVStore = getter.GetKVStoreWithBatch()
+		}
 	}
 
-	return c.flushCommitInfo(targetVersion, cInfo)
+	if oldKVStore != nil && newKVStore != nil {
+		iter, err := oldKVStore.Iterator(nil, nil)
+		if err != nil {
+			return err
+		}
+		defer func() {
+			_ = iter.Close()
+		}()
+
+		batch := newKVStore.NewBatch()
+		for ; iter.Valid(); iter.Next() {
+			if err := batch.Set(iter.Key(), iter.Value()); err != nil {
+				return err
+			}
+			bs, err := batch.GetByteSize()
+			if err != nil {
+				return err
+			}
+			if bs > batchFlushThreshold {
+				if err := batch.Write(); err != nil {
+					return err
+				}
+				if err := batch.Close(); err != nil {
+					return err
+				}
+				batch = newKVStore.NewBatch()
+			}
+		}
+		if err := batch.Write(); err != nil {
+			return err
+		}
+		if err := batch.Close(); err != nil {
+			return err
+		}
+	}
+
+	return nil
 }
 
 func (c *CommitStore) GetCommitInfo(version uint64) (*proof.CommitInfo, error) {
@@ -490,7 +624,18 @@
 		}
 	}
 
-	return snapshotItem, c.LoadVersion(version)
+	return snapshotItem, c.LoadVersion(version, nil)
+}
+
+// GetKVStoreWithBatch implements store.KVStoreGetter.
+func (c *CommitStore) GetKVStoreWithBatch(storeKey string) corestore.KVStoreWithBatch {
+	if tree, ok := c.multiTrees[storeKey]; ok {
+		if kvGetter, ok := tree.(KVStoreGetter); ok {
+			return kvGetter.GetKVStoreWithBatch()
+		}
+	}
+
+	return nil
 }
 
 func (c *CommitStore) Close() (ferr error) {

store/v2/commitment/store_test_suite.go

@@ -19,6 +19,7 @@ import (
 const (
 	storeKey1 = "store1"
 	storeKey2 = "store2"
+	storeKey3 = "store3"
 )
 
 // CommitStoreTestSuite is a test suite to be used for all tree backends.
@@ -166,3 +167,63 @@ func (s *CommitStoreTestSuite) TestStore_Pruning() {
 		}
 	}
 }
+
+func (s *CommitStoreTestSuite) TestStore_Upgrades() {
+	storeKeys := []string{storeKey1, storeKey2, storeKey3}
+	commitDB := dbm.NewMemDB()
+	commitStore, err := s.NewStore(commitDB, storeKeys, log.NewNopLogger())
+	s.Require().NoError(err)
+
+	latestVersion := uint64(10)
+	kvCount := 10
+	for i := uint64(1); i <= latestVersion; i++ {
+		kvPairs := make(map[string]corestore.KVPairs)
+		for _, storeKey := range storeKeys {
+			kvPairs[storeKey] = corestore.KVPairs{}
+			for j := 0; j < kvCount; j++ {
+				key := []byte(fmt.Sprintf("key-%d-%d", i, j))
+				value := []byte(fmt.Sprintf("value-%d-%d", i, j))
+				kvPairs[storeKey] = append(kvPairs[storeKey], corestore.KVPair{Key: key, Value: value})
+			}
+		}
+		s.Require().NoError(commitStore.WriteChangeset(corestore.NewChangesetWithPairs(kvPairs)))
+		_, err = commitStore.Commit(i)
+		s.Require().NoError(err)
+	}
+
+	// create a new commitment store with upgrades
+	upgrades := &corestore.StoreUpgrades{
+		Added: []string{"newStore1", "newStore2"},
+		Renamed: []corestore.StoreRename{
+			{OldKey: storeKey1, NewKey: "renamedStore1"},
+		},
+		Deleted: []string{storeKey3},
+	}
+	newStoreKeys := []string{"store1", "store2", "store3", "renamedStore1", "newStore1", "newStore2"}
+	realStoreKeys := []string{"renamedStore1", "store2", "newStore1", "newStore2"}
+	commitStore, err = s.NewStore(commitDB, newStoreKeys, log.NewNopLogger())
+	s.Require().NoError(err)
+	err = commitStore.LoadVersion(latestVersion, upgrades)
+	s.Require().NoError(err)
+
+	// apply the changeset again
+	for i := latestVersion; i < latestVersion*2; i++ {
+		kvPairs := make(map[string]corestore.KVPairs)
+		for _, storeKey := range realStoreKeys {
+			kvPairs[storeKey] = corestore.KVPairs{}
+			for j := 0; j < kvCount; j++ {
+				key := []byte(fmt.Sprintf("key-%d-%d", i, j))
+				value := []byte(fmt.Sprintf("value-%d-%d", i, j))
+				kvPairs[storeKey] = append(kvPairs[storeKey], corestore.KVPair{Key: key, Value: value})
+			}
+		}
+		s.Require().NoError(commitStore.WriteChangeset(corestore.NewChangesetWithPairs(kvPairs)))
+		commitInfo, err := commitStore.Commit(i)
+		s.Require().NoError(err)
+		s.Require().NotNil(commitInfo)
+		s.Require().Equal(len(realStoreKeys), len(commitInfo.StoreInfos))
+		for _, storeKey := range realStoreKeys {
+			s.Require().NotNil(commitInfo.GetStoreCommitID([]byte(storeKey)))
+		}
+	}
+}

store/v2/commitment/tree.go

@@ -6,6 +6,7 @@ import (
 
 	ics23 "github.com/cosmos/ics23/go"
 
+	corestore "cosmossdk.io/core/store"
 	snapshotstypes "cosmossdk.io/store/v2/snapshots/types"
 )
 
@@ -56,3 +57,10 @@ type Importer interface {
 
 	io.Closer
 }
+
+// KVStoreGetter is an interface that allows getting the KVStoreWithBatch from a Tree.
+// It is used to allow the commitment Tree to be used as a KVStore which can be used
+// in the upgrade module.
+type KVStoreGetter interface {
+	GetKVStoreWithBatch() corestore.KVStoreWithBatch
+}