diff --git a/cpu/cpu_test.go b/cpu/cpu_test.go index 4d9794b..9295dcd 100644 --- a/cpu/cpu_test.go +++ b/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) { diff --git a/cpu/cpu_windows.go b/cpu/cpu_windows.go index de86c04..a5e4785 100644 --- a/cpu/cpu_windows.go +++ b/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() }