db: EFOS and excise race bug fix

Fixes #5390

Author: sumeerbhola

db.go

@@ -501,6 +501,12 @@ type DB struct {
 			// sstables.
 			cumulativePinnedCount uint64
 			cumulativePinnedSize  uint64
+
+			// The ongoing excises. These are under snapshots since this is only
+			// needed to coordinate with EventuallyFileOnlySnapshot creation.
+			ongoingExcises map[SeqNum]KeyRange
+			// Signalled when an ongoingExcise is removed.
+			ongoingExcisesRemovedCond *sync.Cond
 		}
 
 		tableStats struct {
@@ -2812,3 +2818,14 @@ func (d *DB) DebugCurrentVersion() *manifest.Version {
 	defer d.mu.Unlock()
 	return d.mu.versions.currentVersion()
 }
+
+func (d *DB) removeFromOngoingExcises(seqNum base.SeqNum) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	_, ok := d.mu.snapshots.ongoingExcises[seqNum]
+	if !ok {
+		panic(errors.AssertionFailedf("pebble: no ongoing excise for seqnum %d", seqNum))
+	}
+	delete(d.mu.snapshots.ongoingExcises, seqNum)
+	d.mu.snapshots.ongoingExcisesRemovedCond.Broadcast()
+}

ingest.go

