feat(store/v2): route to the commitment during migration (#22290)

Author: cool-develope

store/v2/commitment/iavl/tree.go

@@ -14,6 +14,7 @@ import (
 
 var (
 	_ commitment.Tree      = (*IavlTree)(nil)
+	_ commitment.Reader    = (*IavlTree)(nil)
 	_ store.PausablePruner = (*IavlTree)(nil)
 )
 
@@ -76,6 +77,7 @@ func (t *IavlTree) GetProof(version uint64, key []byte) (*ics23.CommitmentProof,
 	return immutableTree.GetProof(key)
 }
 
+// Get implements the Reader interface.
 func (t *IavlTree) Get(version uint64, key []byte) ([]byte, error) {
 	immutableTree, err := t.tree.GetImmutable(int64(version))
 	if err != nil {
@@ -85,6 +87,16 @@ func (t *IavlTree) Get(version uint64, key []byte) ([]byte, error) {
 	return immutableTree.Get(key)
 }
 
+// Iterator implements the Reader interface.
+func (t *IavlTree) Iterator(version uint64, start, end []byte, ascending bool) (corestore.Iterator, error) {
+	immutableTree, err := t.tree.GetImmutable(int64(version))
+	if err != nil {
+		return nil, fmt.Errorf("failed to get immutable tree at version %d: %w", version, err)
+	}
+
+	return immutableTree.Iterator(start, end, ascending)
+}
+
 // GetLatestVersion returns the latest version of the tree.
 func (t *IavlTree) GetLatestVersion() (uint64, error) {
 	v, err := t.tree.GetLatestVersion()

store/v2/commitment/iavl/tree_test.go

@@ -133,3 +133,82 @@ func TestIavlTree(t *testing.T) {
 	// close the db
 	require.NoError(t, tree.Close())
 }
+
+func TestIavlTreeIterator(t *testing.T) {
+	// generate a new tree
+	tree := generateTree()
+	require.NotNil(t, tree)
+
+	// write a batch of version 1
+	require.NoError(t, tree.Set([]byte("key1"), []byte("value1")))
+	require.NoError(t, tree.Set([]byte("key2"), []byte("value2")))
+	require.NoError(t, tree.Set([]byte("key3"), []byte("value3")))
+
+	// commit the batch
+	_, _, err := tree.Commit()
+	require.NoError(t, err)
+
+	// write a batch of version 2
+	require.NoError(t, tree.Set([]byte("key4"), []byte("value4")))
+	require.NoError(t, tree.Set([]byte("key5"), []byte("value5")))
+	require.NoError(t, tree.Set([]byte("key6"), []byte("value6")))
+	require.NoError(t, tree.Remove([]byte("key1"))) // delete key1
+	_, _, err = tree.Commit()
+	require.NoError(t, err)
+
+	// write a batch of version 3
+	require.NoError(t, tree.Set([]byte("key7"), []byte("value7")))
+	require.NoError(t, tree.Set([]byte("key8"), []byte("value8")))
+	_, _, err = tree.Commit()
+	require.NoError(t, err)
+
+	// iterate over all keys
+	iter, err := tree.Iterator(3, nil, nil, true)
+	require.NoError(t, err)
+	// expect all keys to be iterated over
+	expectedKeys := []string{"key2", "key3", "key4", "key5", "key6", "key7", "key8"}
+	count := 0
+	for i := 0; iter.Valid(); i++ {
+		require.Equal(t, expectedKeys[i], string(iter.Key()))
+		iter.Next()
+		count++
+	}
+	require.Equal(t, len(expectedKeys), count)
+	require.NoError(t, iter.Close())
+
+	// iterate over all keys in reverse
+	iter, err = tree.Iterator(3, nil, nil, false)
+	require.NoError(t, err)
+	expectedKeys = []string{"key8", "key7", "key6", "key5", "key4", "key3", "key2"}
+	for i := 0; iter.Valid(); i++ {
+		require.Equal(t, expectedKeys[i], string(iter.Key()))
+		iter.Next()
+	}
+	require.NoError(t, iter.Close())
+
+	// iterate over keys with version 1
+	iter, err = tree.Iterator(1, nil, nil, true)
+	require.NoError(t, err)
+	expectedKeys = []string{"key1", "key2", "key3"}
+	count = 0
+	for i := 0; iter.Valid(); i++ {
+		require.Equal(t, expectedKeys[i], string(iter.Key()))
+		iter.Next()
+		count++
+	}
+	require.Equal(t, len(expectedKeys), count)
+	require.NoError(t, iter.Close())
+
+	// iterate over keys with version 2
+	iter, err = tree.Iterator(2, nil, nil, false)
+	require.NoError(t, err)
+	expectedKeys = []string{"key6", "key5", "key4", "key3", "key2"}
+	count = 0
+	for i := 0; iter.Valid(); i++ {
+		require.Equal(t, expectedKeys[i], string(iter.Key()))
+		iter.Next()
+		count++
+	}
+	require.Equal(t, len(expectedKeys), count)
+	require.NoError(t, iter.Close())
+}

store/v2/commitment/store.go

@@ -25,6 +25,11 @@ var (
 	_ store.UpgradeableStore      = (*CommitStore)(nil)
 	_ snapshots.CommitSnapshotter = (*CommitStore)(nil)
 	_ store.PausablePruner        = (*CommitStore)(nil)
+
+	// NOTE: It is not recommended to use the CommitStore as a reader. This is only used
+	// during the migration process. Generally, the SC layer does not provide a reader
+	// in the store/v2.
+	_ store.VersionedReader = (*CommitStore)(nil)
 )
 
 // MountTreeFn is a function that mounts a tree given a store key.
@@ -275,21 +280,71 @@ func (c *CommitStore) GetProof(storeKey []byte, version uint64, key []byte) ([]p
 	return []proof.CommitmentOp{commitOp, *storeCommitmentOp}, nil
 }
 
-// Get implements store.VersionedReader.
-func (c *CommitStore) Get(storeKey []byte, version uint64, key []byte) ([]byte, error) {
-	tree, ok := c.multiTrees[conv.UnsafeBytesToStr(storeKey)]
+// getReader returns a reader for the given store key. It will return an error if the
+// store key does not exist or the tree does not implement the Reader interface.
+// WARNING: This function is only used during the migration process. The SC layer
+// generally does not provide a reader for the CommitStore.
+func (c *CommitStore) getReader(storeKey string) (Reader, error) {
+	tree, ok := c.multiTrees[storeKey]
 	if !ok {
 		return nil, fmt.Errorf("store %s not found", storeKey)
 	}
 
-	bz, err := tree.Get(version, key)
+	reader, ok := tree.(Reader)
+	if !ok {
+		return nil, fmt.Errorf("tree for store %s does not implement Reader", storeKey)
+	}
+
+	return reader, nil
+}
+
+// VersionExists implements store.VersionedReader.
+func (c *CommitStore) VersionExists(version uint64) (bool, error) {
+	ci, err := c.metadata.GetCommitInfo(version)
+	return ci != nil, err
+}
+
+// Get implements store.VersionedReader.
+func (c *CommitStore) Get(storeKey []byte, version uint64, key []byte) ([]byte, error) {
+	reader, err := c.getReader(conv.UnsafeBytesToStr(storeKey))
+	if err != nil {
+		return nil, err
+	}
+
+	bz, err := reader.Get(version, key)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get key %s from store %s: %w", key, storeKey, err)
 	}
 
 	return bz, nil
 }
 
+// Has implements store.VersionedReader.
+func (c *CommitStore) Has(storeKey []byte, version uint64, key []byte) (bool, error) {
+	val, err := c.Get(storeKey, version, key)
+	return val != nil, err
+}
+
+// Iterator implements store.VersionedReader.
+func (c *CommitStore) Iterator(storeKey []byte, version uint64, start, end []byte) (corestore.Iterator, error) {
+	reader, err := c.getReader(conv.UnsafeBytesToStr(storeKey))
+	if err != nil {
+		return nil, err
+	}
+
+	return reader.Iterator(version, start, end, true)
+}
+
+// ReverseIterator implements store.VersionedReader.
+func (c *CommitStore) ReverseIterator(storeKey []byte, version uint64, start, end []byte) (corestore.Iterator, error) {
+	reader, err := c.getReader(conv.UnsafeBytesToStr(storeKey))
+	if err != nil {
+		return nil, err
+	}
+
+	return reader.Iterator(version, start, end, false)
+}
+
 // Prune implements store.Pruner.
 func (c *CommitStore) Prune(version uint64) error {
 	// prune the metadata

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"
 )
 
@@ -29,19 +30,20 @@ type Tree interface {
 	SetInitialVersion(version uint64) error
 	GetProof(version uint64, key []byte) (*ics23.CommitmentProof, error)
 
-	// Get attempts to retrieve a value from the tree for a given version.
-	//
-	// NOTE: This method only exists to support migration from IAVL v0/v1 to v2.
-	// Once migration is complete, this method should be removed and/or not used.
-	Get(version uint64, key []byte) ([]byte, error)
-
 	Prune(version uint64) error
 	Export(version uint64) (Exporter, error)
 	Import(version uint64) (Importer, error)
 
 	io.Closer
 }
 
+// Reader is the optional interface that is only used to read data from the tree
+// during the migration process.
+type Reader interface {
+	Get(version uint64, key []byte) ([]byte, error)
+	Iterator(version uint64, start, end []byte, ascending bool) (corestore.Iterator, error)
+}
+
 // Exporter is the interface that wraps the basic Export methods.
 type Exporter interface {
 	Next() (*snapshotstypes.SnapshotIAVLItem, error)

store/v2/database.go

@@ -29,10 +29,6 @@ type VersionedReader interface {
 
 	Iterator(storeKey []byte, version uint64, start, end []byte) (corestore.Iterator, error)
 	ReverseIterator(storeKey []byte, version uint64, start, end []byte) (corestore.Iterator, error)
-
-	// Close releases associated resources. It should NOT be idempotent. It must
-	// only be called once and any call after may panic.
-	io.Closer
 }
 
 // UpgradableDatabase defines an API for a versioned database that allows pruning

store/v2/root/reader.go

@@ -14,38 +14,38 @@ var (
 // operations. This is useful for exposing a read-only view of the RootStore at
 // a specific version in history, which could also be the latest state.
 type ReaderMap struct {
-	rootStore store.RootStore
-	version   uint64
+	vReader store.VersionedReader
+	version uint64
 }
 
-func NewReaderMap(v uint64, rs store.RootStore) *ReaderMap {
+func NewReaderMap(v uint64, vr store.VersionedReader) *ReaderMap {
 	return &ReaderMap{
-		rootStore: rs,
-		version:   v,
+		vReader: vr,
+		version: v,
 	}
 }
 
-func (roa *ReaderMap) GetReader(actor []byte) (corestore.Reader, error) {
-	return NewReader(roa.version, roa.rootStore, actor), nil
+func (rm *ReaderMap) GetReader(actor []byte) (corestore.Reader, error) {
+	return NewReader(rm.version, rm.vReader, actor), nil
 }
 
 // Reader represents a read-only adapter for accessing data from the root store.
 type Reader struct {
-	version   uint64          // The version of the data.
-	rootStore store.RootStore // The root store to read data from.
-	actor     []byte          // The actor associated with the data.
+	version uint64                // The version of the data.
+	vReader store.VersionedReader // The root store to read data from.
+	actor   []byte                // The actor associated with the data.
 }
 
-func NewReader(v uint64, rs store.RootStore, actor []byte) *Reader {
+func NewReader(v uint64, vr store.VersionedReader, actor []byte) *Reader {
 	return &Reader{
-		version:   v,
-		rootStore: rs,
-		actor:     actor,
+		version: v,
+		vReader: vr,
+		actor:   actor,
 	}
 }
 
 func (roa *Reader) Has(key []byte) (bool, error) {
-	val, err := roa.rootStore.GetStateStorage().Has(roa.actor, roa.version, key)
+	val, err := roa.vReader.Has(roa.actor, roa.version, key)
 	if err != nil {
 		return false, err
 	}
@@ -54,13 +54,13 @@ func (roa *Reader) Has(key []byte) (bool, error) {
 }
 
 func (roa *Reader) Get(key []byte) ([]byte, error) {
-	return roa.rootStore.GetStateStorage().Get(roa.actor, roa.version, key)
+	return roa.vReader.Get(roa.actor, roa.version, key)
 }
 
 func (roa *Reader) Iterator(start, end []byte) (corestore.Iterator, error) {
-	return roa.rootStore.GetStateStorage().Iterator(roa.actor, roa.version, start, end)
+	return roa.vReader.Iterator(roa.actor, roa.version, start, end)
 }
 
 func (roa *Reader) ReverseIterator(start, end []byte) (corestore.Iterator, error) {
-	return roa.rootStore.GetStateStorage().ReverseIterator(roa.actor, roa.version, start, end)
+	return roa.vReader.ReverseIterator(roa.actor, roa.version, start, end)
 }

store/v2/root/store.go

@@ -118,27 +118,51 @@ func (s *Store) SetInitialVersion(v uint64) error {
 	return s.stateCommitment.SetInitialVersion(v)
 }
 
-func (s *Store) StateLatest() (uint64, corestore.ReaderMap, error) {
-	v, err := s.GetLatestVersion()
+// getVersionedReader returns a VersionedReader based on the given version. If the
+// version exists in the state storage, it returns the state storage.
+// If not, it checks if the state commitment implements the VersionedReader interface
+// and the version exists in the state commitment, since the state storage will be
+// synced during migration.
+func (s *Store) getVersionedReader(version uint64) (store.VersionedReader, error) {
+	isExist, err := s.stateStorage.VersionExists(version)
 	if err != nil {
-		return 0, nil, err
+		return nil, err
+	}
+	if isExist {
+		return s.stateStorage, nil
 	}
 
-	return v, NewReaderMap(v, s), nil
+	if vReader, ok := s.stateCommitment.(store.VersionedReader); ok {
+		isExist, err := vReader.VersionExists(version)
+		if err != nil {
+			return nil, err
+		}
+		if isExist {
+			return vReader, nil
+		}
+	}
+
+	return nil, fmt.Errorf("version %d does not exist", version)
 }
 
-// StateAt checks if the requested version is present in ss.
-func (s *Store) StateAt(v uint64) (corestore.ReaderMap, error) {
-	// check if version is present in state storage
-	isExist, err := s.stateStorage.VersionExists(v)
+func (s *Store) StateLatest() (uint64, corestore.ReaderMap, error) {
+	v, err := s.GetLatestVersion()
 	if err != nil {
-		return nil, err
+		return 0, nil, err
 	}
-	if !isExist {
-		return nil, fmt.Errorf("version %d does not exist", v)
+
+	vReader, err := s.getVersionedReader(v)
+	if err != nil {
+		return 0, nil, err
 	}
 
-	return NewReaderMap(v, s), nil
+	return v, NewReaderMap(v, vReader), nil
+}
+
+// StateAt returns a read-only view of the state at a given version.
+func (s *Store) StateAt(v uint64) (corestore.ReaderMap, error) {
+	vReader, err := s.getVersionedReader(v)
+	return NewReaderMap(v, vReader), err
 }
 
 func (s *Store) GetStateStorage() store.VersionedWriter {