From 446d3f22f2c7b800ff13808113feffa563c4d144 Mon Sep 17 00:00:00 2001
From: Diptanu Choudhury <diptanuc@gmail.com>
Date: Sat, 11 Jun 2016 01:40:52 +0200
Subject: [PATCH] Pruning out pids which are no longer present

---
 client/driver/executor/executor.go       | 20 ++++++++++++++------
 client/driver/executor/executor_basic.go |  2 +-
 client/driver/executor/executor_linux.go |  8 ++++----
 3 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/client/driver/executor/executor.go b/client/driver/executor/executor.go
index 20cb0e80a..77fd66b24 100644
--- a/client/driver/executor/executor.go
+++ b/client/driver/executor/executor.go
@@ -713,9 +713,17 @@ func (e *UniversalExecutor) collectPids() {
 				e.logger.Printf("[DEBUG] executor: error collecting pids: %v", err)
 			}
 			e.pidLock.Lock()
-			for _, pid := range pids {
-				if _, ok := e.pids[pid.pid]; !ok {
-					e.pids[pid.pid] = pid
+
+			// Adding pids which are not being tracked
+			for pid, np := range pids {
+				if _, ok := e.pids[pid]; !ok {
+					e.pids[pid] = np
+				}
+			}
+			// Removing pids which are no longer present
+			for pid := range e.pids {
+				if _, ok := pids[pid]; !ok {
+					delete(e.pids, pid)
 				}
 			}
 			e.pidLock.Unlock()
@@ -728,7 +736,7 @@ 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(parentPid int, allPids []ps.Process) ([]*nomadPid, error) {
+func (e *UniversalExecutor) scanPids(parentPid int, allPids []ps.Process) (map[int]*nomadPid, error) {
 	processFamily := make(map[int]struct{})
 	processFamily[parentPid] = struct{}{}
 
@@ -758,7 +766,7 @@ func (e *UniversalExecutor) scanPids(parentPid int, allPids []ps.Process) ([]*no
 			break
 		}
 	}
-	res := make([]*nomadPid, 0, len(processFamily))
+	res := make(map[int]*nomadPid)
 	for pid := range processFamily {
 		np := nomadPid{
 			pid:           pid,
@@ -766,7 +774,7 @@ func (e *UniversalExecutor) scanPids(parentPid int, allPids []ps.Process) ([]*no
 			cpuStatsUser:  stats.NewCpuStats(),
 			cpuStatsSys:   stats.NewCpuStats(),
 		}
-		res = append(res, &np)
+		res[pid] = &np
 	}
 	return res, nil
 }
diff --git a/client/driver/executor/executor_basic.go b/client/driver/executor/executor_basic.go
index d66cfb078..d945721fb 100644
--- a/client/driver/executor/executor_basic.go
+++ b/client/driver/executor/executor_basic.go
@@ -43,7 +43,7 @@ func (e *UniversalExecutor) Stats() (*cstructs.TaskResourceUsage, error) {
 	return e.aggregatedResourceUsage(pidStats), nil
 }
 
-func (e *UniversalExecutor) getAllPids() ([]*nomadPid, error) {
+func (e *UniversalExecutor) getAllPids() (map[int]*nomadPid, error) {
 	allProcesses, err := ps.Processes()
 	if err != nil {
 		return nil, err
diff --git a/client/driver/executor/executor_linux.go b/client/driver/executor/executor_linux.go
index d30ad620e..812d1a001 100644
--- a/client/driver/executor/executor_linux.go
+++ b/client/driver/executor/executor_linux.go
@@ -263,16 +263,16 @@ func (e *UniversalExecutor) removeChrootMounts() error {
 // use the libcontainer apis to get the pids when the user is using cgroup
 // isolation and we scan the entire process table if the user is not using any
 // isolation
-func (e *UniversalExecutor) getAllPids() ([]*nomadPid, error) {
+func (e *UniversalExecutor) getAllPids() (map[int]*nomadPid, error) {
 	if e.command.ResourceLimits {
 		manager := getCgroupManager(e.groups, e.cgPaths)
 		pids, err := manager.GetAllPids()
 		if err != nil {
 			return nil, err
 		}
-		np := make([]*nomadPid, len(pids))
-		for idx, pid := range pids {
-			np[idx] = &nomadPid{
+		np := make(map[int]*nomadPid, len(pids))
+		for _, pid := range pids {
+			np[pid] = &nomadPid{
 				pid:           pid,
 				cpuStatsTotal: stats.NewCpuStats(),
 				cpuStatsSys:   stats.NewCpuStats(),