From 6da051503c2ff5850d0a921e7736b505349ad97b Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Thu, 21 Jul 2016 13:41:01 -0700
Subject: [PATCH] Retrieve task runners in helper

---
 client/alloc_runner.go | 40 +++++++++++++++++++++-------------------
 1 file changed, 21 insertions(+), 19 deletions(-)

diff --git a/client/alloc_runner.go b/client/alloc_runner.go
index 5ecbbf2f1064..3f338c8c4c40 100644
--- a/client/alloc_runner.go
+++ b/client/alloc_runner.go
@@ -164,10 +164,9 @@ func (r *AllocRunner) SaveState() error {
 	}
 
 	// Save state for each task
-	r.taskLock.RLock()
-	defer r.taskLock.RUnlock()
+	runners := r.getTaskRunners()
 	var mErr multierror.Error
-	for _, tr := range r.tasks {
+	for _, tr := range runners {
 		if err := r.saveTaskRunnerState(tr); err != nil {
 			mErr.Errors = append(mErr.Errors, err)
 		}
@@ -431,29 +430,25 @@ OUTER:
 			}
 
 			// Update the task groups
-			r.taskLock.RLock()
-			for _, task := range tg.Tasks {
-				tr := r.tasks[task.Name]
+			runners := r.getTaskRunners()
+			for _, tr := range runners {
 				tr.Update(update)
 			}
-			r.taskLock.RUnlock()
-
 		case <-r.destroyCh:
 			break OUTER
 		}
 	}
 
 	// Destroy each sub-task
-	r.taskLock.Lock()
-	for _, tr := range r.tasks {
+	runners := r.getTaskRunners()
+	for _, tr := range runners {
 		tr.Destroy()
 	}
 
 	// Wait for termination of the task runners
-	for _, tr := range r.tasks {
+	for _, tr := range runners {
 		<-tr.WaitCh()
 	}
-	r.taskLock.Unlock()
 
 	// Final state sync
 	r.syncStatus()
@@ -494,6 +489,19 @@ func (r *AllocRunner) StatsReporter() AllocStatsReporter {
 	return r
 }
 
+// getTaskRunners is a helper that returns a copy of the task runners list using
+// the taskLock.
+func (r *AllocRunner) getTaskRunners() []*TaskRunner {
+	// Get the task runners
+	r.taskLock.RLock()
+	defer r.taskLock.RUnlock()
+	runners := make([]*TaskRunner, 0, len(r.tasks))
+	for _, tr := range r.tasks {
+		runners = append(runners, tr)
+	}
+	return runners
+}
+
 // LatestAllocStats returns the latest allocation stats. If the optional taskFilter is set
 // the allocation stats will only include the given task.
 func (r *AllocRunner) LatestAllocStats(taskFilter string) (*cstructs.AllocResourceUsage, error) {
@@ -517,13 +525,7 @@ func (r *AllocRunner) LatestAllocStats(taskFilter string) (*cstructs.AllocResour
 		}
 	} else {
 		// Get the task runners
-		r.taskLock.RLock()
-		runners := make([]*TaskRunner, 0, len(r.tasks))
-		for _, tr := range r.tasks {
-			runners = append(runners, tr)
-		}
-		r.taskLock.RUnlock()
-
+		runners := r.getTaskRunners()
 		for _, tr := range runners {
 			l := tr.LatestResourceUsage()
 			if l != nil {