|
|
|
@ -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)
|
|
|
|
|
}
|
|
|
|
|