Add per-CPU Map Support to Batch Operations

Makefile

@@ -103,14 +103,6 @@ testdata/loader-%-eb.elf: testdata/loader.c
 	$(STRIP) -g $@
 
 .PHONY: update-kernel-deps
-update-kernel-deps: KERNEL_VERSION?=6.6
+update-kernel-deps: export KERNEL_VERSION?=6.6
 update-kernel-deps:
-	$(eval TMP := $(shell mktemp -d))
-	curl -fL https://git.kernel.org/pub/scm/linux/kernel/git/stable/linux.git/plain/tools/lib/bpf/libbpf.c?h=v$(KERNEL_VERSION) -o "$(TMP)/libbpf.c"
-	"$(CURDIR)/internal/cmd/gensections.awk" "$(TMP)/libbpf.c" | gofmt > "$(CURDIR)/elf_sections.go"
-	curl -fL "$(CI_KERNEL_URL)/linux-$(KERNEL_VERSION)-amd64.tgz" -o "$(TMP)/linux.tgz"
-	tar xvf "$(TMP)/linux.tgz" -C "$(TMP)" --strip-components=2 ./boot/vmlinuz ./lib/modules
-	/lib/modules/$(shell uname -r)/build/scripts/extract-vmlinux "$(TMP)/vmlinuz" > "$(TMP)/vmlinux"
-	$(OBJCOPY) --dump-section .BTF=/dev/stdout "$(TMP)/vmlinux" /dev/null | gzip > "btf/testdata/vmlinux.btf.gz"
-	find "$(TMP)/modules" -type f -name bpf_testmod.ko -exec $(OBJCOPY) --dump-section .BTF="btf/testdata/btf_testmod.btf" {} /dev/null \;
-	$(RM) -r "$(TMP)"
+	./testdata/sh/update-kernel-deps.sh

internal/sysenc/buffer.go

@@ -54,6 +54,11 @@ func (b Buffer) CopyTo(dst []byte) int {
 	return copy(dst, b.unsafeBytes())
 }
 
