cpupick: Move AvoidCPUs from util into cpupick as Avoid

This commit is contained in:
JackDoan committed 2026-10-15 12:00:28 -05:00
1 parent 73d26dd87d
commit 36579bae03
6 files changed
+35 -29

No files matched your search

@@ -1,6 +1,6 @@
//go:build linux && !android && !e2e_testing //go:build linux && !android && !e2e_testing
package util package cpupick
import ( import (
"errors" "errors"
@@ -12,11 +12,12 @@ import (
"golang.org/x/sys/unix" "golang.org/x/sys/unix"
) )
// AvoidCPUs removes cpus from the affinity of every thread of the process, for use at startup. New threads inherit // 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 // 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. // 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. // This affects the entire process, including other users of an embedded Nebula. Run it before Default, which then
func AvoidCPUs(cpus []int, apply bool) error { // picks from what remains.
func Avoid(cpus []int, apply bool) error {
const capacity = int(unsafe.Sizeof(unix.CPUSet{})) * 8 const capacity = int(unsafe.Sizeof(unix.CPUSet{})) * 8
for _, cpu := range cpus { for _, cpu := range cpus {
if cpu < 0 || cpu >= capacity { if cpu < 0 || cpu >= capacity {
@@ -1,6 +1,6 @@
//go:build linux && !android && !e2e_testing //go:build linux && !android && !e2e_testing
package util package cpupick
import ( import (
"fmt" "fmt"
@@ -12,22 +12,23 @@ import (
"testing" "testing"
"unsafe" "unsafe"
"github.com/slackhq/nebula/util"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
"golang.org/x/sys/unix" "golang.org/x/sys/unix"
) )
// TestAvoidCPUs changes the affinity of the whole process, so it runs in a child, leaving the test runner untouched. // TestAvoid changes the affinity of the whole process, so it runs in a child, leaving the test runner untouched.
func TestAvoidCPUs(t *testing.T) { func TestAvoid(t *testing.T) {
if os.Getenv("NEBULA_AVOID_CPU_TEST") != "1" { if os.Getenv("NEBULA_CPUPICK_AVOID_TEST") != "1" {
exe, err := os.Executable() exe, err := os.Executable()
require.NoError(t, err) require.NoError(t, err)
cmd := exec.Command(exe, "-test.run=^TestAvoidCPUs$", "-test.v") cmd := exec.Command(exe, "-test.run=^TestAvoid$", "-test.v")
cmd.Env = append(os.Environ(), "NEBULA_AVOID_CPU_TEST=1") cmd.Env = append(os.Environ(), "NEBULA_CPUPICK_AVOID_TEST=1")
out, err := cmd.CombinedOutput() out, err := cmd.CombinedOutput()
require.NoError(t, err, "%s", out) require.NoError(t, err, "%s", out)
return return
} }
allowed, err := AllowedCPUs() allowed, err := util.AllowedCPUs()
require.NoError(t, err) require.NoError(t, err)
if len(allowed) < 2 { if len(allowed) < 2 {
t.Skip("requires at least two allowed CPUs") t.Skip("requires at least two allowed CPUs")
@@ -50,12 +51,12 @@ func TestAvoidCPUs(t *testing.T) {
} }
lockThreads() lockThreads()
require.NoError(t, AvoidCPUs([]int{cpu}, false)) require.NoError(t, Avoid([]int{cpu}, false))
after, err := AllowedCPUs() after, err := util.AllowedCPUs()
require.NoError(t, err) require.NoError(t, err)
require.Equal(t, allowed, after, "configuration check must not mutate affinity") require.Equal(t, allowed, after, "configuration check must not mutate affinity")
require.Error(t, AvoidCPUs(allowed, true), "must reject an empty mask before applying") require.Error(t, Avoid(allowed, true), "must reject an empty mask before applying")
require.NoError(t, AvoidCPUs([]int{cpu, cpu}, true)) require.NoError(t, Avoid([]int{cpu, cpu}, true))
requireAvoided(t, cpu, allowed) requireAvoided(t, cpu, allowed)
// Existing threads stay locked while new ones inherit the restriction // Existing threads stay locked while new ones inherit the restriction
lockThreads() lockThreads()
@@ -88,10 +89,10 @@ func requireAvoided(t *testing.T, cpu int, allowed []int) {
} }
} }
// TestAvoidCPUsRange checks the bounds of unix.CPUSet, whose words are 32 bits wide on 32-bit platforms. // TestAvoidRange checks the bounds of unix.CPUSet, whose words are 32 bits wide on 32-bit platforms.
func TestAvoidCPUsRange(t *testing.T) { func TestAvoidRange(t *testing.T) {
capacity := int(unsafe.Sizeof(unix.CPUSet{})) * 8 capacity := int(unsafe.Sizeof(unix.CPUSet{})) * 8
require.NoError(t, AvoidCPUs([]int{capacity - 1}, false)) require.NoError(t, Avoid([]int{capacity - 1}, false))
require.EqualError(t, AvoidCPUs([]int{capacity}, false), fmt.Sprintf("CPU %d is out of range", capacity)) require.EqualError(t, Avoid([]int{capacity}, false), fmt.Sprintf("CPU %d is out of range", capacity))
require.EqualError(t, AvoidCPUs([]int{-1}, false), "CPU -1 is out of range") require.EqualError(t, Avoid([]int{-1}, false), "CPU -1 is out of range")
} }
+8
View File
@@ -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 }
+4
View File
@@ -16,6 +16,10 @@
// across distinct physical cores when the topology permits, CPU 0's physical // 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 // core demoted to last resort, and the order rotated by a stable per-instance
// key so co-located instances spread instead of stacking. // 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 package cpupick
import ( import (
+1 -1
View File
@@ -443,7 +443,7 @@ func configureAvoidCPUs(c *config.C, l *slog.Logger, configTest bool) error {
if len(cpus) == 0 { if len(cpus) == 0 {
return nil return nil
} }
if err := util.AvoidCPUs(cpus, !configTest); err != nil { if err := cpupick.Avoid(cpus, !configTest); err != nil {
if errors.Is(err, errors.ErrUnsupported) { if errors.Is(err, errors.ErrUnsupported) {
l.Warn("tun.avoid_cpus is only supported on Linux; ignoring") l.Warn("tun.avoid_cpus is only supported on Linux; ignoring")
return nil return nil
-8
View File
@@ -1,8 +0,0 @@
//go:build !linux || android || e2e_testing
package util
import "errors"
// AvoidCPUs is unsupported outside Linux: there is no thread affinity API to restrict.
func AvoidCPUs(_ []int, _ bool) error { return errors.ErrUnsupported }