1. docker
  2. Commit 3707a76921b412cb42e6134beaa0bddbb5e08f25
Browse code

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>

Yong Tang authored on 2016/06/26 14:14:09
Showing 4 changed files
pkg/sysinfo/numcpu.go
History View file @ 3707a76
... ... 
@@ -1,4 +1,4 @@
1 
-// +build !linux
1 
+// +build !linux,!windows
2 2
3 3 
 package sysinfo
4 4
pkg/sysinfo/numcpu_linux.go
History View file @ 3707a76
... ... 
@@ -10,16 +10,6 @@ import (
10 10 
 	"golang.org/x/sys/unix"
11 11 
 )
12 12
13 
-// Returns bit count of 1
14 
-func popcnt(x uint64) (n byte) {
15 
-	x -= (x >> 1) & 0x5555555555555555
16 
-	x = (x>>2)&0x3333333333333333 + x&0x3333333333333333
17 
-	x += x >> 4
18 
-	x &= 0x0f0f0f0f0f0f0f0f
19 
-	x *= 0x0101010101010101
20 
-	return byte(x >> 56)
21 
-}
22 
-
23 13 
 // numCPU queries the system for the count of threads available
24 14 
 // for use to this process.
25 15 
 //
pkg/sysinfo/numcpu_windows.go
History View file @ 3707a76
26 16 
new file mode 100644
... ... 
@@ -0,0 +1,36 @@
0 
+// +build windows
1 
+
2 
+package sysinfo
3 
+
4 
+import (
5 
+	"runtime"
6 
+	"syscall"
7 
+	"unsafe"
8 
+)
9 
+
10 
+var (
11 
+	kernel32               = syscall.NewLazyDLL("kernel32.dll")
12 
+	getCurrentProcess      = kernel32.NewProc("GetCurrentProcess")
13 
+	getProcessAffinityMask = kernel32.NewProc("GetProcessAffinityMask")
14 
+)
15 
+
16 
+func numCPU() int {
17 
+	// Gets the affinity mask for a process
18 
+	var mask, sysmask uintptr
19 
+	currentProcess, _, _ := getCurrentProcess.Call()
20 
+	ret, _, _ := getProcessAffinityMask.Call(currentProcess, uintptr(unsafe.Pointer(&mask)), uintptr(unsafe.Pointer(&sysmask)))
21 
+	if ret == 0 {
22 
+		return 0
23 
+	}
24 
+	// For every available thread a bit is set in the mask.
25 
+	ncpu := int(popcnt(uint64(mask)))
26 
+	return ncpu
27 
+}
28 
+
29 
+// NumCPU returns the number of CPUs which are currently online
30 
+func NumCPU() int {
31 
+	if ncpu := numCPU(); ncpu > 0 {
32 
+		return ncpu
33 
+	}
34 
+	return runtime.NumCPU()
35 
+}
pkg/sysinfo/sysinfo.go
History View file @ 3707a76
... ... 
@@ -126,3 +126,13 @@ func isCpusetListAvailable(provided, available string) (bool, error) {
126 126 
 	}
127 127 
 	return true, nil
128 128 
 }
129 
+
130 
+// Returns bit count of 1, used by NumCPU
131 
+func popcnt(x uint64) (n byte) {
132 
+	x -= (x >> 1) & 0x5555555555555555
133 
+	x = (x>>2)&0x3333333333333333 + x&0x3333333333333333
134 
+	x += x >> 4
135 
+	x &= 0x0f0f0f0f0f0f0f0f
136 
+	x *= 0x0101010101010101
137 
+	return byte(x >> 56)
138 
+}