+// AppendTo appends the buffer onto dst.
+func (b Buffer) AppendTo(dst []byte) []byte {
+	return append(dst, b.unsafeBytes()...)
+}
+
 // Pointer returns the location where a syscall should write.
 func (b Buffer) Pointer() sys.Pointer {
 	// NB: This deliberately ignores b.length to support zero-copy

map.go

@@ -1005,6 +1005,53 @@ type BatchCursor struct {
 }
 
 func (m *Map) batchLookup(cmd sys.Cmd, cursor *BatchCursor, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
+	if m.typ.hasPerCPUValue() {
+		return m.batchLookupPerCPU(cmd, cursor, keysOut, valuesOut, opts)
+	}
+
+	count, err := batchCount(keysOut, valuesOut)
+	if err != nil {
+		return 0, err
+	}
+
+	valueBuf := sysenc.SyscallOutput(valuesOut, count*int(m.fullValueSize))
+
+	n, err := m.batchLookupCmd(cmd, cursor, count, keysOut, valueBuf.Pointer(), opts)
+	if err != nil {
+		return n, err
+	}
+
+	err = valueBuf.Unmarshal(valuesOut)
+	if err != nil {
+		return 0, err
+	}
+
+	return n, nil
+}
+
+func (m *Map) batchLookupPerCPU(cmd sys.Cmd, cursor *BatchCursor, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
+	count, err := sliceLen(keysOut)
+	if err != nil {
+		return 0, fmt.Errorf("keys: %w", err)
+	}
+
+	valueBuf := make([]byte, count*int(m.fullValueSize))
+	valuePtr := sys.NewSlicePointer(valueBuf)
+
+	n, sysErr := m.batchLookupCmd(cmd, cursor, count, keysOut, valuePtr, opts)
+	if sysErr != nil && !errors.Is(sysErr, unix.ENOENT) {
+		return 0, err
+	}
+
+	err = unmarshalBatchPerCPUValue(valuesOut, count, int(m.valueSize), valueBuf)
+	if err != nil {
+		return 0, err
+	}
+
+	return n, sysErr
+}
+
+func (m *Map) batchLookupCmd(cmd sys.Cmd, cursor *BatchCursor, count int, keysOut any, valuePtr sys.Pointer, opts *BatchOptions) (int, error) {
 	cursorLen := int(m.keySize)
 	if cursorLen < 4 {
 		// * generic_map_lookup_batch requires that batch_out is key_size bytes.
@@ -1033,29 +1080,13 @@ func (m *Map) batchLookup(cmd sys.Cmd, cursor *BatchCursor, keysOut, valuesOut i
 	if err := haveBatchAPI(); err != nil {
 		return 0, err
 	}
-	if m.typ.hasPerCPUValue() {
-		return 0, ErrNotSupported
-	}
-	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 := sysenc.SyscallOutput(keysOut, count*int(m.keySize))
-	valueBuf := sysenc.SyscallOutput(valuesOut, count*int(m.fullValueSize))
 
 	attr := sys.MapLookupBatchAttr{
 		MapFd:    m.fd.Uint(),
 		Keys:     keyBuf.Pointer(),
-		Values:   valueBuf.Pointer(),
+		Values:   valuePtr,
 		Count:    uint32(count),
 		InBatch:  sys.NewSlicePointer(inBatch),
 		OutBatch: sys.NewSlicePointer(cursor.opaque),
@@ -1075,9 +1106,6 @@ func (m *Map) batchLookup(cmd sys.Cmd, cursor *BatchCursor, keysOut, valuesOut i
 	if err := keyBuf.Unmarshal(keysOut); err != nil {
 		return 0, err
 	}
-	if err := valueBuf.Unmarshal(valuesOut); err != nil {
-		return 0, err
-	}
 
 	return int(attr.Count), sysErr
 }
@@ -1088,29 +1116,24 @@ func (m *Map) batchLookup(cmd sys.Cmd, cursor *BatchCursor, keysOut, valuesOut i
 // to a slice or buffer will not work.
 func (m *Map) BatchUpdate(keys, values interface{}, opts *BatchOptions) (int, error) {
 	if m.typ.hasPerCPUValue() {
-		return 0, ErrNotSupported
+		return m.batchUpdatePerCPU(keys, values, opts)
 	}
-	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")
-	}
-	var (
-		count    = keysValue.Len()
-		valuePtr sys.Pointer
-		err      error
-	)
-	if count != valuesValue.Len() {
-		return 0, fmt.Errorf("keys and values must be the same length")
+
+	count, err := batchCount(keys, values)
+	if err != nil {
+		return 0, err
 	}
-	keyPtr, err := marshalMapSyscallInput(keys, count*int(m.keySize))
+
+	valuePtr, err := marshalMapSyscallInput(values, count*int(m.valueSize))
 	if err != nil {
 		return 0, err
 	}
-	valuePtr, err = marshalMapSyscallInput(values, count*int(m.valueSize))
+
+	return m.batchUpdate(count, keys, valuePtr, opts)
+}
+
+func (m *Map) batchUpdate(count int, keys any, valuePtr sys.Pointer, opts *BatchOptions) (int, error) {
+	keyPtr, err := marshalMapSyscallInput(keys, count*int(m.keySize))
 	if err != nil {
 		return 0, err
 	}
@@ -1137,17 +1160,28 @@ func (m *Map) BatchUpdate(keys, values interface{}, opts *BatchOptions) (int, er
 	return int(attr.Count), nil
 }
 
+func (m *Map) batchUpdatePerCPU(keys, values any, opts *BatchOptions) (int, error) {
+	count, err := sliceLen(keys)
+	if err != nil {
+		return 0, fmt.Errorf("keys: %w", err)
+	}
+
+	valueBuf, err := marshalBatchPerCPUValue(values, count, int(m.valueSize))
+	if err != nil {
+		return 0, err
+	}
+
+	return m.batchUpdate(count, keys, sys.NewSlicePointer(valueBuf), opts)
+}
+
 // BatchDelete batch deletes entries in the map by keys.
 // "keys" must be of type slice, a pointer to a slice or buffer will not work.
 func (m *Map) BatchDelete(keys interface{}, opts *BatchOptions) (int, error) {
-	if m.typ.hasPerCPUValue() {
-		return 0, ErrNotSupported
-	}
-	keysValue := reflect.ValueOf(keys)
-	if keysValue.Kind() != reflect.Slice {
-		return 0, fmt.Errorf("keys must be a slice")
+	count, err := sliceLen(keys)
+	if err != nil {
+		return 0, fmt.Errorf("keys: %w", err)
 	}
-	count := keysValue.Len()
+
 	keyPtr, err := marshalMapSyscallInput(keys, count*int(m.keySize))
 	if err != nil {
 		return 0, fmt.Errorf("cannot marshal keys: %v", err)
@@ -1174,6 +1208,24 @@ func (m *Map) BatchDelete(keys interface{}, opts *BatchOptions) (int, error) {
 	return int(attr.Count), nil
 }
 
+func batchCount(keys, values any) (int, error) {
+	keysLen, err := sliceLen(keys)
+	if err != nil {
+		return 0, fmt.Errorf("keys: %w", err)
+	}
+
+	valuesLen, err := sliceLen(values)
+	if err != nil {
+		return 0, fmt.Errorf("values: %w", err)
+	}
+
+	if keysLen != valuesLen {
+		return 0, fmt.Errorf("keys and values must have the same length")
+	}
+
+	return keysLen, nil
+}
+
 // Iterate traverses a map.
 //
 // It's safe to create multiple iterators at the same time.
@@ -1552,3 +1604,12 @@ func NewMapFromID(id MapID) (*Map, error) {
 
 	return newMapFromFD(fd)
 }
+
+// sliceLen returns the length if the value is a slice or an error otherwise.
+func sliceLen(slice any) (int, error) {
+	sliceValue := reflect.ValueOf(slice)
+	if sliceValue.Kind() != reflect.Slice {
+		return 0, fmt.Errorf("%T is not a slice", slice)
+	}
+	return sliceValue.Len(), nil
+}