Repository: petermattis/pebble
Pull request #57 (2019-03-07 14:13:42)

Lock-free updates of the cached tombstone fragments
Replace the use of a `RWMutex` with atomic operations to update the
cached slice of tombstone fragments.

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

petermattis: Noticed this possibility in passing while looking at the RocksDB PR to support caching of fragmented range tombstones. 
ajkr: What do you think about adding an pending invalidation atomic counter that works like this:

- The writer increments the value when an invalidation happens
- If a reader sees zero, then it can load from `m.tombstones.vals` and use those tombstones
- If a reader sees nonzero, it does fragmentation work creating its own `newPtr`. It then installs `newPtr` in the same way, CAS'ing it in only if it's longer than what's there currently. Then we decrement the atomic counter although not letting it go below zero.
ajkr: No, wait, I don't think that'll work because the pending invalidation counter may be decremented to zero by re-fragmentations that all began before the most recent invalidation happened.
petermattis: > What if there are two parallel writers, and a single reader that fragments tombstones before the second writer invalidates but updates the pointer after the second writer invalidates? Will think a bit more.

Yeah, I think something isn't quite right here.
ajkr: I think we could use load-link and store-conditional if it were available in Go. Something like:

- Writers store 0x0, same as now.
- A reader uses LL to read the pointer.
  - If it's non-0x0, just return it
  - Otherwise, fragment, then use SC to store the result. Whether it fails or succeeds, return the result of our fragmentation work. (If it fails that means either some other invalidation happened or another reader installed its fragmentation result. In the former case we can rely on a future reader to fragment again. In the latter case we know the pointer holds the fragmentation result for an invalidation at least as new as the one we're working on, so we don't need to do conflict checking.)
petermattis: @ajkr I'm not seeing how your proposal differs from the code in this PR. When the reader stores the fragmentation result, we need to be careful not to clear an invalidation that occurred after the reader produced the fragmentation.

I think we need some sort of invalidation sequence number. I think another level of indirection could accomplish this. `memTable.tombstones.vals` could be a `*tombstoneFrags`:

```
type tombstoneFrags struct {
  seqNum uint32 // count of range tombstones in the memTable
  vals unsafe.Pointer // *[]rangedel.Tombstones
}
```

Whenever a range tombstone is added, we replace atomically update `memtable.tombstones.vals` with a new `tombstoneFrags` if the previous frags struct contained a lower sequence number. In `memTable.newRangeDelIter` we retrieve the cached tombstones, populating `tombstoneFrags.vals` if it is `nil`.

In your scenario of 2 writers and 1 reader, the invalidation by the 2nd writer won't get lost by a reader constructing the fragments. There is an extra per-invalidation allocation in this approach, but that doesn't seem horrible. Another thought is to use `sync.Once` instead of atomically updating `tombstoneFrags.vals`.
ajkr: Yeah I think that should work. There should be no ABA problem between checking `seqNum` and CAS'ing `memtable.tombstones.vals`, right? For example, if during that gap another invalidation deleted the existing `tombstoneFrags` and created a new `tombstoneFrags` at the same memory address, then our CAS could mistakenly succeed. But I think we are holding a reference to the old `tombstoneFrags` throughout so its memory address shouldn't be reused.

Patch:
diff --git a/internal/arenaskl/arena.go b/internal/arenaskl/arena.go
index 39bfdf4..304b582 100644
--- a/internal/arenaskl/arena.go
+++ b/internal/arenaskl/arena.go
@@ -19,16 +19,15 @@ package arenaskl
 
 import (
 	"errors"
-	"math"
 	"sync/atomic"
 	"unsafe"
 
 	"github.com/petermattis/pebble/internal/rawalloc"
 )
 
-// Arena is lock-free.
+// Arena should be lock-free.
 type Arena struct {
-	n   uint64
+	n   uint32
 	buf []byte
 }
 
