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

@@ -1,4 +1,4 @@
-// +build !linux
+// +build !linux,!windows
 
 package sysinfo
 

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.
 //

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()
+}

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)
+}