diff --git a/cpu/cpu_windows.go b/cpu/cpu_windows.go index e86377c..97c0e34 100644 --- a/cpu/cpu_windows.go +++ b/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 }