bpf: Add Batch Methdods

As of kernel v5.6 batch methods allow for the fast lookup,
deletion, and updating of bpf maps so that the syscall
overhead (repeatedly calling into any of these methods)
can be avoided.

The batch methods are as follows:
 * BatchUpdate
 * BatchLookup
 * BatchLookupAndDelete
 * BatchDelete

Only the "array" and "hash" types currently support
batch operations, and the "array" type does not support
batch deletion.

Tests are in place to test every scenario and helper
functions have been written to catch errors that
normally the kernel would give to helpful to users
of the library.

Signed-off-by: Nate Sweet <nathanjsweet@pm.me>

map.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"path/filepath"
+	"reflect"
 	"strings"
 
 	"github.com/cilium/ebpf/internal"
@@ -17,6 +18,7 @@ var (
 	ErrKeyNotExist      = errors.New("key does not exist")
 	ErrKeyExist         = errors.New("key already exists")
 	ErrIterationAborted = errors.New("iteration aborted")
+	ErrBatchOpNotSup    = errors.New("batch operations not supported for this map type")
 )
 
 // MapOptions control loading a map into the kernel.
@@ -579,6 +581,127 @@ func (m *Map) nextKey(key interface{}, nextKeyOut internal.Pointer) error {
 	return nil
 }
 
+// BatchLookup looks up many elements in a map at once
+// with the startKey being the first element to start
+// from.
+func (m *Map) BatchLookup(startKey, nextKey, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
+	return m.batchLookup(false, startKey, nextKey, keysOut, valuesOut, opts)
+}
+
+// BatchLookup looks up many elements in a map at once
+// with the startKey being the first element to start
+// from. It then deletes all those elements.
+func (m *Map) BatchLookupAndDelete(startKey, nextKey, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
+	return m.batchLookup(true, startKey, nextKey, keysOut, valuesOut, opts)
+}
+
+func (m *Map) batchLookup(del bool, startKey, nextKey, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
+	if !m.typ.canBatch() || (del && !m.typ.canBatchDelete()) {
+		return 0, ErrBatchOpNotSup
+	}
+	keysValue := reflect.ValueOf(keysOut)
+	if keysValue.Kind() != reflect.Slice {
+		return 0, fmt.Errorf("keys must be a slice")
+	}
+	valuesValue := reflect.ValueOf(valuesOut)
+	if valuesValue.Kind() != reflect.Slice {
+		return 0, fmt.Errorf("valuesOut must be a slice")
+	}
+	count := keysValue.Len()
+	if count != valuesValue.Len() {
+		return 0, fmt.Errorf("keysOut and valuesOut must be the same length")
+	}
+	keyBuf := make([]byte, count*int(m.keySize))
+	keyPtr := internal.NewSlicePointer(keyBuf)
+	valueBuf := make([]byte, count*int(m.valueSize))
+	valuePtr := internal.NewSlicePointer(valueBuf)
+
+	var (
+		startPtr internal.Pointer
+		err      error
+	)
+	if startKey != nil {
+		startPtr, err = marshalPtr(startKey, int(m.keySize))
+		if err != nil {
+			return 0, err
+		}
+	} else {
+		startPtr = internal.NewPointer(nil)
+	}
+
+	nextPtr, nextBuf := makeBuffer(nextKey, int(m.keySize))
+	ct := uint32(count)
+	if del {
+		err = bpfMapBatchLookupAndDelete(m.fd, startPtr, nextPtr, keyPtr, valuePtr, &ct, opts)
+	} else {
+		err = bpfMapBatchLookup(m.fd, startPtr, nextPtr, keyPtr, valuePtr, &ct, opts)
+	}
+	if err != nil && !errors.Is(err, ErrKeyNotExist) {
+		return 0, err
+	}
+
+	err = m.unmarshalKey(nextKey, nextBuf)
+	if err != nil {
+		return 0, err
+	}
+	err = unmarshalBytes(keysOut, keyBuf)
+	if err != nil {
+		return 0, err
+	}
+	count = int(ct)
+	return count, unmarshalBytes(valuesOut, valueBuf)
+}
+
+// BatchUpdate updates the map with multiple keys and values
+// simultaneously.
+func (m *Map) BatchUpdate(keys, values interface{}, opts *BatchOptions) (int, error) {
+	if !m.typ.canBatch() {
+		return 0, ErrBatchOpNotSup
+	}
+	keysValue := reflect.ValueOf(keys)
+	if keysValue.Kind() != reflect.Slice {
+		return 0, fmt.Errorf("keys must be a slice")
+	}
+	valuesValue := reflect.ValueOf(values)
+	if valuesValue.Kind() != reflect.Slice {
+		return 0, fmt.Errorf("values must be a slice")
+	}
+	count := keysValue.Len()
+	if count != valuesValue.Len() {
+		return 0, fmt.Errorf("keys and values must be the same length")
+	}
+	keyPtr, err := marshalPtr(keys, count*int(m.keySize))
+	if err != nil {
+		return 0, fmt.Errorf("cannot marshal keys: %v", err)
+	}
+	valuePtr, err := marshalPtr(values, count*int(m.valueSize))
+	if err != nil {
+		return 0, fmt.Errorf("cannot marshal keys: %v", err)
+	}
+	ct := uint32(count)
+	err = bpfMapBatchUpdate(m.fd, keyPtr, valuePtr, &ct, opts)
+	return int(ct), err
+}
+
+// BatchDelete batch deletes entries in the map by keys
+func (m *Map) BatchDelete(keys, opts *BatchOptions) (int, error) {
+	if !m.typ.canBatchDelete() {
+		return 0, ErrBatchOpNotSup
+	}
+	keysValue := reflect.ValueOf(keys)
+	if keysValue.Kind() != reflect.Slice {
+		return 0, fmt.Errorf("keys must be a slice")
+	}
+	count := keysValue.Len()
+	keyPtr, err := marshalPtr(keys, count*int(m.keySize))
+	if err != nil {
+		return 0, fmt.Errorf("cannot marshal keys: %v", err)
+	}
+	ct := uint32(count)
+	err = bpfMapBatchDelete(m.fd, keyPtr, &ct, opts)
+	return int(ct), err
+}
+
 // Iterate traverses a map.
 //
 // It's safe to create multiple iterators at the same time.

