util/cpu_pin_bsd.go

@@ -2,6 +2,6 @@
 
 package util
 
-func PinTo(int) error {
+func PinTo(...int) error {
 	return nil
 }

util/cpu_pin_linux.go

@@ -8,11 +8,13 @@ import (
 	"golang.org/x/sys/unix"
 )
 
-func PinTo(cpuID int) error {
-	cpuset := &unix.CPUSet{}
-	cpuset.Set(cpuID)
+func PinTo(cpus ...int) error {
+	set := &unix.CPUSet{}
+	for _, cpu := range cpus {
+		set.Set(cpu)
+	}
 
-	err := unix.SchedSetaffinity(0, cpuset)
+	err := unix.SchedSetaffinity(0, set)
 	if err != nil {
 		return err
 	}
@@ -23,8 +25,13 @@ func PinTo(cpuID int) error {
 		return err
 	}
 
-	if verify.Count() != 1 {
-		return fmt.Errorf("could not pin to CPU %d", cpuID)
+	if verify.Count() != len(cpus) {
+		return fmt.Errorf("could not pin to CPUs %v", cpus)
+	}
+	for _, cpu := range cpus {
+		if !verify.IsSet(cpu) {
+			return fmt.Errorf("could not pin to CPUs %v", cpus)
+		}
 	}
 
 	return nil

util/cpu_pin_test.go

@@ -1,9 +1,44 @@
 package util
 
-import "testing"
+import (
+	"log"
+	"os"
+	"sort"
+	"strconv"
+	"strings"
+	"testing"
+)
+
+func getAllCPUs() ([]int, bool) {
+	entries, err := os.ReadDir("/sys/devices/system/cpu")
+	if err == nil {
+		var cpus []int
+		for _, entry := range entries {
+			cpuid, ok := strings.CutPrefix(entry.Name(), "cpu")
+			if ok {
+				id, err := strconv.ParseInt(cpuid, 10, 64)
+				if err == nil {
+					cpus = append(cpus, int(id))
+				}
+			}
+		}
+		sort.Slice(cpus, func(i, j int) bool {
+			return cpus[i] < cpus[j]
+		})
+		return cpus, true
+	}
+	return nil, false
+}
 
 func TestCPUPin(t *testing.T) {
 	if err := PinTo(0); err != nil {
-		t.Fatal(err)
+		t.Fatalf("could not pin to cpu 0 err=%v", err)
+	}
+
+	if cpus, ok := getAllCPUs(); ok {
+		log.Printf("pinning to all cpus: %v", cpus)
+		if err := PinTo(cpus...); err != nil {
+			t.Fatalf("could not pin to cpus %v err=%v", cpus, err)
+		}
 	}
 }