Merge pull request #1839 from OrbintSoft/remove-wmi

Remove wmi for Get physical CPU core count

cpu/cpu_test.go

@@ -59,16 +59,17 @@ 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)
+	assert.NotZerof(t, logicalCount, "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)
+	assert.NotZerof(t, physicalCount, "could not get physical CPU counts: %v", physicalCount)
+	t.Logf("physical cores: %d", physicalCount)
+	assert.GreaterOrEqualf(t, logicalCount, physicalCount, "logical CPU count should be greater than or equal to physical CPU count: %v >= %v", logicalCount, physicalCount)
 }
 
 func TestTimeStat_String(t *testing.T) {

cpu/cpu_windows.go

@@ -5,6 +5,7 @@ package cpu
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"strconv"
 	"unsafe"
@@ -15,7 +16,10 @@ import (
 	"github.com/shirou/gopsutil/v4/internal/common"
 )
 
-var procGetNativeSystemInfo = common.Modkernel32.NewProc("GetNativeSystemInfo")
+var (
+	procGetNativeSystemInfo              = common.Modkernel32.NewProc("GetNativeSystemInfo")
+	procGetLogicalProcessorInformationEx = common.Modkernel32.NewProc("GetLogicalProcessorInformationEx")
+)
 
 type win32_Processor struct { //nolint:revive //FIXME
 	Family                    uint16
@@ -200,13 +204,70 @@ type systemInfo struct {
 	wProcessorRevision          uint16
 }
 
-func CountsWithContext(ctx context.Context, logical bool) (int, error) {
+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 uint32
+	Size         uint32
+	Processor    processorRelationship
+}
+
+func getPhysicalCoreCount() (int, error) {
+	var length uint32
+	const relationAll = 0xffff
+	const relationProcessorCore = 0x0
+
+	// First call to determine the required buffer size
+	_, _, err := procGetLogicalProcessorInformationEx.Call(uintptr(relationAll), 0, uintptr(unsafe.Pointer(&length)))
+	if err != nil && !errors.Is(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 = procGetLogicalProcessorInformationEx.Call(uintptr(relationAll), uintptr(unsafe.Pointer(&buffer[0])), uintptr(unsafe.Pointer(&length)))
+	if err != nil && !errors.Is(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) {
+		info := (*systemLogicalProcessorInformationEx)(unsafe.Pointer(uintptr(unsafe.Pointer(&buffer[0])) + offset))
+		if info.Relationship == relationProcessorCore {
+			ncpus++
+		}
+		offset += uintptr(info.Size)
+	}
+
+	return ncpus, nil
+}
+
+func CountsWithContext(_ context.Context, logical bool) (int, error) {
 	if logical {
-		// https://github.com/giampaolo/psutil/blob/d01a9eaa35a8aadf6c519839e987a49d8be2d891/psutil/_psutil_windows.c#L97
+		// Get logical processor count https://github.com/giampaolo/psutil/blob/d01a9eaa35a8aadf6c519839e987a49d8be2d891/psutil/_psutil_windows.c#L97
 		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 +275,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 https://github.com/giampaolo/psutil/blob/d01a9eaa35a8aadf6c519839e987a49d8be2d891/psutil/_psutil_windows.c#L499
+	return getPhysicalCoreCount()
 }