From 4b5413396ad1ab208b0a28fd870402e0fefb3be3 Mon Sep 17 00:00:00 2001 From: Yong Tang Date: Sat, 25 Jun 2016 07:26:00 -0700 Subject: [PATCH 1/2] Fix wrong CPU count after CPU hot-plugging This fix tries to address issues raised in #23768 where the CPU count is not updated after cpu ho-plugging. This fix follows the suggestion from #23768 and replace go's `runtime.NumCPU()` with `sysconf(_SC_NPROCESSORS_ONLN)` so that correct CPU count could be obtained even after CPU hot-plugging. This fix is tested manually, as is suggested in #23768. This fix fixes #23768. The NumCPU() in Linux is based on @wmark 's implementation. Signed-off-by: Yong Tang --- sysinfo/numcpu.go | 12 +++++++++ sysinfo/numcpu_linux.go | 55 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+) create mode 100644 sysinfo/numcpu.go create mode 100644 sysinfo/numcpu_linux.go diff --git a/sysinfo/numcpu.go b/sysinfo/numcpu.go new file mode 100644 index 0000000..6cb12cd --- /dev/null +++ b/sysinfo/numcpu.go @@ -0,0 +1,12 @@ +// +build !linux + +package sysinfo + +import ( + "runtime" +) + +// NumCPU returns the number of CPUs +func NumCPU() int { + return runtime.NumCPU() +} diff --git a/sysinfo/numcpu_linux.go b/sysinfo/numcpu_linux.go new file mode 100644 index 0000000..b3c2529 --- /dev/null +++ b/sysinfo/numcpu_linux.go @@ -0,0 +1,55 @@ +// +build linux + +package sysinfo + +import ( + "runtime" + "syscall" + "unsafe" + + "golang.org/x/sys/unix" +) + +// Returns bit count of 1 +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) +} + +// 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() +} From 6c046817ae4fc49d94879a0713b4960c0bb9148e Mon Sep 17 00:00:00 2001 From: Yong Tang Date: Sat, 25 Jun 2016 22:14:09 -0700 Subject: [PATCH 2/2] Fix wrong CPU count after CPU hot-plugging on Windows This fix tries to fix wrong CPU count after CPU hot-plugging. On windows, GetProcessAffinityMask has been used to probe the number of CPUs in real time. Signed-off-by: Yong Tang --- sysinfo/numcpu.go | 2 +- sysinfo/numcpu_linux.go | 10 ---------- sysinfo/numcpu_windows.go | 36 ++++++++++++++++++++++++++++++++++++ sysinfo/sysinfo.go | 10 ++++++++++ 4 files changed, 47 insertions(+), 11 deletions(-) create mode 100644 sysinfo/numcpu_windows.go diff --git a/sysinfo/numcpu.go b/sysinfo/numcpu.go index 6cb12cd..aeb1a3a 100644 --- a/sysinfo/numcpu.go +++ b/sysinfo/numcpu.go @@ -1,4 +1,4 @@ -// +build !linux +// +build !linux,!windows package sysinfo diff --git a/sysinfo/numcpu_linux.go b/sysinfo/numcpu_linux.go index b3c2529..5274696 100644 --- a/sysinfo/numcpu_linux.go +++ b/sysinfo/numcpu_linux.go @@ -10,16 +10,6 @@ import ( "golang.org/x/sys/unix" ) -// Returns bit count of 1 -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) -} - // numCPU queries the system for the count of threads available // for use to this process. // diff --git a/sysinfo/numcpu_windows.go b/sysinfo/numcpu_windows.go new file mode 100644 index 0000000..5077af7 --- /dev/null +++ b/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() +} diff --git a/sysinfo/sysinfo.go b/sysinfo/sysinfo.go index cbd0099..dc71dbf 100644 --- a/sysinfo/sysinfo.go +++ b/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) +}