db: periodically recheck the memtable write stall condition

The period is set to 1s, by writing at that interval to the condition
variable.

A future PR will make Options.MemTableStopWritesThreshold dynamic.

Author: sumeerbhola

db.go

@@ -2341,6 +2341,7 @@ func (d *DB) maybeInduceWriteStall(b *Batch) {
 	// This function will call EventListener.WriteStallBegin at most once.  If
 	// it does call it, it will call EventListener.WriteStallEnd once before
 	// returning.
+	var timer *time.Timer
 	for {
 		var size uint64
 		for i := range d.mu.mem.queue {
@@ -2357,7 +2358,18 @@ func (d *DB) maybeInduceWriteStall(b *Batch) {
 				})
 			}
 			beforeWait := crtime.NowMono()
+			// NB: In a rare case, we can start a write stall, and then the system
+			// may detect WAL failover, resulting in
+			// ElevateWriteStallThresholdForFailover returning true. So we want to
+			// recheck the predicate periodically, which we do by signaling the
+			// condition variable.
+			if timer == nil {
+				timer = time.AfterFunc(time.Second, d.mu.compact.cond.Broadcast)
+			} else {
+				timer.Reset(time.Second)
+			}
 			d.mu.compact.cond.Wait()
+			timer.Stop()
 			if b != nil {
 				b.commitStats.MemTableWriteStallDuration += beforeWait.Elapsed()
 			}

db_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/cockroachdb/datadriven"
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/pebble/internal/base"
+	"github.com/cockroachdb/pebble/internal/buildtags"
 	"github.com/cockroachdb/pebble/internal/cache"
 	"github.com/cockroachdb/pebble/internal/testkeys"
 	"github.com/cockroachdb/pebble/internal/testutils"
@@ -2160,6 +2161,8 @@ func TestRecycleLogs(t *testing.T) {
 
 type sstAndLogFileBlockingFS struct {
 	vfs.FS
+	blockWAL  bool
+	blockSST  bool
 	unblocker sync.WaitGroup
 }
 
@@ -2168,7 +2171,8 @@ var _ vfs.FS = &sstAndLogFileBlockingFS{}
 func (fs *sstAndLogFileBlockingFS) Create(
 	name string, category vfs.DiskWriteCategory,
 ) (vfs.File, error) {
-	if strings.HasSuffix(name, ".log") || strings.HasSuffix(name, ".sst") {
+	if (strings.HasSuffix(name, ".log") && fs.blockWAL) ||
+		(strings.HasSuffix(name, ".sst") && fs.blockSST) {
 		fs.unblocker.Wait()
 	}
 	return fs.FS.Create(name, category)
@@ -2178,16 +2182,16 @@ func (fs *sstAndLogFileBlockingFS) unblock() {
 	fs.unblocker.Done()
 }
 
-func newBlockingFS(fs vfs.FS) *sstAndLogFileBlockingFS {
-	lfbfs := &sstAndLogFileBlockingFS{FS: fs}
+func newBlockingFS(fs vfs.FS, blockWAL, blockSST bool) *sstAndLogFileBlockingFS {
+	lfbfs := &sstAndLogFileBlockingFS{FS: fs, blockWAL: blockWAL, blockSST: blockSST}
 	lfbfs.unblocker.Add(1)
 	return lfbfs
 }
 
 func TestWALFailoverAvoidsWriteStall(t *testing.T) {
 	mem := vfs.NewMem()
 	// All sst and log creation is blocked.
-	primaryFS := newBlockingFS(mem)
+	primaryFS := newBlockingFS(mem, true /*blockWAL*/, true /*blockSST*/)
 	// Secondary for WAL failover can do log creation.
 	secondary := wal.Dir{FS: mem, Dirname: "secondary"}
 	walFailover := &WALFailoverOptions{Secondary: secondary, FailoverOptions: wal.FailoverOptions{
@@ -2221,6 +2225,73 @@ func TestWALFailoverAvoidsWriteStall(t *testing.T) {
 	primaryFS.unblock()
 }
 
+type testLogManager struct {
+	wal.Manager
+	elevateWriteStallThreshold atomic.Bool
+}
+
+func (tlm *testLogManager) ElevateWriteStallThresholdForFailover() bool {
+	return tlm.elevateWriteStallThreshold.Load()
+}
+
+func TestElevateThresholdAfterWriteStallUnblocksStall(t *testing.T) {
+	mem := vfs.NewMem()
+	// All sst writes are blocked.
+	blockingFS := newBlockingFS(mem, false /*blockWAL*/, true /*blockSST*/)
+	writeStallBeginCh := make(chan struct{}, 1)
+	el := EventListener{
+		WriteStallBegin: func(_ WriteStallBeginInfo) {
+			writeStallBeginCh <- struct{}{}
+		},
+	}
+	o := &Options{
+		FS:                          blockingFS,
+		MemTableSize:                4 << 20,
+		MemTableStopWritesThreshold: 2,
+		Logger:                      testutils.Logger{T: t},
+		EventListener:               &el,
+	}
+	d, err := Open("", o)
+	// Replace the log manager with one that can elevate the write stall threshold.
+	d.mu.Lock()
+	testWALManager := &testLogManager{Manager: d.mu.log.manager}
+	d.mu.log.manager = testWALManager
+	d.mu.Unlock()
+	require.NoError(t, err)
+	value := make([]byte, 1<<20)
+	for i := range value {
+		value[i] = byte(rand.Uint32())
+	}
+	go func() {
+		// Wait for write stall to begin.
+		<-writeStallBeginCh
+		t.Logf("write stall has begun")
+		testWALManager.elevateWriteStallThreshold.Store(true)
+	}()
+	done := make(chan struct{})
+	go func() {
+		// After ~8 writes, the default write stall threshold is exceeded.
+		// It is observed by the above goroutine, which removes the stall.
+		for i := 0; i < 200; i++ {
+			require.NoError(t, d.Set([]byte(fmt.Sprintf("%d", i)), value, nil))
+		}
+		done <- struct{}{}
+	}()
+	timeout := 15 * time.Second
+	if buildtags.SlowBuild {
+		timeout = time.Minute
+	}
+	select {
+	case <-time.After(timeout):
+		t.Fatalf("write stall did not terminate")
+	case <-done:
+	}
+	require.True(t, testWALManager.elevateWriteStallThreshold.Load())
+	// Unblock the writes to allow the DB to close.
+	blockingFS.unblock()
+	require.NoError(t, d.Close())
+}
+
 // TestDeterminism is a datadriven test intended to validate determinism of
 // operations in the face of concurrency or randomizing of operations. The test
 // data defines a sequence of commands run sequentially. Then the test may