Repository: petermattis/pebble
Pull request #211 (2019-08-06 12:48:45)

Fix a triple of synchronization bugs
The previous reworking of the `commitQueue` implementation accidentally
omitted an important invariant: `commitQueue.dequeue()` should only
return "applied" batches. The lack of this check was allowing the
visible sequence number to be ratcheted for unapplied batches. Added an
assertion that this invariant is not violated.

The change to not hold `DB.mu` during log writes failed to adequately
protected `DB.mu.log.LogWriter` allowing that field to be modified
concurrently with a call to `LogWriter.SyncRecord`. The concurrent
modification was occurring because `DB.mu` no longer protects that
field, yet some callers of `DB.makeRoomForWrite` were not synchronized.
Specifically, the calls in `DB.{Compact,Flush,AsyncFlush}`.

The change to not use `DB.mu` for commit pipeline synchronization also
led to a new synchronization requirement on `versionSet.logSeqNum`. That
field is incremented while holding `commitPipeline.mu`, but read while
holding `DB.mu`. Neither mutex protects it! The solution is to increment
and read that field atomically.

Added a test which stresses concurrent commit, flush, and compact
operations. All of the above problems were easily visible by running
this test under the race detector.

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/211)
<!-- Reviewable:end -->

petermattis: Some real 🤦‍♂ action going on here. All my fault, too.
petermattis: PS The first commit is #210.

Patch:
diff --git a/commit.go b/commit.go
index 828c4c3..109c4a0 100644
--- a/commit.go
+++ b/commit.go
@@ -79,7 +79,7 @@ func (q *commitQueue) enqueue(b *Batch) {
 		}
 
 		// The head slot is free, so we own it.
-		*slot = unsafe.Pointer(b)
+		atomic.StorePointer(slot, unsafe.Pointer(b))
 
 		// Increment head. This passes ownership of slot to dequeue and acts as a
 		// store barrier for writing the slot.
@@ -89,7 +89,6 @@ func (q *commitQueue) enqueue(b *Batch) {
 }
 
 func (q *commitQueue) dequeue() *Batch {
-	var slot *unsafe.Pointer
 	for {
 		ptrs := atomic.LoadUint64(&q.headTail)
 		head, tail := q.unpack(ptrs)
@@ -98,26 +97,28 @@ func (q *commitQueue) dequeue() *Batch {
 			return nil
 		}
 
+		slot := &q.slots[tail&uint32(len(q.slots)-1)]
+		b := (*Batch)(atomic.LoadPointer(slot))
+		if b == nil || atomic.LoadUint32(&b.applied) == 0 {
+			// The batch is not ready to be dequeued, or another goroutine has
+			// already dequeued it.
+			return nil
+		}
+
 		// Confirm head and tail (for our speculative check above) and increment
 		// tail. If this succeeds, then we own the slot at tail.
 		ptrs2 := q.pack(head, tail+1)
 		if atomic.CompareAndSwapUint64(&q.headTail, ptrs, ptrs2) {
-			// Success.
-			slot = &q.slots[tail&uint32(len(q.slots)-1)]
-			break
+			// We now own slot.
+			//
+			// Tell enqueue that we're done with this slot. Zeroing the slot is also
+			// important so we don't leave behind references that could keep this object
+			// live longer than necessary.
+			atomic.StorePointer(slot, nil)
+			// At this point enqueue owns the slot.
+			return b
 		}
 	}
-
-	// We now own slot.
-	b := (*Batch)(*slot)
-
-	// Tell enqueue that we're done with this slot. Zeroing the slot is also
-	// important so we don't leave behind references that could keep this object
-	// live longer than necessary.
-	atomic.StorePointer(slot, nil)
-	// At this point enqueue owns the slot.
-
-	return b
 }
 
 // commitEnv contains the environment that a commitPipeline interacts
