mdbx: Cursor.GetBatch — batched cursor reads (one cgo call per N records, ~3.5x faster scans)

mdbx/batch.go

@@ -0,0 +1,119 @@
+package mdbx
+
+/*
+#include <stdlib.h>
+#include "mdbxgo.h"
+*/
+import "C"
+import "unsafe"
+
+// GetBatchBuffer is a reusable buffer of MDBX_val (key, value) pairs for
+// Cursor.GetBatch. It lives in C memory so cgo does not scan it for Go
+// pointers on every call. Reusable across cursors and transactions; not safe
+// for concurrent use. Close releases the C allocation.
+type GetBatchBuffer struct {
+	ptr  *C.MDBX_val
+	size int
+	n    int // pairs filled by the most recent GetBatch
+}
+
+// maxBatchPairs bounds GetBatchBuffer sizes to 1Mi pairs (a 32MiB buffer on
+// 64-bit). It keeps the calloc byte count (2 * numPairs * sizeof(MDBX_val))
+// at or below 32MiB even with 32-bit size_t/int, so neither multiplication
+// can overflow on any platform.
+const maxBatchPairs = 1024 * 1024
+
+// NewGetBatchBuffer allocates a buffer holding numPairs key/value pairs
+// (64-512 is a reasonable range).
+func NewGetBatchBuffer(numPairs int) *GetBatchBuffer {
+	if numPairs < 1 || numPairs > maxBatchPairs {
+		panic("mdbx: NewGetBatchBuffer: number of pairs out of range")
+	}
+	p := C.calloc(C.size_t(2*numPairs), C.size_t(unsafe.Sizeof(C.MDBX_val{})))
+	if p == nil {
+		panic("mdbx: NewGetBatchBuffer: OOM")
+	}
+	return &GetBatchBuffer{ptr: (*C.MDBX_val)(p), size: numPairs}
+}
+
+// Close releases the C allocation. No-op if already closed.
+func (b *GetBatchBuffer) Close() {
+	if b.ptr != nil {
+		C.free(unsafe.Pointer(b.ptr))
+		b.ptr = nil
+		b.size = 0
+	}
+}
+
+// Cap returns the buffer capacity in key/value pairs.
+func (b *GetBatchBuffer) Cap() int { return b.size }
+
+func (b *GetBatchBuffer) at(i int) *C.MDBX_val {
+	// Bound by the last fill count, not the capacity: entries past b.n hold
+	// stale pointers from earlier fills (possibly of an already-ended txn).
+	if b.ptr == nil || i < 0 || i >= 2*b.n {
+		panic("mdbx: GetBatchBuffer: index out of range of the last GetBatch fill")
+	}
+	return (*C.MDBX_val)(unsafe.Add(unsafe.Pointer(b.ptr), uintptr(i)*unsafe.Sizeof(C.MDBX_val{})))
+}
+
+// Key returns the i-th key of the most recent GetBatch (i < its pair count).
+// Zero-copy view: read-only, invalid once the txn ends, the buffer is
+// refilled, or (in a write txn) a later Put/Del moves the page.
+func (b *GetBatchBuffer) Key(i int) []byte { return castToBytes(b.at(2 * i)) }
+
+// Val returns the i-th value of the most recent GetBatch. See Key.
+func (b *GetBatchBuffer) Val(i int) []byte { return castToBytes(b.at(2*i + 1)) }
+
+// GetBatch fetches up to buf.Cap() key/value pairs in one cgo call: the first
+// record with opFirst, the rest with opNext. It amortizes cgo overhead over
+// large scans.
+//
+// There is no way to pass a search key or value, so ops that need one (Set,
+// SetRange, GetBoth, ...) cannot be used. For a ranged scan, position the
+// cursor with Get first and batch with (GetCurrent, Next). The *_MULTIPLE
+// ops work but have page granularity: each stored pair is (key, packed page
+// of fixed-size values), so n counts pages, not records.
+//
+// n is the number of pairs stored (read via buf.Key/Val). The first n pairs
+// are valid even when err != nil (the error came from the step after them).
+// eof is true when iteration was exhausted before the buffer filled;
+// otherwise continue with c.GetBatch(buf, opNext, opNext).
+//
+//	buf := mdbx.NewGetBatchBuffer(256)
+//	defer buf.Close()
+//	for opFirst := uint(mdbx.First); ; opFirst = mdbx.Next {
+//		n, eof, err := cur.GetBatch(buf, opFirst, mdbx.Next)
+//		for i := 0; i < n; i++ {
+//			handle(buf.Key(i), buf.Val(i))
+//		}
+//		if err != nil {
+//			return err
+//		}
+//		if eof {
+//			break
+//		}
+//	}
+func (c *Cursor) GetBatch(buf *GetBatchBuffer, opFirst, opNext uint) (n int, eof bool, err error) {
+	if buf != nil {
+		// An errored call must not leave the buffer reporting entries from a
+		// previous fill (possibly views of an already-ended transaction).
+		buf.n = 0
+	}
+	if c._c == nil || buf == nil || buf.ptr == nil || buf.size <= 0 {
+		return 0, false, operrno("mdbx_cursor_get", C.MDBX_EINVAL)
+	}
+	r := C.mdbxgo_cursor_get_batch(
+		c._c, buf.ptr, C.size_t(buf.size),
+		C.MDBX_cursor_op(opFirst), C.MDBX_cursor_op(opNext),
+	)
+	n = int(r.val)
+	buf.n = n
+	if r.err == C.MDBX_NOTFOUND {
+		return n, true, nil
+	}
+	if r.err != success {
+		return n, false, operrno("mdbx_cursor_get", r.err)
+	}
+	return n, false, nil
+}

