From ce6ed4c574066004ab884122d989748388233cad Mon Sep 17 00:00:00 2001
From: Tobias Klauser <tklauser@distanz.ch>
Date: Tue, 17 Aug 2021 10:43:27 +0200
Subject: [PATCH] Use sched_getaffinity syscall to get online CPUs on Linux

The sched_getaffinity syscall returns the mask of online CPUs, so prefer
that over reading sysfs files on Linux. Only if the syscall returns an
errors, fall back to the existing method.

Running `strace getconf _NPROCESSORS_ONLN` this also seems to be the
method glibc's `getconf` uses.
---
 numcpus_linux.go      | 13 +++++++++++++
 numcpus_linux_test.go | 21 +++++++++++++++++++--
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/numcpus_linux.go b/numcpus_linux.go
index a8d96d6..1a30525 100644
--- a/numcpus_linux.go
+++ b/numcpus_linux.go
@@ -20,10 +20,20 @@ import (
 	"path/filepath"
 	"strconv"
 	"strings"
+
+	"golang.org/x/sys/unix"
 )
 
 const sysfsCPUBasePath = "/sys/devices/system/cpu"
 
+func getFromCPUAffinity() (int, error) {
+	var cpuSet unix.CPUSet
+	if err := unix.SchedGetaffinity(0, &cpuSet); err != nil {
+		return 0, err
+	}
+	return cpuSet.Count(), nil
+}
+
 func readCPURange(file string) (int, error) {
 	buf, err := ioutil.ReadFile(filepath.Join(sysfsCPUBasePath, file))
 	if err != nil {
@@ -95,6 +105,9 @@ func getOffline() (int, error) {
 }
 
 func getOnline() (int, error) {
+	if n, err := getFromCPUAffinity(); err == nil {
+		return n, nil
+	}
 	return readCPURange("online")
 }
 
diff --git a/numcpus_linux_test.go b/numcpus_linux_test.go
index a4487f6..0ba9929 100644
--- a/numcpus_linux_test.go
+++ b/numcpus_linux_test.go
@@ -40,13 +40,30 @@ func TestParseCPURange(t *testing.T) {
 		}
 
 		if n != tc.n {
-			t.Errorf("parseCPURange(%s) = %d, expected %d", tc.cpus, n, tc.n)
+			t.Errorf("parseCPURange(%q) = %d, expected %d", tc.cpus, n, tc.n)
 		}
 	}
 
 	str := "invalid"
 	_, err := parseCPURange(str)
 	if err == nil {
-		t.Errorf("parseCPURange(%s) unexpectedly succeeded", str)
+		t.Errorf("parseCPURange(%q) unexpectedly succeeded", str)
+	}
+}
+
+func TestGetFromCPUAffinity(t *testing.T) {
+	nAffinity, err := getFromCPUAffinity()
+	if err != nil {
+		t.Fatalf("getFromCPUAffinity: %v", err)
+	}
+
+	cpus := "online"
+	nSysfs, err := readCPURange(cpus)
+	if err != nil {
+		t.Fatalf("readCPURange(%q): %v", cpus, err)
+	}
+
+	if nAffinity != nSysfs {
+		t.Errorf("getFromCPUAffinity() = %d, readCPURange(%q) = %d, want the same return value", nAffinity, cpus, nSysfs)
 	}
 }