Repository: cilium/ebpf
Pull request #1220 (2023-11-13 01:37:02)

map: allow pre-allocating per-CPU values on lookup
Prequel to #1192 note that this *will* conflict with #1219 and need a rebase.. oh for a stacked diff workflow on GH.

lmb: > oh for a stacked diff workflow on GH.

Agreed, but it's also acceptable to just push multiple commits next time ;)
lmb: FYI I'm sending you an invite to @cilium/ebpf-go-contributors which should allow you to request re-review and the like. No pressure to accept the invitation of course :)

Patch:
diff --git a/map.go b/map.go
index fec0adf..dc25245 100644
--- a/map.go
+++ b/map.go
@@ -642,11 +642,15 @@ func (m *Map) LookupBytes(key interface{}) ([]byte, error) {
 }
 
 func (m *Map) lookupPerCPU(key, valueOut any, flags MapLookupFlags) error {
+	slice, err := ensurePerCPUSlice(valueOut, int(m.valueSize))
+	if err != nil {
+		return err
+	}
 	valueBytes := make([]byte, m.fullValueSize)
 	if err := m.lookup(key, sys.NewSlicePointer(valueBytes), flags); err != nil {
 		return err
 	}
-	return unmarshalPerCPUValue(valueOut, int(m.valueSize), valueBytes)
+	return unmarshalPerCPUValue(slice, int(m.valueSize), valueBytes)
 }
 
 func (m *Map) lookup(key interface{}, valueOut sys.Pointer, flags MapLookupFlags) error {
@@ -669,11 +673,53 @@ func (m *Map) lookup(key interface{}, valueOut sys.Pointer, flags MapLookupFlags
 }
 
 func (m *Map) lookupAndDeletePerCPU(key, valueOut any, flags MapLookupFlags) error {
+	slice, err := ensurePerCPUSlice(valueOut, int(m.valueSize))
+	if err != nil {
+		return err
+	}
 	valueBytes := make([]byte, m.fullValueSize)
 	if err := m.lookupAndDelete(key, sys.NewSlicePointer(valueBytes), flags); err != nil {
 		return err
 	}
-	return unmarshalPerCPUValue(valueOut, int(m.valueSize), valueBytes)
+	return unmarshalPerCPUValue(slice, int(m.valueSize), valueBytes)
+}
+
+// ensurePerCPUSlice allocates a slice for a per-CPU value if necessary.
+func ensurePerCPUSlice(sliceOrPtr any, elemLength int) (any, error) {
+	sliceOrPtrType := reflect.TypeOf(sliceOrPtr)
+	if sliceOrPtrType.Kind() == reflect.Slice {
+		// The target is a slice, the caller is responsible for ensuring that
+		// size is correct.
+		return sliceOrPtr, nil
+	}
+
+	slicePtrType := sliceOrPtrType
+	if slicePtrType.Kind() != reflect.Ptr || slicePtrType.Elem().Kind() != reflect.Slice {
+		return nil, fmt.Errorf("per-cpu value requires a slice or a pointer to slice")
+	}
+
+	possibleCPUs, err := PossibleCPU()
+	if err != nil {
+		return nil, err
+	}
+
+	sliceType := slicePtrType.Elem()
+	slice := reflect.MakeSlice(sliceType, possibleCPUs, possibleCPUs)
+
+	sliceElemType := sliceType.Elem()
+	sliceElemIsPointer := sliceElemType.Kind() == reflect.Ptr
+	reflect.ValueOf(sliceOrPtr).Elem().Set(slice)
+	if !sliceElemIsPointer {
+		return slice.Interface(), nil
+	}
+	sliceElemType = sliceElemType.Elem()
+
+	for i := 0; i < possibleCPUs; i++ {
+		newElem := reflect.New(sliceElemType)
+		slice.Index(i).Set(newElem)
+	}
+
+	return slice.Interface(), nil
 }
 
 func (m *Map) lookupAndDelete(key any, valuePtr sys.Pointer, flags MapLookupFlags) error {
diff --git a/marshalers.go b/marshalers.go
index 00911ad..d77a5fb 100644
--- a/marshalers.go
+++ b/marshalers.go
@@ -84,11 +84,11 @@ func marshalPerCPUValue(slice any, elemLength int) (sys.Pointer, error) {
 // unmarshalPerCPUValue decodes a buffer into a slice containing one value per
 // possible CPU.
 //
-// slicePtr must be a pointer to a slice.
-func unmarshalPerCPUValue(slicePtr any, elemLength int, buf []byte) error {
-	slicePtrType := reflect.TypeOf(slicePtr)
-	if slicePtrType.Kind() != reflect.Ptr || slicePtrType.Elem().Kind() != reflect.Slice {
-		return fmt.Errorf("per-cpu value requires pointer to slice")
+// slice must be a literal slice and not a pointer.
+func unmarshalPerCPUValue(slice any, elemLength int, buf []byte) error {
+	sliceType := reflect.TypeOf(slice)
+	if sliceType.Kind() != reflect.Slice {
+		return fmt.Errorf("per-CPU value requires a slice")
 	}
 
 	possibleCPUs, err := PossibleCPU()
@@ -96,26 +96,26 @@ func unmarshalPerCPUValue(slicePtr any, elemLength int, buf []byte) error {
 		return err
 	}
 
-	sliceType := slicePtrType.Elem()
-	slice := reflect.MakeSlice(sliceType, possibleCPUs, possibleCPUs)
+	sliceValue := reflect.ValueOf(slice)
+	if sliceValue.Len() != possibleCPUs {
+		return fmt.Errorf("per-CPU slice has incorrect length, expected %d, got %d",
+			possibleCPUs, sliceValue.Len())
+	}
 
 	sliceElemType := sliceType.Elem()
 	sliceElemIsPointer := sliceElemType.Kind() == reflect.Ptr
-	if sliceElemIsPointer {
-		sliceElemType = sliceElemType.Elem()
-	}
-
 	stride := internal.Align(elemLength, 8)
 	for i := 0; i < possibleCPUs; i++ {
 		var elem any
+		v := sliceValue.Index(i)
 		if sliceElemIsPointer {
-			newElem := reflect.New(sliceElemType)
-			slice.Index(i).Set(newElem)
-			elem = newElem.Interface()
+			if !v.Elem().CanAddr() {
+				return fmt.Errorf("per-CPU slice elements cannot be nil")
+			}
+			elem = v.Elem().Addr().Interface()
 		} else {
-			elem = slice.Index(i).Addr().Interface()
+			elem = v.Addr().Interface()
 		}
-
 		err := sysenc.Unmarshal(elem, buf[:elemLength])
 		if err != nil {
 			return fmt.Errorf("cpu %d: %w", i, err)
@@ -124,6 +124,5 @@ func unmarshalPerCPUValue(slicePtr any, elemLength int, buf []byte) error {
 		buf = buf[stride:]
 	}
 
-	reflect.ValueOf(slicePtr).Elem().Set(slice)
 	return nil
 }


Test patch:
diff --git a/map_test.go b/map_test.go
index 8c9710d..e18a1ac 100644
--- a/map_test.go
+++ b/map_test.go
@@ -75,6 +75,10 @@ func TestMap(t *testing.T) {
 		t.Error("Want value 42, got", v)
 	}
 
+	sliceVal := make([]uint32, 1)
+	qt.Assert(t, m.Lookup(uint32(0), sliceVal), qt.IsNil)
+	qt.Assert(t, sliceVal, qt.DeepEquals, []uint32{42})
+
 	var slice []byte
 	qt.Assert(t, m.Lookup(uint32(0), &slice), qt.IsNil)
 	qt.Assert(t, slice, qt.DeepEquals, internal.NativeEndian.AppendUint32(nil, 42))
@@ -1345,6 +1349,16 @@ func TestPerCPUMarshaling(t *testing.T) {
 			}
 
 			// Make sure unmarshaling works on slices containing pointers
+			retrievedVal := make([]*customEncoding, numCPU)
+			if err := arr.Lookup(uint32(0), retrievedVal); err == nil {
+				t.Fatal("Slices with nil values should generate error")
+			}
+			for i := range retrievedVal {
+				retrievedVal[i] = &customEncoding{}
+			}
+			if err := arr.Lookup(uint32(0), retrievedVal); err != nil {
+				t.Fatal("Can't retrieve key 0:", err)
+			}
 			var retrieved []*customEncoding
 			if err := arr.Lookup(uint32(0), &retrieved); err != nil {
 				t.Fatal("Can't retrieve key 0:", err)
diff --git a/marshalers_test.go b/marshalers_test.go
new file mode 100644
index 0000000..6695791
--- /dev/null
+++ b/marshalers_test.go
@@ -0,0 +1,37 @@
+package ebpf
+
+import (
+	"testing"
+
+	"github.com/cilium/ebpf/internal"
+
+	qt "github.com/frankban/quicktest"
+)
+
+func TestUnmarshalPerCPUValue(t *testing.T) {
+	possibleCPUs := MustPossibleCPU()
+	expected := make([]uint32, possibleCPUs)
+	for i := 0; i < possibleCPUs; i++ {
+		expected[i] = uint32(1021 * (i + 1))
+	}
+	elemLength := 4
+
+	buf := make([]byte, possibleCPUs*internal.Align(elemLength, 8))
+	b := buf
+	for _, elem := range expected {
+		internal.NativeEndian.PutUint32(b, elem)
+		b = b[8:]
+	}
+	slice := make([]uint32, possibleCPUs)
+	err := unmarshalPerCPUValue(slice, elemLength, buf)
+	if err != nil {
+		t.Fatal(err)
+	}
+	qt.Assert(t, slice, qt.DeepEquals, expected)
+
+	smallSlice := make([]uint32, possibleCPUs-1)
+	qt.Assert(t, unmarshalPerCPUValue(smallSlice, elemLength, buf), qt.IsNotNil)
+
+	nilElemSlice := make([]*uint32, possibleCPUs)
+	qt.Assert(t, unmarshalPerCPUValue(nilElemSlice, elemLength, buf), qt.IsNotNil)
+}