Problem: memiavl snapshot rewrite could rollback state

CHANGELOG.md

@@ -33,6 +33,7 @@
 - [#1087](https://github.com/crypto-org-chain/cronos/pull/1087) memiavl fix LastCommitID when memiavl db not loaded.
 - [#1088](https://github.com/crypto-org-chain/cronos/pull/1088) memiavl fix empty value in write-ahead-log replaying.
 - [#1102](https://github.com/crypto-org-chain/cronos/pull/1102) avoid duplicate cache events emitted from ibc and gravity hook.
+- [#1123](https://github.com/crypto-org-chain/cronos/pull/1123) Fix memiavl snapshot switching
 
 ### Features
 

memiavl/db.go

@@ -8,6 +8,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"time"
 
 	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	"github.com/tendermint/tendermint/libs/log"
@@ -329,6 +330,18 @@ func (db *DB) checkAsyncCommit() error {
 	return nil
 }
 
+// CommittedVersion returns the latest version written in wal, or snapshot version if wal is empty.
+func (db *DB) CommittedVersion() (int64, error) {
+	lastIndex, err := db.wal.LastIndex()
+	if err != nil {
+		return 0, err
+	}
+	if lastIndex == 0 {
+		return db.SnapshotVersion(), nil
+	}
+	return walVersion(lastIndex, db.initialVersion), nil
+}
+
 // checkBackgroundSnapshotRewrite check the result of background snapshot rewrite, cleans up the old snapshots and switches to a new multitree
 func (db *DB) checkBackgroundSnapshotRewrite() error {
 	// check the completeness of background snapshot rewriting
@@ -341,6 +354,19 @@ func (db *DB) checkBackgroundSnapshotRewrite() error {
 			return fmt.Errorf("background snapshot rewriting failed: %w", result.err)
 		}
 
+		// wait for potential pending wal writings to finish, to make sure we catch up to latest state.
+		// in real world, block execution should be slower than wal writing, so this should not block for long.
+		for {
+			committedVersion, err := db.CommittedVersion()
+			if err != nil {
+				return fmt.Errorf("get wal version failed: %w", err)
+			}
+			if db.lastCommitInfo.Version == committedVersion {
+				break
+			}
+			time.Sleep(time.Nanosecond)
+		}
+
 		// catchup the remaining wal
 		if err := result.mtree.CatchupWAL(db.wal, 0); err != nil {
 			return fmt.Errorf("catchup failed: %w", err)

memiavl/db_test.go

@@ -393,3 +393,26 @@ func TestExclusiveLock(t *testing.T) {
 	_, err = Load(dir, Options{})
 	require.NoError(t, err)
 }
+
+func TestFastCommit(t *testing.T) {
+	dir := t.TempDir()
+
+	db, err := Load(dir, Options{CreateIfMissing: true, InitialStores: []string{"test"}, SnapshotInterval: 3, AsyncCommitBuffer: 10})
+	require.NoError(t, err)
+
+	cs := iavl.ChangeSet{
+		Pairs: []*iavl.KVPair{
+			{Key: []byte("hello1"), Value: make([]byte, 1024*1024)},
+		},
+	}
+
+	// the bug reproduce when the wal writing is slower than commit, that happens when wal segment is full and create a new one, the wal writing will slow down a little bit,
+	// segment size is 20m, each change set is 1m, so we need a bit more than 20 commits to reproduce.
+	for i := 0; i < 30; i++ {
+		_, _, err := db.Commit([]*NamedChangeSet{{Name: "test", Changeset: cs}})
+		require.NoError(t, err)
+	}
+
+	<-db.snapshotRewriteChan
+	require.NoError(t, db.Close())
+}