Merge pull request #826 from Lomanic/issue250iocounters

[process][windows] Use win32 API in process.IOCounters() instead of slow WMI call #250

process/process_windows.go

@@ -29,6 +29,7 @@ var (
 
 	procQueryFullProcessImageNameW = common.Modkernel32.NewProc("QueryFullProcessImageNameW")
 	procGetPriorityClass           = common.Modkernel32.NewProc("GetPriorityClass")
+	procGetProcessIoCounters       = common.Modkernel32.NewProc("GetProcessIoCounters")
 )
 
 type SystemProcessInformation struct {
@@ -53,6 +54,17 @@ type MemoryInfoExStat struct {
 type MemoryMapsStat struct {
 }
 
+// ioCounters is an equivalent representation of IO_COUNTERS in the Windows API.
+// https://docs.microsoft.com/windows/win32/api/winnt/ns-winnt-io_counters
+type ioCounters struct {
+	ReadOperationCount  uint64
+	WriteOperationCount uint64
+	OtherOperationCount uint64
+	ReadTransferCount   uint64
+	WriteTransferCount  uint64
+	OtherTransferCount  uint64
+}
+
 type Win32_Process struct {
 	Name                  string
 	ExecutablePath        *string
@@ -479,18 +491,24 @@ func (p *Process) IOCounters() (*IOCountersStat, error) {
 }
 
 func (p *Process) IOCountersWithContext(ctx context.Context) (*IOCountersStat, error) {
-	dst, err := GetWin32ProcWithContext(ctx, p.Pid)
+	c, err := windows.OpenProcess(windows.PROCESS_QUERY_LIMITED_INFORMATION, false, uint32(p.Pid))
-	if err != nil || len(dst) == 0 {
+	if err != nil {
-		return nil, fmt.Errorf("could not get Win32Proc: %s", err)
+		return nil, err
+	}
+	defer windows.CloseHandle(c)
+	var ioCounters ioCounters
+	ret, _, err := procGetProcessIoCounters.Call(uintptr(c), uintptr(unsafe.Pointer(&ioCounters)))
+	if ret == 0 {
+		return nil, err
 	}
-	ret := &IOCountersStat{
+	stats := &IOCountersStat{
-		ReadCount:  uint64(dst[0].ReadOperationCount),
+		ReadCount:  ioCounters.ReadOperationCount,
-		ReadBytes:  uint64(dst[0].ReadTransferCount),
+		ReadBytes:  ioCounters.ReadTransferCount,
-		WriteCount: uint64(dst[0].WriteOperationCount),
+		WriteCount: ioCounters.WriteOperationCount,
-		WriteBytes: uint64(dst[0].WriteTransferCount),
+		WriteBytes: ioCounters.WriteTransferCount,
 	}
 
-	return ret, nil
+	return stats, nil
 }
 func (p *Process) NumCtxSwitches() (*NumCtxSwitchesStat, error) {
 	return p.NumCtxSwitchesWithContext(context.Background())