Merge pull request #23955 from yongtang/23768-cpu-count

Fix wrong CPU count after CPU hot-plugging

sysinfo/numcpu.go

@@ -0,0 +1,12 @@
+// +build !linux,!windows
+
+package sysinfo
+
+import (
+	"runtime"
+)
+
+// NumCPU returns the number of CPUs
+func NumCPU() int {
+	return runtime.NumCPU()
+}

sysinfo/numcpu_linux.go

@@ -0,0 +1,45 @@
+// +build linux
+
+package sysinfo
+
+import (
+	"runtime"
+	"syscall"
+	"unsafe"
+
+	"golang.org/x/sys/unix"
+)
+
+// numCPU queries the system for the count of threads available
+// for use to this process.
+//
+// Issues two syscalls.
+// 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, _, _ := syscall.RawSyscall(unix.SYS_GETPID, 0, 0, 0)
+
+	var mask [1024 / 64]uintptr
+	_, _, err := syscall.RawSyscall(unix.SYS_SCHED_GETAFFINITY, pid, uintptr(len(mask)*8), uintptr(unsafe.Pointer(&mask[0])))
+	if err != 0 {
+		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
+}
+
+// NumCPU returns the number of CPUs which are currently online
+func NumCPU() int {
+	if ncpu := numCPU(); ncpu > 0 {
+		return ncpu
+	}
+	return runtime.NumCPU()
+}

sysinfo/numcpu_windows.go

@@ -0,0 +1,36 @@
+// +build windows
+
+package sysinfo
+
+import (
+	"runtime"
+	"syscall"
+	"unsafe"
+)
+
+var (
+	kernel32               = syscall.NewLazyDLL("kernel32.dll")
+	getCurrentProcess      = kernel32.NewProc("GetCurrentProcess")
+	getProcessAffinityMask = kernel32.NewProc("GetProcessAffinityMask")
+)
+
+func numCPU() int {
+	// Gets the affinity mask for a process
+	var mask, sysmask uintptr
+	currentProcess, _, _ := getCurrentProcess.Call()
+	ret, _, _ := getProcessAffinityMask.Call(currentProcess, uintptr(unsafe.Pointer(&mask)), uintptr(unsafe.Pointer(&sysmask)))
+	if ret == 0 {
+		return 0
+	}
+	// For every available thread a bit is set in the mask.
+	ncpu := int(popcnt(uint64(mask)))
+	return ncpu
+}
+
+// NumCPU returns the number of CPUs which are currently online
+func NumCPU() int {
+	if ncpu := numCPU(); ncpu > 0 {
+		return ncpu
+	}
+	return runtime.NumCPU()
+}

sysinfo/sysinfo.go

@@ -126,3 +126,13 @@ 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)
+}