Repository: petermattis/pebble (go)

Issue:
perf: lock-free reads
`DB.Get` and `DB.NewIter` both currently grab `DB.mu` in order to atomically reference the current memtable, the immutable memtables, and the current version. In RocksDB, the [`SuperVersion`](https://github.com/facebook/rocksdb/blob/master/db/column_family.h#L91) holds a reference to the this state. By bundling this state into the `SuperVersion`, a single reference count can be bumped to ensure that the state is stable for the lifetime of a read operation (`Get` or `Iterator`). The current `SuperVersion` is retrieved from a thread-local variable (see [`ColumnFamilyData::GetThreadLocalSuperVersion`](https://github.com/facebook/rocksdb/blob/master/db/column_family.cc#L1012)). It is unclear how to do something similar in Go given the lack of thread-local variables, but at the very least we could bundle the read state into a single struct and "acquire" it with an atomic operation.

Pull request description:
Acquire state for reading without using DB.mu
Add a `readState` struct which holds the current version and memtable
slice. A `DB.readState` pointer is updated whenever the current version
or memtable slice changes.

```
name                    time/op
ReadState/updates=0-8   49.7ns ± 1%
ReadState/updates=10-8  89.9ns ± 2%
ReadState/updates=50-8   115ns ± 1%
```

Note that the `updates={10,50}` benchmarks are unrealistic. It is
extremely unlikely that 10% or 50% of the operations in a real workload
will be calls to `updateReadStateLocked`.

Fixes #45

Patch:
diff --git a/compaction.go b/compaction.go
index 3b137f8..f22b71c 100644
--- a/compaction.go
+++ b/compaction.go
@@ -480,6 +480,7 @@ func (d *DB) flush1() error {
 			close(d.mu.mem.queue[i].flushed())
 		}
 		d.mu.mem.queue = d.mu.mem.queue[n:]
+		d.updateReadStateLocked()
 		return nil
 	}
 
@@ -506,6 +507,7 @@ func (d *DB) flush1() error {
 		close(d.mu.mem.queue[i].flushed())
 	}
 	d.mu.mem.queue = d.mu.mem.queue[n:]
+	d.updateReadStateLocked()
 
 	// var newDirty int
 	// for _, mem := range d.mu.mem.queue {
@@ -747,6 +749,7 @@ func (d *DB) compact1() (err error) {
 	if err != nil {
 		return err
 	}
+	d.updateReadStateLocked()
 	d.deleteObsoleteFiles(jobID)
 	return nil
 }
