Merge pull request #347 from danielnelson/io-counters-for-names

Add disk.IOCountersForNames function
shirou · Apr 10, 2017 · f5781ca · f5781ca
2 parents e49a95f + ab6db76
commit f5781ca
Show file tree

Hide file tree

Showing 8 changed files with 35 additions and 7 deletions.
diff --git 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_cgo.go b/disk/disk_darwin_cgo.go
@@ -30,9 +30,11 @@ import (
 	"errors"
 	"strings"
 	"unsafe"
+
+	"github.com/shirou/gopsutil/internal/common"
 )
 
-func IOCounters() (map[string]IOCountersStat, error) {
+func IOCountersForNames(names []string) (map[string]IOCountersStat, error) {
 	if C.StartIOCounterFetch() == 0 {
 		return nil, errors.New("Unable to fetch disk list")
 	}
@@ -78,6 +80,10 @@ func IOCounters() (map[string]IOCountersStat, error) {
 			Name:       strings.TrimFunc(C.GoStringN(&di.DiskName[0], C.MAX_DISK_NAME), isRuneNull),
 		}
 
+		if len(names) > 0 && !common.StringsHas(names, d.Name) {
+			continue
+		}
+
 		ret[d.Name] = d
 	}
 

diff --git a/disk/disk_darwin_nocgo.go b/disk/disk_darwin_nocgo.go
@@ -5,6 +5,6 @@ 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_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
@@ -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
@@ -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
@@ -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
@@ -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),