Fix memiavl snapshot race condition

sei-db/state_db/sc/memiavl/snapshot.go

@@ -501,21 +501,23 @@ func (t *Tree) WriteSnapshot(ctx context.Context, snapshotDir string) error {
 // WriteSnapshotWithRateLimit writes snapshot with optional rate limiting.
 // limiter is a shared rate limiter. nil means unlimited.
 func (t *Tree) WriteSnapshotWithRateLimit(ctx context.Context, snapshotDir string, limiter *rate.Limiter) error {
+	root, version := t.snapshotSource()
+
 	// Estimate tree size: root.Size() returns leaf count, total = leaves + branches ≈ 2x
 	treeSize := int64(0)
-	if t.root != nil {
-		treeSize = t.root.Size() * 2 // Total nodes (leaves + branches)
+	if root != nil {
+		treeSize = root.Size() * 2 // Total nodes (leaves + branches)
 	}
 
 	// Use 128MB buffer for all trees (large buffer for better performance)
 	bufSize := bufIOSize
 
-	err := writeSnapshotWithBuffer(ctx, snapshotDir, t.version, bufSize, treeSize, limiter, func(w *snapshotWriter) (uint32, error) {
-		if t.root == nil {
+	err := writeSnapshotWithBuffer(ctx, snapshotDir, version, bufSize, treeSize, limiter, func(w *snapshotWriter) (uint32, error) {
+		if root == nil {
 			return 0, nil
 		}
 
-		if err := w.writePostOrder(t.root); err != nil {
+		if err := w.writePostOrder(root); err != nil {
 			return 0, err
 		}
 		return w.leafCounter, nil
@@ -528,6 +530,20 @@ func (t *Tree) WriteSnapshotWithRateLimit(ctx context.Context, snapshotDir strin
 	return nil
 }
 
+// snapshotSource returns the root to serialize together with the version to
+// record for it, read as a pair so a concurrent SaveVersion cannot land between
+// them and pair a root with the wrong version.
+//
+// The version returned is the last saved one. It describes the root only when
+// the tree has no unsaved writes, which holds for every snapshot writer today
+// because DB.Commit copies after SaveVersion. A snapshot written from a
+// mid-version copy would record a version its nodes are one ahead of.
+func (t *Tree) snapshotSource() (Node, uint32) {
+	t.mtx.RLock()
+	defer t.mtx.RUnlock()
+	return t.root, t.version
+}
+
 // writeSnapshotWithBuffer writes snapshot with specified buffer size and optional rate limiting.
 // limiter is a shared rate limiter. nil means unlimited.
 func writeSnapshotWithBuffer(

sei-db/state_db/sc/memiavl/tree.go

@@ -30,18 +30,20 @@ type Tree struct {
 	// when true, the get and iterator methods could return a slice pointing to mmaped blob files.
 	zeroCopy bool
 
+	// unsavedWrites reports whether Set or Remove has stamped nodes at
+	// version+1 since the last SaveVersion, so maxNodeVersion knows whether the
+	// tree currently holds nodes above version.
+	unsavedWrites bool
+
 	// mtx guards concurrent access to this tree's mutable state (root, version,
 	// cowVersion, snapshot) AND the lazily-populated MemNode.hash caches reachable
 	// from root. Operations that fill those caches in place — RootHash and the
 	// proof builders (GetProof/GetMembership/GetNonMembership) — take the write
 	// lock; pure reads (Get/Has/Iterator) take the read lock. This serialization
 	// only protects a single tree instance: a tree produced by Copy() gets its
-	// own mtx and shares the underlying nodes copy-on-write. Cross-copy hash
-	// consistency therefore relies on (a) the shared nodes already being fully
-	// hashed before the copy is used for hashing/proofs (Copy is taken between
-	// commits, and the commit path hashes via SaveVersion(true)/RootHash), and
-	// (b) cowVersion cloning any shared MemNode before it is structurally mutated,
-	// so a live-tree write never mutates a node another copy is still reading.
+	// own mtx and shares the underlying nodes copy-on-write. Copy is what makes
+	// that safe across instances, by freezing the shared nodes before handing
+	// them over; see its doc comment.
 	mtx *sync.RWMutex
 
 	pendingChanges chan proto.ChangeSet
@@ -108,24 +110,62 @@ func (t *Tree) SetInitialVersion(initialVersion int64) error {
 	return nil
 }
 
-// Copy returns a concurrent-safe snapshot. Acquires the underlying *Snapshot
-// so background rewrites can't unmap it while the copy is live; callers must
-// call Close on the returned tree to release the ref.
+// Copy returns a concurrent-safe snapshot of the tree. Acquires the underlying
+// *Snapshot so background rewrites can't unmap it while the copy is live;
+// callers must call Close on the returned tree to release the ref.
+//
+// The copy shares its nodes with the live tree, so it is only safe to read
+// concurrently if those nodes are frozen: every reachable MemNode is already
+// hashed, and sits at or below cowVersion so a live write clones it rather than
+// mutating it in place. Copy establishes both under the write lock.
 func (t *Tree) Copy() *Tree {
-	t.mtx.RLock()
-	defer t.mtx.RUnlock()
+	t.mtx.Lock()
+	defer t.mtx.Unlock()
+	return t.copyNoLock()
+}
+
+// copyNoLock is Copy without acquiring t.mtx; the caller must already hold the
+// write lock.
+func (t *Tree) copyNoLock() *Tree {
+	// Hash every reachable node now, while we hold the write lock. Otherwise the
+	// copy's reader fills MemNode.hash in place on nodes the live tree also
+	// holds, and the two trees have separate mutexes to serialize with.
+	_ = t.rootHashNoLock()
+
 	if _, ok := t.root.(*MemNode); ok {
-		// protect the existing `MemNode`s from get modified in-place
-		t.cowVersion = t.version
+		// Freeze the shared MemNodes against in-place mutation. The floor is the
+		// highest version present in the tree, not t.version: Set and Remove
+		// stamp nodes at t.version+1, so between a changeset and SaveVersion the
+		// tree holds nodes a t.version floor would leave mutable.
+		t.cowVersion = t.maxNodeVersion()
 	}
+
 	newTree := *t
 	newTree.mtx = &sync.RWMutex{}
+	// The copy is read-only and must not share the live tree's background-write
+	// plumbing: Close would otherwise close the live tree's pendingChanges.
+	newTree.pendingChanges = nil
+	newTree.pendingWg = &sync.WaitGroup{}
 	if newTree.snapshot != nil {
 		newTree.snapshot.Acquire()
 	}
 	return &newTree
 }
 
+// writeVersion returns the version Set and Remove stamp onto the nodes they
+// create or mutate, which is one ahead of the last saved version.
+func (t *Tree) writeVersion() uint32 {
+	return t.version + 1
+}
+
+// maxNodeVersion returns the highest version carried by any node in the tree.
+func (t *Tree) maxNodeVersion() uint32 {
+	if t.unsavedWrites {
+		return t.writeVersion()
+	}
+	return t.version
+}
+
 // ApplyChangeSet apply the change set of a whole version, and update hashes.
 func (t *Tree) ApplyChangeSet(changeSet proto.ChangeSet) {
 	for _, pair := range changeSet.Pairs {
@@ -172,27 +212,37 @@ func (t *Tree) Set(key, value []byte) {
 		// the value could be nil when replaying changes from write-ahead-log because of protobuf decoding
 		value = []byte{}
 	}
-	t.root, _ = setRecursive(t.root, key, value, t.version+1, t.cowVersion)
+	t.root, _ = setRecursive(t.root, key, value, t.writeVersion(), t.cowVersion)
+	t.unsavedWrites = true
 }
 
 func (t *Tree) Remove(key []byte) {
 	t.mtx.Lock()
 	defer t.mtx.Unlock()
-	_, t.root, _ = removeRecursive(t.root, key, t.version+1, t.cowVersion)
+	_, t.root, _ = removeRecursive(t.root, key, t.writeVersion(), t.cowVersion)
+	t.unsavedWrites = true
 }
 
-// SaveVersion increases the version number and optionally updates the hashes
+// SaveVersion increases the version number and optionally updates the hashes.
+//
+// It holds the write lock across both, so a concurrent Copy cannot observe the
+// bumped version before the nodes stamped at the old write version are hashed,
+// which would let it derive a cowVersion that leaves them mutable.
 func (t *Tree) SaveVersion(updateHash bool) ([]byte, int64, error) {
+	t.mtx.Lock()
+	defer t.mtx.Unlock()
+
 	if t.version >= uint32(math.MaxUint32) {
 		return nil, 0, errors.New("version overflows uint32")
 	}
 
 	var hash []byte
 	if updateHash {
-		hash = t.RootHash()
+		hash = t.rootHashNoLock()
 	}
 
 	t.version = nextVersionU32(t.version, t.initialVersion)
+	t.unsavedWrites = false
 	return hash, int64(t.version), nil
 }
 
@@ -376,6 +426,7 @@ func (t *Tree) ReplaceWith(other *Tree) error {
 	t.initialVersion = other.initialVersion
 	t.cowVersion = other.cowVersion
 	t.zeroCopy = other.zeroCopy
+	t.unsavedWrites = other.unsavedWrites
 
 	if snapshot != nil {
 		return snapshot.Close()

sei-db/state_db/sc/memiavl/tree_copy_test.go

@@ -0,0 +1,183 @@
+package memiavl
+
+import (
+	"context"
+	"fmt"
+	"path/filepath"
+	"testing"
+
+	"github.com/sei-protocol/sei-chain/sei-db/proto"
+	"github.com/stretchr/testify/require"
+)
+
+// Copy hands a tree to readers that run concurrently with the live tree: the
+// background snapshot rewrite, which serializes every reachable node, and the
+// DB.Copy consumers that read a copy while consensus keeps committing. For that
+// to be safe the copy must be frozen: every MemNode it can reach is already
+// hashed, and sits at or below cowVersion so a live write clones it instead of
+// mutating it in place. The tests below pin both halves of that invariant.
+
+const copyTestSeedKeys = 3000
+
+// seedTree returns a tree holding keys 0..keys-1, saved and fully hashed.
+func seedTree(t *testing.T, keys int) *Tree {
+	t.Helper()
+	tree := NewEmptyTree(0, 0)
+	tree.ApplyChangeSet(changeSet(0, keys, "seed"))
+	_, _, err := tree.SaveVersion(true)
+	require.NoError(t, err)
+	return tree
+}
+
+// changeSet writes count values over keys offset..offset+count-1, tagged so
+// repeated rounds over the same keys produce distinct values.
+func changeSet(offset, count int, tag string) proto.ChangeSet {
+	cs := proto.ChangeSet{Pairs: make([]*proto.KVPair, 0, count)}
+	for i := offset; i < offset+count; i++ {
+		cs.Pairs = append(cs.Pairs, &proto.KVPair{
+			Key:   []byte(fmt.Sprintf("key%08d", i)),
+			Value: []byte(fmt.Sprintf("%s%08d", tag, i)),
+		})
+	}
+	return cs
+}
+
+// frozenViolations counts the MemNodes reachable from tree's root that break the
+// freeze invariant: unhashed, so a reader fills MemNode.hash in place on a node
+// the live tree also holds; or above cowVersion, so a live write mutates the node
+// rather than cloning it. Either lets a live commit change a node out from under
+// the snapshot writer mid-traversal.
+func frozenViolations(tree *Tree) (unhashed, mutable int) {
+	var walk func(node Node)
+	walk = func(node Node) {
+		// PersistedNode is backed by a read-only mmap and is never mutated.
+		mem, ok := node.(*MemNode)
+		if !ok {
+			return
+		}
+		if mem.hash == nil {
+			unhashed++
+		}
+		if mem.version > tree.cowVersion {
+			mutable++
+		}
+		if !mem.IsLeaf() {
+			walk(mem.left)
+			walk(mem.right)
+		}
+	}
+	if tree.root != nil {
+		walk(tree.root)
+	}
+	return unhashed, mutable
+}
+
+// TestCopyFreezesNodesWrittenThisVersion pins the copy-on-write floor. Set and
+// Remove stamp new nodes at version+1, so between a changeset and SaveVersion the
+// tree holds nodes above t.version. Deriving the floor from t.version alone
+// leaves exactly those nodes mutable in the copy.
+func TestCopyFreezesNodesWrittenThisVersion(t *testing.T) {
+	tree := seedTree(t, copyTestSeedKeys)
+
+	// Mid-version: the changeset is applied but not yet saved.
+	tree.ApplyChangeSet(changeSet(0, 500, "mid"))
+
+	frozen := tree.Copy()
+
+	unhashed, mutable := frozenViolations(frozen)
+	require.Zerof(t, mutable,
+		"copy shares %d MemNode(s) above cowVersion %d; a live write will mutate them in place", mutable, frozen.cowVersion)
+	require.Zerof(t, unhashed,
+		"copy shares %d unhashed MemNode(s); a reader will fill MemNode.hash in place on nodes the live tree also holds", unhashed)
+}
+
+// TestCopyFreezesAfterSaveVersion is the same invariant at the other point a copy
+// is taken, immediately after a commit. This one holds even without the fix, so
+// it guards against a fix that only moves the window rather than closing it.
+func TestCopyFreezesAfterSaveVersion(t *testing.T) {
+	tree := seedTree(t, copyTestSeedKeys)
+	tree.ApplyChangeSet(changeSet(0, 500, "committed"))
+	_, _, err := tree.SaveVersion(true)
+	require.NoError(t, err)
+
+	frozen := tree.Copy()
+
+	unhashed, mutable := frozenViolations(frozen)
+	require.Zero(t, mutable, "copy shares %d MemNode(s) above cowVersion %d", mutable, frozen.cowVersion)
+	require.Zero(t, unhashed, "copy shares %d unhashed MemNode(s)", unhashed)
+}
+
+// TestCopyIsStableWhileLiveTreeAdvances is the behavioural form of the invariant:
+// the tree handed to the snapshot writer must serialize identically no matter how
+// far the live tree advances underneath it. A copy taken mid-version fails this
+// when the live commits mutate its nodes in place.
+func TestCopyIsStableWhileLiveTreeAdvances(t *testing.T) {
+	tree := seedTree(t, copyTestSeedKeys)
+	tree.ApplyChangeSet(changeSet(0, 500, "mid"))
+
+	frozen := tree.Copy()
+	want := frozen.RootHash()
+	wantVersion := frozen.Version()
+
+	// Finish the in-flight version, then keep committing.
+	_, _, err := tree.SaveVersion(true)
+	require.NoError(t, err)
+	for round := 0; round < 20; round++ {
+		tree.ApplyChangeSet(changeSet(round*100, 500, fmt.Sprintf("adv%02d", round)))
+		_, _, err := tree.SaveVersion(true)
+		require.NoError(t, err)
+	}
+
+	require.Equal(t, wantVersion, frozen.Version(), "copy's version moved with the live tree")
+	require.Equal(t, want, frozen.RootHash(), "live-tree commits mutated nodes the copy still references")
+}
+
+// TestSnapshotFromCopyIgnoresLaterCommits carries the freeze invariant through
+// the snapshot writer to what lands on disk: serializing a copy must reproduce
+// the tree as it stood when copied, however far the live tree has moved on.
+//
+// This is the deterministic form of the corruption. The concurrent version needs
+// -race to be caught reliably, because at test scale the writer finishes before
+// the live tree does much damage; here the commits land before the writer runs
+// at all, so an unfrozen copy drifts every time.
+func TestSnapshotFromCopyIgnoresLaterCommits(t *testing.T) {
+	tree := seedTree(t, copyTestSeedKeys)
+	tree.ApplyChangeSet(changeSet(0, 500, "mid"))
+
+	frozen := tree.Copy()
+	want := frozen.RootHash()
+
+	_, _, err := tree.SaveVersion(true)
+	require.NoError(t, err)
+	for round := 0; round < 20; round++ {
+		tree.ApplyChangeSet(changeSet(round*100, 500, fmt.Sprintf("adv%02d", round)))
+		_, _, err := tree.SaveVersion(true)
+		require.NoError(t, err)
+	}
+
+	snapshotDir := filepath.Join(t.TempDir(), "snapshot")
+	require.NoError(t, frozen.WriteSnapshot(context.Background(), snapshotDir))
+
+	snapshot, err := OpenSnapshot(snapshotDir, Options{})
+	require.NoError(t, err)
+	t.Cleanup(func() { require.NoError(t, snapshot.Close()) })
+
+	require.Equal(t, want, snapshot.RootHash(),
+		"snapshot of the copy picked up commits made after Copy returned")
+}
+
+// TestCopyDoesNotInheritBackgroundWriteChannel covers the other way a copy
+// reaches back into the live tree. Close closes pendingChanges, so a copy that
+// inherited the channel closes the live tree's background writer: the next
+// ApplyChangeSetAsync sends on a closed channel and panics, and
+// WaitToCompleteAsyncWrite closes it a second time.
+func TestCopyDoesNotInheritBackgroundWriteChannel(t *testing.T) {
+	tree := seedTree(t, 100)
+	tree.StartBackgroundWrite()
+	defer tree.WaitToCompleteAsyncWrite()
+
+	frozen := tree.Copy()
+
+	require.Nil(t, frozen.pendingChanges, "copy shares the live tree's background-write channel")
+	require.NotSame(t, tree.pendingWg, frozen.pendingWg, "copy shares the live tree's background-write WaitGroup")
+}