map: avoid allocations in MapIterator.Next

Rejig MapIterator.Next to stop allocating nextKey over and over.
This makes Next() not allocate in a steady state:

    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/MapIterator         1.596m ± 0%   1.500m ± 1%  -6.05% (p=0.002 n=6)
    Iterate/MapIteratorDelete   2.862m ± 0%   2.754m ± 0%  -3.75% (p=0.002 n=6)
    geomean                     2.137m        2.032m       -4.91%

                            │   base.txt    │              opt.txt              │
                            │     B/op      │    B/op     vs base               │
    Iterate/MapIterator         56082.00 ± 0%   96.00 ± 0%  -99.83% (p=0.002 n=6)
    Iterate/MapIteratorDelete    56082.0 ± 0%   104.0 ± 0%  -99.81% (p=0.002 n=6)
    geomean                      54.77Ki        99.92       -99.82%

                            │   base.txt    │              opt.txt              │
                            │   allocs/op   │ allocs/op   vs base               │
    Iterate/MapIterator         3004.000 ± 0%   4.000 ± 0%  -99.87% (p=0.002 n=6)
    Iterate/MapIteratorDelete   3004.000 ± 0%   4.000 ± 0%  -99.87% (p=0.002 n=6)
    geomean                       3.004k        4.000       -99.87%

Signed-off-by: Lorenz Bauer <[email protected]>

map.go

@@ -1411,8 +1411,10 @@ func marshalMap(m *Map, length int) ([]byte, error) {
 //
 // See Map.Iterate.
 type MapIterator struct {
-	target            *Map
-	curKey            []byte
+	target *Map
+	// Temporary storage to avoid allocations in Next(). This is any instead
+	// of []byte to avoid allocations.
+	cursor            any
 	count, maxEntries uint32
 	done              bool
 	err               error
@@ -1440,34 +1442,30 @@ func (mi *MapIterator) Next(keyOut, valueOut interface{}) bool {
 		return false
 	}
 
-	// For array-like maps NextKeyBytes returns nil only on after maxEntries
+	// For array-like maps NextKey returns nil only after maxEntries
 	// iterations.
 	for mi.count <= mi.maxEntries {
-		var nextKey []byte
-		if mi.curKey == nil {
-			// Pass nil interface to NextKeyBytes to make sure the Map's first key
+		if mi.cursor == nil {
+			// Pass nil interface to NextKey to make sure the Map's first key
 			// is returned. If we pass an uninitialized []byte instead, it'll see a
 			// non-nil interface and try to marshal it.
-			nextKey, mi.err = mi.target.NextKeyBytes(nil)
-
-			mi.curKey = make([]byte, mi.target.keySize)
+			mi.cursor = make([]byte, mi.target.keySize)
+			mi.err = mi.target.NextKey(nil, mi.cursor)
 		} else {
-			nextKey, mi.err = mi.target.NextKeyBytes(mi.curKey)
-		}
-		if mi.err != nil {
-			mi.err = fmt.Errorf("get next key: %w", mi.err)
-			return false
+			mi.err = mi.target.NextKey(mi.cursor, mi.cursor)
 		}
 
-		if nextKey == nil {
+		if errors.Is(mi.err, ErrKeyNotExist) {
 			mi.done = true
+			mi.err = nil
+			return false
+		} else if mi.err != nil {
+			mi.err = fmt.Errorf("get next key: %w", mi.err)
 			return false
 		}
 
-		mi.curKey = nextKey
-
 		mi.count++
-		mi.err = mi.target.Lookup(nextKey, valueOut)
+		mi.err = mi.target.Lookup(mi.cursor, valueOut)
 		if errors.Is(mi.err, ErrKeyNotExist) {
 			// Even though the key should be valid, we couldn't look up
 			// its value. If we're iterating a hash map this is probably
@@ -1484,10 +1482,11 @@ func (mi *MapIterator) Next(keyOut, valueOut interface{}) bool {
 			return false
 		}
 
+		buf := mi.cursor.([]byte)
 		if ptr, ok := keyOut.(unsafe.Pointer); ok {
-			copy(unsafe.Slice((*byte)(ptr), len(nextKey)), nextKey)
+			copy(unsafe.Slice((*byte)(ptr), len(buf)), buf)
 		} else {
-			mi.err = sysenc.Unmarshal(keyOut, nextKey)
+			mi.err = sysenc.Unmarshal(keyOut, buf)
 		}
 
 		return mi.err == nil

map_test.go

@@ -1103,6 +1103,32 @@ func TestMapIterate(t *testing.T) {
 	}
 }
 
+func TestMapIteratorAllocations(t *testing.T) {
+	arr, err := NewMap(&MapSpec{
+		Type:       Array,
+		KeySize:    4,
+		ValueSize:  4,
+		MaxEntries: 10,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer arr.Close()
+
+	var k, v uint32
+	iter := arr.Iterate()
+	// Allocate any necessary temporary buffers.
+	qt.Assert(t, iter.Next(&k, &v), qt.IsTrue)
+
+	allocs := testing.AllocsPerRun(1, func() {
+		if !iter.Next(&k, &v) {
+			t.Fatal("Next failed")
+		}
+	})
+
+	qt.Assert(t, allocs, qt.Equals, float64(0))
+}
+
 func TestMapIterateHashKeyOneByteFull(t *testing.T) {
 	hash, err := NewMap(&MapSpec{
 		Type:       Hash,