diff --git a/pkg/gpu/stat/stat_windows.go b/pkg/gpu/stat/stat_windows.go index 1fceb6e..7cb566f 100644 --- a/pkg/gpu/stat/stat_windows.go +++ b/pkg/gpu/stat/stat_windows.go @@ -1,31 +1,151 @@ //go:build windows +// Modified from https://github.com/shirou/gopsutil/blob/master/internal/common/common_windows.go +// Original License: BSD-3-Clause + package stat import ( - "os/exec" - "strconv" - "strings" + "errors" + "fmt" + "time" + "unsafe" + + "golang.org/x/sys/windows" ) +const ( + ERROR_SUCCESS = 0 + PDH_FMT_DOUBLE = 0x00000200 + PDH_MORE_DATA = 0x800007d2 + PDH_VAILD_DATA = 0x00000000 + PDH_NEW_DATA = 0x00000001 + PDH_NO_DATA = 0x800007d5 +) + +var ( + modPdh = windows.NewLazySystemDLL("pdh.dll") + + pdhOpenQuery = modPdh.NewProc("PdhOpenQuery") + pdhCollectQueryData = modPdh.NewProc("PdhCollectQueryData") + pdhGetFormattedCounterArrayW = modPdh.NewProc("PdhGetFormattedCounterArrayW") + pdhAddEnglishCounterW = modPdh.NewProc("PdhAddEnglishCounterW") +) + +type PDH_FMT_COUNTERVALUE_DOUBLE struct { + CStatus uint32 + DoubleValue float64 +} + +type PDH_FMT_COUNTERVALUE_ITEM_DOUBLE struct { + SzName *uint16 + FmtValue PDH_FMT_COUNTERVALUE_DOUBLE +} + +// https://github.com/influxdata/telegraf/blob/master/plugins/inputs/win_perf_counters/performance_query.go +func getCounterArrayValue(initialBufSize uint32, counter *win32PerformanceCounter) ([]float64, error) { + for buflen := initialBufSize; buflen <= 100*1024*1024; buflen *= 2 { + time.Sleep(10 * time.Millisecond) // GPU 查询必须设置间隔,否则数据不准 + s, _, err := pdhCollectQueryData.Call(uintptr(counter.Query)) + if s != 0 && err != nil { + if s == PDH_NO_DATA { + return nil, fmt.Errorf("%w: this counter has not data", err) + } + return nil, err + } + buf := make([]byte, buflen) + size := buflen + var itemCount uint32 + r, _, _ := pdhGetFormattedCounterArrayW.Call(uintptr(counter.Counter), PDH_FMT_DOUBLE, uintptr(unsafe.Pointer(&size)), uintptr(unsafe.Pointer(&itemCount)), uintptr(unsafe.Pointer(&buf[0]))) + if r == ERROR_SUCCESS { + items := (*[1 << 20]PDH_FMT_COUNTERVALUE_ITEM_DOUBLE)(unsafe.Pointer(&buf[0]))[:itemCount:itemCount] + values := make([]float64, 0, itemCount) + for _, item := range items { + if item.FmtValue.CStatus == PDH_VAILD_DATA || item.FmtValue.CStatus == PDH_NEW_DATA { + val := item.FmtValue.DoubleValue + values = append(values, val) + } + } + return values, nil + } + if r != PDH_MORE_DATA { + return nil, fmt.Errorf("pdhGetFormattedCounterArrayW failed with status 0x%X", r) + } + } + + return nil, errors.New("buffer limit reached") +} + +func createQuery() (windows.Handle, error) { + var query windows.Handle + r, _, err := pdhOpenQuery.Call(0, 0, uintptr(unsafe.Pointer(&query))) + if r != ERROR_SUCCESS { + return 0, fmt.Errorf("pdhOpenQuery failed with status 0x%X: %v", r, err) + } + return query, nil +} + +type win32PerformanceCounter struct { + PostName string + CounterName string + Query windows.Handle + Counter windows.Handle +} + +func newWin32PerformanceCounter(postName, counterName string) (*win32PerformanceCounter, error) { + query, err := createQuery() + if err != nil { + return nil, err + } + counter := win32PerformanceCounter{ + Query: query, + PostName: postName, + CounterName: counterName, + } + r, _, err := pdhAddEnglishCounterW.Call( + uintptr(counter.Query), + uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(counter.CounterName))), + 0, + uintptr(unsafe.Pointer(&counter.Counter)), + ) + if r != ERROR_SUCCESS { + return nil, fmt.Errorf("pdhAddEnglishCounterW failed with status 0x%X: %v", r, err) + } + return &counter, nil +} + +func getValue(initialBufSize uint32, counter *win32PerformanceCounter) ([]float64, error) { + s, _, err := pdhCollectQueryData.Call(uintptr(counter.Query)) + if s != 0 && err != nil { + if s == PDH_NO_DATA { + return nil, fmt.Errorf("%w: this counter has not data", err) + } + return nil, err + } + + return getCounterArrayValue(initialBufSize, counter) +} + func GetGPUStat() (float64, error) { - shellPath, err := exec.LookPath("powershell.exe") - if err != nil || shellPath == "" { + counter, err := newWin32PerformanceCounter("gpu_utilization", "\\GPU Engine(*engtype_3D)\\Utilization Percentage") + if err != nil { return 0, err } - cmd := exec.Command( - shellPath, - "-Command", - `Write-Output "$([math]::Round((((Get-Counter "\GPU Engine(*engtype_3D)\Utilization Percentage").CounterSamples | where CookedValue).CookedValue | measure -sum).sum,2))"`, - ) - output, err1 := cmd.CombinedOutput() - if err1 != nil { - return 0, err1 + values, err := getValue(1024, counter) + if err != nil { + return 0, err } - t := strings.TrimSpace(string(output)) - gs, err2 := strconv.ParseFloat(t, 64) - if err2 != nil { - return 0, err2 + tot := sumArray(values) + if tot > 100 { + tot = 100 } - return gs, nil + return tot, nil +} + +func sumArray(arr []float64) float64 { + var sum float64 + for _, value := range arr { + sum += value + } + return sum }