feat: Add skipFastStorageUpgrade to MutableTree to control fast storage upgrade

basic_test.go

@@ -470,15 +470,15 @@ func TestPersistence(t *testing.T) {
 	}
 
 	// Construct some tree and save it
-	t1, err := NewMutableTree(db, 0)
+	t1, err := NewMutableTree(db, 0, false)
 	require.NoError(t, err)
 	for key, value := range records {
 		t1.Set([]byte(key), []byte(value))
 	}
 	t1.SaveVersion()
 
 	// Load a tree
-	t2, err := NewMutableTree(db, 0)
+	t2, err := NewMutableTree(db, 0, false)
 	require.NoError(t, err)
 	t2.Load()
 	for key, value := range records {
@@ -524,7 +524,7 @@ func TestProof(t *testing.T) {
 
 func TestTreeProof(t *testing.T) {
 	db := db.NewMemDB()
-	tree, err := NewMutableTree(db, 100)
+	tree, err := NewMutableTree(db, 100, false)
 	require.NoError(t, err)
 	hash, err := tree.Hash()
 	require.NoError(t, err)

benchmarks/bench_test.go

@@ -25,7 +25,7 @@ func randBytes(length int) []byte {
 }
 
 func prepareTree(b *testing.B, db db.DB, size, keyLen, dataLen int) (*iavl.MutableTree, [][]byte) {
-	t, err := iavl.NewMutableTreeWithOpts(db, size, nil)
+	t, err := iavl.NewMutableTreeWithOpts(db, size, nil, false)
 	require.NoError(b, err)
 	keys := make([][]byte, size)
 

benchmarks/cosmos-exim/main.go

@@ -90,7 +90,7 @@ func runExport(dbPath string) (int64, map[string][]*iavl.ExportNode, error) {
 	if err != nil {
 		return 0, nil, err
 	}
-	tree, err := iavl.NewMutableTree(tmdb.NewPrefixDB(ldb, []byte("s/k:main/")), 0)
+	tree, err := iavl.NewMutableTree(tmdb.NewPrefixDB(ldb, []byte("s/k:main/")), 0, false)
 	if err != nil {
 		return 0, nil, err
 	}
@@ -105,7 +105,7 @@ func runExport(dbPath string) (int64, map[string][]*iavl.ExportNode, error) {
 	totalStats := Stats{}
 	for _, name := range stores {
 		db := tmdb.NewPrefixDB(ldb, []byte("s/k:"+name+"/"))
-		tree, err := iavl.NewMutableTree(db, 0)
+		tree, err := iavl.NewMutableTree(db, 0, false)
 		if err != nil {
 			return 0, nil, err
 		}
@@ -170,7 +170,7 @@ func runImport(version int64, exports map[string][]*iavl.ExportNode) error {
 		if err != nil {
 			return err
 		}
-		newTree, err := iavl.NewMutableTree(newDB, 0)
+		newTree, err := iavl.NewMutableTree(newDB, 0, false)
 		if err != nil {
 			return err
 		}

cmd/iaviewer/main.go

@@ -120,7 +120,7 @@ func ReadTree(dir string, version int, prefix []byte) (*iavl.MutableTree, error)
 		db = dbm.NewPrefixDB(db, prefix)
 	}
 
-	tree, err := iavl.NewMutableTree(db, DefaultCacheSize)
+	tree, err := iavl.NewMutableTree(db, DefaultCacheSize, false)
 	if err != nil {
 		return nil, err
 	}

export_test.go

@@ -14,7 +14,7 @@ import (
 // setupExportTreeBasic sets up a basic tree with a handful of
 // create/update/delete operations over a few versions.
 func setupExportTreeBasic(t require.TestingT) *ImmutableTree {
-	tree, err := NewMutableTree(db.NewMemDB(), 0)
+	tree, err := NewMutableTree(db.NewMemDB(), 0, false)
 	require.NoError(t, err)
 
 	tree.Set([]byte("x"), []byte{255})
@@ -59,7 +59,7 @@ func setupExportTreeRandom(t *testing.T) *ImmutableTree {
 	)
 
 	r := rand.New(rand.NewSource(randSeed))
-	tree, err := NewMutableTree(db.NewMemDB(), 0)
+	tree, err := NewMutableTree(db.NewMemDB(), 0, false)
 	require.NoError(t, err)
 
 	var version int64
@@ -119,7 +119,7 @@ func setupExportTreeSized(t require.TestingT, treeSize int) *ImmutableTree {
 	)
 
 	r := rand.New(rand.NewSource(randSeed))
-	tree, err := NewMutableTree(db.NewMemDB(), 0)
+	tree, err := NewMutableTree(db.NewMemDB(), 0, false)
 	require.NoError(t, err)
 
 	for i := 0; i < treeSize; i++ {
@@ -176,7 +176,7 @@ func TestExporter(t *testing.T) {
 
 func TestExporter_Import(t *testing.T) {
 	testcases := map[string]*ImmutableTree{
-		"empty tree": NewImmutableTree(db.NewMemDB(), 0),
+		"empty tree": NewImmutableTree(db.NewMemDB(), 0, false),
 		"basic tree": setupExportTreeBasic(t),
 	}
 	if !testing.Short() {
@@ -192,7 +192,7 @@ func TestExporter_Import(t *testing.T) {
 			exporter := tree.Export()
 			defer exporter.Close()
 
-			newTree, err := NewMutableTree(db.NewMemDB(), 0)
+			newTree, err := NewMutableTree(db.NewMemDB(), 0, false)
 			require.NoError(t, err)
 			importer, err := newTree.Import(tree.Version())
 			require.NoError(t, err)
@@ -256,7 +256,7 @@ func TestExporter_Close(t *testing.T) {
 }
 
 func TestExporter_DeleteVersionErrors(t *testing.T) {
-	tree, err := NewMutableTree(db.NewMemDB(), 0)
+	tree, err := NewMutableTree(db.NewMemDB(), 0, false)
 	require.NoError(t, err)
 
 	tree.Set([]byte("a"), []byte{1})

immutable_tree.go

@@ -14,28 +14,31 @@ import (
 // Returned key/value byte slices must not be modified, since they may point to data located inside
 // IAVL which would also be modified.
 type ImmutableTree struct {
-	root    *Node
-	ndb     *nodeDB
-	version int64
+	root                   *Node
+	ndb                    *nodeDB
+	version                int64
+	skipFastStorageUpgrade bool
 }
 
 // NewImmutableTree creates both in-memory and persistent instances
-func NewImmutableTree(db dbm.DB, cacheSize int) *ImmutableTree {
+func NewImmutableTree(db dbm.DB, cacheSize int, skipFastStorageUpgrade bool) *ImmutableTree {
 	if db == nil {
 		// In-memory Tree.
 		return &ImmutableTree{}
 	}
 	return &ImmutableTree{
 		// NodeDB-backed Tree.
-		ndb: newNodeDB(db, cacheSize, nil),
+		ndb:                    newNodeDB(db, cacheSize, nil),
+		skipFastStorageUpgrade: skipFastStorageUpgrade,
 	}
 }
 
 // NewImmutableTreeWithOpts creates an ImmutableTree with the given options.
-func NewImmutableTreeWithOpts(db dbm.DB, cacheSize int, opts *Options) *ImmutableTree {
+func NewImmutableTreeWithOpts(db dbm.DB, cacheSize int, opts *Options, skipFastStorageUpgrade bool) *ImmutableTree {
 	return &ImmutableTree{
 		// NodeDB-backed Tree.
-		ndb: newNodeDB(db, cacheSize, opts),
+		ndb:                    newNodeDB(db, cacheSize, opts),
+		skipFastStorageUpgrade: skipFastStorageUpgrade,
 	}
 }
 
@@ -172,36 +175,40 @@ func (t *ImmutableTree) GetWithIndex(key []byte) (int64, []byte, error) {
 // Get returns the value of the specified key if it exists, or nil.
 // The returned value must not be modified, since it may point to data stored within IAVL.
 // Get potentially employs a more performant strategy than GetWithIndex for retrieving the value.
+// If tree.skipFastStorageUpgrade is true, this will work almost the same as GetWithIndex.
 func (t *ImmutableTree) Get(key []byte) ([]byte, error) {
 	if t.root == nil {
 		return nil, nil
 	}
 
-	// attempt to get a FastNode directly from db/cache.
-	// if call fails, fall back to the original IAVL logic in place.
-	fastNode, err := t.ndb.GetFastNode(key)
-	if err != nil {
-		_, result, err := t.root.get(t, key)
-		return result, err
-	}
-
-	if fastNode == nil {
-		// If the tree is of the latest version and fast node is not in the tree
-		// then the regular node is not in the tree either because fast node
-		// represents live state.
-		if t.version == t.ndb.latestVersion {
-			return nil, nil
+	if !t.skipFastStorageUpgrade {
+		// attempt to get a FastNode directly from db/cache.
+		// if call fails, fall back to the original IAVL logic in place.
+		fastNode, err := t.ndb.GetFastNode(key)
+		if err != nil {
+			_, result, err := t.root.get(t, key)
+			return result, err
 		}
 
-		_, result, err := t.root.get(t, key)
-		return result, err
-	}
+		if fastNode == nil {
+			// If the tree is of the latest version and fast node is not in the tree
+			// then the regular node is not in the tree either because fast node
+			// represents live state.
+			if t.version == t.ndb.latestVersion {
+				return nil, nil
+			}
+
+			_, result, err := t.root.get(t, key)
+			return result, err
+		}
 
-	if fastNode.GetVersionLastUpdatedAt() <= t.version {
-		return fastNode.GetValue(), nil
+		if fastNode.GetVersionLastUpdatedAt() <= t.version {
+			return fastNode.GetValue(), nil
+		}
 	}
 
-	// Otherwise the cached node was updated later than the current tree. In this case,
+	// otherwise skipFastStorageUpgrade is true or
+	// the cached node was updated later than the current tree. In this case,
 	// we need to use the regular stategy for reading from the current tree to avoid staleness.
 	_, result, err := t.root.get(t, key)
 	return result, err
@@ -239,13 +246,15 @@ func (t *ImmutableTree) Iterate(fn func(key []byte, value []byte) bool) (bool, e
 
 // Iterator returns an iterator over the immutable tree.
 func (t *ImmutableTree) Iterator(start, end []byte, ascending bool) (dbm.Iterator, error) {
-	isFastCacheEnabled, err := t.IsFastCacheEnabled()
-	if err != nil {
-		return nil, err
-	}
+	if !t.skipFastStorageUpgrade {
+		isFastCacheEnabled, err := t.IsFastCacheEnabled()
+		if err != nil {
+			return nil, err
+		}
 
-	if isFastCacheEnabled {
-		return NewFastIterator(start, end, ascending, t.ndb), nil
+		if isFastCacheEnabled {
+			return NewFastIterator(start, end, ascending, t.ndb), nil
+		}
 	}
 	return NewIterator(start, end, ascending, t), nil
 }

import_test.go

@@ -10,7 +10,7 @@ import (
 )
 
 func ExampleImporter() {
-	tree, err := NewMutableTree(db.NewMemDB(), 0)
+	tree, err := NewMutableTree(db.NewMemDB(), 0, false)
 	if err != nil {
 		// handle err
 	}
@@ -41,7 +41,7 @@ func ExampleImporter() {
 		exported = append(exported, node)
 	}
 
-	newTree, err := NewMutableTree(db.NewMemDB(), 0)
+	newTree, err := NewMutableTree(db.NewMemDB(), 0, false)
 	if err != nil {
 		// handle err
 	}
@@ -63,14 +63,14 @@ func ExampleImporter() {
 }
 
 func TestImporter_NegativeVersion(t *testing.T) {
-	tree, err := NewMutableTree(db.NewMemDB(), 0)
+	tree, err := NewMutableTree(db.NewMemDB(), 0, false)
 	require.NoError(t, err)
 	_, err = tree.Import(-1)
 	require.Error(t, err)
 }
 
 func TestImporter_NotEmpty(t *testing.T) {
-	tree, err := NewMutableTree(db.NewMemDB(), 0)
+	tree, err := NewMutableTree(db.NewMemDB(), 0, false)
 	require.NoError(t, err)
 	tree.Set([]byte("a"), []byte{1})
 	_, _, err = tree.SaveVersion()
@@ -83,13 +83,13 @@ func TestImporter_NotEmpty(t *testing.T) {
 func TestImporter_NotEmptyDatabase(t *testing.T) {
 	db := db.NewMemDB()
 
-	tree, err := NewMutableTree(db, 0)
+	tree, err := NewMutableTree(db, 0, false)
 	require.NoError(t, err)
 	tree.Set([]byte("a"), []byte{1})
 	_, _, err = tree.SaveVersion()
 	require.NoError(t, err)
 
-	tree, err = NewMutableTree(db, 0)
+	tree, err = NewMutableTree(db, 0, false)
 	require.NoError(t, err)
 	_, err = tree.Load()
 	require.NoError(t, err)
@@ -99,7 +99,7 @@ func TestImporter_NotEmptyDatabase(t *testing.T) {
 }
 
 func TestImporter_NotEmptyUnsaved(t *testing.T) {
-	tree, err := NewMutableTree(db.NewMemDB(), 0)
+	tree, err := NewMutableTree(db.NewMemDB(), 0, false)
 	require.NoError(t, err)
 	tree.Set([]byte("a"), []byte{1})
 
@@ -126,7 +126,7 @@ func TestImporter_Add(t *testing.T) {
 	for desc, tc := range testcases {
 		tc := tc // appease scopelint
 		t.Run(desc, func(t *testing.T) {
-			tree, err := NewMutableTree(db.NewMemDB(), 0)
+			tree, err := NewMutableTree(db.NewMemDB(), 0, false)
 			require.NoError(t, err)
 			importer, err := tree.Import(1)
 			require.NoError(t, err)
@@ -143,7 +143,7 @@ func TestImporter_Add(t *testing.T) {
 }
 
 func TestImporter_Add_Closed(t *testing.T) {
-	tree, err := NewMutableTree(db.NewMemDB(), 0)
+	tree, err := NewMutableTree(db.NewMemDB(), 0, false)
 	require.NoError(t, err)
 	importer, err := tree.Import(1)
 	require.NoError(t, err)
@@ -155,7 +155,7 @@ func TestImporter_Add_Closed(t *testing.T) {
 }
 
 func TestImporter_Close(t *testing.T) {
-	tree, err := NewMutableTree(db.NewMemDB(), 0)
+	tree, err := NewMutableTree(db.NewMemDB(), 0, false)
 	require.NoError(t, err)
 	importer, err := tree.Import(1)
 	require.NoError(t, err)
@@ -172,7 +172,7 @@ func TestImporter_Close(t *testing.T) {
 }
 
 func TestImporter_Commit(t *testing.T) {
-	tree, err := NewMutableTree(db.NewMemDB(), 0)
+	tree, err := NewMutableTree(db.NewMemDB(), 0, false)
 	require.NoError(t, err)
 	importer, err := tree.Import(1)
 	require.NoError(t, err)
@@ -188,7 +188,7 @@ func TestImporter_Commit(t *testing.T) {
 }
 
 func TestImporter_Commit_Closed(t *testing.T) {
-	tree, err := NewMutableTree(db.NewMemDB(), 0)
+	tree, err := NewMutableTree(db.NewMemDB(), 0, false)
 	require.NoError(t, err)
 	importer, err := tree.Import(1)
 	require.NoError(t, err)
@@ -203,7 +203,7 @@ func TestImporter_Commit_Closed(t *testing.T) {
 }
 
 func TestImporter_Commit_Empty(t *testing.T) {
-	tree, err := NewMutableTree(db.NewMemDB(), 0)
+	tree, err := NewMutableTree(db.NewMemDB(), 0, false)
 	require.NoError(t, err)
 	importer, err := tree.Import(3)
 	require.NoError(t, err)
@@ -232,7 +232,7 @@ func BenchmarkImport(b *testing.B) {
 	b.StartTimer()
 
 	for n := 0; n < b.N; n++ {
-		newTree, err := NewMutableTree(db.NewMemDB(), 0)
+		newTree, err := NewMutableTree(db.NewMemDB(), 0, false)
 		require.NoError(b, err)
 		importer, err := newTree.Import(tree.Version())
 		require.NoError(b, err)