opt: using go-ps to implement pstree

Author: YYCoder

go.mod

@@ -4,8 +4,8 @@ go 1.14
 
 require (
 	github.com/howeyc/fsnotify v0.9.0
+	github.com/mitchellh/go-ps v1.0.0
 	github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect
-	github.com/sbinet/pstree v0.3.0
 	github.com/silenceper/log v0.0.0-20171204144354-e5ac7fa8a76a
 	gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f // indirect
 	gopkg.in/yaml.v1 v1.0.0-20140924161607-9f9df34309c0

go.sum

@@ -3,10 +3,10 @@ github.com/howeyc/fsnotify v0.9.0/go.mod h1:41HzSPxBGeFRQKEEwgh49TRw/nKBsYZ2cF1O
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
+github.com/mitchellh/go-ps v1.0.0 h1:i6ampVEEF4wQFF+bkYfwYgY+F/uYJDktmvLPf7qIgjc=
+github.com/mitchellh/go-ps v1.0.0/go.mod h1:J4lOc8z8yJs6vUwklHw2XEIiT4z4C40KtWVN3nvg8Pg=
 github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e h1:fD57ERR4JtEqsWbfPhv4DMiApHyliiK5xCTNVSPiaAs=
 github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
-github.com/sbinet/pstree v0.3.0 h1:XZLdGCbAg/wg8TzoW/060WtRoAQdECnsTrn+RYgyJoQ=
-github.com/sbinet/pstree v0.3.0/go.mod h1:G208WfJOi4oxq4++w97Y4AeuydVuoOz7tPKCEm8y1oE=
 github.com/silenceper/log v0.0.0-20171204144354-e5ac7fa8a76a h1:COf2KvPmardI1M8p2fhHsXlFS2EXSQygbGgcDYBI9Wc=
 github.com/silenceper/log v0.0.0-20171204144354-e5ac7fa8a76a/go.mod h1:nyN/YUSK3CgJjtNzm6dVTkcou+RYXNMP+XLSlzQu0m0=
 gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f h1:BLraFXnmrev5lT+xlilqcH8XK9/i0At2xKjWk4p6zsU=

gowatch.go

@@ -13,7 +13,7 @@ import (
 	"time"
 
 	"github.com/howeyc/fsnotify"
-	"github.com/sbinet/pstree"
+	"github.com/mitchellh/go-ps"
 	"github.com/silenceper/log"
 )
 
@@ -185,7 +185,7 @@ func Kill() {
 func killAllProcesses(pid int) (err error) {
 	hasAllKilled := make(chan bool)
 	go func() {
-		pids, err := getAllProcesses(pid)
+		pids, err := psTree(pid)
 		if err != nil {
 			log.Fatalf("getting all sub processes error: %v\n", err)
 			return
@@ -232,28 +232,42 @@ func killProcess(pid int) (err error) {
 	return
 }
 
-func getAllProcesses(pid int) (pids []int, err error) {
-	tree, err := pstree.New()
+// implement pstree based on the cross-platform ps utility in go, go-ps
+func psTree(rootPid int) (res []int, err error) {
+	pidOfInterest := map[int]struct{}{rootPid: {}}
+	pss, err := ps.Processes()
 	if err != nil {
-		log.Fatalf("error: %v\n", err)
+		fmt.Println("ERROR: ", err)
 		return
 	}
 
-	traverseChildren(pid, tree, &pids)
-	// we must kill the sub processes from smallest to largest
-	sort.Ints(pids)
-	return
-}
+	// we must sort the ps by ppid && pid first, otherwise we probably will miss some sub-processes
+	// of the root process during for-range searching
+	sort.Slice(pss, func(i, j int) bool {
+		ppidLess := pss[i].PPid() < pss[j].PPid()
+		pidLess := pss[i].PPid() == pss[j].PPid() && pss[i].Pid() < pss[j].Pid()
+
+		return ppidLess || pidLess
+	})
+
+	for _, ps := range pss {
+		ppid := ps.PPid()
+		if _, exists := pidOfInterest[ppid]; exists {
+			pidOfInterest[ps.Pid()] = struct{}{}
+		}
+	}
 
-func traverseChildren(pid int, tree *pstree.Tree, curPids *[]int) {
-	for _, cid := range tree.Procs[pid].Children {
-		*curPids = append(*curPids, cid)
-		traverseChildren(cid, tree, curPids)
+	for pid, _ := range pidOfInterest {
+		if pid != rootPid {
+			res = append(res, pid)
+		}
 	}
+
+	return
 }
 
 func waitForProcess(pid int, hasAllKilled chan bool) (err error) {
-	pids, _ := getAllProcesses(pid)
+	pids, _ := psTree(pid)
 	if len(pids) == 0 {
 		hasAllKilled <- true
 		return