Merge pull request #749 from Lomanic/issue628windows

[cpu][windows] Add support for logical arg in Counts #640 #628

cpu/cpu_windows.go

@@ -5,7 +5,6 @@ package cpu
 import (
 	"context"
 	"fmt"
-	"runtime"
 	"unsafe"
 
 	"github.com/StackExchange/wmi"
@@ -13,12 +12,18 @@ import (
 	"golang.org/x/sys/windows"
 )
 
+var (
+	procGetActiveProcessorCount = common.Modkernel32.NewProc("GetActiveProcessorCount")
+	procGetNativeSystemInfo     = common.Modkernel32.NewProc("GetNativeSystemInfo")
+)
+
 type Win32_Processor struct {
 	LoadPercentage            *uint16
 	Family                    uint16
 	Manufacturer              string
 	Name                      string
 	NumberOfLogicalProcessors uint32
+	NumberOfCores             uint32
 	ProcessorID               *string
 	Stepping                  *string
 	MaxClockSpeed             uint32
@@ -201,6 +206,50 @@ func perfInfo() ([]win32_SystemProcessorPerformanceInformation, error) {
 	return resultBuffer, nil
 }
 
+// SystemInfo is an equivalent representation of SYSTEM_INFO in the Windows API.
+// https://msdn.microsoft.com/en-us/library/ms724958%28VS.85%29.aspx?f=255&MSPPError=-2147217396
+// https://github.com/elastic/go-windows/blob/bb1581babc04d5cb29a2bfa7a9ac6781c730c8dd/kernel32.go#L43
+type systemInfo struct {
+	wProcessorArchitecture      uint16
+	wReserved                   uint16
+	dwPageSize                  uint32
+	lpMinimumApplicationAddress uintptr
+	lpMaximumApplicationAddress uintptr
+	dwActiveProcessorMask       uintptr
+	dwNumberOfProcessors        uint32
+	dwProcessorType             uint32
+	dwAllocationGranularity     uint32
+	wProcessorLevel             uint16
+	wProcessorRevision          uint16
+}
+
 func CountsWithContext(ctx context.Context, logical bool) (int, error) {
-	return runtime.NumCPU(), nil
+	if logical {
+		// https://github.com/giampaolo/psutil/blob/d01a9eaa35a8aadf6c519839e987a49d8be2d891/psutil/_psutil_windows.c#L97
+		err := procGetActiveProcessorCount.Find()
+		if err == nil { // Win7+
+			ret, _, _ := procGetActiveProcessorCount.Call(uintptr(0xffff)) // ALL_PROCESSOR_GROUPS is 0xffff according to Rust's winapi lib https://docs.rs/winapi/*/x86_64-pc-windows-msvc/src/winapi/shared/ntdef.rs.html#120
+			if ret != 0 {
+				return int(ret), nil
+			}
+		}
+		var systemInfo systemInfo
+		_, _, err = procGetNativeSystemInfo.Call(uintptr(unsafe.Pointer(&systemInfo)))
+		if systemInfo.dwNumberOfProcessors == 0 {
+			return 0, err
+		}
+		return int(systemInfo.dwNumberOfProcessors), nil
+	}
+	// physical cores https://github.com/giampaolo/psutil/blob/d01a9eaa35a8aadf6c519839e987a49d8be2d891/psutil/_psutil_windows.c#L499
+	// for the time being, try with unreliable and slow WMI call…
+	var dst []Win32_Processor
+	q := wmi.CreateQuery(&dst, "")
+	if err := common.WMIQueryWithContext(ctx, q, &dst); err != nil {
+		return 0, err
+	}
+	var count uint32
+	for _, d := range dst {
+		count += d.NumberOfCores
+	}
+	return int(count), nil
 }