Merge pull request #977 from AtakanColak/add-load-for-windows

[load][windows] Implement load.Avg() for Windows

internal/common/common_windows.go

@@ -4,6 +4,7 @@ package common
 
 import (
 	"context"
+	"fmt"
 	"path/filepath"
 	"strings"
 	"syscall"
@@ -69,13 +70,13 @@ var (
 	ProcNtWow64QueryInformationProcess64 = ModNt.NewProc("NtWow64QueryInformationProcess64")
 	ProcNtWow64ReadVirtualMemory64       = ModNt.NewProc("NtWow64ReadVirtualMemory64")
 
-	PdhOpenQuery                         = ModPdh.NewProc("PdhOpenQuery")
-	PdhAddCounter                        = ModPdh.NewProc("PdhAddCounterW")
-	PdhCollectQueryData                  = ModPdh.NewProc("PdhCollectQueryData")
-	PdhGetFormattedCounterValue          = ModPdh.NewProc("PdhGetFormattedCounterValue")
-	PdhCloseQuery                        = ModPdh.NewProc("PdhCloseQuery")
+	PdhOpenQuery                = ModPdh.NewProc("PdhOpenQuery")
+	PdhAddCounter               = ModPdh.NewProc("PdhAddCounterW")
+	PdhCollectQueryData         = ModPdh.NewProc("PdhCollectQueryData")
+	PdhGetFormattedCounterValue = ModPdh.NewProc("PdhGetFormattedCounterValue")
+	PdhCloseQuery               = ModPdh.NewProc("PdhCloseQuery")
 
-	procQueryDosDeviceW                  = Modkernel32.NewProc("QueryDosDeviceW")
+	procQueryDosDeviceW = Modkernel32.NewProc("QueryDosDeviceW")
 )
 
 type FILETIME struct {
@@ -93,7 +94,7 @@ func BytePtrToString(p *uint8) string {
 	return string(a[:i])
 }
 
-// CounterInfo
+// CounterInfo struct is used to track a windows performance counter
 // copied from https://github.com/mackerelio/mackerel-agent/
 type CounterInfo struct {
 	PostName    string
@@ -101,7 +102,7 @@ type CounterInfo struct {
 	Counter     windows.Handle
 }
 
-// CreateQuery XXX
+// CreateQuery with a PdhOpenQuery call
 // copied from https://github.com/mackerelio/mackerel-agent/
 func CreateQuery() (windows.Handle, error) {
 	var query windows.Handle
@@ -112,7 +113,7 @@ func CreateQuery() (windows.Handle, error) {
 	return query, nil
 }
 
-// CreateCounter XXX
+// CreateCounter with a PdhAddCounter call
 func CreateCounter(query windows.Handle, pname, cname string) (*CounterInfo, error) {
 	var counter windows.Handle
 	r, _, err := PdhAddCounter.Call(
@@ -130,6 +131,62 @@ func CreateCounter(query windows.Handle, pname, cname string) (*CounterInfo, err
 	}, nil
 }
 
+// GetCounterValue get counter value from handle
+// adapted from https://github.com/mackerelio/mackerel-agent/
+func GetCounterValue(counter windows.Handle) (float64, error) {
+	var value PDH_FMT_COUNTERVALUE_DOUBLE
+	r, _, err := PdhGetFormattedCounterValue.Call(uintptr(counter), PDH_FMT_DOUBLE, uintptr(0), uintptr(unsafe.Pointer(&value)))
+	if r != 0 && r != PDH_INVALID_DATA {
+		return 0.0, err
+	}
+	return value.DoubleValue, 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
+	}
+	var counter = Win32PerformanceCounter{
+		Query:       query,
+		PostName:    postName,
+		CounterName: counterName,
+	}
+	r, _, err := PdhAddCounter.Call(
+		uintptr(counter.Query),
+		uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(counter.CounterName))),
+		0,
+		uintptr(unsafe.Pointer(&counter.Counter)),
+	)
+	if r != 0 {
+		return nil, err
+	}
+	return &counter, nil
+}
+
+func (w *Win32PerformanceCounter) GetValue() (float64, error) {
+	r, _, err := PdhCollectQueryData.Call(uintptr(w.Query))
+	if r != 0 && err != nil {
+		if r == PDH_NO_DATA {
+			return 0.0, fmt.Errorf("%w: this counter has not data", err)
+		}
+		return 0.0, err
+	}
+
+	return GetCounterValue(w.Counter)
+}
+
+func ProcessorQueueLengthCounter() (*Win32PerformanceCounter, error) {
+	return NewWin32PerformanceCounter("processor_queue_length", `\System\Processor Queue Length`)
+}
+
 // WMIQueryWithContext - wraps wmi.Query with a timed-out context to avoid hanging
 func WMIQueryWithContext(ctx context.Context, query string, dst interface{}, connectServerArgs ...interface{}) error {
 	if _, ok := ctx.Deadline(); !ok {

load/load_test.go

@@ -7,7 +7,7 @@ import (
 	"github.com/shirou/gopsutil/internal/common"
 )
 
-func skipIfNotImplementedErr(t *testing.T, err error) {
+func skipIfNotImplementedErr(t testing.TB, err error) {
 	if err == common.ErrNotImplementedError {
 		t.Skip("not implemented")
 	}
@@ -67,3 +67,28 @@ func TestMiscStatString(t *testing.T) {
 	}
 	t.Log(e)
 }
+
+func BenchmarkLoad(b *testing.B) {
+
+	loadAvg := func(t testing.TB) {
+		v, err := Avg()
+		skipIfNotImplementedErr(t, err)
+		if err != nil {
+			t.Errorf("error %v", err)
+		}
+		empty := &AvgStat{}
+		if v == empty {
+			t.Errorf("error load: %v", v)
+		}
+	}
+
+	b.Run("FirstCall", func(b *testing.B) {
+		loadAvg(b)
+	})
+
+	b.Run("SubsequentCalls", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			loadAvg(b)
+		}
+	})
+}

load/load_windows.go

@@ -4,18 +4,73 @@ package load
 
 import (
 	"context"
+	"log"
+	"math"
+	"sync"
+	"time"
 
 	"github.com/shirou/gopsutil/internal/common"
 )
 
+var (
+	loadErr              error
+	loadAvg1M            float64 = 0.0
+	loadAvg5M            float64 = 0.0
+	loadAvg15M           float64 = 0.0
+	loadAvgMutex         sync.RWMutex
+	loadAvgGoroutineOnce sync.Once
+)
+
+// loadAvgGoroutine updates avg data by fetching current load by interval
+// TODO instead of this goroutine, we can register a Win32 counter just as psutil does
+// see https://psutil.readthedocs.io/en/latest/#psutil.getloadavg
+// code https://github.com/giampaolo/psutil/blob/8415355c8badc9c94418b19bdf26e622f06f0cce/psutil/arch/windows/wmi.c
+func loadAvgGoroutine() {
+	var (
+		samplingFrequency time.Duration = 5 * time.Second
+		loadAvgFactor1M   float64       = 1 / math.Exp(samplingFrequency.Seconds()/time.Minute.Seconds())
+		loadAvgFactor5M   float64       = 1 / math.Exp(samplingFrequency.Seconds()/(5*time.Minute).Seconds())
+		loadAvgFactor15M  float64       = 1 / math.Exp(samplingFrequency.Seconds()/(15*time.Minute).Seconds())
+		currentLoad       float64
+	)
+
+	counter, err := common.ProcessorQueueLengthCounter()
+	if err != nil || counter == nil {
+		log.Println("gopsutil: unexpected processor queue length counter error, please file an issue on github")
+		return
+	}
+
+	tick := time.NewTicker(samplingFrequency).C
+	for {
+		currentLoad, err = counter.GetValue()
+		loadAvgMutex.Lock()
+		loadErr = err
+		loadAvg1M = loadAvg1M*loadAvgFactor1M + currentLoad*(1-loadAvgFactor1M)
+		loadAvg5M = loadAvg5M*loadAvgFactor5M + currentLoad*(1-loadAvgFactor5M)
+		loadAvg15M = loadAvg15M*loadAvgFactor15M + currentLoad*(1-loadAvgFactor15M)
+		loadAvgMutex.Unlock()
+		<-tick
+	}
+}
+
+// Avg for Windows may return 0 values for the first few 5 second intervals
 func Avg() (*AvgStat, error) {
 	return AvgWithContext(context.Background())
 }
 
 func AvgWithContext(ctx context.Context) (*AvgStat, error) {
-	ret := AvgStat{}
+	loadAvgGoroutineOnce.Do(func() {
+		go loadAvgGoroutine()
+	})
+	loadAvgMutex.RLock()
+	defer loadAvgMutex.RUnlock()
+	ret := AvgStat{
+		Load1:  loadAvg1M,
+		Load5:  loadAvg5M,
+		Load15: loadAvg15M,
+	}
 
-	return &ret, common.ErrNotImplementedError
+	return &ret, loadErr
 }
 
 func Misc() (*MiscStat, error) {