Repository: petermattis/pebble
Pull request #209 (2019-08-05 16:13:36)

Add `EventListener` functions for write stalls
Call `WriteStallBegin()` when entering write stall mode and
`WriteStallEnd()` when leaving write stall mode.

petermattis: 

<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.png" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/petermattis/pebble/209)
<!-- Reviewable:end -->

ajkr: ^also should investigate why the first flush took 12 seconds.
Ryanfsdf: > ^also should investigate why the first flush took 12 seconds.

I wonder if this has to do with my memtable flushing mechanism (though the mechanism shouldn't throttle flush rate when at the memtable limit)
ajkr: > I'm not sure if the write stall is ended at this point. I think you need a loop:

Good point, what I was thinking.

> That's just my idea for a test. Feel free to suggest something better.

Not sure if better, but I think we could do this entirely with a custom `EventListener` by configuring its `FlushBegin()` callback to sleep until `WriteStallBegin()` is called. Then flush can be released to proceed and we verify `WriteStallEnd()` is eventually called. Will try it out.
petermattis: > Not sure if better, but I think we could do this entirely with a custom EventListener by configuring its FlushBegin() callback to sleep until WriteStallBegin() is called. Then flush can be released to proceed and we verify WriteStallEnd() is eventually called. Will try it out.

I think that should work. That is a better idea. 
ajkr: > > Not sure if better, but I think we could do this entirely with a custom EventListener by configuring its FlushBegin() callback to sleep until WriteStallBegin() is called. Then flush can be released to proceed and we verify WriteStallEnd() is eventually called. Will try it out.
> 
> I think that should work. That is a better idea.

Hm, it's getting kind of complicated because the DB mutex is held while `FlushBegin()` runs. That'd prevent writes to the second memtable or initiating the second flush which prevents us from hitting stall mode. Delaying in the manifest write sounds better now.
ajkr: > Sorry, though, I ended up slightly confused about both ideas. For the first one, would FlushBegin pause writes to the manifest the same way we're doing above?

Actually I think I got your first idea now. Will try it out.

Patch:
diff --git a/db.go b/db.go
index 156186d..1b576d2 100644
--- a/db.go
+++ b/db.go
@@ -757,37 +757,33 @@ func (d *DB) manualCompact(manual *manualCompaction) error {
 
 // Flush the memtable to stable storage.
 func (d *DB) Flush() error {
-	if atomic.LoadInt32(&d.closed) != 0 {
-		panic(ErrClosed)
-	}
-
-	d.commit.mu.Lock()
-	d.mu.Lock()
-	mem := d.mu.mem.mutable
-	err := d.makeRoomForWrite(nil)
-	d.mu.Unlock()
-	d.commit.mu.Unlock()
+	flushDone, err := d.AsyncFlush()
 	if err != nil {
 		return err
 	}
-	<-mem.flushed()
+	<-flushDone
 	return nil
 }
 
 // AsyncFlush asynchronously flushes the memtable to stable storage.
 //
-// TODO(peter): untested
-func (d *DB) AsyncFlush() error {
+// If no error is returned, the caller can receive from the returned channel in
+// order to wait for the flush to complete.
+func (d *DB) AsyncFlush() (<-chan struct{}, error) {
 	if atomic.LoadInt32(&d.closed) != 0 {
 		panic(ErrClosed)
 	}
 
 	d.commit.mu.Lock()
 	d.mu.Lock()
+	mem := d.mu.mem.mutable
 	err := d.makeRoomForWrite(nil)
 	d.mu.Unlock()
 	d.commit.mu.Unlock()
-	return err
+	if err != nil {
+		return nil, err
+	}
+	return mem.flushed(), nil
 }
 
 // Metrics returns metrics about the database.
@@ -839,6 +835,7 @@ func (d *DB) walPreallocateSize() int {
 // may be released and reacquired.
 func (d *DB) makeRoomForWrite(b *Batch) error {
 	force := b == nil || b.flushable != nil
+	stalled := false
 	for {
 		if d.mu.mem.switching {
 			d.mu.mem.cond.Wait()
@@ -846,25 +843,48 @@ func (d *DB) makeRoomForWrite(b *Batch) error {
 		}
 		if b != nil && b.flushable == nil {
 			err := d.mu.mem.mutable.prepare(b)
-			if err == nil {
-				return nil
-			}
 			if err != arenaskl.ErrArenaFull {
+				if stalled {
+					stalled = false
+					if d.opts.EventListener.WriteStallEnd != nil {
+						d.opts.EventListener.WriteStallEnd()
+					}
+				}
 				return err
 			}
 		} else if !force {
+			if stalled {
+				stalled = false
+				if d.opts.EventListener.WriteStallEnd != nil {
+					d.opts.EventListener.WriteStallEnd()
+				}
+			}
 			return nil
 		}
 		if len(d.mu.mem.queue) >= d.opts.MemTableStopWritesThreshold {
 			// We have filled up the current memtable, but the previous one is still
 			// being compacted, so we wait.
-			// fmt.Printf("memtable stop writes threshold\n")
+			if !stalled {
+				stalled = true
+				if d.opts.EventListener.WriteStallBegin != nil {
+					d.opts.EventListener.WriteStallBegin(WriteStallBeginInfo{
+						Reason: "memtable count limit reached",
+					})
+				}
+			}
 			d.mu.compact.cond.Wait()
 			continue
 		}
 		if len(d.mu.versions.currentVersion().files[0]) > d.opts.L0StopWritesThreshold {
 			// There are too many level-0 files, so we wait.
-			// fmt.Printf("L0 stop writes threshold\n")
+			if !stalled {
+				stalled = true
+				if d.opts.EventListener.WriteStallBegin != nil {
+					d.opts.EventListener.WriteStallBegin(WriteStallBeginInfo{
+						Reason: "L0 file count limit exceeded",
+					})
+				}
+			}
 			d.mu.compact.cond.Wait()
 			continue
 		}
diff --git a/event.go b/event.go
index 9fcd295..dc3fb6c 100644
--- a/event.go
+++ b/event.go
@@ -33,6 +33,9 @@ type WALCreateInfo = base.WALCreateInfo
 // WALDeleteInfo exports the base.WALDeleteInfo type.
 type WALDeleteInfo = base.WALDeleteInfo
 
+// WriteStallBeginInfo exports the base.WriteStallBeginInfo type.
+type WriteStallBeginInfo = base.WriteStallBeginInfo
+
 // EventListener exports the base.EventListener type.
 type EventListener = base.EventListener
 
diff --git a/internal/base/event.go b/internal/base/event.go
index ffd74d4..94612e8 100644
--- a/internal/base/event.go
+++ b/internal/base/event.go
@@ -230,6 +230,14 @@ func (i WALDeleteInfo) String() string {
 	return fmt.Sprintf("[JOB %d] WAL deleted %06d", i.JobID, i.FileNum)
 }
 
+type WriteStallBeginInfo struct {
+	Reason string
+}
+
+func (i WriteStallBeginInfo) String() string {
+	return fmt.Sprintf("write stall beginning: %s", i.Reason)
+}
+
 // EventListener contains a set of functions that will be invoked when various
 // significant DB events occur. Note that the functions should not run for an
 // excessive amount of time as they are invokved synchronously by the DB and
@@ -274,6 +282,12 @@ type EventListener struct {
 
 	// WALDeleted is invoked after a WAL has been deleted.
 	WALDeleted func(WALDeleteInfo)
+
+	// WriteStallBegin is invoked when writes are intentionally delayed.
+	WriteStallBegin func(WriteStallBeginInfo)
+
+	// WriteStallEnd is invoked when delayed writes are released.
+	WriteStallEnd func()
 }
 
 // EnsureDefaults ensures that background error events are logged to the
@@ -328,5 +342,11 @@ func MakeLoggingEventListener(logger Logger) EventListener {
 		WALDeleted: func(info WALDeleteInfo) {
 			logger.Infof("%s", info.String())
 		},
+		WriteStallBegin: func(info WriteStallBeginInfo) {
+			logger.Infof("%s", info.String())
+		},
+		WriteStallEnd: func() {
+			logger.Infof("write stall ending")
+		},
 	}
 }


Test patch:
diff --git a/db_test.go b/db_test.go
index e281a08..b48e271 100644
--- a/db_test.go
+++ b/db_test.go
@@ -716,7 +716,7 @@ func TestDBClosed(t *testing.T) {
 
 	require.EqualValues(t, ErrClosed, catch(func() { _ = d.Compact(nil, nil) }))
 	require.EqualValues(t, ErrClosed, catch(func() { _ = d.Flush() }))
-	require.EqualValues(t, ErrClosed, catch(func() { _ = d.AsyncFlush() }))
+	require.EqualValues(t, ErrClosed, catch(func() { _, _ = d.AsyncFlush() }))
 
 	require.EqualValues(t, ErrClosed, catch(func() { _, _ = d.Get(nil) }))
 	require.EqualValues(t, ErrClosed, catch(func() { _ = d.Delete(nil, nil) }))
@@ -757,7 +757,7 @@ func TestDBConcurrentCommitCompactFlush(t *testing.T) {
 			case 1:
 				err = d.Flush()
 			case 2:
-				err = d.AsyncFlush()
+				_, err = d.AsyncFlush()
 			}
 			if err != nil {
 				t.Fatal(err)
diff --git a/event_listener_test.go b/event_listener_test.go
index ae631b7..cfe4a50 100644
--- a/event_listener_test.go
+++ b/event_listener_test.go
@@ -8,13 +8,16 @@ import (
 	"bytes"
 	"fmt"
 	"io"
+	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 
 	"github.com/petermattis/pebble/internal/base"
 	"github.com/petermattis/pebble/internal/datadriven"
 	"github.com/petermattis/pebble/sstable"
 	"github.com/petermattis/pebble/vfs"
+	"github.com/stretchr/testify/require"
 )
 
 type syncedBuffer struct {
@@ -173,3 +176,130 @@ func TestEventListener(t *testing.T) {
 		}
 	})
 }
+
+type createDelayingFS struct {
+	vfs.FS
+	index          int32         // index of sstable whose creation will be delayed
+	createReleased chan struct{} // when closed, the delayed sstable creation can proceed
+	createDelayed  chan struct{} // if non-nil, will be closed before delaying sstable creation
+}
+
+func (fs *createDelayingFS) Create(name string) (vfs.File, error) {
+	if strings.HasSuffix(name, ".sst") && atomic.AddInt32(&fs.index, -1) == -1 {
+		if fs.createDelayed != nil {
+			close(fs.createDelayed)
+		}
+		<-fs.createReleased
+	}
+	return fs.FS.Create(name)
+}
+
+// Verify `WriteStallBegin()` and `WriteStallEnd()` events for stalls triggered
+// for reaching memtable count limit.
+func TestWriteStallForMemTableLimit(t *testing.T) {
+	const memTableLimitReason = "memtable count limit reached"
+	const memTableLimit = 2
+
+	stallBegan := make(chan struct{})
+	stallEnded := make(chan struct{})
+	mem := vfs.NewMem()
+	// `customFS` delays the sstable with index 0, which is the first flushed sstable.
+	// Delay until writes stall, which we trigger by calling `AsyncFlush()` again while
+	// the first flush is still delayed.
+	customFS := createDelayingFS{
+		FS:             mem,
+		index:          0,
+		createReleased: stallBegan,
+		createDelayed:  nil,
+	}
+	listener := EventListener{
+		WriteStallBegin: func(info WriteStallBeginInfo) {
+			require.EqualValues(t, memTableLimitReason, info.Reason)
+			close(stallBegan)
+		},
+		WriteStallEnd: func() {
+			close(stallEnded)
+		},
+	}
+	d, err := Open("db", &Options{
+		EventListener:               listener,
+		FS:                          &customFS,
+		MemTableStopWritesThreshold: memTableLimit,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	for i := 0; i < memTableLimit; i++ {
+		if err := d.Set([]byte("a"), nil, NoSync); err != nil {
+			t.Fatal(err)
+		}
+		_, err := d.AsyncFlush()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+	<-stallBegan
+	<-stallEnded
+}
+
+// Verify `WriteStallBegin()` and `WriteStallEnd()` events for stalls triggered
+// for reaching L0 file count limit.
+func TestWriteStallForL0FileLimit(t *testing.T) {
+	const l0FileLimitReason = "L0 file count limit exceeded"
+	const (
+		l0FileCompactionTrigger = 2
+		l0FileStopTrigger       = 4
+	)
+
+	stallBegan := make(chan struct{})
+	stallEnded := make(chan struct{})
+	createDelayed := make(chan struct{})
+	mem := vfs.NewMem()
+	// The indexes of the first few sstables are as follows:
+	//
+	// * index 0: first flushed sstable
+	// * index 1: second flushed sstable
+	// * index 2: first sstable created by L0->Lbase compaction
+	//
+	// `customFS` delays the sstable with index 2, i.e., an Lbase file created during
+	// compaction. Delay until writes stall, which we trigger by calling `Flush()`
+	// repeatedly while the L0->Lbase compaction is still delayed.
+	customFS := createDelayingFS{mem, l0FileCompactionTrigger /* index */, stallBegan, /* createReleased */
+		createDelayed}
+	listener := EventListener{
+		WriteStallBegin: func(info WriteStallBeginInfo) {
+			require.EqualValues(t, l0FileLimitReason, info.Reason)
+			close(stallBegan)
+		},
+		WriteStallEnd: func() {
+			close(stallEnded)
+		},
+	}
+	d, err := Open("db", &Options{
+		EventListener:         listener,
+		FS:                    &customFS,
+		L0CompactionThreshold: l0FileCompactionTrigger,
+		L0StopWritesThreshold: l0FileStopTrigger,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// TODO(ajkr): Are the two files pending L0->Lbase not counted towards the
+	// limit? Shouldn't they be?
+	for i := 0; i < l0FileStopTrigger+l0FileCompactionTrigger; i++ {
+		if err := d.Set([]byte("a"), nil, NoSync); err != nil {
+			t.Fatal(err)
+		}
+		d.Flush()
+		if i == l0FileCompactionTrigger-1 {
+			// make sure compaction gets started on the Lbase file before
+			// we proceed to flushing the next sstable (we wouldn't want the
+			// delay to be applied on an L0 file accidentally).
+			<-createDelayed
+		}
+	}
+	<-stallBegan
+	<-stallEnded
+}
diff --git a/flush_test.go b/flush_test.go
index 7a3c10b..8f30ad8 100644
--- a/flush_test.go
+++ b/flush_test.go
@@ -46,7 +46,7 @@ func TestManualFlush(t *testing.T) {
 			cur := d.mu.versions.currentVersion()
 			d.mu.Unlock()
 
-			if err := d.AsyncFlush(); err != nil {
+			if _, err := d.AsyncFlush(); err != nil {
 				return err.Error()
 			}