@@ -51,12 +50,7 @@ func NewArena(size, extValueThreshold uint32) *Arena {
 }
 
 func (a *Arena) Size() uint32 {
-	s := atomic.LoadUint64(&a.n)
-	if s > math.MaxUint32 {
-		// Saturate at MaxUint32.
-		return math.MaxUint32
-	}
-	return uint32(s)
+	return atomic.LoadUint32(&a.n)
 }
 
 func (a *Arena) Capacity() uint32 {
@@ -64,22 +58,16 @@ func (a *Arena) Capacity() uint32 {
 }
 
 func (a *Arena) alloc(size, align uint32) (uint32, error) {
-	// Verify that the arena isn't already full.
-	origSize := atomic.LoadUint64(&a.n)
-	if int(origSize) > len(a.buf) {
-		return 0, ErrArenaFull
-	}
-
 	// Pad the allocation with enough bytes to ensure the requested alignment.
 	padded := uint32(size) + align
 
-	newSize := atomic.AddUint64(&a.n, uint64(padded))
+	newSize := atomic.AddUint32(&a.n, padded)
 	if int(newSize) > len(a.buf) {
 		return 0, ErrArenaFull
 	}
 
 	// Return the aligned offset.
-	offset := (uint32(newSize) - padded + uint32(align)) & ^uint32(align)
+	offset := (newSize - padded + uint32(align)) & ^uint32(align)
 	return offset, nil
 }
 
diff --git a/mem_table.go b/mem_table.go
index ed283d9..616644a 100644
--- a/mem_table.go
+++ b/mem_table.go
@@ -7,6 +7,7 @@ package pebble
 import (
 	"sync"
 	"sync/atomic"
+	"unsafe"
 
 	"github.com/petermattis/pebble/db"
 	"github.com/petermattis/pebble/internal/arenaskl"
@@ -34,12 +35,7 @@ type memTable struct {
 	reserved    uint32
 	refs        int32
 	flushedCh   chan struct{}
-
-	tombstones struct {
-		count uint32
-		sync.RWMutex
-		vals []rangedel.Tombstone
-	}
+	tombstones  rangeTombstoneCache
 }
 
 // newMemTable returns a new MemTable.
@@ -123,8 +119,8 @@ func (m *memTable) prepare(batch *Batch) error {
 
 func (m *memTable) apply(batch *Batch, seqNum uint64) error {
 	var ins arenaskl.Inserter
+	var tombstoneCount uint32
 	startSeqNum := seqNum
-	invalidateTombstones := false
 	for iter := batch.iter(); ; seqNum++ {
 		kind, ukey, value, ok := iter.next()
 		if !ok {
@@ -134,8 +130,7 @@ func (m *memTable) apply(batch *Batch, seqNum uint64) error {
 		ikey := db.MakeInternalKey(ukey, seqNum, kind)
 		if kind == db.InternalKeyKindRangeDelete {
 			err = m.rangeDelSkl.Add(ikey, value)
-			atomic.AddUint32(&m.tombstones.count, 1)
-			invalidateTombstones = true
+			tombstoneCount++
 		} else {
 			err = ins.Add(&m.skl, ikey, value)
 		}
@@ -146,10 +141,8 @@ func (m *memTable) apply(batch *Batch, seqNum uint64) error {
 	if seqNum != startSeqNum+uint64(batch.count()) {
 		panic("pebble: inconsistent batch count")
 	}
-	if invalidateTombstones {
-		m.tombstones.Lock()
-		m.tombstones.vals = nil
-		m.tombstones.Unlock()
+	if tombstoneCount != 0 {
+		m.tombstones.invalidate(tombstoneCount)
 	}
 	return nil
 }
@@ -163,47 +156,96 @@ func (m *memTable) newIter(*db.IterOptions) internalIterator {
 }
 
 func (m *memTable) newRangeDelIter(*db.IterOptions) internalIterator {
-	if atomic.LoadUint32(&m.tombstones.count) == 0 {
+	tombstones := m.tombstones.get(m)
+	if tombstones == nil {
 		return nil
 	}
+	return rangedel.NewIter(m.cmp, tombstones)
+}
 
-	// We need to return an iterator over fragmented tombstones. The fragmented
-	// tombstones are cached and the cache is invalided whenever a new range
-	// tombstone is added to the memtable.
-	m.tombstones.RLock()
-	tombstones := m.tombstones.vals
-	m.tombstones.RUnlock()
+func (m *memTable) close() error {
+	return nil
+}
 
-	if tombstones == nil {
+// empty returns whether the MemTable has no key/value pairs.
+func (m *memTable) empty() bool {
+	return m.skl.Size() == m.emptySize
+}
+
+// A rangeTombstoneFrags holds a set of fragmented range tombstones generated
+// at a particular "sequence number" for a memtable. Rather than use actual
+// sequence numbers, this cache uses a count of the number of range tombstones
+// in the memTable. Note that the count of range tombstones in a memTable only
+// ever increases, which provides a monotonically increasing sequence.
+type rangeTombstoneFrags struct {
+	count      uint32
+	once       sync.Once
+	tombstones []rangedel.Tombstone
+}
+
+// get retrieves the fragmented tombstones, populating them if necessary. Note
+// that the populated tombstone fragments may be built from more than f.count
+// memTable range tombstones, but that is ok for correctness. All we're
+// requiring is that the memTable contains at least f.count range
+// tombstones. This situation can occur if there are multiple concurrent
+// additions of range tombstones and a concurrent reader. The reader can load a
+// tombstoneFrags and populate it even though is has been invalidated
+// (i.e. replaced with a newer tombstoneFrags).
+func (f *rangeTombstoneFrags) get(m *memTable) []rangedel.Tombstone {
+	f.once.Do(func() {
 		frag := &rangedel.Fragmenter{
 			Cmp: m.cmp,
 			Emit: func(fragmented []rangedel.Tombstone) {
-				tombstones = append(tombstones, fragmented...)
+				f.tombstones = append(f.tombstones, fragmented...)
 			},
 		}
 		for it := m.rangeDelSkl.NewIter(); it.Next(); {
 			frag.Add(it.Key(), it.Value())
 		}
 		frag.Finish()
+	})
+	return f.tombstones
+}
 
-		m.tombstones.Lock()
-		// Adding tombstones to the memtable can only increase the number of
-		// fragmented tombstones. If there was already a cached value, only
-		// overwrite it if we generated more tombstones.
-		if len(m.tombstones.vals) < len(tombstones) {
-			m.tombstones.vals = tombstones
-		}
-		m.tombstones.Unlock()
-	}
-
-	return rangedel.NewIter(m.cmp, tombstones)
+// A rangeTombstoneCache is used to cache a set of fragmented tombstones. The
+// cache is invalidated whenever a tombstone is added to a memTable, and
+// populated when empty when a range-del iterator is created.
+type rangeTombstoneCache struct {
+	count uint32
+	frags unsafe.Pointer
 }
 
-func (m *memTable) close() error {
-	return nil
+// Invalidate the current set of cached tombstones, indicating the number of
+// tombstones that were added.
+func (c *rangeTombstoneCache) invalidate(count uint32) {
+	newCount := atomic.AddUint32(&c.count, count)
+	var frags *rangeTombstoneFrags
+
+	for {
+		oldPtr := atomic.LoadPointer(&c.frags)
+		if oldPtr != nil {
+			oldFrags := (*rangeTombstoneFrags)(oldPtr)
+			if oldFrags.count >= newCount {
+				// Someone else invalidated the cache before us and their invalidation
+				// subsumes ours.
+				break
+			}
+		}
+		if frags == nil {
+			frags = &rangeTombstoneFrags{count: newCount}
+		}
+		if atomic.CompareAndSwapPointer(&c.frags, oldPtr, unsafe.Pointer(frags)) {
+			// We successfully invalidated the cache.
+			break
+		}
+		// Someone else invalidated the cache. Loop and try again.
+	}
 }
 
-// empty returns whether the MemTable has no key/value pairs.
-func (m *memTable) empty() bool {
-	return m.skl.Size() == m.emptySize
+func (c *rangeTombstoneCache) get(m *memTable) []rangedel.Tombstone {
+	frags := (*rangeTombstoneFrags)(atomic.LoadPointer(&c.frags))
+	if frags == nil {
+		return nil
+	}
+	return frags.get(m)
 }


Test patch:
diff --git a/internal/arenaskl/arena_test.go b/internal/arenaskl/arena_test.go
deleted file mode 100644
index 39ff3e1..0000000
--- a/internal/arenaskl/arena_test.go
+++ /dev/null
@@ -1,48 +0,0 @@
-/*
- * Copyright 2017 Dgraph Labs, Inc. and Contributors
- * Modifications copyright (C) 2017 Andy Kimball and Contributors
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package arenaskl
-
-import (
-	"math"
-	"testing"
-
-	"github.com/stretchr/testify/require"
-)
-
-// TestArenaSizeOverflow tests that large allocations do not cause Arena's
-// internal size accounting to overflow and produce incorrect results.
-func TestArenaSizeOverflow(t *testing.T) {
-	a := NewArena(math.MaxUint32, 0)
-
-	// Allocating under the limit throws no error.
-	offset, err := a.alloc(math.MaxUint16, 0)
-	require.Nil(t, err)
-	require.Equal(t, uint32(1), offset)
-	require.Equal(t, uint32(math.MaxUint16)+1, a.Size())
-
-	// Allocating over the limit could cause an accounting
-	// overflow if 32-bit arithmetic was used. It shouldn't.
-	_, err = a.alloc(math.MaxUint32, 0)
-	require.Equal(t, ErrArenaFull, err)
-	require.Equal(t, uint32(math.MaxUint32), a.Size())
-
-	// Continuing to allocate continues to throw an error.
-	_, err = a.alloc(math.MaxUint16, 0)
-	require.Equal(t, ErrArenaFull, err)
-	require.Equal(t, uint32(math.MaxUint32), a.Size())
-}
diff --git a/mem_table_test.go b/mem_table_test.go
index a6ee58f..6cf31c4 100644
--- a/mem_table_test.go
+++ b/mem_table_test.go
@@ -10,6 +10,7 @@ import (
 	"math/rand"
 	"strconv"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -27,10 +28,7 @@ func (m *memTable) set(key db.InternalKey, value []byte) error {
 		if err := m.rangeDelSkl.Add(key, value); err != nil {
 			return err
 		}
-		atomic.AddUint32(&m.tombstones.count, 1)
-		m.tombstones.Lock()
-		m.tombstones.vals = nil
-		m.tombstones.Unlock()
+		m.tombstones.invalidate(1)
 		return nil
 	}
 	return m.skl.Add(key, value)
@@ -283,6 +281,48 @@ func TestMemTableDeleteRange(t *testing.T) {
 	})
 }
 
+func TestMemTableConcurrentDeleteRange(t *testing.T) {
+	// Concurrently write and read range tombstones. Workers add range
+	// tombstones, and then immediately retrieve them verifying that the
+	// tombstones they've added are all present.
+
+	m := newMemTable(&db.Options{MemTableSize: 64 << 20})
+
+	const workers = 10
+	var wg sync.WaitGroup
+	wg.Add(workers)
+	seqNum := uint64(1)
+	for i := 0; i < workers; i++ {
+		go func(i int) {
+			defer wg.Done()
+			start := ([]byte)(fmt.Sprintf("%03d", i))
+			end := ([]byte)(fmt.Sprintf("%03d", i+1))
+			for j := 0; j < 100; j++ {
+				b := newBatch(nil)
+				b.DeleteRange(start, end, nil)
+				n := atomic.AddUint64(&seqNum, 1) - 1
+				if err := m.apply(b, n); err != nil {
+					t.Fatal(err)
+				}
+				b.release()
+
+				var count int
+				it := m.newRangeDelIter(nil)
+				for valid := it.SeekGE(start); valid; valid = it.Next() {
+					if m.cmp(it.Key().UserKey, end) >= 0 {
+						break
+					}
+					count++
+				}
+				if j+1 != count {
+					t.Fatalf("%d: expected %d tombstones, but found %d", i, j+1, count)
+				}
+			}
+		}(i)
+	}
+	wg.Wait()
+}
+
 func buildMemTable(b *testing.B) (*memTable, [][]byte) {
 	m := newMemTable(nil)
 	var keys [][]byte