mdbx/batch_test.go

@@ -0,0 +1,264 @@
+package mdbx
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+)
+
+func fillBatchDB(tb testing.TB, env *Env, name string, numItems int) DBI {
+	tb.Helper()
+	var db DBI
+	err := env.Update(func(txn *Txn) (err error) {
+		db, err = txn.OpenDBISimple(name, Create)
+		if err != nil {
+			return err
+		}
+		for i := range numItems {
+			k := []byte(fmt.Sprintf("key-%08d", i))
+			v := []byte(fmt.Sprintf("val-%08d", i))
+			if err := txn.Put(db, k, v, Append); err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+	if err != nil {
+		tb.Fatal(err)
+	}
+	return db
+}
+
+func TestCursor_GetBatch(t *testing.T) {
+	env, _ := setup(t)
+	const numItems = 1000
+	db := fillBatchDB(t, env, "testbatch", numItems)
+
+	buf := NewGetBatchBuffer(64)
+	defer buf.Close()
+
+	err := env.View(func(txn *Txn) error {
+		cur, err := txn.OpenCursor(db)
+		if err != nil {
+			return err
+		}
+		defer cur.Close()
+
+		seen := 0
+		for opFirst := uint(First); ; opFirst = Next {
+			n, eof, err := cur.GetBatch(buf, opFirst, Next)
+			if err != nil {
+				return err
+			}
+			for i := range n {
+				wantK := fmt.Sprintf("key-%08d", seen)
+				wantV := fmt.Sprintf("val-%08d", seen)
+				if !bytes.Equal(buf.Key(i), []byte(wantK)) {
+					t.Fatalf("pair %d: key = %q, want %q", seen, buf.Key(i), wantK)
+				}
+				if !bytes.Equal(buf.Val(i), []byte(wantV)) {
+					t.Fatalf("pair %d: val = %q, want %q", seen, buf.Val(i), wantV)
+				}
+				seen++
+			}
+			if eof {
+				break
+			}
+		}
+		if seen != numItems {
+			t.Errorf("scanned %d items, want %d", seen, numItems)
+		}
+		return nil
+	})
+	if err != nil {
+		t.Error(err)
+	}
+}
+
+func TestCursor_GetBatch_EmptyDB(t *testing.T) {
+	env, _ := setup(t)
+	db := fillBatchDB(t, env, "testbatchempty", 0)
+
+	buf := NewGetBatchBuffer(8)
+	defer buf.Close()
+
+	err := env.View(func(txn *Txn) error {
+		cur, err := txn.OpenCursor(db)
+		if err != nil {
+			return err
+		}
+		defer cur.Close()
+
+		n, eof, err := cur.GetBatch(buf, First, Next)
+		if err != nil {
+			return err
+		}
+		if n != 0 || !eof {
+			t.Errorf("GetBatch on empty table: n=%d eof=%v, want 0/true", n, eof)
+		}
+		return nil
+	})
+	if err != nil {
+		t.Error(err)
+	}
+}
+
+func TestCursor_GetBatch_NoAllocs(t *testing.T) {
+	env, _ := setup(t)
+	db := fillBatchDB(t, env, "testbatchnoalloc", 100)
+
+	buf := NewGetBatchBuffer(32)
+	defer buf.Close()
+
+	err := env.View(func(txn *Txn) error {
+		cur, err := txn.OpenCursor(db)
+		if err != nil {
+			return err
+		}
+		defer cur.Close()
+
+		assertNoAllocs(t, "Cursor.GetBatch", func() { _, _, _ = cur.GetBatch(buf, First, Next) })
+		return nil
+	})
+	if err != nil {
+		t.Error(err)
+	}
+}
+
+// BenchmarkCursorScan compares a full forward scan performed with one cgo
+// call per record (Get/Next) against batched retrieval (GetBatch).
+func BenchmarkCursorScan(b *testing.B) {
+	env, _ := setup(b)
+	const numItems = 100_000
+	db := fillBatchDB(b, env, "benchscan", numItems)
+
+	txn, err := env.BeginTxn(nil, Readonly)
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer txn.Abort()
+	cur, err := txn.OpenCursor(db)
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer cur.Close()
+
+	b.Run("Get_Next", func(b *testing.B) {
+		b.ResetTimer()
+		var total int
+		for range b.N {
+			count := 0
+			for _, _, err := cur.Get(nil, nil, First); err == nil; _, _, err = cur.Get(nil, nil, Next) {
+				count++
+			}
+			total = count
+		}
+		if total != numItems {
+			b.Fatalf("scanned %d, want %d", total, numItems)
+		}
+		b.ReportMetric(float64(b.Elapsed().Nanoseconds())/float64(b.N)/numItems, "ns/record")
+	})
+
+	for _, batch := range []int{64, 256, 1024} {
+		b.Run(fmt.Sprintf("GetBatch_%d", batch), func(b *testing.B) {
+			buf := NewGetBatchBuffer(batch)
+			defer buf.Close()
+			b.ResetTimer()
+			var total int
+			for range b.N {
+				count := 0
+				for opFirst := uint(First); ; opFirst = Next {
+					n, eof, err := cur.GetBatch(buf, opFirst, Next)
+					if err != nil {
+						b.Fatal(err)
+					}
+					count += n
+					if eof {
+						break
+					}
+				}
+				total = count
+			}
+			if total != numItems {
+				b.Fatalf("scanned %d, want %d", total, numItems)
+			}
+			b.ReportMetric(float64(b.Elapsed().Nanoseconds())/float64(b.N)/numItems, "ns/record")
+		})
+	}
+}
+
+// Exact-fill boundary: the batch that fills completely on the table's last
+// record reports eof=false; the follow-up call must return (0, true, nil).
+func TestCursor_GetBatch_ExactFill(t *testing.T) {
+	env, _ := setup(t)
+	db := fillBatchDB(t, env, "testbatchexact", 128)
+
+	buf := NewGetBatchBuffer(64)
+	defer buf.Close()
+
+	err := env.View(func(txn *Txn) error {
+		cur, err := txn.OpenCursor(db)
+		if err != nil {
+			return err
+		}
+		defer cur.Close()
+
+		for i, want := range []struct {
+			n   int
+			eof bool
+		}{{64, false}, {64, false}, {0, true}} {
+			op := uint(Next)
+			if i == 0 {
+				op = First
+			}
+			n, eof, err := cur.GetBatch(buf, op, Next)
+			if err != nil {
+				return err
+			}
+			if n != want.n || eof != want.eof {
+				t.Errorf("batch %d: n=%d eof=%v, want n=%d eof=%v", i, n, eof, want.n, want.eof)
+			}
+		}
+		return nil
+	})
+	if err != nil {
+		t.Error(err)
+	}
+}
+
+// A failing step mid-batch must leave the pairs fetched before it valid.
+func TestCursor_GetBatch_PartialOnError(t *testing.T) {
+	env, _ := setup(t)
+	db := fillBatchDB(t, env, "testbatchpartial", 10)
+
+	buf := NewGetBatchBuffer(8)
+	defer buf.Close()
+
+	err := env.View(func(txn *Txn) error {
+		cur, err := txn.OpenCursor(db)
+		if err != nil {
+			return err
+		}
+		defer cur.Close()
+
+		// First succeeds; GetMultiple on a non-DupFixed table fails with
+		// MDBX_INCOMPATIBLE on the second step.
+		n, eof, err := cur.GetBatch(buf, First, GetMultiple)
+		if err == nil {
+			t.Fatal("GetBatch(First, GetMultiple): expected error, got nil")
+		}
+		if eof {
+			t.Error("eof must be false on error")
+		}
+		if n != 1 {
+			t.Fatalf("n = %d, want 1 pair fetched before the failing step", n)
+		}
+		if got := string(buf.Key(0)); got != "key-00000000" {
+			t.Errorf("Key(0) = %q, want key-00000000", got)
+		}
+		return nil
+	})
+	if err != nil {
+		t.Error(err)
+	}
+}