Repository: petermattis/pebble
Pull request #231 (2019-08-19 20:55:18)

batch: Don't increment count on LogData(), add misc helpers
RocksDB batches contain a count field in their byte representation
which is usually NOT incremented for any operation of kind LogData
that exists in the WAL but isn't applied to the memtable. As
Cockroach moves to using pebble.Batch instead of its own batch builder,
reconcile this difference so RocksDB doesn't error out on pebble.Batch
produced batches.

Also adds two exported helper functions - MaybeInit and NewBatchFromRepr,
which will be used in Cockroach.

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


Patch:
diff --git a/batch.go b/batch.go
index 9073369..addd97e 100644
--- a/batch.go
+++ b/batch.go
@@ -78,6 +78,36 @@ func (s *batchStorage) Compare(a []byte, b uint32) int {
 	return 1
 }
 
+// DeferredBatchOp represents a batch operation (eg. set, merge, delete) that is
+// being inserted into the batch. Indexing is not performed on the specified key
+// until Finish is called, hence the name deferred. This struct lets the caller
+// copy or encode keys/values directly into the batch representation instead of
+// copying into an intermediary buffer then having pebble.Batch copy off of it.
+type DeferredBatchOp struct {
+	batch *Batch
+
+	// Key and Value point to parts of the binary batch representation where
+	// keys and values should be encoded/copied into. len(Key) and len(Value)
+	// bytes must be copied into these slices respectively before calling
+	// Finish(). Changing where these slices point to is not allowed.
+	Key, Value []byte
+	offset uint32
+}
+
+// Finish completes the addition of this batch operation, and adds it to the
+// index if necessary. Must be called once (and exactly once) keys/values
+// have been filled into Key and Value. Not calling Finish or not
+// copying/encoding keys will result in an incomplete index, and calling Finish
+// twice may result in a panic.
+func (d DeferredBatchOp) Finish() {
+	if d.batch.index != nil {
+		if err := d.batch.index.Add(d.offset); err != nil {
+			// We never add duplicate entries, so an error should never occur.
+			panic(err)
+		}
+	}
+}
+
 // A Batch is a sequence of Sets, Merges, Deletes, and/or DeleteRanges that are
 // applied atomically. Batch implements the Reader interface, but only an
 // indexed batch supports reading (without error) via Get or NewIter. A
