From 3f35f001da07342bc08834f96f3f9007196c1a75 Mon Sep 17 00:00:00 2001
From: Daniel Nelson <daniel.nelson@influxdb.com>
Date: Thu, 6 Apr 2017 17:54:50 -0700
Subject: [PATCH] Add disk.IOCountersForNames function

Operates like disk.IOCounters, but accepts an array of names to limit
the results.
---
 disk/disk.go          | 4 ++++
 disk/disk_darwin.go   | 4 ++++
 disk/disk_fallback.go | 2 +-
 disk/disk_freebsd.go  | 6 +++++-
 disk/disk_linux.go    | 7 ++++++-
 disk/disk_openbsd.go  | 6 +++++-
 disk/disk_windows.go  | 7 ++++++-
 7 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/disk/disk.go b/disk/disk.go
index a2c4720..b2c9419 100644
--- a/disk/disk.go
+++ b/disk/disk.go
@@ -62,3 +62,7 @@ func (d IOCountersStat) String() string {
 	s, _ := json.Marshal(d)
 	return string(s)
 }
+
+func IOCounters() (map[string]IOCountersStat, error) {
+	return IOCountersForNames([]string{})
+}
diff --git a/disk/disk_darwin.go b/disk/disk_darwin.go
index dc642df..21ec491 100644
--- a/disk/disk_darwin.go
+++ b/disk/disk_darwin.go
@@ -87,6 +87,10 @@ func Partitions(all bool) ([]PartitionStat, error) {
 	return ret, nil
 }
 
+func IOCountersForNames(names []string) (map[string]IOCountersStat, error) {
+	return nil, common.ErrNotImplementedError
+}
+
 func Getfsstat(buf []Statfs_t, flags int) (n int, err error) {
 	var _p0 unsafe.Pointer
 	var bufsize uintptr
diff --git a/disk/disk_fallback.go b/disk/disk_fallback.go
index 6fb01a9..8bff71c 100644
--- a/disk/disk_fallback.go
+++ b/disk/disk_fallback.go
@@ -4,7 +4,7 @@ package disk
 
 import "github.com/shirou/gopsutil/internal/common"
 
-func IOCounters() (map[string]IOCountersStat, error) {
+func IOCountersForNames(names []string) (map[string]IOCountersStat, error) {
 	return nil, common.ErrNotImplementedError
 }
 
diff --git a/disk/disk_freebsd.go b/disk/disk_freebsd.go
index 6e76f31..7a84a2c 100644
--- a/disk/disk_freebsd.go
+++ b/disk/disk_freebsd.go
@@ -94,7 +94,7 @@ func Partitions(all bool) ([]PartitionStat, error) {
 	return ret, nil
 }
 
-func IOCounters() (map[string]IOCountersStat, error) {
+func IOCountersForNames(names []string) (map[string]IOCountersStat, error) {
 	// statinfo->devinfo->devstat
 	// /usr/include/devinfo.h
 	ret := make(map[string]IOCountersStat)
@@ -119,6 +119,10 @@ func IOCounters() (map[string]IOCountersStat, error) {
 		un := strconv.Itoa(int(d.Unit_number))
 		name := common.IntToString(d.Device_name[:]) + un
 
+		if len(names) > 0 && !common.StringsHas(names, name) {
+			continue
+		}
+
 		ds := IOCountersStat{
 			ReadCount:  d.Operations[DEVSTAT_READ],
 			WriteCount: d.Operations[DEVSTAT_WRITE],
diff --git a/disk/disk_linux.go b/disk/disk_linux.go
index 51f17cd..a50e66f 100644
--- a/disk/disk_linux.go
+++ b/disk/disk_linux.go
@@ -272,7 +272,7 @@ func getFileSystems() ([]string, error) {
 	return ret, nil
 }
 
-func IOCounters() (map[string]IOCountersStat, error) {
+func IOCountersForNames(names []string) (map[string]IOCountersStat, error) {
 	filename := common.HostProc("diskstats")
 	lines, err := common.ReadLines(filename)
 	if err != nil {
@@ -288,6 +288,11 @@ func IOCounters() (map[string]IOCountersStat, error) {
 			continue
 		}
 		name := fields[2]
+
+		if len(names) > 0 && !common.StringsHas(names, name) {
+			continue
+		}
+
 		reads, err := strconv.ParseUint((fields[3]), 10, 64)
 		if err != nil {
 			return ret, err
diff --git a/disk/disk_openbsd.go b/disk/disk_openbsd.go
index 2129b2b..b836fa3 100644
--- a/disk/disk_openbsd.go
+++ b/disk/disk_openbsd.go
@@ -63,7 +63,7 @@ func Partitions(all bool) ([]PartitionStat, error) {
 	return ret, nil
 }
 
-func IOCounters() (map[string]IOCountersStat, error) {
+func IOCountersForNames(names []string) (map[string]IOCountersStat, error) {
 	ret := make(map[string]IOCountersStat)
 
 	r, err := syscall.Sysctl("hw.diskstats")
@@ -84,6 +84,10 @@ func IOCounters() (map[string]IOCountersStat, error) {
 		}
 		name := common.IntToString(d.Name[:])
 
+		if len(names) > 0 && !common.StringsHas(names, name) {
+			continue
+		}
+
 		ds := IOCountersStat{
 			ReadCount:  d.Rxfer,
 			WriteCount: d.Wxfer,
diff --git a/disk/disk_windows.go b/disk/disk_windows.go
index b3a30d6..dc79b8b 100644
--- a/disk/disk_windows.go
+++ b/disk/disk_windows.go
@@ -129,7 +129,7 @@ func Partitions(all bool) ([]PartitionStat, error) {
 	return ret, nil
 }
 
-func IOCounters() (map[string]IOCountersStat, error) {
+func IOCountersForNames(names []string) (map[string]IOCountersStat, error) {
 	ret := make(map[string]IOCountersStat, 0)
 	var dst []Win32_PerfFormattedData
 
@@ -141,6 +141,11 @@ func IOCounters() (map[string]IOCountersStat, error) {
 		if len(d.Name) > 3 { // not get _Total or Harddrive
 			continue
 		}
+
+		if len(names) > 0 && !common.StringsHas(names, name) {
+			continue
+		}
+
 		ret[d.Name] = IOCountersStat{
 			Name:       d.Name,
 			ReadCount:  uint64(d.AvgDiskReadQueueLength),