From 7ee4a4c6ff23179775d0365ec3593583d49c0b95 Mon Sep 17 00:00:00 2001
From: Leonid Podolny <leonid@digitalocean.com>
Date: Sun, 1 Oct 2017 21:27:31 -0400
Subject: [PATCH] Fix Threads() signature to better match python original

---
 process/process_darwin.go   |  4 ++--
 process/process_fallback.go |  2 +-
 process/process_freebsd.go  |  4 ++--
 process/process_linux.go    | 37 ++++++++++++++++++++++++++++++-------
 process/process_openbsd.go  |  4 ++--
 process/process_test.go     | 20 ++++++++++++++++++++
 process/process_windows.go  |  4 ++--
 7 files changed, 59 insertions(+), 16 deletions(-)

diff --git a/process/process_darwin.go b/process/process_darwin.go
index 1a9ed94..aa17d1d 100644
--- a/process/process_darwin.go
+++ b/process/process_darwin.go
@@ -268,8 +268,8 @@ func (p *Process) NumThreads() (int32, error) {
 	}
 	return int32(len(r)), nil
 }
-func (p *Process) Threads() ([]int32, error) {
-	ret := make([]int32, 0)
+func (p *Process) Threads() (map[int32]*cpu.TimesStat, error) {
+	ret := make(map[int32]*cpu.TimesStat)
 	return ret, common.ErrNotImplementedError
 }
 
diff --git a/process/process_fallback.go b/process/process_fallback.go
index 7c7d71d..12bd3a6 100644
--- a/process/process_fallback.go
+++ b/process/process_fallback.go
@@ -95,7 +95,7 @@ func (p *Process) NumFDs() (int32, error) {
 func (p *Process) NumThreads() (int32, error) {
 	return 0, common.ErrNotImplementedError
 }
-func (p *Process) Threads() ([]int32, error) {
+func (p *Process) Threads() (map[int32]*cpu.TimesStat, error) {
 	return nil, common.ErrNotImplementedError
 }
 func (p *Process) Times() (*cpu.TimesStat, error) {
diff --git a/process/process_freebsd.go b/process/process_freebsd.go
index a6b978d..aa3e85c 100644
--- a/process/process_freebsd.go
+++ b/process/process_freebsd.go
@@ -204,8 +204,8 @@ func (p *Process) NumThreads() (int32, error) {
 
 	return k.Numthreads, nil
 }
-func (p *Process) Threads() ([]int32, error) {
-	ret := make([]int32, 0)
+func (p *Process) Threads() (map[int32]*cpu.TimesStat, error) {
+	ret := make(map[int32]*cpu.TimesStat)
 	return ret, common.ErrNotImplementedError
 }
 func (p *Process) Times() (*cpu.TimesStat, error) {
diff --git a/process/process_linux.go b/process/process_linux.go
index dcf4d7a..1839eb2 100644
--- a/process/process_linux.go
+++ b/process/process_linux.go
@@ -290,12 +290,24 @@ func (p *Process) NumThreads() (int32, error) {
 	return p.numThreads, nil
 }
 
-// Threads returns a map of threads
-//
-// Notice: Not implemented yet. always returns empty map.
-func (p *Process) Threads() ([]int32, error) {
+func (p *Process) Threads() (map[int32]*cpu.TimesStat, error) {
+	ret := make(map[int32]*cpu.TimesStat)
 	taskPath := common.HostProc(strconv.Itoa(int(p.Pid)), "task")
-	return readPidsFromDir(taskPath)
+
+	tids, err := readPidsFromDir(taskPath)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, tid := range tids {
+		_, _, cpuTimes, _, _, _, err := p.fillFromTIDStat(tid)
+		if err != nil {
+			return nil, err
+		}
+		ret[tid] = cpuTimes
+	}
+
+	return ret, nil
 }
 
 // Times returns CPU times of the process.
@@ -922,9 +934,16 @@ func (p *Process) fillFromStatus() error {
 	return nil
 }
 
-func (p *Process) fillFromStat() (string, int32, *cpu.TimesStat, int64, uint32, int32, error) {
+func (p *Process) fillFromTIDStat(tid int32) (string, int32, *cpu.TimesStat, int64, uint32, int32, error) {
 	pid := p.Pid
-	statPath := common.HostProc(strconv.Itoa(int(pid)), "stat")
+	var statPath string
+
+	if tid == -1 {
+		statPath = common.HostProc(strconv.Itoa(int(pid)), "stat")
+	} else {
+		statPath = common.HostProc(strconv.Itoa(int(pid)), "task", strconv.Itoa(int(tid)), "stat")
+	}
+
 	contents, err := ioutil.ReadFile(statPath)
 	if err != nil {
 		return "", 0, nil, 0, 0, 0, err
@@ -989,6 +1008,10 @@ func (p *Process) fillFromStat() (string, int32, *cpu.TimesStat, int64, uint32,
 	return terminal, int32(ppid), cpuTimes, createTime, uint32(rtpriority), nice, nil
 }
 
+func (p *Process) fillFromStat() (string, int32, *cpu.TimesStat, int64, uint32, int32, error) {
+	return p.fillFromTIDStat(-1)
+}
+
 // Pids returns a slice of process ID list which are running now.
 func Pids() ([]int32, error) {
 	return readPidsFromDir(common.HostProc())
diff --git a/process/process_openbsd.go b/process/process_openbsd.go
index b8365d5..fb553a6 100644
--- a/process/process_openbsd.go
+++ b/process/process_openbsd.go
@@ -194,8 +194,8 @@ func (p *Process) NumThreads() (int32, error) {
 	/* not supported, just return 1 */
 	return 1, nil
 }
-func (p *Process) Threads() ([]int32, error) {
-	ret := make([]int32, 0)
+func (p *Process) Threads() (map[int32]*cpu.TimesStat, error) {
+	ret := make(map[int32]*cpu.TimesStat)
 	return ret, common.ErrNotImplementedError
 }
 func (p *Process) Times() (*cpu.TimesStat, error) {
diff --git a/process/process_test.go b/process/process_test.go
index 2f33c27..9947dea 100644
--- a/process/process_test.go
+++ b/process/process_test.go
@@ -219,6 +219,26 @@ func Test_Process_NumThread(t *testing.T) {
 	}
 }
 
+func Test_Process_Threads(t *testing.T) {
+	p := testGetProcess()
+
+	n, err := p.NumThreads()
+	if err != nil {
+		t.Errorf("geting NumThread error %v", err)
+	}
+	if n < 0 {
+		t.Errorf("invalid NumThread: %d", n)
+	}
+
+	ts, err := p.Threads()
+	if err != nil {
+		t.Errorf("geting Threads error %v", err)
+	}
+	if len(ts) != int(n) {
+		t.Errorf("unexpected number of threads: %v vs %v", len(ts), n)
+	}
+}
+
 func Test_Process_Name(t *testing.T) {
 	p := testGetProcess()
 
diff --git a/process/process_windows.go b/process/process_windows.go
index ff55274..e02765e 100644
--- a/process/process_windows.go
+++ b/process/process_windows.go
@@ -273,8 +273,8 @@ func (p *Process) NumThreads() (int32, error) {
 	}
 	return int32(dst[0].ThreadCount), nil
 }
-func (p *Process) Threads() ([]int32, error) {
-	ret := make([]int32, 0)
+func (p *Process) Threads() (map[int32]*cpu.TimesStat, error) {
+	ret := make(map[int32]*cpu.TimesStat)
 	return ret, common.ErrNotImplementedError
 }
 func (p *Process) Times() (*cpu.TimesStat, error) {