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

mdbx/batch.go

@@ -0,0 +1,94 @@
+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
+}
+
+// NewGetBatchBuffer allocates a buffer holding numPairs key/value pairs
+// (64-512 is a reasonable range).
+func NewGetBatchBuffer(numPairs int) *GetBatchBuffer {
+	if numPairs < 1 {
+		panic("mdbx: NewGetBatchBuffer requires a positive number of pairs")
+	}
+	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 {
+	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 after the txn ends or the buffer is
+// refilled.
+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.
+//
+// n is the number of pairs stored (read via buf.Key/Val). eof is true when
+// iteration was exhausted before the buffer filled; otherwise continue with
+// 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)
+//		if err != nil {
+//			return err
+//		}
+//		for i := 0; i < n; i++ {
+//			handle(buf.Key(i), buf.Val(i))
+//		}
+//		if eof {
+//			break
+//		}
+//	}
+func (c *Cursor) GetBatch(buf *GetBatchBuffer, opFirst, opNext uint) (n int, eof bool, err error) {
+	if c._c == nil || buf == nil || buf.ptr == nil {
+		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)
+	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,186 @@
+package mdbx
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+)
+
+func fillBatchDB(t testing.TB, env *Env, name string, numItems int) DBI {
+	t.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 := 0; i < numItems; i++ {
+			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 {
+		t.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 := 0; i < n; i++ {
+				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) {
+		var total int
+		for i := 0; i < b.N; i++ {
+			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()
+			var total int
+			for i := 0; i < b.N; i++ {
+				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")
+		})
+	}
+}

mdbx/mdbxgo.c

@@ -114,6 +114,23 @@ mdbxgo_val_result mdbxgo_cursor_get_val(MDBX_cursor *cur, char *kdata, size_t kn
     return r;
 }
 
+mdbxgo_size_result mdbxgo_cursor_get_batch(MDBX_cursor *cur, MDBX_val *pairs, size_t max_pairs,
+                                           MDBX_cursor_op op_first, MDBX_cursor_op op_next) {
+    mdbxgo_size_result r = {0};
+    MDBX_cursor_op op = op_first;
+    while (r.val < max_pairs) {
+        MDBX_val key = {0}, val = {0};
+        r.err = mdbx_cursor_get(cur, &key, &val, op);
+        if (r.err != MDBX_SUCCESS)
+            break;
+        pairs[2 * r.val] = key;
+        pairs[2 * r.val + 1] = val;
+        r.val++;
+        op = op_next;
+    }
+    return r;
+}
+
 /* Compare two items lexically */
 // static int __hot cmp_lexical(const MDBX_val *a, const MDBX_val *b) {
 //   if (a->iov_len == b->iov_len)

mdbx/mdbxgo.h

@@ -79,6 +79,14 @@ typedef struct { int err; char *kbase; size_t klen; char *vbase; size_t vlen; }
 mdbxgo_val_result        mdbxgo_cursor_get_empty(MDBX_cursor *cur, MDBX_cursor_op op);
 mdbxgo_val_result        mdbxgo_cursor_get_val(MDBX_cursor *cur, char *kdata, size_t kn, char *vdata, size_t vn, MDBX_cursor_op op);
 
+/* mdbxgo_cursor_get_batch fills pairs[0..2*max_pairs) with (key, value)
+ * result pairs: the first mdbx_cursor_get step uses op_first, the rest
+ * op_next. r.val holds the filled pair count; r.err is MDBX_SUCCESS when the
+ * buffer was filled, MDBX_NOTFOUND when iteration was exhausted first.
+ * Amortizes cgo call overhead: one Go->C call retrieves max_pairs records.
+ * */
+mdbxgo_size_result       mdbxgo_cursor_get_batch(MDBX_cursor *cur, MDBX_val *pairs, size_t max_pairs, MDBX_cursor_op op_first, MDBX_cursor_op op_next);
+
 typedef struct {
     int err;
     uint64_t pages_allocated;