diff --git a/sysinfo/numcpu.go b/sysinfo/numcpu.go new file mode 100644 index 0000000..aeb1a3a --- /dev/null +++ b/sysinfo/numcpu.go @@ -0,0 +1,12 @@ +// +build !linux,!windows + +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..5274696 --- /dev/null +++ b/sysinfo/numcpu_linux.go @@ -0,0 +1,45 @@ +// +build linux + +package sysinfo + +import ( + "runtime" + "syscall" + "unsafe" + + "golang.org/x/sys/unix" +) + +// 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() +} 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) +}