Просмотр исходного кода

cpupick: Move AvoidCPUs from util into cpupick as Avoid

JackDoan 5 дней назад
Родитель
Сommit
36579bae03
6 измененных файлов с 35 добавлено и 29 удалено
  1. 5 4
      cpupick/avoid_linux.go
  2. 17 16
      cpupick/avoid_linux_test.go
  3. 8 0
      cpupick/avoid_other.go
  4. 4 0
      cpupick/cpupick.go
  5. 1 1
      main.go
  6. 0 8
      util/cpuavoid_other.go

+ 5 - 4
util/cpuavoid_linux.go → cpupick/avoid_linux.go

@@ -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.
-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
 	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 {

+ 17 - 16
util/cpuavoid_linux_test.go → cpupick/avoid_linux_test.go

@@ -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.
-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()
 		exe, err := os.Executable()
 		require.NoError(t, err)
 		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()
 		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))
-	after, err := AllowedCPUs()
+	require.NoError(t, Avoid([]int{cpu}, false))
+	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.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)
 	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.
-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
 	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 - 0
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 }

+ 4 - 0
cpupick/cpupick.go

@@ -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
main.go

@@ -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

+ 0 - 8
util/cpuavoid_other.go

@@ -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 }