feat: speed up rollback command

CHANGELOG.md

@@ -7,6 +7,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 `offlineRollback` flag 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,33 @@ 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)
+}
+
+// LoadVersionForOverwritingWithMode call LoadVersionForOverwriting with offlineRollback
+// to allow rollback in a quick and dirty way.
+func (tree *MutableTree) LoadVersionForOverwritingWithMode(targetVersion int64, offlineRollback bool) (int64, error) {
+	var latestVersion int64
+	var err error
+	if !offlineRollback {
+		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, offlineRollback); err != nil {
 		return latestVersion, err
 	}
 
-	if !tree.skipFastStorageUpgrade {
+	if !tree.skipFastStorageUpgrade && !offlineRollback {
 		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

@@ -419,8 +419,8 @@ func (ndb *nodeDB) DeleteVersion(version int64, checkLatestVersion bool) error {
 	return err
 }
 
-// DeleteVersionsFrom permanently deletes all tree versions from the given version upwards.
-func (ndb *nodeDB) DeleteVersionsFrom(version int64) error {
+// DeleteVersionsFrom permanently deletes all tree versions from the given version upwards with offlineRollback.
+func (ndb *nodeDB) DeleteVersionsFrom(version int64, offlineRollback bool) error {
 	latest, err := ndb.getLatestVersion()
 	if err != nil {
 		return err
@@ -444,18 +444,9 @@ 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 {
+	traverseInnerFunc := 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
@@ -470,8 +461,20 @@ func (ndb *nodeDB) DeleteVersionsFrom(version int64) error {
 			}
 		}
 		return nil
-	})
-
+	}
+	if !offlineRollback {
+		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(traverseInnerFunc)
+	} else {
+		// toVersion in orphan records is current version-1
+		err = ndb.traverseOrphansVersion(version-1, traverseInnerFunc)
+	}
 	if err != nil {
 		return err
 	}
@@ -489,24 +492,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 offlineRollback
+	// with the assumption that the rollback happens offline
+	// since fast nodes will be reinforced when next start up
+	if !offlineRollback {
+		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,101 @@ func makeAndPopulateMutableTree(tb testing.TB) *MutableTree {
 	require.Nil(tb, err, "Expected .SaveVersion to succeed")
 	return tree
 }
+
+type customBatch struct {
+	db.Batch
+	DeleteCount int
+}
+
+func (b *customBatch) Delete(key []byte) error {
+	b.DeleteCount += 1
+	return b.Batch.Delete(key)
+}
+
+type customMemDB struct {
+	*db.MemDB
+	LastBatch *customBatch
+}
+
+func (db *customMemDB) NewBatch() db.Batch {
+	if db.LastBatch != nil {
+		db.LastBatch.Batch = db.MemDB.NewBatch()
+	} else {
+		db.LastBatch = &customBatch{db.MemDB.NewBatch(), 0}
+	}
+	return db.LastBatch
+}
+
+func TestDeleteVersion(t *testing.T) {
+	v := []byte("value")
+
+	var version int64 = 100
+	testCases := []struct {
+		name            string
+		v               int64
+		offlineRollback bool
+		delCount        int
+	}{
+		{
+			"delete from version without offline rollback",
+			version,
+			false,
+			2,
+		},
+		{
+			"delete from version -1 without offline rollback",
+			version - 1,
+			false,
+			4,
+		},
+		{
+			"enable offline rollback",
+			version - 1,
+			true,
+			1,
+		},
+	}
+	for _, tc := range testCases {
+		db := &customMemDB{
+			MemDB:     db.NewMemDB(),
+			LastBatch: nil,
+		}
+		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.offlineRollback)
+		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.offlineRollback {
+			if tc.v <= version {
+				require.Empty(t, bz)
+			}
+			if tc.v < version {
+				require.Empty(t, leftBz)
+			} else {
+				require.NotEmpty(t, leftBz)
+			}
+		} else {
+			require.NotEmpty(t, bz)
+			require.NotEmpty(t, leftBz)
+		}
+		require.Equal(t, tc.delCount, db.LastBatch.DeleteCount, "Delete call count mismatch")
+	}
+}