diff --git a/db.go b/db.go
index 616f070..db1bd0b 100644
--- a/db.go
+++ b/db.go
@@ -137,6 +137,13 @@ type DB struct {
 	largeBatchThreshold int
 	optionsFileNum      uint64
 
+	// readState provides access to the state needed for reading without needing
+	// to acquire DB.mu.
+	readState struct {
+		sync.RWMutex
+		val *readState
+	}
+
 	// TODO(peter): describe exactly what this mutex protects. So far: every
 	// field in the struct.
 	mu struct {
@@ -192,14 +199,10 @@ func (d *DB) Get(key []byte) ([]byte, error) {
 }
 
 func (d *DB) getInternal(key []byte, b *Batch, s *Snapshot) ([]byte, error) {
-	d.mu.Lock()
-	// 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.
-	current := d.mu.versions.currentVersion()
-	current.ref()
-	memtables := d.mu.mem.queue
-	d.mu.Unlock()
+	// Grab and reference the current readState. This prevents the underlying
+	// files in the associated version from being deleted if there is a current
+	// compaction. The readState is unref'd by Iterator.Close().
+	readState := d.loadReadState()
 
 	// Determine the seqnum to read at after grabbing the read state (current and
 	// memtables) above.
@@ -222,16 +225,16 @@ func (d *DB) getInternal(key []byte, b *Batch, s *Snapshot) ([]byte, error) {
 	get.snapshot = seqNum
 	get.key = key
 	get.batch = b
-	get.mem = memtables
-	get.l0 = current.files[0]
-	get.version = current
+	get.mem = readState.memtables
+	get.l0 = readState.current.files[0]
+	get.version = readState.current
 
 	i := &buf.dbi
 	i.cmp = d.cmp
 	i.equal = d.equal
 	i.merge = d.merge
 	i.iter = get
-	i.version = current
+	i.readState = readState
 
 	defer i.Close()
 	if !i.Next() {
@@ -379,14 +382,10 @@ func (d *DB) newIterInternal(
 	s *Snapshot,
 	o *db.IterOptions,
 ) *Iterator {
-	d.mu.Lock()
-	// 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.
-	current := d.mu.versions.currentVersion()
-	current.ref()
-	memtables := d.mu.mem.queue
-	d.mu.Unlock()
+	// Grab and reference the current readState. This prevents the underlying
+	// files in the associated version from being deleted if there is a current
+	// compaction. The readState is unref'd by Iterator.Close().
+	readState := d.loadReadState()
 
 	// Determine the seqnum to read at after grabbing the read state (current and
 	// memtables) above.
@@ -413,7 +412,7 @@ func (d *DB) newIterInternal(
 	dbi.cmp = d.cmp
 	dbi.equal = d.equal
 	dbi.merge = d.merge
-	dbi.version = current
+	dbi.readState = readState
 
 	iters := buf.iters[:0]
 	rangeDelIters := buf.rangeDelIters[:0]
@@ -427,6 +426,7 @@ func (d *DB) newIterInternal(
 	// TODO(peter): We only need to add memtables which contain sequence numbers
 	// older than seqNum. Unfortunately, memtables don't track their oldest
 	// sequence number currently.
+	memtables := readState.memtables
 	for i := len(memtables) - 1; i >= 0; i-- {
 		mem := memtables[i]
 		iters = append(iters, mem.newIter(o))
@@ -435,6 +435,7 @@ func (d *DB) newIterInternal(
 	}
 
 	// The level 0 files need to be added from newest to oldest.
+	current := readState.current
 	for i := len(current.files[0]) - 1; i >= 0; i-- {
 		f := &current.files[0][i]
 		iter, rangeDelIter, err := d.newIters(f)
@@ -557,6 +558,8 @@ func (d *DB) Close() error {
 	d.mu.closed = true
 
 	if err == nil {
+		d.readState.val.unrefLocked()
+
 		current := d.mu.versions.currentVersion()
 		for v := d.mu.versions.versions.front(); true; v = v.next {
 			refs := atomic.LoadInt32(&v.refs)
@@ -788,6 +791,7 @@ func (d *DB) makeRoomForWrite(b *Batch) error {
 			d.mu.mem.mutable = newMemTable(d.opts)
 		}
 		d.mu.mem.queue = append(d.mu.mem.queue, d.mu.mem.mutable)
+		d.updateReadStateLocked()
 		if (imm != nil && imm.unref()) || scheduleFlush {
 			d.maybeScheduleFlush()
 		}
diff --git a/ingest.go b/ingest.go
index 009ffea..54796bf 100644
--- a/ingest.go
+++ b/ingest.go
@@ -357,5 +357,6 @@ func (d *DB) ingestApply(meta []*fileMetadata) (*versionEdit, error) {
 	if err := d.mu.versions.logAndApply(ve); err != nil {
 		return nil, err
 	}
+	d.updateReadStateLocked()
 	return ve, nil
 }
diff --git a/iterator.go b/iterator.go
index 4dbfdb6..391c3a0 100644
--- a/iterator.go
+++ b/iterator.go
@@ -36,7 +36,7 @@ type Iterator struct {
 	equal     db.Equal
 	merge     db.Merge
 	iter      internalIterator
-	version   *version
+	readState *readState
 	err       error
 	key       []byte
 	keyBuf    []byte
@@ -373,9 +373,9 @@ func (i *Iterator) Error() error {
 // It is valid to call Close multiple times. Other methods should not be
 // called after the iterator has been closed.
 func (i *Iterator) Close() error {
-	if i.version != nil {
-		i.version.unref()
-		i.version = nil
+	if i.readState != nil {
+		i.readState.unref()
+		i.readState = nil
 	}
 	if err := i.iter.Close(); err != nil && i.err != nil {
 		i.err = err
diff --git a/open.go b/open.go
index 57b6f13..dae415c 100644
--- a/open.go
+++ b/open.go
@@ -174,6 +174,7 @@ func Open(dirname string, opts *db.Options) (*DB, error) {
 	if err := d.mu.versions.logAndApply(&ve); err != nil {
 		return nil, err
 	}
+	d.updateReadStateLocked()
 
 	// Write the current options to disk.
 	d.optionsFileNum = d.mu.versions.nextFileNum()
diff --git a/read_state.go b/read_state.go
new file mode 100644
index 0000000..e507a3b
--- /dev/null
+++ b/read_state.go
@@ -0,0 +1,82 @@
+// Copyright 2019 The LevelDB-Go and Pebble Authors. All rights reserved. Use
+// of this source code is governed by a BSD-style license that can be found in
+// the LICENSE file.
+
+package pebble
+
+import (
+	"sync/atomic"
+)
+
+// readState encapsulates the state needed for reading (the current version and
+// list of memtables). Loading the readState is done without grabbing
+// DB.mu. Instead, a separate DB.readState.RWMutex is used for
+// synchronization. This mutex solely covers the current readState object which
+// means it is rarely or ever contended.
+//
+// Note that various fancy lock-free mechanisms can be imagined for loading the
+// readState, but benchmarking showed the ones considered to purely be
+// pessimizations. The RWMutex version is a single atomic increment for the
+// RLock and an atomic decrement for the RUnlock. It is difficult to do better
+// than that without something like thread-local storage which isn't available
+// in Go.
+type readState struct {
+	refcnt    int32
+	current   *version
+	memtables []flushable
+}
+
+// ref adds a reference to the readState.
+func (s *readState) ref() {
+	atomic.AddInt32(&s.refcnt, 1)
+}
+
+// unref removes a reference to the readState. If this was the last reference,
+// the reference the readState holds on the version is released. Requires DB.mu
+// is NOT held as version.unref() will acquire it. See unrefLocked() if DB.mu
+// is held by the caller.
+func (s *readState) unref() {
+	if atomic.AddInt32(&s.refcnt, -1) == 0 {
+		s.current.unref()
+	}
+}
+
+// unrefLocked removes a reference to the readState. If this was the last
+// reference, the reference the readState holds on the version is
+// released. Requires DB.mu is held as version.unrefLocked() requires it. See
+// unref() if DB.mu is NOT held by the caller.
+func (s *readState) unrefLocked() {
+	if atomic.AddInt32(&s.refcnt, -1) == 0 {
+		s.current.unrefLocked()
+	}
+}
+
+// loadReadState returns the current readState. The returned readState must be
+// unreferenced when the caller is finished with it.
+func (d *DB) loadReadState() *readState {
+	d.readState.RLock()
+	state := d.readState.val
+	state.ref()
+	d.readState.RUnlock()
+	return state
+}
+
+// updateReadStateLocked creates a new readState from the current version and
+// list of memtables. Requires DB.mu is held.
+func (d *DB) updateReadStateLocked() {
+	s := &readState{
+		refcnt:    1,
+		current:   d.mu.versions.currentVersion(),
+		memtables: d.mu.mem.queue,
+	}
+	s.current.ref()
+
+	d.readState.Lock()
+	old := d.readState.val
+	d.readState.val = s
+	d.readState.Unlock()
+
+	if old != nil {
+		old.unrefLocked()
+	}
+}


Test patch:
diff --git a/data_test.go b/data_test.go
index dac8627..da48407 100644
--- a/data_test.go
+++ b/data_test.go
@@ -273,6 +273,7 @@ func runDBDefineCmd(td *datadriven.TestData) (*DB, error) {
 				if !d.mu.mem.mutable.empty() {
 					d.mu.mem.mutable = newMemTable(d.opts)
 					d.mu.mem.queue = append(d.mu.mem.queue, d.mu.mem.mutable)
+					d.updateReadStateLocked()
 				}
 				mem = d.mu.mem.mutable
 				fields = fields[1:]
@@ -307,6 +308,7 @@ func runDBDefineCmd(td *datadriven.TestData) (*DB, error) {
 		if err := d.mu.versions.logAndApply(ve); err != nil {
 			return nil, err
 		}
+		d.updateReadStateLocked()
 		for i := range ve.newFiles {
 			meta := &ve.newFiles[i].meta
 			delete(d.mu.compact.pendingOutputs, meta.fileNum)
diff --git a/read_state_test.go b/read_state_test.go
new file mode 100644
index 0000000..d694093
--- /dev/null
+++ b/read_state_test.go
@@ -0,0 +1,47 @@
+// Copyright 2019 The LevelDB-Go and Pebble Authors. All rights reserved. Use
+// of this source code is governed by a BSD-style license that can be found in
+// the LICENSE file.
+
+package pebble
+
+import (
+	"fmt"
+	"math/rand"
+	"testing"
+	"time"
+
+	"github.com/petermattis/pebble/db"
+	"github.com/petermattis/pebble/storage"
+)
+
+func BenchmarkReadState(b *testing.B) {
+	d, err := Open("", &db.Options{
+		Storage: storage.NewMem(),
+	})
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	for _, updateFrac := range []float32{0, 0.1, 0.5} {
+		b.Run(fmt.Sprintf("updates=%.0f", updateFrac*100), func(b *testing.B) {
+			b.RunParallel(func(pb *testing.PB) {
+				rng := rand.New(rand.NewSource(time.Now().UnixNano()))
+
+				for pb.Next() {
+					if rng.Float32() < updateFrac {
+						d.mu.Lock()
+						d.updateReadStateLocked()
+						d.mu.Unlock()
+					} else {
+						s := d.loadReadState()
+						s.unref()
+					}
+				}
+			})
+		})
+	}
+
+	if err := d.Close(); err != nil {
+		b.Fatal(err)
+	}
+}