diff --git a/sei-db/state_db/sc/memiavl/snapshot.go b/sei-db/state_db/sc/memiavl/snapshot.go index d5434a5f07..97417f1b19 100644 --- a/sei-db/state_db/sc/memiavl/snapshot.go +++ b/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( diff --git a/sei-db/state_db/sc/memiavl/tree.go b/sei-db/state_db/sc/memiavl/tree.go index 599681ca1a..d943683bae 100644 --- a/sei-db/state_db/sc/memiavl/tree.go +++ b/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() diff --git a/sei-db/state_db/sc/memiavl/tree_copy_test.go b/sei-db/state_db/sc/memiavl/tree_copy_test.go new file mode 100644 index 0000000000..4ff184d4db --- /dev/null +++ b/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") +} diff --git a/sei-db/state_db/sc/memiavl/tree_race_test.go b/sei-db/state_db/sc/memiavl/tree_race_test.go index 8846826a02..b820392639 100644 --- a/sei-db/state_db/sc/memiavl/tree_race_test.go +++ b/sei-db/state_db/sc/memiavl/tree_race_test.go @@ -1,7 +1,10 @@ package memiavl import ( + "context" "fmt" + "os" + "path/filepath" "sync" "testing" @@ -100,3 +103,150 @@ func TestTreeConcurrentRootHash(t *testing.T) { require.Equal(t, hashes[0], hashes[i], "concurrent RootHash produced divergent hashes") } } + +// TestCopyReadRaceWithLiveCommits covers the cross-copy hole that the per-tree +// write lock does not close. RootHash and the proof builders take the write lock +// because MemNode.Hash fills the hash cache in place (Immunefi 83246), but Copy +// gives each tree its own mutex, so a copy and the tree it came from serialize +// against nothing. Reading a copy whose nodes the live tree is still free to +// mutate is the same unsynchronized access, one lock away from the fix for it. +// +// This is the trace-snapshot path: SnapshotSCStore hands EndBlock a DB.Copy, and +// ApplyChangeSets has already released db.mtx by then, so the copy can share +// nodes stamped at version+1 that the old cowVersion = t.version floor left +// mutable. Its consequence is a wrong trace or proof rather than a bad snapshot +// file, which is why it is separate from the writer test below. +func TestCopyReadRaceWithLiveCommits(t *testing.T) { + const ( + rounds = 4 + blocksPerRound = 20 + keysPerChangeSet = 200 + ) + + tree := seedTree(t, copyTestSeedKeys) + + for round := 0; round < rounds; round++ { + // Copy mid-version, where DB.Copy lands between ApplyChangeSets and Commit. + tree.ApplyChangeSet(changeSet(round*53, keysPerChangeSet, fmt.Sprintf("pre%02d", round))) + leased := tree.Copy() + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < blocksPerRound; i++ { + _ = leased.RootHash() + _ = leased.Get([]byte(fmt.Sprintf("key%08d", i))) + _ = leased.GetProof([]byte(fmt.Sprintf("key%08d", i))) + } + }() + + _, _, err := tree.SaveVersion(true) + require.NoError(t, err) + for block := 0; block < blocksPerRound; block++ { + tree.ApplyChangeSet(changeSet(block*17, keysPerChangeSet, fmt.Sprintf("r%02db%02d", round, block))) + _, _, err := tree.SaveVersion(true) + require.NoError(t, err) + } + wg.Wait() + } +} + +// TestSnapshotWriteRaceWithLiveCommits reproduces the corrupted-snapshot failure +// that surfaces at startup as "leaves file size N is not a multiple of 48". The +// background rewrite serializes a Copy of the tree while consensus keeps +// committing. If the copy is not frozen, those commits mutate nodes the writer is +// traversing: Mutate clears MemNode.hash under the writer, and writeLeafDirect +// emits the 16-byte header and the hash as two writes, so a torn read of the +// cleared slice header writes a leaf record with no hash and leaves the file 32 +// bytes short. +// +// The copy is taken mid-version, between a changeset and SaveVersion, which is +// what DB.Copy and CommitStore.Copy produce: ApplyChangeSets releases db.mtx on +// return, so a copy taken before the following Commit sees nodes stamped at +// version+1. That is the state the old cowVersion = t.version floor left +// mutable. The background rewrite, by contrast, copies inside Commit after +// SaveVersion, where that floor already covered every node. +// +// Driving the snapshot writer over such a copy composes the two: the mid-version +// copy is the state a real caller produces, and the writer is the reader that +// touches every reachable node, which makes it the strongest probe of the freeze. +// +// Under -race the unsynchronized MemNode access is the signal. The assertions +// then cover the on-disk result: well-formedness, and that each snapshot holds +// the tree as it stood when copied rather than some blend of it and later blocks. +func TestSnapshotWriteRaceWithLiveCommits(t *testing.T) { + const ( + rounds = 6 + blocksPerRound = 30 + keysPerChangeSet = 400 + ) + + dir := t.TempDir() + tree := seedTree(t, copyTestSeedKeys) + roundDir := func(round int) string { + return filepath.Join(dir, fmt.Sprintf("snapshot-%d", round)) + } + + var ( + wg sync.WaitGroup + errMtx sync.Mutex + writeErr []error + ) + + // Root hash of each copy at the moment it was handed to the writer. + wantHash := make([][]byte, rounds) + + for round := 0; round < rounds; round++ { + tree.ApplyChangeSet(changeSet(round*137, keysPerChangeSet, fmt.Sprintf("pre%02d", round))) + frozen := tree.Copy() + wantHash[round] = frozen.RootHash() + + snapshotDir := roundDir(round) + wg.Add(1) + go func() { + defer wg.Done() + if err := frozen.WriteSnapshot(context.Background(), snapshotDir); err != nil { + errMtx.Lock() + writeErr = append(writeErr, err) + errMtx.Unlock() + } + }() + + _, _, err := tree.SaveVersion(true) + require.NoError(t, err) + for block := 0; block < blocksPerRound; block++ { + tree.ApplyChangeSet(changeSet(block*29, keysPerChangeSet, fmt.Sprintf("r%02db%02d", round, block))) + _, _, err := tree.SaveVersion(true) + require.NoError(t, err) + } + } + wg.Wait() + require.Empty(t, writeErr, "snapshot writer failed while the live tree committed") + + for round := 0; round < rounds; round++ { + snapshotDir := roundDir(round) + + info, err := os.Stat(filepath.Join(snapshotDir, FileNameLeaves)) + require.NoError(t, err) + require.Zerof(t, info.Size()%int64(SizeLeaf), + "round %d: leaves file size %d is not a multiple of %d", round, info.Size(), SizeLeaf) + + info, err = os.Stat(filepath.Join(snapshotDir, FileNameNodes)) + require.NoError(t, err) + require.Zerof(t, info.Size()%int64(SizeNode), + "round %d: nodes file size %d is not a multiple of %d", round, info.Size(), SizeNode) + + // The published snapshot must reopen, which is the check that panics the + // node on restart when a rewrite raced a commit. + snapshot, err := OpenSnapshot(snapshotDir, Options{}) + require.NoErrorf(t, err, "round %d: snapshot written during live commits does not reopen", round) + + // A well-formed snapshot holding the wrong nodes reopens fine, so compare + // contents: the serialized tree must be the one that was copied. + require.Equalf(t, wantHash[round], snapshot.RootHash(), + "round %d: snapshot contents drifted from the copy the writer was given", round) + + require.NoError(t, snapshot.Close()) + } +}