@@ -238,15 +268,22 @@ func (b *Batch) release() {
 	// NB: This is ugly, but necessary so that we can use atomic.StoreUint32 for
 	// the Batch.applied field. Without using an atomic to clear that field the
 	// Go race detector complains.
-	b.reset()
+	b.Reset()
 	b.storage.cmp = nil
 	b.storage.abbreviatedKey = nil
 	b.memTableSize = 0
-	b.db = nil
+
 	b.flushable = nil
 	b.commit = sync.WaitGroup{}
 	atomic.StoreUint32(&b.applied, 0)
 
+	if b.db == nil {
+		// Batch not created using newBatch or newIndexedBatch, so don't put it
+		// back in the pool.
+		return
+	}
+	b.db = nil
+
 	if b.index == nil {
 		batchPool.Put(b)
 	} else {
@@ -337,6 +374,29 @@ func (b *Batch) encodeKeyValue(key, value []byte, kind InternalKeyKind) uint32 {
 	return offset
 }
 
+func (b *Batch) encodeKeyValuePlaceholders(
+	keyLen, valueLen int, kind InternalKeyKind,
+) (keyDest, valueDest []byte, offset uint32) {
+	pos := len(b.storage.data)
+	offset = uint32(pos)
+	b.grow(1 + 2*maxVarintLen32 + keyLen + valueLen)
+	b.storage.data[pos] = byte(kind)
+	pos++
+
+	varlen1 := putUvarint32(b.storage.data[pos:], uint32(keyLen))
+	pos += varlen1
+	keyDest = b.storage.data[pos:pos+keyLen]
+	pos += keyLen
+
+	varlen2 := putUvarint32(b.storage.data[pos:], uint32(valueLen))
+	pos += varlen2
+	valueDest = b.storage.data[pos:pos+valueLen]
+	pos += valueLen
+	b.storage.data = b.storage.data[:len(b.storage.data)-(2*maxVarintLen32-varlen1-varlen2)]
+
+	return
+}
+
 // Set adds an action to the batch that sets the key to map to the value.
 //
 // It is safe to modify the contents of the arguments after Set returns.
@@ -360,6 +420,33 @@ func (b *Batch) Set(key, value []byte, _ *WriteOptions) error {
 	return nil
 }
 
+// SetDeferred is similar to Set in that it adds a set operation to the batch,
+// except it only takes in key/value lengths insted of complete slices, letting
+// the caller encode into those objects and then call Finish() on the returned
+// object.
+func (b *Batch) SetDeferred(keyLen, valueLen int, _ *WriteOptions) (DeferredBatchOp, error) {
+	// Code duplication between Set and SetDeferred lets us preserve the fast
+	// path where the entire byte slices are available (in the Set case).
+	if len(b.storage.data) == 0 {
+		b.init(keyLen + valueLen + 2*binary.MaxVarintLen64 + batchHeaderLen)
+	}
+	if !b.increment() {
+		return DeferredBatchOp{}, ErrInvalidBatch
+	}
+
+	keyDest, valueDest, offset := b.encodeKeyValuePlaceholders(keyLen, valueLen, InternalKeyKindSet)
+
+	deferredOp := DeferredBatchOp{
+		batch:  b,
+		Key:    keyDest,
+		Value:  valueDest,
+		offset: offset,
+	}
+
+	b.memTableSize += memTableEntrySize(keyLen, valueLen)
+	return deferredOp, nil
+}
+
 // Merge adds an action to the batch that merges the value at key with the new
 // value. The details of the merge are dependent upon the configured merge
 // operator.
@@ -385,6 +472,33 @@ func (b *Batch) Merge(key, value []byte, _ *WriteOptions) error {
 	return nil
 }
 
+// MergeDeferred is similar to Merge in that it adds a merge operation to the
+// batch, except it only takes in key/value lengths insted of complete slices,
+// letting the caller encode into those objects and then call Finish() on the
+// returned object.
+func (b *Batch) MergeDeferred(keyLen, valueLen int, _ *WriteOptions) (DeferredBatchOp, error) {
+	// Code duplication with Merge is so that the Merge case (where byte slices
+	// are provided) can preserve the fast path.
+	if len(b.storage.data) == 0 {
+		b.init(keyLen + valueLen + 2*binary.MaxVarintLen64 + batchHeaderLen)
+	}
+	if !b.increment() {
+		return DeferredBatchOp{}, ErrInvalidBatch
+	}
+
+	keyDest, valueDest, offset := b.encodeKeyValuePlaceholders(keyLen, valueLen, InternalKeyKindMerge)
+
+	deferredOp := DeferredBatchOp{
+		batch:  b,
+		Key:    keyDest,
+		Value:  valueDest,
+		offset: offset,
+	}
+
+	b.memTableSize += memTableEntrySize(keyLen, valueLen)
+	return deferredOp, nil
+}
+
 // Delete adds an action to the batch that deletes the entry for key.
 //
 // It is safe to modify the contents of the arguments after Delete returns.
@@ -413,6 +527,39 @@ func (b *Batch) Delete(key []byte, _ *WriteOptions) error {
 	return nil
 }
 
+// DeleteDeferred is similar to Delete in that it adds a delete operation to the
+// batch, except it only takes in key/value lengths insted of complete slices,
+// letting the caller encode into those objects and then call Finish() on the
+// returned object.
+func (b *Batch) DeleteDeferred(keyLen int, _ *WriteOptions) (DeferredBatchOp, error) {
+	// Code duplication with Delete is so that the Delete case (where byte
+	// slices are provided) can preserve the fast path.
+	if len(b.storage.data) == 0 {
+		b.init(keyLen + binary.MaxVarintLen64 + batchHeaderLen)
+	}
+	if !b.increment() {
+		return DeferredBatchOp{}, ErrInvalidBatch
+	}
+
+	deferredOp := DeferredBatchOp{
+		batch:     b,
+	}
+
+	pos := len(b.storage.data)
+	deferredOp.offset = uint32(pos)
+	b.grow(1 + maxVarintLen32 +  keyLen)
+	b.storage.data[pos] = byte(InternalKeyKindDelete)
+	pos++
+	varlen1 := putUvarint32(b.storage.data[pos:], uint32(keyLen))
+	pos += varlen1
+	deferredOp.Key = b.storage.data[pos:pos+keyLen]
+
+	b.storage.data = b.storage.data[:len(b.storage.data)-(maxVarintLen32-varlen1)]
+
+	b.memTableSize += memTableEntrySize(keyLen,0)
+	return deferredOp, nil
+}
+
 // DeleteRange deletes all of the keys (and values) in the range [start,end)
 // (inclusive on start, exclusive on end).
 //
@@ -452,9 +599,9 @@ func (b *Batch) LogData(data []byte, _ *WriteOptions) error {
 	if len(b.storage.data) == 0 {
 		b.init(len(data) + binary.MaxVarintLen64 + batchHeaderLen)
 	}
-	if !b.increment() {
-		return ErrInvalidBatch
-	}
+	// Since LogData only writes to the WAL and does not affect the memtable,
+	// we don't increment b.count here. b.count only tracks operations that
+	// are applied to the memtable.
 
 	pos := len(b.storage.data)
 	b.grow(1 + maxVarintLen32 + len(data))
@@ -467,6 +614,9 @@ func (b *Batch) LogData(data []byte, _ *WriteOptions) error {
 // Repr returns the underlying batch representation. It is not safe to modify
 // the contents.
 func (b *Batch) Repr() []byte {
+	if len(b.storage.data) == 0 {
+		b.init(batchHeaderLen)
+	}
 	return b.storage.data
 }
 
@@ -569,7 +719,11 @@ func (b *Batch) init(cap int) {
 	b.storage.data = b.storage.data[:batchHeaderLen]
 }
 
-func (b *Batch) reset() {
+// Reset clears the underlying byte slice and effectively empties the batch for
+// reuse. Used in cases where Batch is only being used to build a batch, and
+// where the end result is a Repr() call, not a Commit call or a Close call.
+// Commits and Closes take care of releasing resources when appropriate.
+func (b *Batch) Reset() {
 	if b.storage.data != nil {
 		if cap(b.storage.data) > batchMaxRetainedSize {
 			// If the capacity of the buffer is larger than our maximum
@@ -659,8 +813,12 @@ func (b *Batch) setCount(v uint32) {
 	binary.LittleEndian.PutUint32(b.countData(), v)
 }
 
-// Count returns the number of records in the batch.
+// Count returns the count of memtable-modifying operations in this batch. All
+// operations with the except of LogData increment this count.
 func (b *Batch) Count() uint32 {
+	if len(b.storage.data) == 0 {
+		return 0
+	}
 	return binary.LittleEndian.Uint32(b.countData())
 }
 


Test patch:
diff --git a/batch_test.go b/batch_test.go
index d3608c5..36c062d 100644
--- a/batch_test.go
+++ b/batch_test.go
@@ -18,10 +18,31 @@ import (
 )
 
 func TestBatch(t *testing.T) {
-	testCases := []struct {
+	type testCase struct {
 		kind       InternalKeyKind
 		key, value string
-	}{
+	}
+
+	verifyTestCases := func(b Batch, testCases []testCase) {
+		r := b.Reader()
+
+		for _, tc := range testCases {
+			kind, k, v, ok := r.Next()
+			if !ok {
+				t.Fatalf("next returned !ok: test case = %v", tc)
+			}
+			key, value := string(k), string(v)
+			if kind != tc.kind || key != tc.key || value != tc.value {
+				t.Errorf("got (%d, %q, %q), want (%d, %q, %q)",
+					kind, key, value, tc.kind, tc.key, tc.value)
+			}
+		}
+		if len(r) != 0 {
+			t.Errorf("reader was not exhausted: remaining bytes = %q", r)
+		}
+	}
+
+	testCases := []testCase{
 		{InternalKeyKindSet, "roses", "red"},
 		{InternalKeyKindSet, "violets", "blue"},
 		{InternalKeyKindDelete, "roses", ""},
@@ -57,21 +78,37 @@ func TestBatch(t *testing.T) {
 			_ = b.LogData([]byte(tc.key), nil)
 		}
 	}
-	r := b.Reader()
+	verifyTestCases(b, testCases)
+
+	b.Reset()
+	// Run the same operations, this time using the Deferred variants of each
+	// operation (eg. SetDeferred).
 	for _, tc := range testCases {
-		kind, k, v, ok := r.Next()
-		if !ok {
-			t.Fatalf("next returned !ok: test case = %v", tc)
-		}
-		key, value := string(k), string(v)
-		if kind != tc.kind || key != tc.key || value != tc.value {
-			t.Errorf("got (%d, %q, %q), want (%d, %q, %q)",
-				kind, key, value, tc.kind, tc.key, tc.value)
+		key := []byte(tc.key)
+		value := []byte(tc.value)
+		switch tc.kind {
+		case InternalKeyKindSet:
+			d, _ := b.SetDeferred(len(key), len(value), nil)
+			copy(d.Key, key)
+			copy(d.Value, value)
+			d.Finish()
+		case InternalKeyKindMerge:
+			d, _ := b.MergeDeferred(len(key), len(value), nil)
+			copy(d.Key, key)
+			copy(d.Value, value)
+			d.Finish()
+		case InternalKeyKindDelete:
+			d, _ := b.DeleteDeferred(len(key), nil)
+			copy(d.Key, key)
+			copy(d.Value, value)
+			d.Finish()
+		case InternalKeyKindRangeDelete:
+			_ = b.DeleteRange([]byte(tc.key), []byte(tc.value), nil)
+		case InternalKeyKindLogData:
+			_ = b.LogData([]byte(tc.key), nil)
 		}
 	}
-	if len(r) != 0 {
-		t.Errorf("reader was not exhausted: remaining bytes = %q", r)
-	}
+	verifyTestCases(b, testCases)
 }
 
 func TestBatchIncrement(t *testing.T) {
@@ -121,6 +158,48 @@ func TestBatchIncrement(t *testing.T) {
 	}
 }
 
+func TestBatchOpDoesIncrement(t *testing.T) {
+	var b Batch
+	key := []byte("foo")
+	value := []byte("bar")
+
+	if b.Count() != 0 {
+		t.Fatalf("new batch has a nonzero count: %d", b.Count())
+	}
+
+	// Should increment count by 1
+	_ = b.Set(key, value, nil)
+	if b.Count() != 1 {
+		t.Fatalf("expected count: %d, got %d", 1, b.Count())
+	}
+
+	var b2 Batch
+	// Should increment count by 1 each
+	_ = b2.Set(key, value, nil)
+	_ = b2.Delete(key, nil)
+	if b2.Count() != 2 {
+		t.Fatalf("expected count: %d, got %d", 2, b2.Count())
+	}
+
+	// Should increment count by b2.count()
+	_ = b.Apply(&b2, nil)
+	if b.Count() != 3 {
+		t.Fatalf("expected count: %d, got %d", 3, b.Count())
+	}
+
+	// Should increment count by 1
+	_ = b.Merge(key, value, nil)
+	if b.Count() != 4 {
+		t.Fatalf("expected count: %d, got %d", 4, b.Count())
+	}
+
+	// Should NOT increment count.
+	_ = b.LogData([]byte("foobarbaz"), nil)
+	if b.Count() != 4 {
+		t.Fatalf("expected count: %d, got %d", 4, b.Count())
+	}
+}
+
 func TestBatchGet(t *testing.T) {
 	for _, method := range []string{"build", "apply"} {
 		t.Run(method, func(t *testing.T) {
@@ -445,6 +524,7 @@ func BenchmarkBatchSet(b *testing.B) {
 		value[i] = byte(i)
 	}
 	key := make([]byte, 8)
+	batch := newBatch(nil)
 
 	b.ResetTimer()
 
@@ -455,12 +535,11 @@ func BenchmarkBatchSet(b *testing.B) {
 			end = b.N
 		}
 
-		batch := newBatch(nil)
 		for j := i; j < end; j++ {
 			binary.BigEndian.PutUint64(key, uint64(j))
 			batch.Set(key, value, nil)
 		}
-		batch.release()
+		batch.Reset()
 	}
 
 	b.StopTimer()
@@ -472,6 +551,7 @@ func BenchmarkIndexedBatchSet(b *testing.B) {
 		value[i] = byte(i)
 	}
 	key := make([]byte, 8)
+	batch := newIndexedBatch(nil, DefaultComparer)
 
 	b.ResetTimer()
 
@@ -482,12 +562,81 @@ func BenchmarkIndexedBatchSet(b *testing.B) {
 			end = b.N
 		}
 
-		batch := newIndexedBatch(nil, DefaultComparer)
 		for j := i; j < end; j++ {
 			binary.BigEndian.PutUint64(key, uint64(j))
 			batch.Set(key, value, nil)
 		}
-		batch.release()
+		batch.Reset()
+	}
+
+	b.StopTimer()
+}
+
+func BenchmarkBatchSetDeferred(b *testing.B) {
+	value := make([]byte, 10)
+	for i := range value {
+		value[i] = byte(i)
+	}
+	key := make([]byte, 8)
+	batch := newBatch(nil)
+
+	b.ResetTimer()
+
+	const batchSize = 1000
+	for i := 0; i < b.N; i += batchSize {
+		end := i + batchSize
+		if end > b.N {
+			end = b.N
+		}
+
+		for j := i; j < end; j++ {
+			binary.BigEndian.PutUint64(key, uint64(j))
+			deferredOp, err := batch.SetDeferred(len(key), len(value), nil)
+			if err != nil {
+				b.Fatal(err)
+			}
+
+			copy(deferredOp.Key, key)
+			copy(deferredOp.Value, value)
+
+			deferredOp.Finish()
+		}
+		batch.Reset()
+	}
+
+	b.StopTimer()
+}
+
+func BenchmarkIndexedBatchSetDeferred(b *testing.B) {
+	value := make([]byte, 10)
+	for i := range value {
+		value[i] = byte(i)
+	}
+	key := make([]byte, 8)
+	batch := newIndexedBatch(nil, DefaultComparer)
+
+	b.ResetTimer()
+
+	const batchSize = 1000
+	for i := 0; i < b.N; i += batchSize {
+		end := i + batchSize
+		if end > b.N {
+			end = b.N
+		}
+
+		for j := i; j < end; j++ {
+			binary.BigEndian.PutUint64(key, uint64(j))
+			deferredOp, err := batch.SetDeferred(len(key), len(value), nil)
+			if err != nil {
+				b.Fatal(err)
+			}
+
+			copy(deferredOp.Key, key)
+			copy(deferredOp.Value, value)
+
+			deferredOp.Finish()
+		}
+		batch.Reset()
 	}
 
 	b.StopTimer()