From ba8a15694ae8036ec6589e468e25f87eae190a7c Mon Sep 17 00:00:00 2001
From: Tobias Klauser <tklauser@distanz.ch>
Date: Tue, 18 Jun 2019 19:26:56 +0200
Subject: [PATCH] Use functions from x/sys/unix to get number of CPUs on Linux

Use Getpid and SchedGetaffinity from golang.org/x/sys/unix to get the
number of CPUs in numCPU on Linux.

Signed-off-by: Tobias Klauser <tklauser@distanz.ch>
---
 pkg/sysinfo/numcpu_linux.go   | 19 +++++--------------
 pkg/sysinfo/numcpu_windows.go | 10 ++++++++++
 pkg/sysinfo/sysinfo.go        | 10 ----------
 3 files changed, 15 insertions(+), 24 deletions(-)

diff --git a/pkg/sysinfo/numcpu_linux.go b/pkg/sysinfo/numcpu_linux.go
index 5f6c6df8c4..3d9d4cceb3 100644
--- a/pkg/sysinfo/numcpu_linux.go
+++ b/pkg/sysinfo/numcpu_linux.go
@@ -2,7 +2,6 @@ package sysinfo // import "github.com/docker/docker/pkg/sysinfo"
 
 import (
 	"runtime"
-	"unsafe"
 
 	"golang.org/x/sys/unix"
 )
@@ -14,23 +13,15 @@ import (
 // Returns 0 on errors. Use |runtime.NumCPU| in that case.
 func numCPU() int {
 	// Gets the affinity mask for a process: The very one invoking this function.
-	pid, _, _ := unix.RawSyscall(unix.SYS_GETPID, 0, 0, 0)
+	pid := unix.Getpid()
 
-	var mask [1024 / 64]uintptr
-	_, _, err := unix.RawSyscall(unix.SYS_SCHED_GETAFFINITY, pid, uintptr(len(mask)*8), uintptr(unsafe.Pointer(&mask[0])))
-	if err != 0 {
+	var mask unix.CPUSet
+	err := unix.SchedGetaffinity(pid, &mask)
+	if err != nil {
 		return 0
 	}
 
-	// For every available thread a bit is set in the mask.
-	ncpu := 0
-	for _, e := range mask {
-		if e == 0 {
-			continue
-		}
-		ncpu += int(popcnt(uint64(e)))
-	}
-	return ncpu
+	return mask.Count()
 }
 
 // NumCPU returns the number of CPUs which are currently online
diff --git a/pkg/sysinfo/numcpu_windows.go b/pkg/sysinfo/numcpu_windows.go
index 13523f671f..4135c2a2c3 100644
--- a/pkg/sysinfo/numcpu_windows.go
+++ b/pkg/sysinfo/numcpu_windows.go
@@ -13,6 +13,16 @@ var (
 	getProcessAffinityMask = kernel32.NewProc("GetProcessAffinityMask")
 )
 
+// Returns bit count of 1, used by NumCPU
+func popcnt(x uint64) (n byte) {
+	x -= (x >> 1) & 0x5555555555555555
+	x = (x>>2)&0x3333333333333333 + x&0x3333333333333333
+	x += x >> 4
+	x &= 0x0f0f0f0f0f0f0f0f
+	x *= 0x0101010101010101
+	return byte(x >> 56)
+}
+
 func numCPU() int {
 	// Gets the affinity mask for a process
 	var mask, sysmask uintptr
diff --git a/pkg/sysinfo/sysinfo.go b/pkg/sysinfo/sysinfo.go
index de9d736f5d..c0d29fb9e2 100644
--- a/pkg/sysinfo/sysinfo.go
+++ b/pkg/sysinfo/sysinfo.go
@@ -146,13 +146,3 @@ func isCpusetListAvailable(provided, available string) (bool, error) {
 	}
 	return true, nil
 }
-
-// Returns bit count of 1, used by NumCPU
-func popcnt(x uint64) (n byte) {
-	x -= (x >> 1) & 0x5555555555555555
-	x = (x>>2)&0x3333333333333333 + x&0x3333333333333333
-	x += x >> 4
-	x &= 0x0f0f0f0f0f0f0f0f
-	x *= 0x0101010101010101
-	return byte(x >> 56)
-}