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

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

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