Merge pull request #11 from hungrybirder/optimize_on_linux

Code Optimization

go.mod

@@ -0,0 +1,3 @@
+module github.com/struCoder/pidusage
+
+go 1.13

pidusage.go

@@ -2,6 +2,7 @@ package pidusage
 
 import (
 	"errors"
+	"fmt"
 	"io/ioutil"
 	"math"
 	"os/exec"
@@ -12,6 +13,11 @@ import (
 	"sync"
 )
 
+const (
+	statTypePS   = "ps"
+	statTypeProc = "proc"
+)
+
 // SysInfo will record cpu and memory data
 type SysInfo struct {
 	CPU    float64
@@ -37,11 +43,10 @@ var history map[int]Stat
 var historyLock sync.Mutex
 var eol string
 
-func wrapper(statType string) func(pid int) (*SysInfo, error) {
-	return func(pid int) (*SysInfo, error) {
-		return stat(pid, statType)
-	}
-}
+// Linux platform
+var clkTck float64 = 100    // default
+var pageSize float64 = 4096 // default
+
 func init() {
 	platform = runtime.GOOS
 	if eol = "\n"; strings.Index(platform, "win") == 0 {
@@ -57,7 +62,31 @@ func init() {
 	fnMap["linux"] = wrapper("proc")
 	fnMap["netbsd"] = wrapper("proc")
 	fnMap["win"] = wrapper("win")
+
+	if platform == "linux" || platform == "netbsd" {
+		initProc()
+	}
 }
+
+func initProc() {
+	clkTckStdout, err := exec.Command("getconf", "CLK_TCK").Output()
+	if err == nil {
+		clkTck = parseFloat(formatStdOut(clkTckStdout, 0)[0])
+	}
+
+	pageSizeStdout, err := exec.Command("getconf", "PAGESIZE").Output()
+	if err == nil {
+		pageSize = parseFloat(formatStdOut(pageSizeStdout, 0)[0])
+	}
+
+}
+
+func wrapper(statType string) func(pid int) (*SysInfo, error) {
+	return func(pid int) (*SysInfo, error) {
+		return stat(pid, statType)
+	}
+}
+
 func formatStdOut(stdout []byte, userfulIndex int) []string {
 	infoArr := strings.Split(string(stdout), eol)[userfulIndex]
 	ret := strings.Fields(infoArr)
@@ -69,86 +98,93 @@ func parseFloat(val string) float64 {
 	return floatVal
 }
 
-func stat(pid int, statType string) (*SysInfo, error) {
+func statFromPS(pid int) (*SysInfo, error) {
+	sysInfo := &SysInfo{}
+	args := "-o pcpu,rss -p"
+	if platform == "aix" {
+		args = "-o pcpu,rssize -p"
+	}
+	stdout, _ := exec.Command("ps", args, strconv.Itoa(pid)).Output()
+	ret := formatStdOut(stdout, 1)
+	if len(ret) == 0 {
+		return sysInfo, errors.New("Can't find process with this PID: " + strconv.Itoa(pid))
+	}
+	sysInfo.CPU = parseFloat(ret[0])
+	sysInfo.Memory = parseFloat(ret[1]) * 1024
+	return sysInfo, nil
+}
+
+func statFromProc(pid int) (*SysInfo, error) {
 	sysInfo := &SysInfo{}
+	uptimeFileBytes, err := ioutil.ReadFile(path.Join("/proc", "uptime"))
+	if err != nil {
+		return nil, err
+	}
+	uptime := parseFloat(strings.Split(string(uptimeFileBytes), " ")[0])
+
+	procStatFileBytes, err := ioutil.ReadFile(path.Join("/proc", strconv.Itoa(pid), "stat"))
+	if err != nil {
+		return nil, err
+	}
+	splitAfter := strings.SplitAfter(string(procStatFileBytes), ")")
+
+	if len(splitAfter) == 0 || len(splitAfter) == 1 {
+		return sysInfo, errors.New("Can't find process with this PID: " + strconv.Itoa(pid))
+	}
+	infos := strings.Split(splitAfter[1], " ")
+	stat := &Stat{
+		utime:  parseFloat(infos[12]),
+		stime:  parseFloat(infos[13]),
+		cutime: parseFloat(infos[14]),
+		cstime: parseFloat(infos[15]),
+		start:  parseFloat(infos[20]) / clkTck,
+		rss:    parseFloat(infos[22]),
+		uptime: uptime,
+	}
+
+	_stime := 0.0
+	_utime := 0.0
+
+	historyLock.Lock()
+	defer historyLock.Unlock()
+
 	_history := history[pid]
-	if statType == "ps" {
-		args := "-o pcpu,rss -p"
-		if platform == "aix" {
-			args = "-o pcpu,rssize -p"
-		}
-		stdout, _ := exec.Command("ps", args, strconv.Itoa(pid)).Output()
-		ret := formatStdOut(stdout, 1)
-		if len(ret) == 0 {
-			return sysInfo, errors.New("Can't find process with this PID: " + strconv.Itoa(pid))
-		}
-		sysInfo.CPU = parseFloat(ret[0])
-		sysInfo.Memory = parseFloat(ret[1]) * 1024
-	} else if statType == "proc" {
-		// default clkTck and pageSize
-		var clkTck float64 = 100
-		var pageSize float64 = 4096
-
-		uptimeFileBytes, err := ioutil.ReadFile(path.Join("/proc", "uptime"))
-		uptime := parseFloat(strings.Split(string(uptimeFileBytes), " ")[0])
-
-		clkTckStdout, err := exec.Command("getconf", "CLK_TCK").Output()
-		if err == nil {
-			clkTck = parseFloat(formatStdOut(clkTckStdout, 0)[0])
-		}
-
-		pageSizeStdout, err := exec.Command("getconf", "PAGESIZE").Output()
-		if err == nil {
-			pageSize = parseFloat(formatStdOut(pageSizeStdout, 0)[0])
-		}
-
-		procStatFileBytes, err := ioutil.ReadFile(path.Join("/proc", strconv.Itoa(pid), "stat"))
-		splitAfter := strings.SplitAfter(string(procStatFileBytes), ")")
-
-		if len(splitAfter) == 0 || len(splitAfter) == 1 {
-			return sysInfo, errors.New("Can't find process with this PID: " + strconv.Itoa(pid))
-		}
-		infos := strings.Split(splitAfter[1], " ")
-		stat := &Stat{
-			utime:  parseFloat(infos[12]),
-			stime:  parseFloat(infos[13]),
-			cutime: parseFloat(infos[14]),
-			cstime: parseFloat(infos[15]),
-			start:  parseFloat(infos[20]) / clkTck,
-			rss:    parseFloat(infos[22]),
-			uptime: uptime,
-		}
-
-		_stime := 0.0
-		_utime := 0.0
-		if _history.stime != 0 {
-			_stime = _history.stime
-		}
-
-		if _history.utime != 0 {
-			_utime = _history.utime
-		}
-		total := stat.stime - _stime + stat.utime - _utime
-		total = total / clkTck
-
-		seconds := stat.start - uptime
-		if _history.uptime != 0 {
-			seconds = uptime - _history.uptime
-		}
-
-		seconds = math.Abs(seconds)
-		if seconds == 0 {
-			seconds = 1
-		}
-
-		historyLock.Lock()
-		history[pid] = *stat
-		historyLock.Unlock()
-		sysInfo.CPU = (total / seconds) * 100
-		sysInfo.Memory = stat.rss * pageSize
+
+	if _history.stime != 0 {
+		_stime = _history.stime
+	}
+
+	if _history.utime != 0 {
+		_utime = _history.utime
+	}
+	total := stat.stime - _stime + stat.utime - _utime
+	total = total / clkTck
+
+	seconds := stat.start - uptime
+	if _history.uptime != 0 {
+		seconds = uptime - _history.uptime
 	}
+
+	seconds = math.Abs(seconds)
+	if seconds == 0 {
+		seconds = 1
+	}
+
+	history[pid] = *stat
+	sysInfo.CPU = (total / seconds) * 100
+	sysInfo.Memory = stat.rss * pageSize
 	return sysInfo, nil
+}
 
+func stat(pid int, statType string) (*SysInfo, error) {
+	switch statType {
+	case statTypePS:
+		return statFromPS(pid)
+	case statTypeProc:
+		return statFromProc(pid)
+	default:
+		return nil, fmt.Errorf("Unsupported OS %s", runtime.GOOS)
+	}
 }
 
 // GetStat will return current system CPU and memory data

pidusage_test.go

@@ -1 +1,34 @@
 package pidusage
+
+import (
+	"os"
+	"testing"
+)
+
+var pid = os.Getpid()
+
+func BenchmarkGetStat(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		GetStat(pid)
+	}
+}
+
+// Before optimize
+// $ go clean -testcache && go test -test.v -bench=BenchmarkGetStat
+// goos: linux
+// goarch: amd64
+// pkg: github.com/struCoder/pidusage
+// BenchmarkGetStat
+// BenchmarkGetStat-12    	     470	   2690727 ns/op
+// PASS
+// ok  	github.com/struCoder/pidusage	1.533s
+
+// After optimize
+// $ go clean -testcache && go test -test.v -bench=BenchmarkGetStat
+// goos: linux
+// goarch: amd64
+// pkg: github.com/struCoder/pidusage
+// BenchmarkGetStat
+// BenchmarkGetStat-12    	   28416	     36234 ns/op
+// PASS
+// ok  	github.com/struCoder/pidusage	1.472s