Repository: cilium/ebpf
Pull request #1223 (2023-11-17 06:30:51)

map: Introduce BatchCursor abstraction
See commit msgs.

* Introduce BatchCursor abstraction
* Add unit test for batch lookup with chunking
* Add warning to Batch* API operations

---

Main commit for convenience: 

> Introduce BatchCursor abstraction
> 
> The batch API in the kernel is quite tricky to get right. In order to
> pass the arguments in correctly, especially regarding the "in_batch" and
> "out_batch" arguments, it requires knowledge of how the batch map
> iteration is implemented in the kernel.
> 
> In order to avoid pushing the cognitive overhead onto users of the
> library, abstract away the details of batch map iteration into an opaque
> type, BatchCursor.
> 
> Users are expected to pass in a reference to a BatchCursor and the
> library will handle allocating the underlying buffer when needed.
> Specifically, the first time that a batch API is called using the
> cursor, the underlying buffer will be nil, which signfifies that the
> "in_batch" (aka "prevKey") should be nil to indicate the starting of a
> batching operation, and "out_batch" (aka "nextKey") is set to an
> newly-allocated buffer. Previously, users were expected to carry the
> "out_batch" reference in order to pass it in as "in_batch" upon the next
> call to the batching operation. Instead, with this abstraction, the
> cursor handles it for the user, by using the underlying buffer for both
> "in_batch" and "out_batch" in subsequent calls.
> 
> Signed-off-by: Chris Tarazi <chris@isovalent.com>
> Suggested-by: Lorenz Bauer <lmb@isovalent.com>


christarazi: cc @lorenz @ti-mo 
lmb: Sorry this took so long. I needed to make some time to dig into the kernel sources what this parameter actually does. From the selftests it looks like the type of the batch is either uint32 or uint64:
* https://elixir.bootlin.com/linux/v6.6.3/source/tools/testing/selftests/bpf/map_tests/array_map_batch_ops.c#L75
* https://elixir.bootlin.com/linux/v6.6.3/source/tools/testing/selftests/bpf/map_tests/htab_map_batch_ops.c#L79
* https://elixir.bootlin.com/linux/v6.6.3/source/tools/testing/selftests/bpf/map_tests/lpm_trie_map_batch_ops.c#L73
* https://elixir.bootlin.com/linux/v6.6.3/source/tools/testing/selftests/bpf/map_tests/map_in_map_batch_ops.c#L126

It turns out that things are even more complicated though!

