From 4cc5e8950a4ade803d771075766f5ff958b69515 Mon Sep 17 00:00:00 2001
From: Lorenz Bauer <lmb@isovalent.com>
Date: Tue, 5 Dec 2023 17:31:55 +0000
Subject: [PATCH] map: avoid allocations during batch lookup of common types
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Extend the optimization from 4609dc73 ("map: zero-allocation operations for
common types") to BatchLookup. This means that types without implicit padding
will not incur allocations.

    core: 1
    goos: linux
    goarch: amd64
    pkg: github.com/cilium/ebpf
    cpu: 12th Gen Intel(R) Core(TM) i7-1260P
                        │  base.txt   │              opt.txt               │
                        │   sec/op    │   sec/op     vs base               │
    Iterate/BatchLookup   5.997µ ± 2%   2.204µ ± 0%  -63.26% (p=0.002 n=6)

                        │   base.txt    │              opt.txt              │
                        │     B/op      │    B/op     vs base               │
    Iterate/BatchLookup   16440.00 ± 0%   56.00 ± 0%  -99.66% (p=0.002 n=6)

                        │  base.txt  │              opt.txt              │
                        │ allocs/op  │ allocs/op   vs base               │
    Iterate/BatchLookup   5.000 ± 0%   3.000 ± 0%  -40.00% (p=0.002 n=6)

Signed-off-by: Lorenz Bauer <lmb@isovalent.com>
---
 internal/sysenc/buffer.go |  1 +
 map.go                    | 17 +++++++----------
 map_test.go               | 32 ++++++++++++++++++++++++++++++--
 3 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/internal/sysenc/buffer.go b/internal/sysenc/buffer.go
index c6959d9cc..97ac473c8 100644
--- a/internal/sysenc/buffer.go
+++ b/internal/sysenc/buffer.go
@@ -32,6 +32,7 @@ func UnsafeBuffer(ptr unsafe.Pointer) Buffer {
 
 // SyscallOutput prepares a Buffer for a syscall to write into.
 //
+// size is the length of the desired buffer in bytes.
 // The buffer may point at the underlying memory of dst, in which case [Unmarshal]
 // becomes a no-op.
 //
diff --git a/map.go b/map.go
index 87a91d133..5e71ccfc2 100644
--- a/map.go
+++ b/map.go
@@ -1048,15 +1048,14 @@ func (m *Map) batchLookup(cmd sys.Cmd, cursor *BatchCursor, keysOut, valuesOut i
 	if count != valuesValue.Len() {
 		return 0, fmt.Errorf("keysOut and valuesOut must be the same length")
 	}
-	keyBuf := make([]byte, count*int(m.keySize))
-	keyPtr := sys.NewSlicePointer(keyBuf)
-	valueBuf := make([]byte, count*int(m.fullValueSize))
-	valuePtr := sys.NewSlicePointer(valueBuf)
+
+	keyBuf := sysenc.SyscallOutput(keysOut, count*int(m.keySize))
+	valueBuf := sysenc.SyscallOutput(valuesOut, count*int(m.fullValueSize))
 
 	attr := sys.MapLookupBatchAttr{
 		MapFd:    m.fd.Uint(),
-		Keys:     keyPtr,
-		Values:   valuePtr,
+		Keys:     keyBuf.Pointer(),
+		Values:   valueBuf.Pointer(),
 		Count:    uint32(count),
 		InBatch:  sys.NewSlicePointer(inBatch),
 		OutBatch: sys.NewSlicePointer(cursor.opaque),
@@ -1073,12 +1072,10 @@ func (m *Map) batchLookup(cmd sys.Cmd, cursor *BatchCursor, keysOut, valuesOut i
 		return 0, sysErr
 	}
 
-	err := sysenc.Unmarshal(keysOut, keyBuf)
-	if err != nil {
+	if err := keyBuf.Unmarshal(keysOut); err != nil {
 		return 0, err
 	}
-	err = sysenc.Unmarshal(valuesOut, valueBuf)
-	if err != nil {
+	if err := valueBuf.Unmarshal(valuesOut); err != nil {
 		return 0, err
 	}
 
diff --git a/map_test.go b/map_test.go
index 415953916..13e52b2a3 100644
--- a/map_test.go
+++ b/map_test.go
@@ -1093,9 +1093,8 @@ func TestMapIteratorAllocations(t *testing.T) {
 
 	var k, v uint32
 	iter := arr.Iterate()
-	// Allocate any necessary temporary buffers.
-	qt.Assert(t, qt.IsTrue(iter.Next(&k, &v)))
 
+	// AllocsPerRun warms up the function for us.
 	allocs := testing.AllocsPerRun(1, func() {
 		if !iter.Next(&k, &v) {
 			t.Fatal("Next failed")
@@ -1105,6 +1104,35 @@ func TestMapIteratorAllocations(t *testing.T) {
 	qt.Assert(t, qt.Equals(allocs, float64(0)))
 }
 
+func TestMapBatchLookupAllocations(t *testing.T) {
+	testutils.SkipIfNotSupported(t, haveBatchAPI())
+
+	arr, err := NewMap(&MapSpec{
+		Type:       Array,
+		KeySize:    4,
+		ValueSize:  4,
+		MaxEntries: 10,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer arr.Close()
+
+	var cursor BatchCursor
+	tmp := make([]uint32, 2)
+	input := any(tmp)
+
+	// AllocsPerRun warms up the function for us.
+	allocs := testing.AllocsPerRun(1, func() {
+		_, err := arr.BatchLookup(&cursor, input, input, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+	})
+
+	qt.Assert(t, qt.Equals(allocs, 0))
+}
+
 func TestMapIterateHashKeyOneByteFull(t *testing.T) {
 	hash, err := NewMap(&MapSpec{
 		Type:       Hash,