disk: implements DiskIoCounterInfo on windows. Note: only read/write counts during 500 msec.

Thanks for mackerel-agent!

common/common_windows.go

@@ -10,12 +10,54 @@ import (
 	"unsafe"
 )
 
+// for double values
+type PDH_FMT_COUNTERVALUE_DOUBLE struct {
+	CStatus     uint32
+	DoubleValue float64
+}
+
+// for 64 bit integer values
+type PDH_FMT_COUNTERVALUE_LARGE struct {
+	CStatus    uint32
+	LargeValue int64
+}
+
+// for long values
+type PDH_FMT_COUNTERVALUE_LONG struct {
+	CStatus   uint32
+	LongValue int32
+	padding   [4]byte
+}
+
+// windows system const
+const (
+	ERROR_SUCCESS        = 0
+	ERROR_FILE_NOT_FOUND = 2
+	DRIVE_REMOVABLE      = 2
+	DRIVE_FIXED          = 3
+	HKEY_LOCAL_MACHINE   = 0x80000002
+	RRF_RT_REG_SZ        = 0x00000002
+	RRF_RT_REG_DWORD     = 0x00000010
+	PDH_FMT_LONG         = 0x00000100
+	PDH_FMT_DOUBLE       = 0x00000200
+	PDH_FMT_LARGE        = 0x00000400
+	PDH_INVALID_DATA     = 0xc0000bc6
+	PDH_INVALID_HANDLE   = 0xC0000bbc
+	PDH_NO_DATA          = 0x800007d5
+)
+
 var (
 	Modkernel32 = syscall.NewLazyDLL("kernel32.dll")
 	ModNt       = syscall.NewLazyDLL("ntdll.dll")
+	ModPdh      = syscall.NewLazyDLL("pdh.dll")
 
 	ProcGetSystemTimes           = Modkernel32.NewProc("GetSystemTimes")
 	ProcNtQuerySystemInformation = ModNt.NewProc("NtQuerySystemInformation")
+	PdhOpenQuery                 = ModPdh.NewProc("PdhOpenQuery")
+	PdhAddCounter                = ModPdh.NewProc("PdhAddCounterW")
+	PdhCollectQueryData          = ModPdh.NewProc("PdhCollectQueryData")
+	PdhGetFormattedCounterValue  = ModPdh.NewProc("PdhGetFormattedCounterValue")
+	PdhCloseQuery                = ModPdh.NewProc("PdhCloseQuery")
 )
 
 type FILETIME struct {
@@ -47,3 +89,40 @@ func GetWmic(target string, query string) ([]string, error) {
 	// skip first two line
 	return lines[2:], nil
 }
+
+// CounterInfo
+// copied from https://github.com/mackerelio/mackerel-agent/
+type CounterInfo struct {
+	PostName    string
+	CounterName string
+	Counter     syscall.Handle
+}
+
+// CreateQuery XXX
+// copied from https://github.com/mackerelio/mackerel-agent/
+func CreateQuery() (syscall.Handle, error) {
+	var query syscall.Handle
+	r, _, err := PdhOpenQuery.Call(0, 0, uintptr(unsafe.Pointer(&query)))
+	if r != 0 {
+		return 0, err
+	}
+	return query, nil
+}
+
+// CreateCounter XXX
+func CreateCounter(query syscall.Handle, pname, cname string) (*CounterInfo, error) {
+	var counter syscall.Handle
+	r, _, err := PdhAddCounter.Call(
+		uintptr(query),
+		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(cname))),
+		0,
+		uintptr(unsafe.Pointer(&counter)))
+	if r != 0 {
+		return nil, err
+	}
+	return &CounterInfo{
+		PostName:    pname,
+		CounterName: cname,
+		Counter:     counter,
+	}, nil
+}

disk/disk_windows.go

@@ -4,7 +4,9 @@ package disk
 
 import (
 	"bytes"
+	"fmt"
 	"syscall"
+	"time"
 	"unsafe"
 
 	common "github.com/shirou/gopsutil/common"
@@ -22,6 +24,8 @@ var (
 	FileReadOnlyVolume  = int64(524288) // 0x00080000
 )
 
+const WaitMSec = 500
+
 func DiskUsage(path string) (DiskUsageStat, error) {
 	ret := DiskUsageStat{}
 
@@ -115,5 +119,79 @@ func DiskPartitions(all bool) ([]DiskPartitionStat, error) {
 
 func DiskIOCounters() (map[string]DiskIOCountersStat, error) {
 	ret := make(map[string]DiskIOCountersStat, 0)
-	return ret, common.NotImplementedError
+	query, err := common.CreateQuery()
+	if err != nil {
+		return ret, err
+	}
+
+	drivebuf := make([]byte, 256)
+	r, _, err := procGetLogicalDriveStringsW.Call(
+		uintptr(len(drivebuf)),
+		uintptr(unsafe.Pointer(&drivebuf[0])))
+
+	if r == 0 {
+		return ret, err
+	}
+
+	drivemap := make(map[string][]*common.CounterInfo, 0)
+	for _, v := range drivebuf {
+		if v >= 65 && v <= 90 {
+			drive := string(v)
+			r, _, err = procGetDriveType.Call(uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(drive + `:\`))))
+			if r != common.DRIVE_FIXED {
+				continue
+			}
+			drivemap[drive] = make([]*common.CounterInfo, 0, 2)
+			var counter *common.CounterInfo
+
+			counter, err = common.CreateCounter(query,
+				"read",
+				fmt.Sprintf(`\PhysicalDisk(0 %s:)\Disk Reads/sec`, drive))
+			if err != nil {
+				return nil, err
+			}
+			drivemap[drive] = append(drivemap[drive], counter)
+			counter, err = common.CreateCounter(query,
+				"write",
+				fmt.Sprintf(`\PhysicalDisk(0 %s:)\Disk Writes/sec`, drive))
+			if err != nil {
+				return nil, err
+			}
+			drivemap[drive] = append(drivemap[drive], counter)
+		}
+	}
+	r, _, err = common.PdhCollectQueryData.Call(uintptr(query))
+	if r != 0 && err != nil {
+		return nil, err
+	}
+	time.Sleep(time.Duration(WaitMSec) * time.Millisecond)
+	r, _, err = common.PdhCollectQueryData.Call(uintptr(query))
+	if r != 0 && err != nil {
+		return nil, err
+	}
+
+	for drive, counters := range drivemap {
+		stat := DiskIOCountersStat{}
+		for _, v := range counters {
+			var fmtValue common.PDH_FMT_COUNTERVALUE_LARGE
+			r, _, err := common.PdhGetFormattedCounterValue.Call(uintptr(v.Counter), common.PDH_FMT_LARGE, uintptr(0), uintptr(unsafe.Pointer(&fmtValue)))
+			if r != 0 && r != common.PDH_INVALID_DATA {
+				return nil, err
+			}
+
+			switch v.PostName {
+			case "read":
+				stat.ReadCount = uint64(fmtValue.LargeValue)
+			case "write":
+				stat.WriteCount = uint64(fmtValue.LargeValue)
+			default:
+				return ret, fmt.Errorf("unknown postname: %s", v.PostName)
+			}
+			stat.Name = drive
+		}
+		ret[drive] = stat
+	}
+	fmt.Println(ret)
+
+	return ret, nil
 }