From 007cd557919a8d0b66d6e57cc30bbea48086d6b6 Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Mon, 26 Dec 2022 10:50:08 +0100
Subject: [PATCH] feat: speed up rollback command (backport #636) (#648)

Co-authored-by: yihuang <huang@crypto.com>
Co-authored-by: Marko Baricevic <marbar3778@yahoo.com>
---
 CHANGELOG.md    |  1 +
 mutable_tree.go | 19 +++++++------------
 nodedb.go       | 36 ++++++++++++------------------------
 3 files changed, 20 insertions(+), 36 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3f5d3e246..0cd91992d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@
 ## Unreleased
 
 - [#640](https://github.com/cosmos/iavl/pull/640) commit `NodeDB` batch in `LoadVersionForOverwriting`.
+- [#636](https://github.com/cosmos/iavl/pull/636) Speed up rollback method: `LoadVersionForOverwriting`.
 
 ## 0.19.4 (October 28, 2022)
 
diff --git a/mutable_tree.go b/mutable_tree.go
index f704b375f..831a142e4 100644
--- a/mutable_tree.go
+++ b/mutable_tree.go
@@ -630,17 +630,18 @@ func (tree *MutableTree) LoadVersionForOverwriting(targetVersion int64) (int64,
 		return latestVersion, err
 	}
 
+	tree.mtx.Lock()
+	defer tree.mtx.Unlock()
+
+	tree.ndb.resetLatestVersion(latestVersion)
+
 	if !tree.skipFastStorageUpgrade {
-		if err := tree.enableFastStorageAndCommitLocked(); err != nil {
+		// it'll repopulates the fast node index because of version mismatch.
+		if _, err := tree.enableFastStorageAndCommitIfNotEnabled(); err != nil {
 			return latestVersion, err
 		}
 	}
 
-	tree.ndb.resetLatestVersion(latestVersion)
-
-	tree.mtx.Lock()
-	defer tree.mtx.Unlock()
-
 	for v := range tree.versions {
 		if v > targetVersion {
 			delete(tree.versions, v)
@@ -706,12 +707,6 @@ func (tree *MutableTree) enableFastStorageAndCommitIfNotEnabled() (bool, error)
 	return true, nil
 }
 
-func (tree *MutableTree) enableFastStorageAndCommitLocked() error {
-	tree.mtx.Lock()
-	defer tree.mtx.Unlock()
-	return tree.enableFastStorageAndCommit()
-}
-
 func (tree *MutableTree) enableFastStorageAndCommit() error {
 	var err error
 
diff --git a/nodedb.go b/nodedb.go
index 62c56d1ec..bc8064a07 100644
--- a/nodedb.go
+++ b/nodedb.go
@@ -32,6 +32,7 @@ const (
 	defaultStorageVersionValue = "1.0.0"
 	fastStorageVersionValue    = "1.1.0"
 	fastNodeCacheSize          = 100000
+	maxVersion                 = int64(math.MaxInt64)
 )
 
 var (
@@ -458,7 +459,7 @@ func (ndb *nodeDB) DeleteVersionsFrom(version int64) error {
 	// Next, delete orphans:
 	// - Delete orphan entries *and referred nodes* with fromVersion >= version
 	// - Delete orphan entries with toVersion >= version-1 (since orphans at latest are not orphans)
-	err = ndb.traverseOrphans(func(key, hash []byte) error {
+	err = ndb.traverseRange(orphanKeyFormat.Key(version-1), orphanKeyFormat.Key(maxVersion), func(key, hash []byte) error {
 		var fromVersion, toVersion int64
 		orphanKeyFormat.Scan(key, &toVersion, &fromVersion)
 
@@ -483,7 +484,7 @@ func (ndb *nodeDB) DeleteVersionsFrom(version int64) error {
 	}
 
 	// Delete the version root entries
-	err = ndb.traverseRange(rootKeyFormat.Key(version), rootKeyFormat.Key(int64(math.MaxInt64)), func(k, v []byte) error {
+	err = ndb.traverseRange(rootKeyFormat.Key(version), rootKeyFormat.Key(maxVersion), func(k, v []byte) error {
 		if err = ndb.batch.Delete(k); err != nil {
 			return err
 		}
@@ -494,27 +495,7 @@ func (ndb *nodeDB) DeleteVersionsFrom(version int64) error {
 		return err
 	}
 
-	// Delete fast node entries
-	err = ndb.traverseFastNodes(func(keyWithPrefix, v []byte) error {
-		key := keyWithPrefix[1:]
-		fastNode, err := DeserializeFastNode(key, v)
-
-		if err != nil {
-			return err
-		}
-
-		if version <= fastNode.versionLastUpdatedAt {
-			if err = ndb.batch.Delete(keyWithPrefix); err != nil {
-				return err
-			}
-			ndb.fastNodeCache.Remove(key)
-		}
-		return nil
-	})
-
-	if err != nil {
-		return err
-	}
+	// NOTICE: we don't touch fast node indexes here, because it'll be rebuilt later because of version mismatch.
 
 	return nil
 }
@@ -612,6 +593,11 @@ func (ndb *nodeDB) deleteNodesFrom(version int64, hash []byte) error {
 		return err
 	}
 
+	if node.version < version {
+		// We can skip the whole sub-tree since children.version <= parent.version.
+		return nil
+	}
+
 	if node.leftHash != nil {
 		if err := ndb.deleteNodesFrom(version, node.leftHash); err != nil {
 			return err
@@ -729,7 +715,7 @@ func (ndb *nodeDB) rootKey(version int64) []byte {
 func (ndb *nodeDB) getLatestVersion() (int64, error) {
 	if ndb.latestVersion == 0 {
 		var err error
-		ndb.latestVersion, err = ndb.getPreviousVersion(1<<63 - 1)
+		ndb.latestVersion, err = ndb.getPreviousVersion(maxVersion)
 		if err != nil {
 			return 0, err
 		}
@@ -793,6 +779,7 @@ func (ndb *nodeDB) traverseOrphans(fn func(keyWithPrefix, v []byte) error) error
 }
 
 // Traverse fast nodes and return error if any, nil otherwise
+// nolint: unused
 func (ndb *nodeDB) traverseFastNodes(fn func(k, v []byte) error) error {
 	return ndb.traversePrefix(fastKeyFormat.Key(), fn)
 }
@@ -1015,6 +1002,7 @@ func (ndb *nodeDB) orphans() ([][]byte, error) {
 // Not efficient.
 // NOTE: DB cannot implement Size() because
 // mutations are not always synchronous.
+//
 //nolint:unused
 func (ndb *nodeDB) size() int {
 	size := 0