From f59ad3da53c80d103178cae62d635c20d6c94f43 Mon Sep 17 00:00:00 2001
From: Diptanu Choudhury <diptanuc@gmail.com>
Date: Wed, 25 May 2016 14:58:32 -0700
Subject: [PATCH] Fixed the logic of scanpids

---
 client/driver/executor/executor.go       | 39 +++++++++++++++------
 client/driver/executor/executor_basic.go | 10 +++++-
 client/driver/executor/executor_test.go  | 44 ++++++++++++++++++++++++
 3 files changed, 81 insertions(+), 12 deletions(-)

diff --git a/client/driver/executor/executor.go b/client/driver/executor/executor.go
index f0e4f98b9..d0ffff68b 100644
--- a/client/driver/executor/executor.go
+++ b/client/driver/executor/executor.go
@@ -706,18 +706,35 @@ func (e *UniversalExecutor) collectPids() {
 
 // scanPids scans all the pids on the machine running the current executor and
 // returns the child processes of the executor.
-func (e *UniversalExecutor) scanPids() ([]*nomadPid, error) {
+func (e *UniversalExecutor) scanPids(parentPid int, allPids []ps.Process) ([]*nomadPid, error) {
 	processFamily := make(map[int]struct{})
-	processFamily[os.Getpid()] = struct{}{}
-	pids, err := ps.Processes()
-	if err != nil {
-		return nil, err
-	}
-	for _, pid := range pids {
-		_, parentPid := processFamily[pid.Pid()]
-		_, childPid := processFamily[pid.PPid()]
-		if parentPid || childPid {
-			processFamily[pid.Pid()] = struct{}{}
+	processFamily[parentPid] = struct{}{}
+
+	// A buffer for holding pids which haven't matched with any parent pid
+	var pidsRemaining []ps.Process
+	for {
+		// flag to indicate if we have found a match
+		foundNewPid := false
+
+		for _, pid := range allPids {
+			_, parentPid := processFamily[pid.Pid()]
+			_, childPid := processFamily[pid.PPid()]
+
+			// checking if the pid is a child of any of the parents
+			if parentPid || childPid {
+				processFamily[pid.Pid()] = struct{}{}
+				foundNewPid = true
+			} else {
+				// if it is not, then we add the pid to the buffer
+				pidsRemaining = append(pidsRemaining, pid)
+			}
+			// scan only the pids which are left in the buffer
+			allPids = pidsRemaining
+		}
+
+		// not scanning anymore if we couldn't find a single match
+		if !foundNewPid {
+			break
 		}
 	}
 	res := make([]*nomadPid, 0, len(processFamily))
diff --git a/client/driver/executor/executor_basic.go b/client/driver/executor/executor_basic.go
index ca0120dc7..eb07a406d 100644
--- a/client/driver/executor/executor_basic.go
+++ b/client/driver/executor/executor_basic.go
@@ -3,7 +3,11 @@
 package executor
 
 import (
+	"os"
+
 	cstructs "github.com/hashicorp/nomad/client/driver/structs"
+
+	"github.com/mitchellh/go-ps"
 	cgroupConfig "github.com/opencontainers/runc/libcontainer/configs"
 )
 
@@ -36,5 +40,9 @@ func (e *UniversalExecutor) Stats() (*cstructs.TaskResourceUsage, error) {
 }
 
 func (e *UniversalExecutor) getAllPids() ([]*nomadPid, error) {
-	return e.scanPids()
+	allProcesses, err := ps.Processes()
+	if err != nil {
+		return nil, err
+	}
+	return e.scanPids(os.Getpid(), allProcesses)
 }
diff --git a/client/driver/executor/executor_test.go b/client/driver/executor/executor_test.go
index 92b26510d..fef7b8e3b 100644
--- a/client/driver/executor/executor_test.go
+++ b/client/driver/executor/executor_test.go
@@ -12,6 +12,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/mitchellh/go-ps"
+
 	"github.com/hashicorp/nomad/client/allocdir"
 	"github.com/hashicorp/nomad/client/driver/env"
 	cstructs "github.com/hashicorp/nomad/client/driver/structs"
@@ -343,3 +345,45 @@ func TestExecutorInterpolateServices(t *testing.T) {
 		t.Fatalf("expected: %v, actual: %v", expectedCheckArgs, task.Services[0].Checks[0].Args)
 	}
 }
+
+func TestScanPids(t *testing.T) {
+	p1 := NewFakeProcess(2, 5)
+	p2 := NewFakeProcess(10, 2)
+	p3 := NewFakeProcess(15, 6)
+	p4 := NewFakeProcess(3, 10)
+	p5 := NewFakeProcess(20, 18)
+
+	// Make a fake exececutor
+	ctx := testExecutorContext(t)
+	defer ctx.AllocDir.Destroy()
+	executor := NewExecutor(log.New(os.Stdout, "", log.LstdFlags)).(*UniversalExecutor)
+
+	nomadPids, err := executor.scanPids(5, []ps.Process{p1, p2, p3, p4, p5})
+	if err != nil {
+		t.Fatalf("error: %v", err)
+	}
+	if len(nomadPids) != 4 {
+		t.Fatalf("expected: 4, actual: %v", len(nomadPids))
+	}
+}
+
+type FakeProcess struct {
+	pid  int
+	ppid int
+}
+
+func (f FakeProcess) Pid() int {
+	return f.pid
+}
+
+func (f FakeProcess) PPid() int {
+	return f.ppid
+}
+
+func (f FakeProcess) Executable() string {
+	return "fake"
+}
+
+func NewFakeProcess(pid int, ppid int) ps.Process {
+	return FakeProcess{pid: pid, ppid: ppid}
+}