map_test.go

@@ -7,6 +7,7 @@ import (
 	"math"
 	"os"
 	"path/filepath"
+	"reflect"
 	"sort"
 	"strings"
 	"testing"
@@ -68,6 +69,148 @@ func TestMap(t *testing.T) {
 	}
 }
 
+func TestBatchAPIArray(t *testing.T) {
+	if err := hasBatchAPI(); err != nil {
+		t.Skipf("batch api not available: %v", err)
+	}
+	m, err := NewMap(&MapSpec{
+		Type:       Array,
+		KeySize:    4,
+		ValueSize:  4,
+		MaxEntries: 10,
+	})
+	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)
+	)
+
+	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 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)
+	}
+
+	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)
+	}
+
+	_, err = m.BatchLookupAndDelete(nil, &nextKey, deleteKeys, deleteValues, nil)
+	if !errors.Is(err, ErrBatchOpNotSup) {
+		t.Fatalf("BatchLookUpDelete: expected error %v, but got %v", ErrBatchOpNotSup, err)
+	}
+}
+
+func TestBatchAPIHash(t *testing.T) {
+	if err := hasBatchAPI(); err != nil {
+		t.Skipf("batch api not available: %v", err)
+	}
+	m, err := NewMap(&MapSpec{
+		Type:       Hash,
+		KeySize:    4,
+		ValueSize:  4,
+		MaxEntries: 10,
+	})
+	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)
+	)
+
+	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 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)
+	}
+
+	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))
+	}
+	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)
+	}
+
+	count, err = m.BatchLookupAndDelete(nil, &nextKey, deleteKeys, deleteValues, nil)
+	if err != nil {
+		t.Fatalf("BatchLookupAndDelete: %v", 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)
+	}
+
+	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)
+	}
+}
+
 func TestMapClose(t *testing.T) {
 	m := createArray(t)
 

run-tests.sh

@@ -2,7 +2,7 @@
 # Test the current package under a different kernel.
 # Requires virtme and qemu to be installed.
 
-set -eu
+set -eux
 set -o pipefail
 
 if [[ "${1:-}" = "--in-vm" ]]; then