From f98100720e5c0c28e1665b404f0cb1340677be6a Mon Sep 17 00:00:00 2001
From: Lomanic <Lomanic@users.noreply.github.com>
Date: Sun, 1 Mar 2020 23:30:05 +0100
Subject: [PATCH] [process][windows] Fix #846 use win32 API in
 disk.IOCounters() instead of slow WMI call returning wrong values

---
 disk/disk_windows.go | 91 ++++++++++++++++++++++++++++++++++------------------
 1 file changed, 60 insertions(+), 31 deletions(-)

diff --git a/disk/disk_windows.go b/disk/disk_windows.go
index 02c965d..f871e9d 100644
--- a/disk/disk_windows.go
+++ b/disk/disk_windows.go
@@ -5,6 +5,8 @@ package disk
 import (
 	"bytes"
 	"context"
+	"fmt"
+	"syscall"
 	"unsafe"
 
 	"github.com/shirou/gopsutil/internal/common"
@@ -23,18 +25,24 @@ var (
 	FileReadOnlyVolume  = int64(524288) // 0x00080000
 )
 
-type Win32_PerfFormattedData struct {
-	Name                    string
-	AvgDiskBytesPerRead     uint64
-	AvgDiskBytesPerWrite    uint64
-	AvgDiskReadQueueLength  uint64
-	AvgDiskWriteQueueLength uint64
-	AvgDisksecPerRead       uint64
-	AvgDisksecPerWrite      uint64
+// diskPerformance is an equivalent representation of DISK_PERFORMANCE in the Windows API.
+// https://docs.microsoft.com/fr-fr/windows/win32/api/winioctl/ns-winioctl-disk_performance
+type diskPerformance struct {
+	BytesRead           int64
+	BytesWritten        int64
+	ReadTime            int64
+	WriteTime           int64
+	IdleTime            int64
+	ReadCount           uint32
+	WriteCount          uint32
+	QueueDepth          uint32
+	SplitCount          uint32
+	QueryTime           int64
+	StorageDeviceNumber uint32
+	StorageManagerName  [8]uint16
+	alignmentPadding    uint32 // necessary for 32bit support, see https://github.com/elastic/beats/pull/16553
 }
 
-const WaitMSec = 500
-
 func Usage(path string) (*UsageStat, error) {
 	return UsageWithContext(context.Background(), path)
 }
@@ -136,31 +144,52 @@ func IOCounters(names ...string) (map[string]IOCountersStat, error) {
 }
 
 func IOCountersWithContext(ctx context.Context, names ...string) (map[string]IOCountersStat, error) {
-	ret := make(map[string]IOCountersStat, 0)
-	var dst []Win32_PerfFormattedData
+	// https://github.com/giampaolo/psutil/blob/544e9daa4f66a9f80d7bf6c7886d693ee42f0a13/psutil/arch/windows/disk.c#L83
+	drivemap := make(map[string]IOCountersStat, 0)
+	var diskPerformance diskPerformance
 
-	err := common.WMIQueryWithContext(ctx, "SELECT * FROM Win32_PerfFormattedData_PerfDisk_LogicalDisk", &dst)
+	lpBuffer := make([]uint16, 254)
+	lpBufferLen, err := windows.GetLogicalDriveStrings(uint32(len(lpBuffer)), &lpBuffer[0])
 	if err != nil {
-		return ret, err
+		return drivemap, err
 	}
-	for _, d := range dst {
-		if len(d.Name) > 3 { // not get _Total or Harddrive
-			continue
-		}
-
-		if len(names) > 0 && !common.StringsHas(names, d.Name) {
-			continue
-		}
+	for _, v := range lpBuffer[:lpBufferLen] {
+		if 'A' <= v && v <= 'Z' {
+			path := string(v) + ":"
+			typepath, _ := windows.UTF16PtrFromString(path)
+			typeret := windows.GetDriveType(typepath)
+			if typeret == 0 {
+				return drivemap, windows.GetLastError()
+			}
+			if typeret != windows.DRIVE_FIXED {
+				continue
+			}
+			szDevice := fmt.Sprintf(`\\.\%s`, path)
+			const IOCTL_DISK_PERFORMANCE = 0x70020
+			h, err := windows.CreateFile(syscall.StringToUTF16Ptr(szDevice), 0, windows.FILE_SHARE_READ|windows.FILE_SHARE_WRITE, nil, windows.OPEN_EXISTING, 0, 0)
+			if err != nil {
+				if err == windows.ERROR_FILE_NOT_FOUND {
+					continue
+				}
+				return drivemap, err
+			}
+			defer windows.CloseHandle(h)
 
-		ret[d.Name] = IOCountersStat{
-			Name:       d.Name,
-			ReadCount:  uint64(d.AvgDiskReadQueueLength),
-			WriteCount: d.AvgDiskWriteQueueLength,
-			ReadBytes:  uint64(d.AvgDiskBytesPerRead),
-			WriteBytes: uint64(d.AvgDiskBytesPerWrite),
-			ReadTime:   d.AvgDisksecPerRead,
-			WriteTime:  d.AvgDisksecPerWrite,
+			var diskPerformanceSize uint32
+			err = windows.DeviceIoControl(h, IOCTL_DISK_PERFORMANCE, nil, 0, (*byte)(unsafe.Pointer(&diskPerformance)), uint32(unsafe.Sizeof(diskPerformance)), &diskPerformanceSize, nil)
+			if err != nil {
+				return drivemap, err
+			}
+			drivemap[path] = IOCountersStat{
+				ReadBytes:  uint64(diskPerformance.BytesRead),
+				WriteBytes: uint64(diskPerformance.BytesWritten),
+				ReadCount:  uint64(diskPerformance.ReadCount),
+				WriteCount: uint64(diskPerformance.WriteCount),
+				ReadTime:   uint64(diskPerformance.ReadTime / 10000 / 1000), // convert to ms: https://github.com/giampaolo/psutil/issues/1012
+				WriteTime:  uint64(diskPerformance.WriteTime / 10000 / 1000),
+				Name:       path,
+			}
 		}
 	}
-	return ret, nil
+	return drivemap, nil
 }