feat(store/v2): handle store upgrades in RootStore (#18277)

Co-authored-by: cool-developer <[email protected]>

store/kv/branch/store.go

@@ -1,6 +1,7 @@
 package branch
 
 import (
+	"fmt"
 	"io"
 	"slices"
 	"sync"
@@ -84,17 +85,16 @@ func (s *Store) GetChangeset() *store.Changeset {
 	return store.NewChangeset(pairs...)
 }
 
-func (s *Store) Reset() error {
+func (s *Store) Reset(toVersion uint64) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
-	latestVersion, err := s.storage.GetLatestVersion()
-	if err != nil {
-		return err
+	if err := s.storage.SetLatestVersion(toVersion); err != nil {
+		return fmt.Errorf("failed to set SS latest version %d: %w", toVersion, err)
 	}
 
 	clear(s.changeset)
-	s.version = latestVersion
+	s.version = toVersion
 
 	return nil
 }

store/kv/branch/store_test.go

@@ -63,7 +63,7 @@ func (s *StoreTestSuite) TestGetChangeset() {
 }
 
 func (s *StoreTestSuite) TestReset() {
-	s.Require().NoError(s.kvStore.Reset())
+	s.Require().NoError(s.kvStore.Reset(1))
 
 	cs := s.kvStore.GetChangeset()
 	s.Require().Zero(cs.Size())

store/kv/gas/store.go

@@ -62,8 +62,8 @@ func (s *Store) GetChangeset() *store.Changeset {
 	return s.parent.GetChangeset()
 }
 
-func (s *Store) Reset() error {
-	return s.parent.Reset()
+func (s *Store) Reset(toVersion uint64) error {
+	return s.parent.Reset(toVersion)
 }
 
 func (s *Store) Write() {

store/kv/gas/store_test.go

@@ -35,7 +35,7 @@ func (s *StoreTestSuite) SetupTest() {
 }
 
 func (s *StoreTestSuite) TearDownTest() {
-	err := s.gasKVStore.Reset()
+	err := s.gasKVStore.Reset(1)
 	s.Require().NoError(err)
 }
 

store/kv/mem/store.go

@@ -88,7 +88,7 @@ func (s *Store) GetChangeset() *store.Changeset {
 	return store.NewChangeset(kvPairs...)
 }
 
-func (s *Store) Reset() error {
+func (s *Store) Reset(_ uint64) error {
 	s.tree.Clear()
 	return nil
 }

store/kv/mem/store_test.go

@@ -44,7 +44,7 @@ func (s *StoreTestSuite) TestGetChangeset() {
 }
 
 func (s *StoreTestSuite) TestReset() {
-	s.Require().NoError(s.kvStore.Reset())
+	s.Require().NoError(s.kvStore.Reset(1))
 
 	cs := s.kvStore.GetChangeset()
 	s.Require().Zero(cs.Size())

store/kv/trace/store.go

@@ -80,8 +80,8 @@ func (s *Store) Delete(key []byte) {
 	s.parent.Delete(key)
 }
 
-func (s *Store) Reset() error {
-	return s.parent.Reset()
+func (s *Store) Reset(toVersion uint64) error {
+	return s.parent.Reset(toVersion)
 }
 
 func (s *Store) Write() {

store/root/store.go

@@ -101,15 +101,6 @@ func (s *Store) GetSCStore(_ string) store.Tree {
 	return s.stateCommitment
 }
 
-func (s *Store) LoadLatestVersion() error {
-	lv, err := s.GetLatestVersion()
-	if err != nil {
-		return err
-	}
-
-	return s.loadVersion(lv, nil)
-}
-
 // LastCommitID returns a CommitID based off of the latest internal CommitInfo.
 // If an internal CommitInfo is not set, a new one will be returned with only the
 // latest version set, which is based off of the SS view.
@@ -173,11 +164,6 @@ func (s *Store) Query(storeKey string, version uint64, key []byte, prove bool) (
 	return result, nil
 }
 
-// LoadVersion loads a specific version returning an error upon failure.
-func (s *Store) LoadVersion(v uint64) (err error) {
-	return s.loadVersion(v, nil)
-}
-
 // GetKVStore returns the store's root KVStore. Any writes to this store without
 // branching will be committed to SC and SS upon Commit(). Branching will create
 // a branched KVStore that allow writes to be discarded and propagated to the
@@ -198,17 +184,37 @@ func (s *Store) GetBranchedKVStore(_ string) store.BranchedKVStore {
 	return s.rootKVStore
 }
 
-func (s *Store) loadVersion(v uint64, upgrades any) error {
+func (s *Store) LoadLatestVersion() error {
+	lv, err := s.GetLatestVersion()
+	if err != nil {
+		return err
+	}
+
+	return s.loadVersion(lv)
+}
+
+func (s *Store) LoadVersion(version uint64) error {
+	return s.loadVersion(version)
+}
+
+func (s *Store) loadVersion(v uint64) error {
 	s.logger.Debug("loading version", "version", v)
 
+	// Reset the root KVStore s.t. the latest version is v. Any writes will
+	// overwrite existing versions.
+	if err := s.rootKVStore.Reset(v); err != nil {
+		return err
+	}
+
 	if err := s.stateCommitment.LoadVersion(v); err != nil {
 		return fmt.Errorf("failed to load SS version %d: %w", v, err)
 	}
 
-	// TODO: Complete this method to handle upgrades. See legacy RMS loadVersion()
-	// for reference.
-	//
-	// Ref: https://github.com/cosmos/cosmos-sdk/issues/17314
+	s.workingHash = nil
+	s.commitHeader = nil
+
+	// set lastCommitInfo explicitly s.t. Commit commits the correct version, i.e. v+1
+	s.lastCommitInfo = &store.CommitInfo{Version: v}
 
 	return nil
 }
@@ -305,7 +311,7 @@ func (s *Store) Commit() ([]byte, error) {
 		s.lastCommitInfo.Timestamp = s.commitHeader.GetTime()
 	}
 
-	if err := s.rootKVStore.Reset(); err != nil {
+	if err := s.rootKVStore.Reset(version); err != nil {
 		return nil, fmt.Errorf("failed to reset root KVStore: %w", err)
 	}
 
@@ -341,16 +347,14 @@ func (s *Store) writeSC() error {
 		version = previousHeight + 1
 	}
 
-	workingHash := s.stateCommitment.WorkingHash()
-
 	s.lastCommitInfo = &store.CommitInfo{
 		Version: version,
 		StoreInfos: []store.StoreInfo{
 			{
 				Name: defaultStoreKey,
 				CommitID: store.CommitID{
 					Version: version,
-					Hash:    workingHash,
+					Hash:    s.stateCommitment.WorkingHash(),
 				},
 			},
 		},

store/root/store_test.go

@@ -128,6 +128,73 @@ func (s *RootStoreTestSuite) TestBranch() {
 	s.Require().Equal([]byte("updated_bar"), bs.Get([]byte("foo")))
 }
 
+func (s *RootStoreTestSuite) TestLoadVersion() {
+	// write and commit a few changesets
+	for v := 1; v <= 5; v++ {
+		bs := s.rootStore.GetBranchedKVStore("")
+		val := fmt.Sprintf("val%03d", v) // val001, val002, ..., val005
+		bs.Set([]byte("key"), []byte(val))
+
+		workingHash, err := s.rootStore.WorkingHash()
+		s.Require().NoError(err)
+		s.Require().NotNil(workingHash)
+
+		commitHash, err := s.rootStore.Commit()
+		s.Require().NoError(err)
+		s.Require().NotNil(commitHash)
+		s.Require().Equal(workingHash, commitHash)
+	}
+
+	// ensure the latest version is correct
+	latest, err := s.rootStore.GetLatestVersion()
+	s.Require().NoError(err)
+	s.Require().Equal(uint64(5), latest)
+
+	// attempt to load a non-existent version
+	err = s.rootStore.LoadVersion(6)
+	s.Require().Error(err)
+
+	// attempt to load a previously committed version
+	err = s.rootStore.LoadVersion(3)
+	s.Require().NoError(err)
+
+	// ensure the latest version is correct
+	latest, err = s.rootStore.GetLatestVersion()
+	s.Require().NoError(err)
+	s.Require().Equal(uint64(3), latest)
+
+	// query state and ensure values returned are based on the loaded version
+	kvStore := s.rootStore.GetKVStore("")
+	val := kvStore.Get([]byte("key"))
+	s.Require().Equal([]byte("val003"), val)
+
+	// attempt to write and commit a few changesets
+	for v := 4; v <= 5; v++ {
+		bs := s.rootStore.GetBranchedKVStore("")
+		val := fmt.Sprintf("overwritten_val%03d", v) // overwritten_val004, overwritten_val005
+		bs.Set([]byte("key"), []byte(val))
+
+		workingHash, err := s.rootStore.WorkingHash()
+		s.Require().NoError(err)
+		s.Require().NotNil(workingHash)
+
+		commitHash, err := s.rootStore.Commit()
+		s.Require().NoError(err)
+		s.Require().NotNil(commitHash)
+		s.Require().Equal(workingHash, commitHash)
+	}
+
+	// ensure the latest version is correct
+	latest, err = s.rootStore.GetLatestVersion()
+	s.Require().NoError(err)
+	s.Require().Equal(uint64(5), latest)
+
+	// query state and ensure values returned are based on the loaded version
+	kvStore = s.rootStore.GetKVStore("")
+	val = kvStore.Get([]byte("key"))
+	s.Require().Equal([]byte("overwritten_val005"), val)
+}
+
 func (s *RootStoreTestSuite) TestMultiBranch() {
 	// write and commit a changeset
 	bs := s.rootStore.GetKVStore("")

store/store.go

@@ -52,7 +52,10 @@ type RootStore interface {
 	// TracingEnabled returns true if tracing is enabled on the RootStore.
 	TracingEnabled() bool
 
+	// LoadVersion loads the RootStore to the given version.
 	LoadVersion(version uint64) error
+	// LoadLatestVersion behaves identically to LoadVersion except it loads the
+	// latest version implicitly.
 	LoadLatestVersion() error
 
 	// GetLatestVersion returns the latest version, i.e. height, committed.
@@ -80,6 +83,20 @@ type RootStore interface {
 	io.Closer
 }
 
+// UpgradeableRootStore extends the RootStore interface to support loading versions
+// with upgrades.
+type UpgradeableRootStore interface {
+	RootStore
+
+	// LoadVersionAndUpgrade behaves identically to LoadVersion except it also
+	// accepts a StoreUpgrades object that defines a series of transformations to
+	// apply to store keys (if any).
+	//
+	// Note, handling StoreUpgrades is optional depending on the underlying RootStore
+	// implementation.
+	LoadVersionAndUpgrade(version uint64, upgrades *StoreUpgrades) error
+}
+
 // BranchedRootStore defines an extension of the RootStore interface that allows
 // for nested branching and flushing of writes. It extends RootStore by allowing
 // a caller to call Branch() which should return a BranchedRootStore that has all
@@ -117,7 +134,7 @@ type KVStore interface {
 	GetChangeset() *Changeset
 
 	// Reset resets the store, which is implementation dependent.
-	Reset() error
+	Reset(toVersion uint64) error
 
 	// Iterator creates a new Iterator over the domain [start, end). Note:
 	//

store/upgrade.go

@@ -0,0 +1,53 @@
+package store
+
+import "golang.org/x/exp/slices"
+
+// StoreUpgrades defines a series of transformations to apply the RootStore upon
+// loading a version.
+type StoreUpgrades struct {
+	Add    []string      `json:"add"`
+	Rename []StoreRename `json:"rename"`
+	Delete []string      `json:"delete"`
+}
+
+// StoreRename defines a change in a store key. All data previously stored under
+// the current store key should be migrated to the new store key, while also
+// deleting the old store key.
+type StoreRename struct {
+	OldKey string `json:"old_key"`
+	NewKey string `json:"new_key"`
+}
+
+// IsAdded returns true if the given key should be added.
+func (s *StoreUpgrades) IsAdded(key string) bool {
+	if s == nil {
+		return false
+	}
+
+	return slices.Contains(s.Add, key)
+}
+
+// IsDeleted returns true if the given key should be deleted.
+func (s *StoreUpgrades) IsDeleted(key string) bool {
+	if s == nil {
+		return false
+	}
+
+	return slices.Contains(s.Delete, key)
+}
+
+// RenamedFrom returns the oldKey if it was renamed. It returns an empty string
+// if it was not renamed.
+func (s *StoreUpgrades) RenamedFrom(key string) string {
+	if s == nil {
+		return ""
+	}
+
+	for _, re := range s.Rename {
+		if re.NewKey == key {
+			return re.OldKey
+		}
+	}
+
+	return ""
+}