From 1ebcd369abb6046f214105dcb7705ffef333302f Mon Sep 17 00:00:00 2001
From: Stefano Balzarotti <stefano.balzarotti@orbintsoft.net>
Date: Sun, 13 Apr 2025 17:47:36 +0200
Subject: [PATCH] implement get physical core api

---
 cpu/cpu_test.go    | 27 ++++++++++++------
 cpu/cpu_windows.go | 80 +++++++++++++++++++++++++++++++++++++++++++++---------
 2 files changed, 86 insertions(+), 21 deletions(-)

diff --git a/cpu/cpu_test.go b/cpu/cpu_test.go
index 4d9794b..b4662f5 100644
--- a/cpu/cpu_test.go
+++ b/cpu/cpu_test.go
@@ -59,16 +59,27 @@ func TestTimes(t *testing.T) {
 }
 
 func TestCounts(t *testing.T) {
-	v, err := Counts(true)
+	logicalCount, err := Counts(true)
 	common.SkipIfNotImplementedErr(t, err)
-	require.NoError(t, err)
-	assert.NotZerof(t, v, "could not get logical CPU counts: %v", v)
-	t.Logf("logical cores: %d", v)
-	v, err = Counts(false)
+	if err != nil {
+		t.Errorf("error %v", err)
+	}
+	if logicalCount == 0 {
+		t.Errorf("could not get logical CPU counts: %v", logicalCount)
+	}
+	t.Logf("logical cores: %d", logicalCount)
+	physicalCount, err := Counts(false)
 	common.SkipIfNotImplementedErr(t, err)
-	require.NoError(t, err)
-	assert.NotZerof(t, v, "could not get physical CPU counts: %v", v)
-	t.Logf("physical cores: %d", v)
+	if err != nil {
+		t.Errorf("error %v", err)
+	}
+	if physicalCount == 0 {
+		t.Errorf("could not get physical CPU counts: %v", physicalCount)
+	}
+	t.Logf("physical cores: %d", physicalCount)
+	if physicalCount > logicalCount {
+		t.Errorf("physical cpu cannot be more than logical cpu: %v > %v", physicalCount, logicalCount)
+	}
 }
 
 func TestTimeStat_String(t *testing.T) {
diff --git a/cpu/cpu_windows.go b/cpu/cpu_windows.go
index 3ca8281..aeed5fb 100644
--- a/cpu/cpu_windows.go
+++ b/cpu/cpu_windows.go
@@ -16,6 +16,7 @@ import (
 )
 
 var procGetNativeSystemInfo = common.Modkernel32.NewProc("GetNativeSystemInfo")
+var getLogicalProcessorInformationEx = common.Modkernel32.NewProc("GetLogicalProcessorInformationEx")
 
 type win32_Processor struct { //nolint:revive //FIXME
 	Family                    uint16
@@ -200,13 +201,75 @@ type systemInfo struct {
 	wProcessorRevision          uint16
 }
 
+type groupAffinity struct {
+	mask     uintptr // https://learn.microsoft.com/it-it/windows-hardware/drivers/kernel/interrupt-affinity-and-priority#about-kaffinity
+	group    uint16
+	reserved [3]uint16
+}
+
+// https://learn.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-processor_relationship
+type processorRelationship struct {
+	flags          byte
+	efficientClass byte
+	reserved       [20]byte
+	groupCount     uint16
+	groupMask      [1]groupAffinity
+}
+
+// https://learn.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-system_logical_processor_information_ex
+type systemLogicalProcessorInformationEx struct {
+	Relationship int
+	Size         uint32
+	Processor    processorRelationship
+}
+
+func getPhysicalCoreCount(ctx context.Context) (int, error) {
+	var length uint32
+	const relationAll = 0xffff
+	const relationProcessorCore = 0x0
+
+	// First call to determine the required buffer size
+	_, _, err := getLogicalProcessorInformationEx.Call(uintptr(relationAll), 0, uintptr(unsafe.Pointer(&length)))
+	if err != nil && err != windows.ERROR_INSUFFICIENT_BUFFER {
+		return 0, fmt.Errorf("failed to get buffer size: %w", err)
+	}
+
+	// Allocate the buffer
+	buffer := make([]byte, length)
+
+	// Second call to retrieve the processor information
+	_, _, err = getLogicalProcessorInformationEx.Call(uintptr(relationAll), uintptr(unsafe.Pointer(&buffer[0])), uintptr(unsafe.Pointer(&length)))
+	if err != nil && err != windows.NTE_OP_OK {
+		return 0, fmt.Errorf("failed to get logical processor information: %w", err)
+	}
+
+	// Iterate through the buffer to count physical cores
+	offset := uintptr(0)
+	ncpus := 0
+	for offset < uintptr(length) {
+		select {
+		case <-ctx.Done():
+			return 0, ctx.Err()
+		default:
+			info := (*systemLogicalProcessorInformationEx)(unsafe.Pointer(uintptr(unsafe.Pointer(&buffer[0])) + offset))
+			if info.Relationship == relationProcessorCore {
+				ncpus++
+			}
+			offset += uintptr(info.Size)
+		}
+	}
+
+	return ncpus, nil
+}
+
 func CountsWithContext(ctx context.Context, logical bool) (int, error) {
 	if logical {
-		// https://github.com/giampaolo/psutil/blob/d01a9eaa35a8aadf6c519839e987a49d8be2d891/psutil/_psutil_windows.c#L97
+		// Get logical processor count
 		ret := windows.GetActiveProcessorCount(windows.ALL_PROCESSOR_GROUPS)
 		if ret != 0 {
 			return int(ret), nil
 		}
+
 		var systemInfo systemInfo
 		_, _, err := procGetNativeSystemInfo.Call(uintptr(unsafe.Pointer(&systemInfo)))
 		if systemInfo.dwNumberOfProcessors == 0 {
@@ -214,16 +277,7 @@ func CountsWithContext(ctx context.Context, logical bool) (int, error) {
 		}
 		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
+
+	// Get physical core count
+	return getPhysicalCoreCount(ctx)
 }