- [generic_map_lookup_batch](https://elixir.bootlin.com/linux/v6.6.3/source/kernel/bpf/syscall.c#L1859) requires that batch_out is key_size bytes. This is used by array and LPM maps.
- [__htab_map_lookup_and_delete_batch](https://elixir.bootlin.com/linux/v6.6.3/source/kernel/bpf/hashtab.c#L1892) requires u64. This is used by the various hash maps.

I don't want an API which requires the user to know how batch map iteration is implemented in the kernel. So at the very least we need to abstract away the job of allocating the batch storage. Minimal change I can think of:

```go
type BatchCursor struct{
    opaque []byte // len() is max(key_size, 8) to be on the safe side
}

func (*Map) BatchLookup(cursor *BatchCursor, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
    if (cursor.opaque == nil) {
        // this is the first time the lookup is called. allocate cursor.opaque and pass it as in_batch and out_batch.
    }
}
```

Usage would be something like the following:

```go
var cursor BatchCursor
count, err := m.Lookup(&cursor, ...)
```

Is that something you want to contribute? Otherwise I can fix it as well.
christarazi: @lmb Thanks for taking a look. I'll take a stab at it. I'll repurpose this PR.
christarazi: @lmb I've pushed "Introduce BatchCursor abstraction" commit which is the minimal implementation of your idea of the BatchCursor. Let me know if that's in line with what your thinking and I'll proceed with tightening up the godocs, etc.
lmb: One thought: the API isn't really for a lookup, since there is no way to specify a key. `BatchIterate`?
christarazi: > One thought: the API isn't really for a lookup, since there is no way to specify a key. `BatchIterate`?

I get the sentiment but my honest opinion is that it's a bit confusing to name it differently than the kernel operation that it's doing. Ultimately, the kernel is issuing a `map_lookup_batch`. My 2 cents.

Patch:
diff --git a/map.go b/map.go
index d4260a2..87a91d1 100644
--- a/map.go
+++ b/map.go
@@ -963,14 +963,19 @@ func (m *Map) guessNonExistentKey() ([]byte, error) {
 //
 // "keysOut" and "valuesOut" must be of type slice, a pointer
 // to a slice or buffer will not work.
-// "prevKey" is the key to start the batch lookup from, it will
-// *not* be included in the results. Use nil to start at the first key.
+// "cursor" is an pointer to an opaque handle. It must be non-nil. Pass
+// "cursor" to subsequent calls of this function to continue the batching
+// operation in the case of chunking.
+//
+// Warning: This API is not very safe to use as the kernel implementation for
+// batching relies on the user to be aware of subtle details with regarding to
+// different map type implementations.
 //
 // ErrKeyNotExist is returned when the batch lookup has reached
 // the end of all possible results, even when partial results
 // are returned. It should be used to evaluate when lookup is "done".
-func (m *Map) BatchLookup(prevKey, nextKeyOut, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
-	return m.batchLookup(sys.BPF_MAP_LOOKUP_BATCH, prevKey, nextKeyOut, keysOut, valuesOut, opts)
+func (m *Map) BatchLookup(cursor *BatchCursor, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
+	return m.batchLookup(sys.BPF_MAP_LOOKUP_BATCH, cursor, keysOut, valuesOut, opts)
 }
 
 // BatchLookupAndDelete looks up many elements in a map at once,
@@ -978,17 +983,53 @@ func (m *Map) BatchLookup(prevKey, nextKeyOut, keysOut, valuesOut interface{}, o
 // It then deletes all those elements.
 // "keysOut" and "valuesOut" must be of type slice, a pointer
 // to a slice or buffer will not work.
-// "prevKey" is the key to start the batch lookup from, it will
-// *not* be included in the results. Use nil to start at the first key.
+// "cursor" is an pointer to an opaque handle. It must be non-nil. Pass
+// "cursor" to subsequent calls of this function to continue the batching
+// operation in the case of chunking.
+//
+// Warning: This API is not very safe to use as the kernel implementation for
+// batching relies on the user to be aware of subtle details with regarding to
+// different map type implementations.
 //
 // ErrKeyNotExist is returned when the batch lookup has reached
 // the end of all possible results, even when partial results
 // are returned. It should be used to evaluate when lookup is "done".
-func (m *Map) BatchLookupAndDelete(prevKey, nextKeyOut, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
-	return m.batchLookup(sys.BPF_MAP_LOOKUP_AND_DELETE_BATCH, prevKey, nextKeyOut, keysOut, valuesOut, opts)
-}
+func (m *Map) BatchLookupAndDelete(cursor *BatchCursor, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
+	return m.batchLookup(sys.BPF_MAP_LOOKUP_AND_DELETE_BATCH, cursor, keysOut, valuesOut, opts)
+}
+
+// BatchCursor represents a starting point for a batch operation.
+type BatchCursor struct {
+	m      *Map
+	opaque []byte
+}
+
+func (m *Map) batchLookup(cmd sys.Cmd, cursor *BatchCursor, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
+	cursorLen := int(m.keySize)
+	if cursorLen < 4 {
+		// * generic_map_lookup_batch requires that batch_out is key_size bytes.
+		//   This is used by array and LPM maps.
+		//
+		// * __htab_map_lookup_and_delete_batch requires u32. This is used by the
+		//   various hash maps.
+		//
+		// Use a minimum of 4 bytes to avoid having to distinguish between the two.
+		cursorLen = 4
+	}
+
+	inBatch := cursor.opaque
+	if inBatch == nil {
+		// This is the first lookup, allocate a buffer to hold the cursor.
+		cursor.opaque = make([]byte, cursorLen)
+		cursor.m = m
+	} else if cursor.m != m {
+		// Prevent reuse of a cursor across maps. First, it's unlikely to work.
+		// Second, the maps may require different cursorLen and cursor.opaque
+		// may therefore be too short. This could lead to the kernel clobbering
+		// user space memory.
+		return 0, errors.New("a cursor may not be reused across maps")
+	}
 
-func (m *Map) batchLookup(cmd sys.Cmd, startKey, nextKeyOut, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
 	if err := haveBatchAPI(); err != nil {
 		return 0, err
 	}
@@ -1011,14 +1052,14 @@ func (m *Map) batchLookup(cmd sys.Cmd, startKey, nextKeyOut, keysOut, valuesOut
 	keyPtr := sys.NewSlicePointer(keyBuf)
 	valueBuf := make([]byte, count*int(m.fullValueSize))
 	valuePtr := sys.NewSlicePointer(valueBuf)
-	nextBuf := makeMapSyscallOutput(nextKeyOut, int(m.keySize))
 
 	attr := sys.MapLookupBatchAttr{
 		MapFd:    m.fd.Uint(),
 		Keys:     keyPtr,
 		Values:   valuePtr,
 		Count:    uint32(count),
-		OutBatch: nextBuf.Pointer(),
+		InBatch:  sys.NewSlicePointer(inBatch),
+		OutBatch: sys.NewSlicePointer(cursor.opaque),
 	}
 
 	if opts != nil {
@@ -1026,25 +1067,13 @@ func (m *Map) batchLookup(cmd sys.Cmd, startKey, nextKeyOut, keysOut, valuesOut
 		attr.Flags = opts.Flags
 	}
 
-	var err error
-	if startKey != nil {
-		attr.InBatch, err = marshalMapSyscallInput(startKey, int(m.keySize))
-		if err != nil {
-			return 0, err
-		}
-	}
-
 	_, sysErr := sys.BPF(cmd, unsafe.Pointer(&attr), unsafe.Sizeof(attr))
 	sysErr = wrapMapError(sysErr)
 	if sysErr != nil && !errors.Is(sysErr, unix.ENOENT) {
 		return 0, sysErr
 	}
 
-	err = nextBuf.Unmarshal(nextKeyOut)
-	if err != nil {
-		return 0, err
-	}
-	err = sysenc.Unmarshal(keysOut, keyBuf)
+	err := sysenc.Unmarshal(keysOut, keyBuf)
 	if err != nil {
 		return 0, err
 	}


Test patch:
diff --git a/map_test.go b/map_test.go
index 196b2b7..1b2bf70 100644
--- a/map_test.go
+++ b/map_test.go
@@ -6,7 +6,6 @@ import (
 	"math"
 	"os"
 	"path/filepath"
-	"reflect"
 	"sort"
 	"testing"
 	"unsafe"
@@ -92,146 +91,130 @@ func TestMap(t *testing.T) {
 	}
 }
 
-func TestBatchAPIArray(t *testing.T) {
+func TestMapBatch(t *testing.T) {
 	if err := haveBatchAPI(); err != nil {
 		t.Skipf("batch api not available: %v", err)
 	}
-	m, err := NewMap(&MapSpec{
+
+	contents := map[uint32]uint32{
+		0: 42, 1: 4242, 2: 23, 3: 2323,
+	}
+
+	hash, err := NewMap(&MapSpec{
+		Type:      Hash,
+		KeySize:   4,
+		ValueSize: 4,
+		// Make the map large enough to avoid ENOSPC.
+		MaxEntries: uint32(len(contents)) * 10,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer hash.Close()
+
+	array, err := NewMap(&MapSpec{
 		Type:       Array,
 		KeySize:    4,
 		ValueSize:  4,
-		MaxEntries: 2,
+		MaxEntries: 4,
 	})
 	if err != nil {
 		t.Fatal(err)
 	}
-	defer m.Close()
-
-	var (
-		nextKey      uint32
-		keys         = []uint32{0, 1}
-		values       = []uint32{42, 4242}
-		lookupKeys   = make([]uint32, 2)
-		lookupValues = make([]uint32, 2)
-		deleteKeys   = make([]uint32, 2)
-		deleteValues = make([]uint32, 2)
-	)
+	defer array.Close()
 
-	count, err := m.BatchUpdate(keys, values, nil)
-	if err != nil {
-		t.Fatalf("BatchUpdate: %v", err)
-	}
-	if count != len(keys) {
-		t.Fatalf("BatchUpdate: expected count, %d, to be %d", count, len(keys))
+	var keys, values []uint32
+	for key, value := range contents {
+		keys = append(keys, key)
+		values = append(values, value)
 	}
 
-	var v uint32
-	if err := m.Lookup(uint32(0), &v); err != nil {
-		t.Fatal("Can't lookup 0:", err)
-	}
-	if v != 42 {
-		t.Error("Want value 42, got", v)
-	}
+	for _, m := range []*Map{array, hash} {
+		t.Run(m.Type().String(), func(t *testing.T) {
+			count, err := m.BatchUpdate(keys, values, nil)
+			qt.Assert(t, err, qt.IsNil)
+			qt.Assert(t, count, qt.Equals, len(contents))
+
+			n := len(contents) / 2 // cut buffer in half
+			lookupKeys := make([]uint32, n)
+			lookupValues := make([]uint32, n)
+
+			var cursor BatchCursor
+			var total int
+			for {
+				count, err = m.BatchLookup(&cursor, lookupKeys, lookupValues, nil)
+				total += count
+				if errors.Is(err, ErrKeyNotExist) {
+					break
+				}
+				qt.Assert(t, err, qt.IsNil)
 
-	count, err = m.BatchLookup(nil, &nextKey, lookupKeys, lookupValues, nil)
-	if err != nil {
-		t.Fatalf("BatchLookup: %v", err)
-	}
-	if count != len(lookupKeys) {
-		t.Fatalf("BatchLookup: returned %d results, expected %d", count, len(lookupKeys))
-	}
-	if nextKey != lookupKeys[1] {
-		t.Fatalf("BatchLookup: expected nextKey, %d, to be the same as the lastKey returned, %d", nextKey, lookupKeys[1])
-	}
-	if !reflect.DeepEqual(keys, lookupKeys) {
-		t.Errorf("BatchUpdate and BatchLookup keys disagree: %v %v", keys, lookupKeys)
-	}
-	if !reflect.DeepEqual(values, lookupValues) {
-		t.Errorf("BatchUpdate and BatchLookup values disagree: %v %v", values, lookupValues)
-	}
+				qt.Assert(t, count, qt.Equals, len(lookupKeys))
+				for i, key := range lookupKeys {
+					value := lookupValues[i]
+					qt.Assert(t, value, qt.Equals, contents[key], qt.Commentf("value for key %d should match", key))
+				}
+			}
+			qt.Assert(t, total, qt.Equals, len(contents))
+
+			if m.Type() == Array {
+				// Array doesn't support batch delete
+				return
+			}
+
+			cursor = BatchCursor{}
+			total = 0
+			for {
+				count, err = m.BatchLookupAndDelete(&cursor, lookupKeys, lookupValues, nil)
+				total += count
+				if errors.Is(err, ErrKeyNotExist) {
+					break
+				}
+				qt.Assert(t, err, qt.IsNil)
+
+				qt.Assert(t, count, qt.Equals, len(lookupKeys))
+				for i, key := range lookupKeys {
+					value := lookupValues[i]
+					qt.Assert(t, value, qt.Equals, contents[key], qt.Commentf("value for key %d should match", key))
+				}
+			}
+			qt.Assert(t, total, qt.Equals, len(contents))
 
-	_, err = m.BatchLookupAndDelete(nil, &nextKey, deleteKeys, deleteValues, nil)
-	if !errors.Is(err, ErrNotSupported) {
-		t.Fatalf("BatchLookUpDelete: expected error %v, but got %v", ErrNotSupported, err)
+			var v uint32
+			qt.Assert(t, m.Lookup(uint32(0), &v), qt.ErrorIs, ErrKeyNotExist)
+		})
 	}
 }
 
-func TestBatchAPIHash(t *testing.T) {
-	if err := haveBatchAPI(); err != nil {
-		t.Skipf("batch api not available: %v", err)
-	}
-	m, err := NewMap(&MapSpec{
-		Type:       Hash,
+func TestMapBatchCursorReuse(t *testing.T) {
+	spec := &MapSpec{
+		Type:       Array,
 		KeySize:    4,
 		ValueSize:  4,
-		MaxEntries: 10,
-	})
-	if err != nil {
-		t.Fatal(err)
+		MaxEntries: 4,
 	}
-	defer m.Close()
-
-	var (
-		nextKey      uint32
-		keys         = []uint32{0, 1}
-		values       = []uint32{42, 4242}
-		lookupKeys   = make([]uint32, 2)
-		lookupValues = make([]uint32, 2)
-		deleteKeys   = make([]uint32, 2)
-		deleteValues = make([]uint32, 2)
-	)
 
-	count, err := m.BatchUpdate(keys, values, nil)
+	arr1, err := NewMap(spec)
 	if err != nil {
-		t.Fatalf("BatchUpdate: %v", err)
-	}
-	if count != len(keys) {
-		t.Fatalf("BatchUpdate: expected count, %d, to be %d", count, len(keys))
+		t.Fatal(err)
 	}
+	defer arr1.Close()
 
-	var v uint32
-	if err := m.Lookup(uint32(0), &v); err != nil {
-		t.Fatal("Can't lookup 0:", err)
-	}
-	if v != 42 {
-		t.Error("Want value 42, got", v)
+	arr2, err := NewMap(spec)
+	if err != nil {
+		t.Fatal(err)
 	}
+	defer arr2.Close()
 
-	count, err = m.BatchLookup(nil, &nextKey, lookupKeys, lookupValues, nil)
-	if !errors.Is(err, ErrKeyNotExist) {
-		t.Fatalf("BatchLookup: expected %v got %v", ErrKeyNotExist, err)
-	}
-	if count != len(lookupKeys) {
-		t.Fatalf("BatchLookup: returned %d results, expected %d", count, len(lookupKeys))
-	}
-	sort.Slice(lookupKeys, func(i, j int) bool { return lookupKeys[i] < lookupKeys[j] })
-	if !reflect.DeepEqual(keys, lookupKeys) {
-		t.Errorf("BatchUpdate and BatchLookup keys disagree: %v %v", keys, lookupKeys)
-	}
-	sort.Slice(lookupValues, func(i, j int) bool { return lookupValues[i] < lookupValues[j] })
-	if !reflect.DeepEqual(values, lookupValues) {
-		t.Errorf("BatchUpdate and BatchLookup values disagree: %v %v", values, lookupValues)
-	}
+	tmp := make([]uint32, 2)
 
-	count, err = m.BatchLookupAndDelete(nil, &nextKey, deleteKeys, deleteValues, nil)
-	if !errors.Is(err, ErrKeyNotExist) {
-		t.Fatalf("BatchLookupAndDelete: expected %v got %v", ErrKeyNotExist, err)
-	}
-	if count != len(deleteKeys) {
-		t.Fatalf("BatchLookupAndDelete: returned %d results, expected %d", count, len(deleteKeys))
-	}
-	sort.Slice(deleteKeys, func(i, j int) bool { return deleteKeys[i] < deleteKeys[j] })
-	if !reflect.DeepEqual(keys, deleteKeys) {
-		t.Errorf("BatchUpdate and BatchLookupAndDelete keys disagree: %v %v", keys, deleteKeys)
-	}
-	sort.Slice(deleteValues, func(i, j int) bool { return deleteValues[i] < deleteValues[j] })
-	if !reflect.DeepEqual(values, deleteValues) {
-		t.Errorf("BatchUpdate and BatchLookupAndDelete values disagree: %v %v", values, deleteValues)
-	}
+	var cursor BatchCursor
+	_, err = arr1.BatchLookup(&cursor, tmp, tmp, nil)
+	testutils.SkipIfNotSupported(t, err)
+	qt.Assert(t, err, qt.IsNil)
 
-	if err := m.Lookup(uint32(0), &v); !errors.Is(err, ErrKeyNotExist) {
-		t.Fatalf("Lookup should have failed with error, %v, instead error is %v", ErrKeyNotExist, err)
-	}
+	_, err = arr2.BatchLookup(&cursor, tmp, tmp, nil)
+	qt.Assert(t, err, qt.IsNotNil)
 }
 
 func TestMapLookupKeyTooSmall(t *testing.T) {
@@ -345,10 +328,10 @@ func TestBatchMapWithLock(t *testing.T) {
 			t.Fatalf("BatchUpdate: expected count, %d, to be %d", count, len(keys))
 		}
 
-		nextKey := uint32(0)
+		var cursor BatchCursor
 		lookupKeys := make([]uint32, 2)
 		lookupValues := make([]spinLockValue, 2)
-		count, err = m.BatchLookup(nil, &nextKey, lookupKeys, lookupValues, &BatchOptions{ElemFlags: uint64(LookupLock)})
+		count, err = m.BatchLookup(&cursor, lookupKeys, lookupValues, &BatchOptions{ElemFlags: uint64(LookupLock)})
 		if !errors.Is(err, ErrKeyNotExist) {
 			t.Fatalf("BatchLookup: %v", err)
 		}
@@ -356,10 +339,10 @@ func TestBatchMapWithLock(t *testing.T) {
 			t.Fatalf("BatchLookup: expected two keys, got %d", count)
 		}
 
-		nextKey = uint32(0)
+		cursor = BatchCursor{}
 		deleteKeys := []uint32{0, 1}
 		deleteValues := make([]spinLockValue, 2)
-		count, err = m.BatchLookupAndDelete(nil, &nextKey, deleteKeys, deleteValues, nil)
+		count, err = m.BatchLookupAndDelete(&cursor, deleteKeys, deleteValues, nil)
 		if !errors.Is(err, ErrKeyNotExist) {
 			t.Fatalf("BatchLookupAndDelete: %v", err)
 		}
@@ -2108,9 +2091,9 @@ func BenchmarkIterate(b *testing.B) {
 
 		b.ReportAllocs()
 
+		var cursor BatchCursor
 		for i := 0; i < b.N; i++ {
-			var next uint32
-			_, err := m.BatchLookup(nil, &next, k, v, nil)
+			_, err := m.BatchLookup(&cursor, k, v, nil)
 			if err != nil && !errors.Is(err, ErrKeyNotExist) {
 				b.Fatal(err)
 			}
@@ -2130,8 +2113,8 @@ func BenchmarkIterate(b *testing.B) {
 			}
 			b.StartTimer()
 
-			var next uint32
-			_, err := m.BatchLookupAndDelete(nil, &next, k, v, nil)
+			var cursor BatchCursor
+			_, err := m.BatchLookupAndDelete(&cursor, k, v, nil)
 			if err != nil && !errors.Is(err, ErrKeyNotExist) {
 				b.Fatal(err)
 			}