cpupick: Move AvoidCPUs from util into cpupick as Avoid

cpupick/avoid_linux.go

@@ -0,0 +1,75 @@
+//go:build linux && !android && !e2e_testing
+
+package cpupick
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"strconv"
+	"unsafe"
+
+	"golang.org/x/sys/unix"
+)
+
+// Avoid removes cpus from the affinity of every thread of the process, for use at startup. New threads inherit
+// their parent's mask, so it repeats until a pass finds nothing left to change; a thread that exits meanwhile is
+// skipped. Every mask is checked before any is changed. With apply false it only validates, for configuration tests.
+// This affects the entire process, including other users of an embedded Nebula. Run it before Default, which then
+// picks from what remains.
+func Avoid(cpus []int, apply bool) error {
+	const capacity = int(unsafe.Sizeof(unix.CPUSet{})) * 8
+	for _, cpu := range cpus {
+		if cpu < 0 || cpu >= capacity {
+			return fmt.Errorf("CPU %d is out of range", cpu)
+		}
+	}
+	if len(cpus) == 0 {
+		return nil
+	}
+	for {
+		masks, err := avoidingMasks(cpus)
+		if err != nil || !apply || len(masks) == 0 {
+			return err
+		}
+		for tid, mask := range masks {
+			if err := unix.SchedSetaffinity(tid, &mask); err != nil && !errors.Is(err, unix.ESRCH) {
+				return fmt.Errorf("set thread %d affinity: %w", tid, err)
+			}
+		}
+	}
+}
+
+// avoidingMasks returns the mask without cpus of every thread whose mask has any of them.
+func avoidingMasks(cpus []int) (map[int]unix.CPUSet, error) {
+	entries, err := os.ReadDir("/proc/self/task")
+	if err != nil {
+		return nil, err
+	}
+	masks := make(map[int]unix.CPUSet, len(entries))
+	for _, entry := range entries {
+		tid, err := strconv.Atoi(entry.Name())
+		if err != nil {
+			return nil, err
+		}
+		var mask unix.CPUSet
+		err = unix.SchedGetaffinity(tid, &mask)
+		if errors.Is(err, unix.ESRCH) {
+			continue
+		}
+		if err != nil {
+			return nil, fmt.Errorf("read thread %d affinity: %w", tid, err)
+		}
+		next := mask
+		for _, cpu := range cpus {
+			next.Clear(cpu)
+		}
+		if next.Count() == 0 {
+			return nil, fmt.Errorf("avoid_cpus leaves thread %d with no allowed CPU", tid)
+		}
+		if next != mask {
+			masks[tid] = next
+		}
+	}
+	return masks, nil
+}

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")
+}

cpupick/avoid_other.go

@@ -0,0 +1,8 @@
+//go:build !linux || android || e2e_testing
+
+package cpupick
+
+import "errors"
+
+// Avoid is unsupported outside Linux: there is no thread affinity API to restrict.
+func Avoid(_ []int, _ bool) error { return errors.ErrUnsupported }

cpupick/cpupick.go

@@ -16,6 +16,10 @@
 // across distinct physical cores when the topology permits, CPU 0's physical
 // core demoted to last resort, and the order rotated by a stable per-instance
 // key so co-located instances spread instead of stacking.
+//
+// Avoid narrows the process-wide allowed set before any of that happens
+// (tun.avoid_cpus), so Default and tun.cpu_affinity validation only ever see
+// the CPUs left over.
 package cpupick
 
 import (