TreeDB/caching: reserve pointer batch index metadata

TreeDB/caching/batch_aux_pool_test.go

@@ -35,16 +35,17 @@ func TestBatchCloseReleasesAuxiliaryIndexSlices(t *testing.T) {
 		eligibleIdxs: make([]int, 1, 4),
 		shardAdds:    make([]int64, 2, 4),
 		shardCnts:    make([]int, 2, 4),
+		ptrValueIdxs: make([]int, 2, 4),
 		shardEntries: [][]batch.Entry{make([]batch.Entry, 0, 2)},
 		shardIdxSets: [][]int{make([]int, 0, 2)},
 	}
 
 	if err := b.Close(); err != nil {
 		t.Fatalf("Close: %v", err)
 	}
-	if b.entries != nil || b.shardIdxs != nil || b.eligibleIdxs != nil || b.shardAdds != nil || b.shardCnts != nil || b.shardEntries != nil || b.shardIdxSets != nil {
-		t.Fatalf("Close retained auxiliary slices: entries=%v shardIdxs=%v eligible=%v shardAdds=%v shardCnts=%v shardEntries=%v shardIdxSets=%v",
-			b.entries, b.shardIdxs, b.eligibleIdxs, b.shardAdds, b.shardCnts, b.shardEntries, b.shardIdxSets)
+	if b.entries != nil || b.shardIdxs != nil || b.eligibleIdxs != nil || b.shardAdds != nil || b.shardCnts != nil || b.ptrValueIdxs != nil || b.shardEntries != nil || b.shardIdxSets != nil {
+		t.Fatalf("Close retained auxiliary slices: entries=%v shardIdxs=%v eligible=%v shardAdds=%v shardCnts=%v ptrValueIdxs=%v shardEntries=%v shardIdxSets=%v",
+			b.entries, b.shardIdxs, b.eligibleIdxs, b.shardAdds, b.shardCnts, b.ptrValueIdxs, b.shardEntries, b.shardIdxSets)
 	}
 
 	// The returned pools must still hand out correctly sized, owned scratch slices.
@@ -59,3 +60,27 @@ func TestBatchCloseReleasesAuxiliaryIndexSlices(t *testing.T) {
 	}
 	db.putBatchInt64Slice(adds)
 }
+
+func TestBatchAppendPtrValueIdxReservesEntryCapacity(t *testing.T) {
+	db := &DB{}
+	b := &Batch{
+		db:      db,
+		entries: db.getBatchEntries(16),
+	}
+	defer b.Close()
+
+	for i := 0; i < 8; i++ {
+		b.appendPtrValueIdx(i)
+	}
+	if got, want := len(b.ptrValueIdxs), 8; got != want {
+		t.Fatalf("len(ptrValueIdxs)=%d want %d", got, want)
+	}
+	if got := cap(b.ptrValueIdxs); got < cap(b.entries) {
+		t.Fatalf("cap(ptrValueIdxs)=%d want >= entries cap %d", got, cap(b.entries))
+	}
+	for i, got := range b.ptrValueIdxs {
+		if got != i {
+			t.Fatalf("ptrValueIdxs[%d]=%d want %d", i, got, i)
+		}
+	}
+}

TreeDB/caching/batch_setwithrevision_bench_test.go

