Repository: petermattis/pebble (go)

Issue:
perf: avoid holding DB.mu during log writes
Mutual exclusion for appending to the log (in memory) is currently achieved by holding `DB.mu`. This is unfortunate as a log write can result in a disk write and thus hold `DB.mu` for a significant amount of time. RocksDB has a separate log-write mutex for providing mutual exclusion for log writes. We should investigate performing a similar separation of locking duties in Pebble.

Pull request description:
Avoid holding DB.mu during log writes
Change `commitPipeline` to use a separate mutex rather than `DB.mu` for
serializing writing to the WAL. The use of `DB.mu` was a potential
performance problem because `DB.mu` is used for reads and holding it
during a log write meant that it was potentially held during disk
IO. There was also a super subtle bug here: `DB.makeRoomForWrite`
releases and then reacquires `DB.mu`. This behavior is necessary when
rotating the log, but the effect was that all of the careful ordering
guarantees in `commitePipeline` could be violated.

Fixes #44

Patch:
diff --git a/commit.go b/commit.go
index 4aa6db4..b1bcf2f 100644
--- a/commit.go
+++ b/commit.go
@@ -110,11 +110,8 @@ func (q *commitQueue) dequeue(cond *sync.Cond) *Batch {
 // with. This allows fine-grained testing of commitPipeline behavior without
 // construction of an entire DB.
 type commitEnv struct {
-	// The mutex to use for synchronizing access to logSeqNum and serializing
-	// calls to write().
-	mu *sync.Mutex
-	// The next sequence number to give to a batch. Mutated atomically by the
-	// current WAL writer.
+	// The next sequence number to give to a batch. Protected by
+	// commitPipeline.mu.
 	logSeqNum *uint64
 	// The visible sequence number at which reads should be performed. Ratcheted
 	// upwards atomically as batches are applied to the memtable.
@@ -126,7 +123,7 @@ type commitEnv struct {
 	sync func() error
 	// Write the batch to the WAL. The data is not persisted until a call to
 	// sync() is performed. Returns the memtable the batch should be applied
-	// to. Called serially.
+	// to. Serial execution enforced by commitPipeline.mu.
 	write func(b *Batch) (*memTable, error)
 }
 
@@ -167,7 +164,7 @@ type commitEnv struct {
 //
 // The full outline of the commit pipeline operation is as follows:
 //
-//   with DB mutex locked:
+//   with commitPipeline mutex locked:
 //     assign batch sequence number
 //     write batch to WAL
 //   (optionally) add batch to WAL sync list
@@ -176,8 +173,8 @@ type commitEnv struct {
 //   ratchet read sequence number
 //   (optionally) wait for the WAL to sync
 //
-// As soon as a batch has been written to the WAL, the DB mutex is released
-// allowing another batch to write to the WAL. Each commit operation
+// As soon as a batch has been written to the WAL, the commitPipeline mutex is
+// released allowing another batch to write to the WAL. Each commit operation
 // individually applies its batch to the memtable providing concurrency. The
 // WAL sync happens concurrently with applying to the memtable (see
 // commitPipeline.syncLoop).
@@ -192,8 +189,8 @@ type commitEnv struct {
 // though it avoids waking a goroutine to perform work that another goroutine
 // can perform. A commitQueue (a single-producer, multiple-consumer queue)
 // holds the ordered list of committing batches. Addition to the queue is done
-// while holding the DB mutex ensuring the same ordering of batches in the
-// queue as the ordering in the WAL. When a batch finishes applying to the
+// while holding commitPipeline.mutex ensuring the same ordering of batches in
+// the queue as the ordering in the WAL. When a batch finishes applying to the
 // memtable, it atomically updates its Batch.applied field. Ratcheting of the
 // visible sequence number is done by commitPipeline.publish which loops
 // dequeueing "applied" batches and ratcheting the visible sequence number. If
@@ -202,6 +199,9 @@ type commitEnv struct {
 // batch in the queue. See commitPipeline.publish for additional commentary.
 type commitPipeline struct {
 	env commitEnv
+	// The mutex to use for synchronizing access to logSeqNum and serializing
+	// calls to commitEnv.write().
+	mu sync.Mutex
 	// Condition var to signal upon changes to the pending queue.
 	cond sync.Cond
 	// Queue of pending batches to commit.
@@ -219,7 +219,7 @@ func newCommitPipeline(env commitEnv) *commitPipeline {
 	p := &commitPipeline{
 		env: env,
 	}
-	p.cond.L = p.env.mu
+	p.cond.L = &p.mu
 	p.pending.init()
 	p.syncer.cond.L = &p.syncer.Mutex
 	go p.syncLoop()
@@ -300,8 +300,8 @@ func (p *commitPipeline) Commit(b *Batch, syncWAL bool) error {
 // number. AllocateSeqNum does not write to the WAL or add entries to the
 // memtable. AllocateSeqNum can be used to sequence an operation such as
 // sstable ingestion within the commit pipeline. The prepare callback is
-// invoked with commitEnv.mu held, making it suitable for flushing the memtable
-// if necessary.
+// invoked with commitPipeline.mu held, but note that DB.mu is not held and
+// must be locked if necessary.
 func (p *commitPipeline) AllocateSeqNum(prepare func(), apply func(seqNum uint64)) {
 	// This method is similar to Commit and prepare. Be careful about trying to
 	// share additional code with those methods because Commit and prepare are
@@ -316,7 +316,7 @@ func (p *commitPipeline) AllocateSeqNum(prepare func(), apply func(seqNum uint64
 	b.setCount(1)
 	b.commit.Add(1)
 
-	p.env.mu.Lock()
+	p.mu.Lock()
 
 	// Enqueue the batch in the pending queue. Note that while the pending queue
 	// is lock-free, we want the order of batches to be the same as the sequence
@@ -324,9 +324,11 @@ func (p *commitPipeline) AllocateSeqNum(prepare func(), apply func(seqNum uint64
 	p.pending.enqueue(b, &p.cond)
 
 	// Assign the batch a sequence number.
-	seqNum := atomic.AddUint64(p.env.logSeqNum, 1) - 1
+	seqNum := *p.env.logSeqNum
+	*p.env.logSeqNum++
 	if seqNum == 0 {
-		seqNum = atomic.AddUint64(p.env.logSeqNum, 1) - 1
+		seqNum = *p.env.logSeqNum
+		*p.env.logSeqNum++
 		b.setCount(2)
 	}
 	b.setSeqNum(seqNum)
@@ -336,7 +338,7 @@ func (p *commitPipeline) AllocateSeqNum(prepare func(), apply func(seqNum uint64
 	// order to allow the commit pipeline to proceed.
 	prepare()
 
-	p.env.mu.Unlock()
+	p.mu.Unlock()
 
 	// Invoke the apply callback.
 	apply(b.seqNum())
@@ -356,7 +358,7 @@ func (p *commitPipeline) prepare(b *Batch, syncWAL bool) (*memTable, error) {
 	}
 	b.commit.Add(count)
 
-	p.env.mu.Lock()
+	p.mu.Lock()
 
 	// Enqueue the batch in the pending queue. Note that while the pending queue
 	// is lock-free, we want the order of batches to be the same as the sequence
@@ -364,12 +366,13 @@ func (p *commitPipeline) prepare(b *Batch, syncWAL bool) (*memTable, error) {
 	p.pending.enqueue(b, &p.cond)
 
 	// Assign the batch a sequence number.
-	b.setSeqNum(atomic.AddUint64(p.env.logSeqNum, n) - n)
+	b.setSeqNum(*p.env.logSeqNum)
+	*p.env.logSeqNum += n
 
 	// Write the data to the WAL.
 	mem, err := p.env.write(b)
 
-	p.env.mu.Unlock()
+	p.mu.Unlock()
 
 	if syncWAL {
 		s := &p.syncer
diff --git a/db.go b/db.go
index ae44f49..616f070 100644
--- a/db.go
+++ b/db.go
@@ -192,13 +192,7 @@ func (d *DB) Get(key []byte) ([]byte, error) {
 }
 
 func (d *DB) getInternal(key []byte, b *Batch, s *Snapshot) ([]byte, error) {
-	var seqNum uint64
 	d.mu.Lock()
-	if s != nil {
-		seqNum = s.seqNum
-	} else {
-		seqNum = atomic.LoadUint64(&d.mu.versions.visibleSeqNum)
-	}
 	// Grab and reference the current version to prevent its underlying files
 	// from being deleted if we have a concurrent compaction. Note that
 	// version.unref() can be called without holding DB.mu.
@@ -207,6 +201,15 @@ func (d *DB) getInternal(key []byte, b *Batch, s *Snapshot) ([]byte, error) {
 	memtables := d.mu.mem.queue
 	d.mu.Unlock()
 
+	// Determine the seqnum to read at after grabbing the read state (current and
+	// memtables) above.
+	var seqNum uint64
+	if s != nil {
+		seqNum = s.seqNum
+	} else {
+		seqNum = atomic.LoadUint64(&d.mu.versions.visibleSeqNum)
+	}
+
 	var buf struct {
 		dbi Iterator
 		get getIter
@@ -340,7 +343,7 @@ func (d *DB) commitSync() error {
 }
 
 func (d *DB) commitWrite(b *Batch) (*memTable, error) {
-	// NB: commitWrite is called with d.mu locked.
+	d.mu.Lock()
 
 	// Throttle writes if there are too many L0 tables.
 	d.throttleWrite()
@@ -349,9 +352,11 @@ func (d *DB) commitWrite(b *Batch) (*memTable, error) {
 		b.flushable.seqNum = b.seqNum()
 	}
 
-	// Switch out the memtable if there was not enough room to store the
-	// batch.
-	if err := d.makeRoomForWrite(b); err != nil {
+	// Switch out the memtable if there was not enough room to store the batch.
+	err := d.makeRoomForWrite(b)
+
+	d.mu.Unlock()
+	if err != nil {
 		return nil, err
 	}
 
@@ -359,7 +364,7 @@ func (d *DB) commitWrite(b *Batch) (*memTable, error) {
 		return d.mu.mem.mutable, nil
 	}
 
-	_, err := d.mu.log.WriteRecord(b.storage.data)
+	_, err = d.mu.log.WriteRecord(b.storage.data)
 	if err != nil {
 		panic(err)
 	}
@@ -374,13 +379,7 @@ func (d *DB) newIterInternal(
 	s *Snapshot,
 	o *db.IterOptions,
 ) *Iterator {
-	var seqNum uint64
 	d.mu.Lock()
-	if s != nil {
-		seqNum = s.seqNum
-	} else {
-		seqNum = atomic.LoadUint64(&d.mu.versions.visibleSeqNum)
-	}
 	// Grab and reference the current version to prevent its underlying files
 	// from being deleted if we have a concurrent compaction. Note that
 	// version.unref() can be called without holding DB.mu.
@@ -389,6 +388,15 @@ func (d *DB) newIterInternal(
 	memtables := d.mu.mem.queue
 	d.mu.Unlock()
 
+	// Determine the seqnum to read at after grabbing the read state (current and
+	// memtables) above.
+	var seqNum uint64
+	if s != nil {
+		seqNum = s.seqNum
+	} else {
+		seqNum = atomic.LoadUint64(&d.mu.versions.visibleSeqNum)
+	}
+
 	// Bundle various structures under a single umbrella in order to allocate
 	// them together.
 	var buf struct {
@@ -518,9 +526,11 @@ func (d *DB) NewIter(o *db.IterOptions) *Iterator {
 // deleted. Instead, a snapshot prevents deletion of sequence numbers
 // referenced by the snapshot.
 func (d *DB) NewSnapshot() *Snapshot {
-	s := &Snapshot{db: d}
+	s := &Snapshot{
+		db:     d,
+		seqNum: atomic.LoadUint64(&d.mu.versions.visibleSeqNum),
+	}
 	d.mu.Lock()
-	s.seqNum = atomic.LoadUint64(&d.mu.versions.visibleSeqNum)
 	d.mu.snapshots.pushBack(s)
 	d.mu.Unlock()
 	return s
diff --git a/ingest.go b/ingest.go
index d7123f8..009ffea 100644
--- a/ingest.go
+++ b/ingest.go
@@ -259,8 +259,9 @@ func (d *DB) Ingest(paths []string) error {
 	}
 
 	var mem flushable
-	prepareLocked := func() {
-		// NB: prepare is called with d.mu locked.
+	prepare := func() {
+		d.mu.Lock()
+		defer d.mu.Unlock()
 
 		// If the mutable memtable contains keys which overlap any of the sstables
 		// then flush the memtable. Note that apply will wait for the flushing to
@@ -286,7 +287,7 @@ func (d *DB) Ingest(paths []string) error {
 	var ve *versionEdit
 	apply := func(seqNum uint64) {
 		if err != nil {
-			// An error occurred during prepareLocked.
+			// An error occurred during prepare.
 			return
 		}
 
@@ -296,8 +297,8 @@ func (d *DB) Ingest(paths []string) error {
 			return
 		}
 
-		// If we flushed the mutable memtable in prepareLocked wait for the flush
-		// to finish.
+		// If we flushed the mutable memtable in prepare wait for the flush to
+		// finish.
 		if mem != nil {
 			<-mem.flushed()
 		}
@@ -307,7 +308,7 @@ func (d *DB) Ingest(paths []string) error {
 		ve, err = d.ingestApply(meta)
 	}
 
-	d.commit.AllocateSeqNum(prepareLocked, apply)
+	d.commit.AllocateSeqNum(prepare, apply)
 
 	if err != nil {
 		if err2 := ingestCleanup(d.opts.Storage, d.dirname, meta); err2 != nil {
diff --git a/open.go b/open.go
index 6852889..57b6f13 100644
--- a/open.go
+++ b/open.go
@@ -73,7 +73,6 @@ func Open(dirname string, opts *db.Options) (*DB, error) {
 	d.tableCache.init(dirname, opts.Storage, d.opts, tableCacheSize)
 	d.newIters = d.tableCache.newIters
 	d.commit = newCommitPipeline(commitEnv{
-		mu:            &d.mu.Mutex,
 		logSeqNum:     &d.mu.versions.logSeqNum,
 		visibleSeqNum: &d.mu.versions.visibleSeqNum,
 		apply:         d.commitApply,


Test patch:
diff --git a/commit_test.go b/commit_test.go
index ec7e66f..3ed3476 100644
--- a/commit_test.go
+++ b/commit_test.go
@@ -19,7 +19,6 @@ import (
 )
 
 type testCommitEnv struct {
-	mu            sync.Mutex
 	logSeqNum     uint64
 	visibleSeqNum uint64
 	writePos      int64
@@ -32,7 +31,6 @@ type testCommitEnv struct {
 
 func (e *testCommitEnv) env() commitEnv {
 	return commitEnv{
-		mu:            &e.mu,
 		logSeqNum:     &e.logSeqNum,
 		visibleSeqNum: &e.visibleSeqNum,
 		apply:         e.apply,
@@ -136,7 +134,6 @@ func BenchmarkCommitPipeline(b *testing.B) {
 			wal := record.NewLogWriter(ioutil.Discard)
 
 			nullCommitEnv := commitEnv{
-				mu:            new(sync.Mutex),
 				logSeqNum:     new(uint64),
 				visibleSeqNum: new(uint64),
 				apply: func(b *Batch, mem *memTable) error {