From bbb2c07fbd9ccaf593c5d463fa14d95808df9686 Mon Sep 17 00:00:00 2001
From: Lorenz Bauer <lmb@isovalent.com>
Date: Tue, 9 Apr 2024 14:09:03 +0200
Subject: [PATCH] perf: fix TestPerfReaderWakeupEvents

The WakeupEvents limit is per ring, one per CPU. And when we execute
BPF_PROG_RUN multiple times, we sometimes execute on different CPUs.
This means that samples don't all go into a single ring.

Fix this by forcing the producer to run on a single CPU.

Fixes https://github.com/cilium/ebpf/issues/1419

Co-developed-by: Dylan Reimerink <dylan@isovalent.com>
Signed-off-by: Lorenz Bauer <lmb@isovalent.com>
---
 internal/testutils/cpu.go    | 32 ++++++++++++++++++++++++++++++++
 internal/unix/types_linux.go |  9 +++++++++
 internal/unix/types_other.go | 12 ++++++++++++
 perf/reader_test.go          |  2 ++
 4 files changed, 55 insertions(+)
 create mode 100644 internal/testutils/cpu.go

diff --git a/internal/testutils/cpu.go b/internal/testutils/cpu.go
new file mode 100644
index 000000000..6ce55eaa5
--- /dev/null
+++ b/internal/testutils/cpu.go
@@ -0,0 +1,32 @@
+package testutils
+
+import (
+	"runtime"
+	"testing"
+
+	"github.com/cilium/ebpf/internal/unix"
+
+	"github.com/go-quicktest/qt"
+)
+
+// LockOSThreadToSingleCPU force the current goroutine to run on a single CPU.
+func LockOSThreadToSingleCPU(tb testing.TB) {
+	tb.Helper()
+
+	runtime.LockOSThread()
+	tb.Cleanup(runtime.UnlockOSThread)
+
+	var old unix.CPUSet
+	err := unix.SchedGetaffinity(0, &old)
+	qt.Assert(tb, qt.IsNil(err))
+
+	// Schedule test to run on only CPU 0
+	var first unix.CPUSet
+	first.Set(0)
+	err = unix.SchedSetaffinity(0, &first)
+	qt.Assert(tb, qt.IsNil(err))
+
+	tb.Cleanup(func() {
+		_ = unix.SchedSetaffinity(0, &old)
+	})
+}
diff --git a/internal/unix/types_linux.go b/internal/unix/types_linux.go
index 2f22b1278..025774f72 100644
--- a/internal/unix/types_linux.go
+++ b/internal/unix/types_linux.go
@@ -100,6 +100,7 @@ type PerfEventMmapPage = linux.PerfEventMmapPage
 type EpollEvent = linux.EpollEvent
 type PerfEventAttr = linux.PerfEventAttr
 type Utsname = linux.Utsname
+type CPUSet = linux.CPUSet
 
 func Syscall(trap, a1, a2, a3 uintptr) (r1, r2 uintptr, err syscall.Errno) {
 	return linux.Syscall(trap, a1, a2, a3)
@@ -204,3 +205,11 @@ func Fstat(fd int, stat *Stat_t) error {
 func SetsockoptInt(fd, level, opt, value int) error {
 	return linux.SetsockoptInt(fd, level, opt, value)
 }
+
+func SchedSetaffinity(pid int, set *CPUSet) error {
+	return linux.SchedSetaffinity(pid, set)
+}
+
+func SchedGetaffinity(pid int, set *CPUSet) error {
+	return linux.SchedGetaffinity(pid, set)
+}
diff --git a/internal/unix/types_other.go b/internal/unix/types_other.go
index e5bad0469..807037181 100644
--- a/internal/unix/types_other.go
+++ b/internal/unix/types_other.go
@@ -296,3 +296,15 @@ func Fstat(fd int, stat *Stat_t) error {
 func SetsockoptInt(fd, level, opt, value int) error {
 	return errNonLinux
 }
+
+type CPUSet struct{}
+
+func (*CPUSet) Set(int) {}
+
+func SchedSetaffinity(pid int, set *CPUSet) error {
+	return errNonLinux
+}
+
+func SchedGetaffinity(pid int, set *CPUSet) error {
+	return errNonLinux
+}
diff --git a/perf/reader_test.go b/perf/reader_test.go
index c2c67da0e..bb2988ac4 100644
--- a/perf/reader_test.go
+++ b/perf/reader_test.go
@@ -490,6 +490,8 @@ func TestPause(t *testing.T) {
 }
 
 func TestPerfReaderWakeupEvents(t *testing.T) {
+	testutils.LockOSThreadToSingleCPU(t)
+
 	events := perfEventArray(t)
 
 	numEvents := 2