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
package util
package cpupick
import (
"errors"
@@ -12,11 +12,12 @@ import (
"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
// 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.
func AvoidCPUs(cpus []int, apply bool) error {
// 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 {
@@ -1,6 +1,6 @@
//go:build linux && !android && !e2e_testing
package util
package cpupick
import (
"fmt"
@@ -12,22 +12,23 @@ import (
"testing"
"unsafe"
"github.com/slackhq/nebula/util"
"github.com/stretchr/testify/require"
"golang.org/x/sys/unix"
)
// TestAvoidCPUs changes the affinity of the whole process, so it runs in a child, leaving the test runner untouched.
func TestAvoidCPUs(t *testing.T) {
if os.Getenv("NEBULA_AVOID_CPU_TEST") != "1" {
// 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=^TestAvoidCPUs$", "-test.v")
cmd.Env = append(os.Environ(), "NEBULA_AVOID_CPU_TEST=1")
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 := AllowedCPUs()
allowed, err := util.AllowedCPUs()
require.NoError(t, err)
if len(allowed) < 2 {
t.Skip("requires at least two allowed CPUs")
@@ -50,12 +51,12 @@ func TestAvoidCPUs(t *testing.T) {
}
lockThreads()
require.NoError(t, AvoidCPUs([]int{cpu}, false))
after, err := AllowedCPUs()
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, AvoidCPUs(allowed, true), "must reject an empty mask before applying")
require.NoError(t, AvoidCPUs([]int{cpu, cpu}, true))
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()
@@ -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.
func TestAvoidCPUsRange(t *testing.T) {
// 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, AvoidCPUs([]int{capacity - 1}, false))
require.EqualError(t, AvoidCPUs([]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.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")
}
+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
// 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 (
+1 -1
View File
@@ -443,7 +443,7 @@ func configureAvoidCPUs(c *config.C, l *slog.Logger, configTest bool) error {
if len(cpus) == 0 {
return nil
}
if err := util.AvoidCPUs(cpus, !configTest); err != nil {
if err := cpupick.Avoid(cpus, !configTest); err != nil {
if errors.Is(err, errors.ErrUnsupported) {
l.Warn("tun.avoid_cpus is only supported on Linux; ignoring")
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 }