diff --git a/common/common_windows.go b/common/common_windows.go index df79d27c43bf84f663965da44f8a3a1e3c8f2075..d148e54038cc39bbc72d1dcda16c9ccfd0df4f2c 100644 --- a/common/common_windows.go +++ b/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 +} diff --git a/disk/disk_windows.go b/disk/disk_windows.go index cd1fe36d71dfaef644f3ba10a0d76f806d123297..51492906c586057db791d35f6d1355753fd864ad 100644 --- a/disk/disk_windows.go +++ b/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 }