From e6aa38bec2904e8135120d37712e2b0bf82a9de6 Mon Sep 17 00:00:00 2001
From: Lomanic <Lomanic@users.noreply.github.com>
Date: Wed, 8 May 2019 20:34:59 +0200
Subject: [PATCH] [cpu][windows] WIP Add support for logical arg in Counts #640
 #628

---
 cpu/cpu_windows.go | 39 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/cpu/cpu_windows.go b/cpu/cpu_windows.go
index e86377c..4807964 100644
--- a/cpu/cpu_windows.go
+++ b/cpu/cpu_windows.go
@@ -13,6 +13,11 @@ 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
@@ -201,6 +206,40 @@ 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) {
+	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
 	return runtime.NumCPU(), nil
 }