@@ -295,12 +296,13 @@ func (p *commitPipeline) AllocateSeqNum(prepare func(), apply func(seqNum uint64
 	// number order.
 	p.pending.enqueue(b)
 
-	// Assign the batch a sequence number.
-	seqNum := *p.env.logSeqNum
-	*p.env.logSeqNum++
+	// Assign the batch a sequence number. Note that logSeqNum is not protected
+	// by commitPipeline.mu or DB.mu and thus needs to be incremented atomically.
+	seqNum := atomic.AddUint64(p.env.logSeqNum, 1) - 1
 	if seqNum == 0 {
-		seqNum = *p.env.logSeqNum
-		*p.env.logSeqNum++
+		// We can't use the value 0 for the global seqnum during ingestion, because
+		// 0 indicates no global seqnum. So allocate another seqnum.
+		seqNum = atomic.AddUint64(p.env.logSeqNum, 1) - 1
 		b.setCount(2)
 	}
 	b.setSeqNum(seqNum)
@@ -344,9 +346,9 @@ func (p *commitPipeline) prepare(b *Batch, syncWAL bool) (*memTable, error) {
 	// number order.
 	p.pending.enqueue(b)
 
-	// Assign the batch a sequence number.
-	b.setSeqNum(*p.env.logSeqNum)
-	*p.env.logSeqNum += n
+	// Assign the batch a sequence number. Note that logSeqNum is not protected
+	// by commitPipeline.mu or DB.mu and thus needs to be incremented atomically.
+	b.setSeqNum(atomic.AddUint64(p.env.logSeqNum, n) - n)
 
 	// Write the data to the WAL.
 	mem, err := p.env.write(b, syncWG)
@@ -375,6 +377,9 @@ func (p *commitPipeline) publish(b *Batch) {
 			b.commit.Wait()
 			break
 		}
+		if atomic.LoadUint32(&t.applied) != 1 {
+			panic("not reached")
+		}
 
 		// We're responsible for publishing the sequence number for batch t, but
 		// another concurrent goroutine might sneak in and publish the sequence
diff --git a/db.go b/db.go
index 12b3607..156186d 100644
--- a/db.go
+++ b/db.go
@@ -195,6 +195,10 @@ type DB struct {
 			queue   []uint64
 			size    uint64
 			bytesIn uint64
+			// The LogWriter is protected by commitPipeline.mu. This allows log
+			// writes to be performed without holding DB.mu, but requires both
+			// commitPipeline.mu and DB.mu to be held when rotating the WAL/memtable
+			// (i.e. makeRoomForWrite).
 			*record.LogWriter
 		}
 
@@ -417,13 +421,18 @@ func (d *DB) commitWrite(b *Batch, wg *sync.WaitGroup) (*memTable, error) {
 		d.mu.log.bytesIn += uint64(len(b.storage.data))
 	}
 
+	// Grab a reference to the memtable while holding DB.mu. Note that
+	// makeRoomForWrite() add a reference to the memtable which will prevent it
+	// from being flushed until we unreference it.
+	mem := d.mu.mem.mutable
+
 	d.mu.Unlock()
 	if err != nil {
 		return nil, err
 	}
 
 	if d.opts.DisableWAL {
-		return d.mu.mem.mutable, nil
+		return mem, nil
 	}
 
 	size, err := d.mu.log.SyncRecord(b.storage.data, wg)
@@ -432,7 +441,7 @@ func (d *DB) commitWrite(b *Batch, wg *sync.WaitGroup) (*memTable, error) {
 	}
 
 	atomic.StoreUint64(&d.mu.log.size, uint64(size))
-	return d.mu.mem.mutable, err
+	return mem, err
 }
 
 type iterAlloc struct {
@@ -683,7 +692,19 @@ func (d *DB) Compact(start, end []byte /* CompactionOptions */) error {
 	mem, err := func() (flushable, error) {
 		if ingestMemtableOverlaps(d.cmp, d.mu.mem.mutable, meta) {
 			mem := d.mu.mem.mutable
-			return mem, d.makeRoomForWrite(nil)
+
+			// We have to hold both commitPipeline.mu and DB.mu when calling
+			// makeRoomForWrite(). Lock order requirements elsewhere force us to
+			// unlock DB.mu in order to grab commitPipeline.mu first.
+			d.mu.Unlock()
+			d.commit.mu.Lock()
+			d.mu.Lock()
+			defer d.commit.mu.Unlock()
+			if mem == d.mu.mem.mutable {
+				// Only flush if the active memtable is unchanged.
+				return mem, d.makeRoomForWrite(nil)
+			}
+			return mem, nil
 		}
 		// Check to see if any files overlap with any of the immutable
 		// memtables. The queue is ordered from oldest to newest. We want to wait
@@ -740,10 +761,12 @@ func (d *DB) Flush() error {
 		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()
 	if err != nil {
 		return err
 	}
@@ -759,9 +782,11 @@ func (d *DB) AsyncFlush() error {
 		panic(ErrClosed)
 	}
 
+	d.commit.mu.Lock()
 	d.mu.Lock()
 	err := d.makeRoomForWrite(nil)
 	d.mu.Unlock()
+	d.commit.mu.Unlock()
 	return err
 }
 
@@ -803,6 +828,15 @@ func (d *DB) walPreallocateSize() int {
 	return size
 }
 
+// makeRoomForWrite ensures that the memtable has room to hold the contents of
+// Batch. It reserves the space in the memtable and adds a reference to the
+// memtable. The caller must later ensure that the memtable is unreferenced. If
+// the memtable is full, or a nil Batch is provided, the current memtable is
+// rotated (marked as immutable) and a new mutable memtable is allocated. This
+// memtable rotation also causes a log rotation.
+//
+// Both DB.mu and commitPipeline.mu must be held by the caller. Note that DB.mu
+// may be released and reacquired.
 func (d *DB) makeRoomForWrite(b *Batch) error {
 	force := b == nil || b.flushable != nil
 	for {


Test patch:
diff --git a/batch_test.go b/batch_test.go
index d3608c5..41d68e0 100644
--- a/batch_test.go
+++ b/batch_test.go
@@ -424,13 +424,7 @@ func TestFlushableBatchBytesIterated(t *testing.T) {
 		var bytesIterated uint64
 		it := fb.newFlushIter(nil, &bytesIterated)
 
-		var prevIterated uint64
-		for it.First(); it.Valid(); it.Next() {
-			if bytesIterated < prevIterated {
-				t.Fatalf("bytesIterated moved backward: %d < %d", bytesIterated, prevIterated)
-			}
-			prevIterated = bytesIterated
-		}
+		for it.First(); it.Valid(); it.Next() {}
 
 		expected := fb.totalBytes()
 		if bytesIterated != expected {
diff --git a/commit_test.go b/commit_test.go
index 8500914..975ad94 100644
--- a/commit_test.go
+++ b/commit_test.go
@@ -52,6 +52,30 @@ func (e *testCommitEnv) write(b *Batch, _ *sync.WaitGroup) (*memTable, error) {
 	return nil, nil
 }
 
+func TestCommitQueue(t *testing.T) {
+	var q commitQueue
+	var batches [16]Batch
+	for i := range batches {
+		q.enqueue(&batches[i])
+	}
+	if b := q.dequeue(); b != nil {
+		t.Fatalf("unexpectedly dequeued batch: %p", b)
+	}
+	atomic.StoreUint32(&batches[1].applied, 1)
+	if b := q.dequeue(); b != nil {
+		t.Fatalf("unexpectedly dequeued batch: %p", b)
+	}
+	for i := range batches {
+		atomic.StoreUint32(&batches[i].applied, 1)
+		if b := q.dequeue(); b != &batches[i] {
+			t.Fatalf("%d: expected batch %p, but found %p", i, &batches[i], b)
+		}
+	}
+	if b := q.dequeue(); b != nil {
+		t.Fatalf("unexpectedly dequeued batch: %p", b)
+	}
+}
+
 func TestCommitPipeline(t *testing.T) {
 	var e testCommitEnv
 	p := newCommitPipeline(e.env())
diff --git a/db_test.go b/db_test.go
index 2ba3a4b..e281a08 100644
--- a/db_test.go
+++ b/db_test.go
@@ -732,3 +732,41 @@ func TestDBClosed(t *testing.T) {
 	require.EqualValues(t, ErrClosed, catch(func() { _ = d.Apply(b, nil) }))
 	require.EqualValues(t, ErrClosed, catch(func() { _ = b.NewIter(nil) }))
 }
+
+func TestDBConcurrentCommitCompactFlush(t *testing.T) {
+	d, err := Open("", &Options{
+		FS: vfs.NewMem(),
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Concurrently commit, compact, and flush in order to stress the locking around
+	// those operations.
+	const n = 1000
+	var wg sync.WaitGroup
+	wg.Add(n)
+	for i := 0; i < n; i++ {
+		go func(i int) {
+			defer wg.Done()
+			_ = d.Set([]byte(fmt.Sprint(i)), nil, nil)
+			var err error
+			switch i % 3 {
+			case 0:
+				err = d.Compact(nil, []byte("\xff"))
+			case 1:
+				err = d.Flush()
+			case 2:
+				err = d.AsyncFlush()
+			}
+			if err != nil {
+				t.Fatal(err)
+			}
+		}(i)
+	}
+	wg.Wait()
+
+	if err := d.Close(); err != nil {
+		t.Fatal(err)
+	}
+}
diff --git a/internal/arenaskl/skl_test.go b/internal/arenaskl/skl_test.go
index 2263296..1c46417 100644
--- a/internal/arenaskl/skl_test.go
+++ b/internal/arenaskl/skl_test.go
@@ -692,7 +692,7 @@ func TestBytesIterated(t *testing.T) {
 	l := NewSkiplist(NewArena(arenaSize, 0), bytes.Compare)
 	emptySize := l.arena.Size()
 	for i := 0; i < 200; i++ {
-		bytesIterated := l.bytesIterated(t)
+		bytesIterated := l.bytesIterated()
 		expected := uint64(l.arena.Size() - emptySize)
 		if bytesIterated != expected {
 			t.Fatalf("bytesIterated: got %d, want %d", bytesIterated, expected)
@@ -702,15 +702,9 @@ func TestBytesIterated(t *testing.T) {
 }
 
 // bytesIterated returns the number of bytes iterated in the skiplist.
-func (s *Skiplist) bytesIterated(t *testing.T) (bytesIterated uint64) {
+func (s *Skiplist) bytesIterated() (bytesIterated uint64) {
 	x := s.NewFlushIter(&bytesIterated)
-	var prevIterated uint64
-	for x.First(); x.Valid(); x.Next() {
-		if bytesIterated < prevIterated {
-			t.Fatalf("bytesIterated moved backward: %d < %d", bytesIterated, prevIterated)
-		}
-		prevIterated = bytesIterated
-	}
+	for x.First(); x.Valid(); x.Next() {}
 	if x.Close() != nil {
 		return 0
 	}
diff --git a/mem_table_test.go b/mem_table_test.go
index ff98721..0e6a693 100644
--- a/mem_table_test.go
+++ b/mem_table_test.go
@@ -47,15 +47,9 @@ func (m *memTable) count() (n int) {
 }
 
 // bytesIterated returns the number of bytes iterated in a DB.
-func (m *memTable) bytesIterated(t *testing.T) (bytesIterated uint64) {
+func (m *memTable) bytesIterated() (bytesIterated uint64) {
 	x := internalIterAdapter{m.newFlushIter(nil, &bytesIterated)}
-	var prevIterated uint64
-	for valid := x.First(); valid; valid = x.Next() {
-		if bytesIterated < prevIterated {
-			t.Fatalf("bytesIterated moved backward: %d < %d", bytesIterated, prevIterated)
-		}
-		prevIterated = bytesIterated
-	}
+	for valid := x.First(); valid; valid = x.Next() {}
 	if x.Close() != nil {
 		return 0
 	}
@@ -134,7 +128,7 @@ func TestMemTableCount(t *testing.T) {
 func TestMemTableBytesIterated(t *testing.T) {
 	m := newMemTable(nil)
 	for i := 0; i < 200; i++ {
-		bytesIterated := m.bytesIterated(t)
+		bytesIterated := m.bytesIterated()
 		expected := m.totalBytes()
 		if bytesIterated != expected {
 			t.Fatalf("bytesIterated: got %d, want %d", bytesIterated, expected)
diff --git a/sstable/reader_test.go b/sstable/reader_test.go
index 38492c1..8ae50e8 100644
--- a/sstable/reader_test.go
+++ b/sstable/reader_test.go
@@ -319,7 +319,7 @@ func buildTestTable(t *testing.T, numEntries uint64, blockSize int, compression
 	var ikey InternalKey
 	for i := uint64(0); i < numEntries; i++ {
 		key := make([]byte, 8+i%3)
-		value := make([]byte, i%100)
+		value := make([]byte, 7+i%5)
 		binary.BigEndian.PutUint64(key, i)
 		ikey.UserKey = key
 		w.Add(ikey, value)