From 3707a76921b412cb42e6134beaa0bddbb5e08f25 Mon Sep 17 00:00:00 2001
From: Yong Tang <yong.tang.github@outlook.com>
Date: Sat, 25 Jun 2016 22:14:09 -0700
Subject: [PATCH] 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 <yong.tang.github@outlook.com>
---
 pkg/sysinfo/numcpu.go         |  2 +-
 pkg/sysinfo/numcpu_linux.go   | 10 ----------
 pkg/sysinfo/numcpu_windows.go | 36 +++++++++++++++++++++++++++++++++++
 pkg/sysinfo/sysinfo.go        | 10 ++++++++++
 4 files changed, 47 insertions(+), 11 deletions(-)
 create mode 100644 pkg/sysinfo/numcpu_windows.go

diff --git a/pkg/sysinfo/numcpu.go b/pkg/sysinfo/numcpu.go
index 6cb12cdcbc..aeb1a3a804 100644
--- a/pkg/sysinfo/numcpu.go
+++ b/pkg/sysinfo/numcpu.go
@@ -1,4 +1,4 @@
-// +build !linux
+// +build !linux,!windows
 
 package sysinfo
 
diff --git a/pkg/sysinfo/numcpu_linux.go b/pkg/sysinfo/numcpu_linux.go
index b3c25294e0..5274696da3 100644
--- a/pkg/sysinfo/numcpu_linux.go
+++ b/pkg/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/pkg/sysinfo/numcpu_windows.go b/pkg/sysinfo/numcpu_windows.go
new file mode 100644
index 0000000000..5077af7298
--- /dev/null
+++ b/pkg/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/pkg/sysinfo/sysinfo.go b/pkg/sysinfo/sysinfo.go
index cbd0099957..dc71dbfb80 100644
--- a/pkg/sysinfo/sysinfo.go
+++ b/pkg/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)
+}