From 4b5413396ad1ab208b0a28fd870402e0fefb3be3 Mon Sep 17 00:00:00 2001
From: Yong Tang <yong.tang.github@outlook.com>
Date: Sat, 25 Jun 2016 07:26:00 -0700
Subject: [PATCH] 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 <yong.tang.github@outlook.com>
---
 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()
+}