cpupick: Move AvoidCPUs from util into cpupick as Avoid

cpupick/avoid_linux_test.go

@@ -0,0 +1,98 @@
+//go:build linux && !android && !e2e_testing
+
+package cpupick
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"runtime"
+	"strconv"
+	"sync"
+	"testing"
+	"unsafe"
+
+	"github.com/slackhq/nebula/util"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sys/unix"
+)
+
+// TestAvoid changes the affinity of the whole process, so it runs in a child, leaving the test runner untouched.
+func TestAvoid(t *testing.T) {
+	if os.Getenv("NEBULA_CPUPICK_AVOID_TEST") != "1" {
+		exe, err := os.Executable()
+		require.NoError(t, err)
+		cmd := exec.Command(exe, "-test.run=^TestAvoid$", "-test.v")
+		cmd.Env = append(os.Environ(), "NEBULA_CPUPICK_AVOID_TEST=1")
+		out, err := cmd.CombinedOutput()
+		require.NoError(t, err, "%s", out)
+		return
+	}
+	allowed, err := util.AllowedCPUs()
+	require.NoError(t, err)
+	if len(allowed) < 2 {
+		t.Skip("requires at least two allowed CPUs")
+	}
+	cpu := allowed[0]
+	release := make(chan struct{})
+	var wg sync.WaitGroup
+	// lockThreads occupies 8 OS threads until release, so the runtime has to start more
+	lockThreads := func() {
+		var ready sync.WaitGroup
+		for range 8 {
+			ready.Add(1)
+			wg.Go(func() {
+				runtime.LockOSThread()
+				ready.Done()
+				<-release
+			})
+		}
+		ready.Wait()
+	}
+	lockThreads()
+
+	require.NoError(t, Avoid([]int{cpu}, false))
+	after, err := util.AllowedCPUs()
+	require.NoError(t, err)
+	require.Equal(t, allowed, after, "configuration check must not mutate affinity")
+	require.Error(t, Avoid(allowed, true), "must reject an empty mask before applying")
+	require.NoError(t, Avoid([]int{cpu, cpu}, true))
+	requireAvoided(t, cpu, allowed)
+	// Existing threads stay locked while new ones inherit the restriction
+	lockThreads()
+	requireAvoided(t, cpu, allowed)
+	close(release)
+	wg.Wait()
+}
+
+// requireAvoided checks no thread may run on cpu, and every thread still may on some CPU of allowed only.
+func requireAvoided(t *testing.T, cpu int, allowed []int) {
+	t.Helper()
+	entries, err := os.ReadDir("/proc/self/task")
+	require.NoError(t, err)
+	for _, entry := range entries {
+		tid, err := strconv.Atoi(entry.Name())
+		require.NoError(t, err)
+		var mask unix.CPUSet
+		err = unix.SchedGetaffinity(tid, &mask)
+		if err == unix.ESRCH {
+			continue
+		}
+		require.NoError(t, err)
+		require.False(t, mask.IsSet(cpu), "thread %d", tid)
+		require.Positive(t, mask.Count())
+		for id := range len(mask) * 64 {
+			if mask.IsSet(id) {
+				require.Contains(t, allowed, id, "must not widen the inherited mask")
+			}
+		}
+	}
+}
+
+// TestAvoidRange checks the bounds of unix.CPUSet, whose words are 32 bits wide on 32-bit platforms.
+func TestAvoidRange(t *testing.T) {
+	capacity := int(unsafe.Sizeof(unix.CPUSet{})) * 8
+	require.NoError(t, Avoid([]int{capacity - 1}, false))
+	require.EqualError(t, Avoid([]int{capacity}, false), fmt.Sprintf("CPU %d is out of range", capacity))
+	require.EqualError(t, Avoid([]int{-1}, false), "CPU -1 is out of range")
+}