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