Backport release/v6.7: fix(memiavl): hold a snapshot reference for an iterator's lifetime

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

@@ -15,6 +15,10 @@ type Iterator struct {
 	ascending  bool
 	zeroCopy   bool
 
+	// snapshot is the snapshot backing the nodes this iterator walks, held for the
+	// iterator's lifetime and released by Close. nil when no snapshot backs them.
+	snapshot *Snapshot
+
 	// cache the next key-value pair
 	key, value []byte
 
@@ -23,13 +27,24 @@ type Iterator struct {
 	stack []Node
 }
 
-func NewIterator(start, end []byte, ascending bool, root Node, zeroCopy bool) *Iterator {
+// NewIterator returns an iterator over the tree rooted at root, backed by
+// snapshot, which is nil when no snapshot backs those nodes. The iterator holds a
+// reference on snapshot until Close.
+func NewIterator(start, end []byte, ascending bool, root Node, zeroCopy bool, snapshot *Snapshot) *Iterator {
+	// A persisted node's key and value are slices into this mmap, and both Key and
+	// Value read them after the tree's read lock is gone, so the reference is what
+	// stops a concurrent rewrite unmapping the region mid-iteration. That unmap
+	// surfaces as a fatal fault in runtime.memmove, which no caller can recover.
+	if snapshot != nil {
+		snapshot.Acquire()
+	}
 	iter := &Iterator{
 		start:     start,
 		end:       end,
 		ascending: ascending,
 		valid:     true,
 		zeroCopy:  zeroCopy,
+		snapshot:  snapshot,
 	}
 
 	if root != nil {
@@ -117,5 +132,16 @@ func (iter *Iterator) Next() {
 func (iter *Iterator) Close() error {
 	iter.valid = false
 	iter.stack = nil
-	return nil
+	// Under zeroCopy these still point into the mapping the release below may
+	// unmap, so drop them with the reference rather than leaving a caller able to
+	// reach freed pages through a closed iterator.
+	iter.key = nil
+	iter.value = nil
+
+	snapshot := iter.snapshot
+	iter.snapshot = nil
+	if snapshot == nil {
+		return nil
+	}
+	return snapshot.Close()
 }

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

@@ -170,3 +170,92 @@ func TestSnapshotDoubleCloseReturnsError(t *testing.T) {
 	require.Error(t, snapshot.Close())
 	require.Panics(t, snapshot.Acquire)
 }
+
+// snapshotBackedDB returns a DB whose "test" tree is backed by a mapped snapshot
+// holding pairs, which is the state an iterator must be opened over to exercise
+// the mmap: an iterator over unpersisted MemNodes reads the heap instead.
+func snapshotBackedDB(t *testing.T, pairs []*proto.KVPair) *DB {
+	t.Helper()
+	db, err := OpenDB(0, Options{
+		Config:          Config{SnapshotKeepRecent: 0},
+		Dir:             t.TempDir(),
+		CreateIfMissing: true,
+		InitialStores:   []string{"test"},
+	})
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = db.Close() })
+
+	require.NoError(t, db.ApplyChangeSets([]*proto.NamedChangeSet{{
+		Name:      "test",
+		Changeset: proto.ChangeSet{Pairs: pairs},
+	}}))
+	_, err = db.Commit()
+	require.NoError(t, err)
+
+	require.NoError(t, db.RewriteSnapshot(context.Background()))
+	require.NoError(t, db.Reload())
+	return db
+}
+
+// Regression: an iterator handed out by Tree.Iterator must stay readable across a
+// snapshot rewrite and reload. Key and Value clone out of the snapshot's mmap
+// after Tree.Iterator has already released the read lock, so before the iterator
+// took a reference the reload's snapshot.Close() unmapped pages it was still
+// reading. That killed the process with a fatal fault in runtime.memmove, which
+// no recover can catch, and it was observed killing validators mid-block.
+func TestTreeIteratorOutlivesSnapshotRewrite(t *testing.T) {
+	db := snapshotBackedDB(t, []*proto.KVPair{
+		{Key: []byte("k1"), Value: []byte("v1")},
+		{Key: []byte("k2"), Value: []byte("v2")},
+	})
+
+	tree := db.TreeByName("test")
+	require.NotNil(t, tree)
+	// Production opens memiavl with ZeroCopy false, which is what routes Key and
+	// Value through utils.Clone and so reads the mapping on every pair.
+	tree.SetZeroCopy(false)
+	iter := tree.Iterator(nil, nil, true)
+
+	// Rotate underneath the open iterator: Reload's ReplaceWith closes the
+	// snapshot the iterator is reading.
+	require.NoError(t, db.ApplyChangeSets([]*proto.NamedChangeSet{{
+		Name: "test",
+		Changeset: proto.ChangeSet{Pairs: []*proto.KVPair{
+			{Key: []byte("k1"), Value: []byte("OVERWRITTEN")},
+		}},
+	}}))
+	_, err := db.Commit()
+	require.NoError(t, err)
+	require.NoError(t, db.RewriteSnapshot(context.Background()))
+	require.NoError(t, db.Reload())
+
+	// The iterator reports the contents it was opened over, not the rewrite's.
+	require.Equal(t, []pair{
+		{key: []byte("k1"), value: []byte("v1")},
+		{key: []byte("k2"), value: []byte("v2")},
+	}, collectIter(iter))
+	require.NoError(t, iter.Close())
+}
+
+// The reference an iterator takes has to come back on Close, or every iterator
+// pins its snapshot's blob files mapped for the life of the process.
+func TestTreeIteratorReleasesSnapshotOnClose(t *testing.T) {
+	db := snapshotBackedDB(t, []*proto.KVPair{{Key: []byte("k"), Value: []byte("v")}})
+
+	tree := db.TreeByName("test")
+	require.NotNil(t, tree)
+	held := tree.snapshot
+	require.NotNil(t, held)
+	before := held.refCount.Load()
+
+	iter := tree.Iterator(nil, nil, true)
+	require.Equal(t, before+1, held.refCount.Load(), "iterator did not take a reference")
+
+	require.NoError(t, iter.Close())
+	require.Equal(t, before, held.refCount.Load(), "Close did not return the reference")
+
+	// Close is idempotent, so a caller that closes twice cannot drive the
+	// refcount below what it took and unmap the snapshot under someone else.
+	require.NoError(t, iter.Close())
+	require.Equal(t, before, held.refCount.Load())
+}

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

@@ -339,10 +339,16 @@ func (t *Tree) hasNoLock(key []byte) bool {
 	return value != nil
 }
 
+// Iterator returns an iterator over the tree's current contents. The caller must
+// Close it: it holds a reference on the tree's snapshot, so an iterator left open
+// keeps that snapshot's blob files mapped.
 func (t *Tree) Iterator(start, end []byte, ascending bool) dbm.Iterator {
+	// The read lock also makes the reference safe to take: ReplaceWith and Close
+	// swap and close t.snapshot under the write lock, so the snapshot read here
+	// cannot already have reached a zero refcount.
 	t.mtx.RLock()
 	defer t.mtx.RUnlock()
-	return NewIterator(start, end, ascending, t.root, t.zeroCopy)
+	return NewIterator(start, end, ascending, t.root, t.zeroCopy, t.snapshot)
 }
 
 // ScanPostOrder scans the tree in post-order, and call the callback function on each node.