feat: speed up rollback command

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Unreleased
 
 - [#586](https://github.com/cosmos/iavl/pull/586) Remove the `RangeProof` and refactor the ics23_proof to use the internal methods.
+- [#620](https://github.com/cosmos/iavl/pull/620) Add fast mode to rollback chain state in quick and dirty way, may leave some orphan nodes in db, not a big deal.
 
 ## 0.19.4 (October 28, 2022)
 

mutable_tree.go

@@ -616,19 +616,31 @@ func (tree *MutableTree) LoadVersion(targetVersion int64) (int64, error) {
 // LoadVersionForOverwriting attempts to load a tree at a previously committed
 // version, or the latest version below it. Any versions greater than targetVersion will be deleted.
 func (tree *MutableTree) LoadVersionForOverwriting(targetVersion int64) (int64, error) {
-	latestVersion, err := tree.LoadVersion(targetVersion)
+	return tree.LoadVersionForOverwritingWithMode(targetVersion, false)
+}
+
+func (tree *MutableTree) LoadVersionForOverwritingWithMode(targetVersion int64, fastMode bool) (int64, error) {
+	var latestVersion int64
+	var err error
+	if !fastMode {
+		latestVersion, err = tree.LoadVersion(targetVersion)
+	} else {
+		latestVersion, err = tree.LazyLoadVersion(targetVersion)
+	}
 	if err != nil {
 		return latestVersion, err
 	}
 
-	if err = tree.ndb.DeleteVersionsFrom(targetVersion + 1); err != nil {
+	if err = tree.ndb.DeleteVersionsFrom(targetVersion+1, fastMode); err != nil {
 		return latestVersion, err
 	}
 
 	if !tree.skipFastStorageUpgrade {
 		if err := tree.enableFastStorageAndCommitLocked(); err != nil {
 			return latestVersion, err
 		}
+	} else if err = tree.ndb.Commit(); err != nil {
+		return latestVersion, err
 	}
 
 	tree.ndb.resetLatestVersion(latestVersion)

nodedb.go

@@ -420,7 +420,7 @@ func (ndb *nodeDB) DeleteVersion(version int64, checkLatestVersion bool) error {
 }
 
 // DeleteVersionsFrom permanently deletes all tree versions from the given version upwards.
-func (ndb *nodeDB) DeleteVersionsFrom(version int64) error {
+func (ndb *nodeDB) DeleteVersionsFrom(version int64, fastMode bool) error {
 	latest, err := ndb.getLatestVersion()
 	if err != nil {
 		return err
@@ -444,34 +444,53 @@ func (ndb *nodeDB) DeleteVersionsFrom(version int64) error {
 
 	// First, delete all active nodes in the current (latest) version whose node version is after
 	// the given version.
-	err = ndb.deleteNodesFrom(version, root)
-	if err != nil {
-		return err
-	}
-
-	// 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 {
-		var fromVersion, toVersion int64
-		orphanKeyFormat.Scan(key, &toVersion, &fromVersion)
-
-		if fromVersion >= version {
-			if err = ndb.batch.Delete(key); err != nil {
-				return err
-			}
-			if err = ndb.batch.Delete(ndb.nodeKey(hash)); err != nil {
-				return err
+	if !fastMode {
+		err = ndb.deleteNodesFrom(version, root)
+		if err != nil {
+			return err
+		}
+		// 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 {
+			var fromVersion, toVersion int64
+			orphanKeyFormat.Scan(key, &toVersion, &fromVersion)
+
+			if fromVersion >= version {
+				if err = ndb.batch.Delete(key); err != nil {
+					return err
+				}
+				if err = ndb.batch.Delete(ndb.nodeKey(hash)); err != nil {
+					return err
+				}
+				ndb.nodeCache.Remove(hash)
+			} else if toVersion >= version-1 {
+				if err = ndb.batch.Delete(key); err != nil {
+					return err
+				}
 			}
-			ndb.nodeCache.Remove(hash)
-		} else if toVersion >= version-1 {
-			if err = ndb.batch.Delete(key); err != nil {
-				return err
+			return nil
+		})
+	} else {
+		err = ndb.traverseOrphansVersion(version-1, func(key, hash []byte) error {
+			var fromVersion, toVersion int64
+			orphanKeyFormat.Scan(key, &toVersion, &fromVersion)
+			if fromVersion >= version {
+				if err = ndb.batch.Delete(key); err != nil {
+					return err
+				}
+				if err = ndb.batch.Delete(ndb.nodeKey(hash)); err != nil {
+					return err
+				}
+				ndb.nodeCache.Remove(hash)
+			} else if toVersion >= version-1 {
+				if err = ndb.batch.Delete(key); err != nil {
+					return err
+				}
 			}
-		}
-		return nil
-	})
-
+			return nil
+		})
+	}
 	if err != nil {
 		return err
 	}
@@ -489,24 +508,29 @@ func (ndb *nodeDB) DeleteVersionsFrom(version int64) error {
 	}
 
 	// Delete fast node entries
-	err = ndb.traverseFastNodes(func(keyWithPrefix, v []byte) error {
-		key := keyWithPrefix[1:]
-		fastNode, err := fastnode.DeserializeNode(key, v)
-		if err != nil {
-			return err
-		}
-
-		if version <= fastNode.GetVersionLastUpdatedAt() {
-			if err = ndb.batch.Delete(keyWithPrefix); err != nil {
+	// Delete step will be skipped with enable fastMode
+	// with the assumption that the rollback happens offline
+	// since fast nodes will be reinforced when next start up
+	if !fastMode {
+		err = ndb.traverseFastNodes(func(keyWithPrefix, v []byte) error {
+			key := keyWithPrefix[1:]
+			fastNode, err := fastnode.DeserializeNode(key, v)
+			if err != nil {
 				return err
 			}
-			ndb.fastNodeCache.Remove(key)
-		}
-		return nil
-	})
 
-	if err != nil {
-		return err
+			if version <= fastNode.GetVersionLastUpdatedAt() {
+				if err = ndb.batch.Delete(keyWithPrefix); err != nil {
+					return err
+				}
+				ndb.fastNodeCache.Remove(key)
+			}
+			return nil
+		})
+
+		if err != nil {
+			return err
+		}
 	}
 
 	return nil

nodedb_test.go

@@ -293,3 +293,66 @@ func makeAndPopulateMutableTree(tb testing.TB) *MutableTree {
 	require.Nil(tb, err, "Expected .SaveVersion to succeed")
 	return tree
 }
+
+func TestDeleteVersion(t *testing.T) {
+	v := []byte("value")
+
+	var version int64 = 100
+	testCases := []struct {
+		name     string
+		v        int64
+		fastMode bool
+	}{
+		{
+			"delete from version without fast mode",
+			version,
+			false,
+		},
+		{
+			"delete from version -1 without fast mode",
+			version - 1,
+			false,
+		},
+		{
+			"enable fast mode",
+			version,
+			true,
+		},
+	}
+	for _, tc := range testCases {
+		db := db.NewMemDB()
+		ndb := newNodeDB(db, 0, nil)
+		leftNode := NewNode([]byte("left_key"), v, version-1)
+		rightNode := NewNode([]byte("right_key"), v, version-1)
+		node := NewNode([]byte("key"), v, version)
+		node.leftNode = leftNode
+		node.rightNode = rightNode
+		node.subtreeHeight = 1
+		node.size = 2
+		hash, err := ndb.SaveBranch(node)
+		require.NoError(t, err)
+		err = ndb.Commit()
+		require.NoError(t, err)
+		key := ndb.rootKey(version)
+		err = ndb.db.Set(key, hash)
+		require.NoError(t, err)
+		err = ndb.DeleteVersionsFrom(tc.v, tc.fastMode)
+		require.NoError(t, err)
+		err = ndb.Commit()
+		require.NoError(t, err)
+		bz, err := ndb.db.Get(ndb.nodeKey(hash))
+		require.NoError(t, err)
+		leftBz, err := ndb.db.Get(ndb.nodeKey(leftNode.hash))
+		require.NoError(t, err)
+		if !tc.fastMode {
+			if tc.v <= version {
+				require.Empty(t, bz)
+			}
+			if tc.v < version {
+				require.Empty(t, leftBz)
+			} else {
+				require.NotEmpty(t, leftBz)
+			}
+		}
+	}
+}