@@ -1526,7 +1526,9 @@ func (d *DB) ingest(ctx context.Context, args ingestArgs) (IngestOperationStats,
 	// asFlushable indicates whether the sstable was ingested as a flushable.
 	var asFlushable bool
 	var waitFlushStart crtime.Mono
+	var assignedSeqNum SeqNum
 	prepare := func(seqNum base.SeqNum) {
+		assignedSeqNum = seqNum
 		// Note that d.commit.mu is held by commitPipeline when calling prepare.
 
 		// Determine the set of bounds we care about for the purpose of checking
@@ -1547,14 +1549,35 @@ func (d *DB) ingest(ctx context.Context, args ingestArgs) (IngestOperationStats,
 		}
 
 		d.mu.Lock()
-		defer d.mu.Unlock()
-
+		defer func() {
+			if args.ExciseSpan.Valid() && !asFlushable {
+				// This can conflict with a concurrent EFOS creation (unlike a
+				// flushable ingest, which does not destructively excise).
+				_, ok := d.mu.snapshots.ongoingExcises[seqNum]
+				if ok {
+					panic(errors.AssertionFailedf("pebble: excise with seqnum %s already in map", seqNum))
+				}
+				d.mu.snapshots.ongoingExcises[seqNum] = args.ExciseSpan
+			}
+			d.mu.Unlock()
+		}()
 		if args.ExciseSpan.Valid() {
 			// Check if any of the currently-open EventuallyFileOnlySnapshots
-			// overlap in key ranges with the excise span. If so, we need to
-			// check for memtable overlaps with all bounds of that
-			// EventuallyFileOnlySnapshot in addition to the ingestion's own
-			// bounds too.
+			// overlap in key ranges with the excise span. If so, we need to check
+			// for memtable overlaps with *all* bounds of that
+			// EventuallyFileOnlySnapshot in addition to the ingestion's own bounds
+			// too. This is needed because excise is a destructive operation that
+			// mutates the LSM in a non-additive way (unlike ingests and writes).
+			//
+			// Say there is an existing EFOS for the key-range [a, d) where a is in
+			// the memtable and visible to the EFOS, and a later excise wants to
+			// excise [c, g) and there is nothing overlapping with it in the
+			// memtable. By adding [a, d) to the overlap bounds, we ensure that the
+			// excise waits until all of that span is flushed from the memtables,
+			// which makes the EFOS transition to a FOS before the excise modifies
+			// the LSM. If the excise were to not wait, the [c, d) state that the
+			// EFOS needs in the LSM when transitioning to FOS will already have
+			// been excised.
 			overlapBounds = append(overlapBounds, exciseOverlapBounds(
 				d.cmp, &d.mu.snapshots.snapshotList, args.ExciseSpan, seqNum)...)
 		}
@@ -1735,6 +1758,13 @@ func (d *DB) ingest(ctx context.Context, args ingestArgs) (IngestOperationStats,
 	d.commit.ingestSem <- struct{}{}
 	d.commit.AllocateSeqNum(seqNumCount, prepare, apply)
 	<-d.commit.ingestSem
+	if args.ExciseSpan.Valid() && !asFlushable {
+		// NB: this must happen after the assignedSeqNum has become visible, so
+		// that any concurrent EFOS creation that acquires d.mu after the removal
+		// of this excise gets a visible seqnum after the excise. The
+		// assignedSeqNum becomes visible in AllocateSeqNum.
+		d.removeFromOngoingExcises(assignedSeqNum)
+	}
 
 	if err != nil {
 		if err2 := ingestCleanup(d.objProvider, loadResult.local); err2 != nil {

open.go

@@ -13,6 +13,7 @@ import (
 	"math"
 	"os"
 	"slices"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -242,7 +243,9 @@ func Open(dirname string, opts *Options) (db *DB, err error) {
 	d.mu.compact.cond.L = &d.mu.Mutex
 	d.mu.compact.inProgress = make(map[compaction]struct{})
 	d.mu.compact.noOngoingFlushStartTime = crtime.NowMono()
-	d.mu.snapshots.init()
+	d.mu.snapshots.snapshotList.init()
+	d.mu.snapshots.ongoingExcises = make(map[SeqNum]KeyRange)
+	d.mu.snapshots.ongoingExcisesRemovedCond = sync.NewCond(&d.mu.Mutex)
 	// logSeqNum is the next sequence number that will be assigned.
 	// Start assigning sequence numbers from base.SeqNumStart to leave
 	// room for reserved sequence numbers (see comments around

snapshot.go

@@ -234,12 +234,83 @@ type EventuallyFileOnlySnapshot struct {
 }
 
 func (d *DB) makeEventuallyFileOnlySnapshot(keyRanges []KeyRange) *EventuallyFileOnlySnapshot {
-	isFileOnly := true
-
+	var snapshotSeqNum, exciseSeqNumToWaitFor base.SeqNum
+	// tryGetSnapshotSeqNum attempts to initialize a snapshotSeqNum. The
+	// snapshotSeqNum should be ignored if exciseSeqNumToWaitFor is > 0. In this
+	// case the caller should wait for this seqnum to become visible and call
+	// this function again.
+	tryGetSnapshotSeqNum := func() {
+		// In AllocateSeqNum, for an ingest-and-excise, some seqnums are
+		// allocated, say starting at N. Then AllocateSeqNum calls prepare, which
+		// acquires DB.mu and grabs all the existing EFOS that have not
+		// transitioned to FOS (which are in d.mu.snapshots.snapshotList) and
+		// overlap with the excise. After releasing DB.mu, in apply, the
+		// ingest-and-excise waits for all the previous EFOSs to transition to FOS
+		// (as a side effect of waiting for a memtable flush to complete). Note,
+		// the visible seqnum is still <= N-1, and will not be bumped to N until
+		// the ingest-and-excise completes. Any EFOS that gets created after
+		// prepare looked at the existing EFOSs, but before the ingest-and-excise
+		// completes, will be created with EventuallyFileOnlySnapshot.seqNum=N-1,
+		// and may not transition to FOS until after the excise, which is
+		// incorrect (the version at the transition to FOS has already experienced
+		// the excise). To avoid this incorrectness, tryGetSnapshotSeqNum is
+		// called in a loop, until there are no such ongoing excises. The loop can
+		// starve EFOS creation if the keyRanges keep overlapping with new ongoing
+		// ingest-and-excises. So EFOS should only be used when ingest-and-excises
+		// are rare over the keyRanges.
+		//
+		// Improving this starvation behavior would require this snapshot to
+		// register itself in a way that blocks future ingest-and-excises. The
+		// blocking would need to be done before allocating the seqnum, since
+		// blocking after can delay (latency sensitive) writes that get a seqnum
+		// later than the ingest-and-excise. Then the problem shifts to not
+		// starving the ingest-and-excise if EFOS creation is frequent. We observe
+		// that in the CockroachDB use case (a) ingest-and-excises are not very
+		// frequent, so EFOS starvation is unlikely, (b) EFOS creation is not
+		// latency sensitive. Hence, we ignore this starvation problem.
+		snapshotSeqNum = d.mu.versions.visibleSeqNum.Load()
+		// Check if any of the keyRanges overlap with an ongoing
+		// ingest-and-excise.
+		//
+		// NB: The zero seqnum cannot occur in practice, since base.SeqNumStart >
+		// 0.
+		exciseSeqNumToWaitFor = base.SeqNum(0)
+		for seqNum, span := range d.mu.snapshots.ongoingExcises {
+			if base.Visible(seqNum, snapshotSeqNum, base.SeqNumMax) {
+				// Skip this excise, since this is visible to the snapshot.
+				continue
+			}
+			// INVARIANT: seqNum >= snapshotSeqNum.
+			if seqNum <= exciseSeqNumToWaitFor {
+				// We are already waiting for a later excise.
+				continue
+			}
+			for i := range keyRanges {
+				if keyRanges[i].OverlapsKeyRange(d.cmp, span) {
+					exciseSeqNumToWaitFor = seqNum
+					break
+				}
+			}
+		}
+	}
 	d.mu.Lock()
 	defer d.mu.Unlock()
-	seqNum := d.mu.versions.visibleSeqNum.Load()
-	// Check if any of the keyRanges overlap with a memtable.
+	for {
+		// This call updates snapshotSeqNum and exciseSeqNumToWaitFor.
+		tryGetSnapshotSeqNum()
+		if exciseSeqNumToWaitFor == 0 {
+			break
+		}
+		for !base.Visible(exciseSeqNumToWaitFor, d.mu.versions.visibleSeqNum.Load(), base.SeqNumMax) {
+			d.mu.snapshots.ongoingExcisesRemovedCond.Wait()
+		}
+	}
+	isFileOnly := true
+	// Check if any of the keyRanges overlap with a memtable. It is possible
+	// (with very low probability) that all these memtable have seqnums later
+	// than seqNum, and the overlap is a false positive. This is harmless, and
+	// this EFOS will transition to FOS, when the false positive memtable is
+	// flushed.
 	for i := range d.mu.mem.queue {
 		d.mu.mem.queue[i].computePossibleOverlaps(func(bounded) shouldContinue {
 			isFileOnly = false
@@ -248,7 +319,7 @@ func (d *DB) makeEventuallyFileOnlySnapshot(keyRanges []KeyRange) *EventuallyFil
 	}
 	es := &EventuallyFileOnlySnapshot{
 		db:              d,
-		seqNum:          seqNum,
+		seqNum:          snapshotSeqNum,
 		protectedRanges: keyRanges,
 		closed:          make(chan struct{}),
 	}
@@ -258,7 +329,7 @@ func (d *DB) makeEventuallyFileOnlySnapshot(keyRanges []KeyRange) *EventuallyFil
 	} else {
 		s := &Snapshot{
 			db:     d,
-			seqNum: seqNum,
+			seqNum: snapshotSeqNum,
 		}
 		s.efos = es
 		es.mu.snap = s

snapshot_test.go

@@ -384,11 +384,8 @@ func TestNewSnapshotRace(t *testing.T) {
 	require.NoError(t, d.Close())
 }
 
-// Test for https://github.com/cockroachdb/pebble/issues/5390.
+// Test for fix to https://github.com/cockroachdb/pebble/issues/5390.
 func TestEFOSAndExciseRace(t *testing.T) {
-	// Skipping since: panic: unexpected excise of an
-	// EventuallyFileOnlySnapshot's bounds.
-	t.Skip("#5390")
 	o := &Options{
 		FS:                          vfs.NewMem(),
 		L0CompactionThreshold:       10,
@@ -414,25 +411,33 @@ func TestEFOSAndExciseRace(t *testing.T) {
 	<-beforeIngestApply
 	// IngestAndExcise is blocked, after waiting for memtable flush.
 	//
-	// Create a EFOS that overlaps with the excise.
-	snap := d.NewEventuallyFileOnlySnapshot([]KeyRange{{Start: []byte("a"), End: []byte("d")}})
-	defer snap.Close()
-	checkIter := func() {
-		iter, err := snap.NewIter(nil)
-		// The iter sees a and c, since it precedes the excise.
-		require.NoError(t, err)
-		require.True(t, iter.First())
-		require.Equal(t, []byte("a"), iter.Key())
-		require.True(t, iter.Next())
-		require.Equal(t, []byte("c"), iter.Key())
-		require.False(t, iter.Next())
-		require.NoError(t, iter.Close())
-	}
-	checkIter()
-	// Unblock the excise.
+	// Create a EFOS that overlaps with the excise. It will block due to the
+	// ongoing excise.
+	snapDone := make(chan struct{})
+	go func() {
+		snap := d.NewEventuallyFileOnlySnapshot([]KeyRange{{Start: []byte("a"), End: []byte("d")}})
+		defer func() {
+			require.NoError(t, snap.Close())
+			snapDone <- struct{}{}
+		}()
+		checkIter := func() {
+			iter, err := snap.NewIter(nil)
+			// The iter only sees a, since it comes after the excise.
+			require.NoError(t, err)
+			require.True(t, iter.First())
+			require.Equal(t, []byte("a"), iter.Key())
+			require.False(t, iter.Next())
+			require.NoError(t, iter.Close())
+		}
+		checkIter()
+		// Flush and transition to FOS.
+		require.NoError(t, d.Flush())
+		require.NoError(t, snap.WaitForFileOnlySnapshot(context.Background(), 0))
+		checkIter()
+	}()
+	time.Sleep(100 * time.Millisecond)
+	// Unblock the excise, which eventually unblocks the snapshot creation.
 	waitForCh <- struct{}{}
-	// Flush and transition to FOS.
-	require.NoError(t, d.Flush())
-	require.NoError(t, snap.WaitForFileOnlySnapshot(context.Background(), 0))
-	checkIter()
+	// Wait for the snapshot checking to finish.
+	<-snapDone
 }