@@ -0,0 +1,112 @@
+package caching
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/snissn/gomap/TreeDB/page"
+)
+
+var batchSetWithRevisionBenchSink int
+
+func BenchmarkBatchSetWithRevisionAllocationShape(b *testing.B) {
+	const (
+		entries   = 64
+		keySize   = 32
+		inlineVal = 96
+		ptrVal    = 2048
+	)
+
+	cases := []struct {
+		name             string
+		valueSize        int
+		pointerThreshold int
+		set              func(*Batch, []byte, []byte, page.EntryRevision) error
+	}{
+		{
+			name:             "SetWithRevision_inline",
+			valueSize:        inlineVal,
+			pointerThreshold: 4096,
+			set: func(batch *Batch, key, value []byte, revision page.EntryRevision) error {
+				return batch.SetWithRevision(key, value, revision)
+			},
+		},
+		{
+			name:             "SetWithRevision_pointer",
+			valueSize:        ptrVal,
+			pointerThreshold: 32,
+			set: func(batch *Batch, key, value []byte, revision page.EntryRevision) error {
+				return batch.SetWithRevision(key, value, revision)
+			},
+		},
+		{
+			name:             "SetViewValidatedWithRevision_pointer",
+			valueSize:        ptrVal,
+			pointerThreshold: 32,
+			set: func(batch *Batch, key, value []byte, revision page.EntryRevision) error {
+				return batch.SetViewValidatedWithRevision(key, value, revision)
+			},
+		},
+	}
+
+	for _, tc := range cases {
+		tc := tc
+		b.Run(fmt.Sprintf("%s_%dx%dx%d", tc.name, entries, keySize, tc.valueSize), func(b *testing.B) {
+			db, err := Open(b.TempDir(), NewMockBackend(), Options{
+				AllowUnsafe:              true,
+				DisableWAL:               true,
+				MemtableMode:             "btree",
+				MemtableShards:           1,
+				FlushThreshold:           1 << 30,
+				ValueLogPointerThreshold: tc.pointerThreshold,
+			})
+			if err != nil {
+				b.Fatalf("open: %v", err)
+			}
+			defer db.Close()
+
+			keys, values := makeBatchSetWithRevisionBenchPayload(entries, keySize, tc.valueSize)
+
+			b.ReportAllocs()
+			b.SetBytes(int64(entries * (keySize + tc.valueSize)))
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				batch := db.NewBatchWithSize(entries)
+				for j := 0; j < entries; j++ {
+					if err := tc.set(batch, keys[j], values[j], page.EntryRevision(j+1)); err != nil {
+						b.Fatalf("set %d: %v", j, err)
+					}
+				}
+				batchSetWithRevisionBenchSink += len(batch.entries) + batch.size + len(batch.ptrValueIdxs)
+				if err := batch.Close(); err != nil {
+					b.Fatalf("close: %v", err)
+				}
+			}
+		})
+	}
+}
+
+func makeBatchSetWithRevisionBenchPayload(entries, keySize, valueSize int) ([][]byte, [][]byte) {
+	keys := make([][]byte, entries)
+	values := make([][]byte, entries)
+	for i := 0; i < entries; i++ {
+		key := make([]byte, keySize)
+		copy(key, "batch-setwithrevision-key:")
+		encodeBenchUint(key[len(key)-8:], uint64(i))
+		keys[i] = key
+
+		value := make([]byte, valueSize)
+		for j := range value {
+			value[j] = byte((i*31 + j*17) & 0xff)
+		}
+		values[i] = value
+	}
+	return keys, values
+}
+
+func encodeBenchUint(dst []byte, value uint64) {
+	for i := len(dst) - 1; i >= 0; i-- {
+		dst[i] = byte(value)
+		value >>= 8
+	}
+}

TreeDB/caching/db.go

@@ -33221,6 +33221,35 @@ func (b *Batch) shouldCopyValueToPtrArena(key, value []byte) bool {
 	return b.db.shouldWriteViaValueLogForKeyValue(key, value)
 }
 
+func (b *Batch) appendPtrValueIdx(idx int) {
+	if len(b.ptrValueIdxs) == cap(b.ptrValueIdxs) {
+		minCap := len(b.ptrValueIdxs) + 1
+		nextCap := cap(b.entries)
+		if nextCap < minCap {
+			nextCap = cap(b.ptrValueIdxs) * 2
+			if nextCap < minCap {
+				nextCap = minCap
+			}
+		}
+		if nextCap > batchIntSlicePoolMaxRetain && minCap <= batchIntSlicePoolMaxRetain {
+			nextCap = batchIntSlicePoolMaxRetain
+		}
+
+		var next []int
+		if b.db != nil {
+			next = b.db.getBatchIntSlice(nextCap)
+		} else {
+			next = make([]int, 0, nextCap)
+		}
+		next = append(next, b.ptrValueIdxs...)
+		if b.db != nil && b.ptrValueIdxs != nil {
+			b.db.putBatchIntSlice(b.ptrValueIdxs)
+		}
+		b.ptrValueIdxs = next
+	}
+	b.ptrValueIdxs = append(b.ptrValueIdxs, idx)
+}
+
 func (b *Batch) assignLegacyPointRevisions(entries []batch.Entry) {
 	if b == nil || b.db == nil || len(entries) == 0 {
 		return
@@ -33272,7 +33301,7 @@ func (b *Batch) SetWithRevision(key, value []byte, revision page.EntryRevision)
 	if b.shouldCopyValueToPtrArena(key, value) {
 		keyCopy = b.cloneKey(key)
 		valCopy = b.clonePtrValue(value)
-		b.ptrValueIdxs = append(b.ptrValueIdxs, idx)
+		b.appendPtrValueIdx(idx)
 	} else {
 		keyCopy, valCopy = b.cloneKeyValue(key, value)
 	}
@@ -35563,6 +35592,9 @@ func (b *Batch) Close() error {
 	if b.db != nil && b.shardCnts != nil {
 		b.db.putBatchIntSlice(b.shardCnts)
 	}
+	if b.db != nil && b.ptrValueIdxs != nil {
+		b.db.putBatchIntSlice(b.ptrValueIdxs)
+	}
 	if b.db != nil && b.shardIdxSets != nil && cap(b.shardIdxSets) > 0 {
 		full := b.shardIdxSets[:cap(b.shardIdxSets)]
 		for i := range full {