From aa2da9ea50384ed7f780a67e760bb13429c76d7f Mon Sep 17 00:00:00 2001
From: Pete Wildsmith <pete@dabapps.com>
Date: Fri, 28 Apr 2017 09:19:15 +0100
Subject: [PATCH] address feedback

---
 client/task_runner.go | 21 ++++++++-------------
 1 file changed, 8 insertions(+), 13 deletions(-)

diff --git a/client/task_runner.go b/client/task_runner.go
index be32215ec..2bb98922f 100644
--- a/client/task_runner.go
+++ b/client/task_runner.go
@@ -918,7 +918,6 @@ func (r *TaskRunner) run() {
 			select {
 			case success := <-prestartResultCh:
 				if !success {
-					r.consulCleanup()
 					r.cleanup()
 					r.setState(structs.TaskStateDead, nil)
 					return
@@ -1026,17 +1025,16 @@ func (r *TaskRunner) run() {
 				r.runningLock.Lock()
 				running := r.running
 				r.runningLock.Unlock()
-
-				// Remove from consul before killing the task so that traffic
-				// can be rerouted
-				r.consulCleanup()
-
 				if !running {
 					r.cleanup()
 					r.setState(structs.TaskStateDead, r.destroyEvent)
 					return
 				}
 
+				// Remove from consul before killing the task so that traffic
+				// can be rerouted
+				r.consul.RemoveTask(r.alloc.ID, r.task)
+
 				// Store the task event that provides context on the task
 				// destroy. The Killed event is set from the alloc_runner and
 				// doesn't add detail
@@ -1065,7 +1063,6 @@ func (r *TaskRunner) run() {
 		// shouldRestart will block if the task should restart after a delay.
 		restart := r.shouldRestart()
 		if !restart {
-			r.consulCleanup()
 			r.cleanup()
 			r.setState(structs.TaskStateDead, nil)
 			return
@@ -1080,14 +1077,12 @@ func (r *TaskRunner) run() {
 	}
 }
 
-func (r *TaskRunner) consulCleanup() {
-	// Remove from Consul
-	r.consul.RemoveTask(r.alloc.ID, r.task)
-}
-
-// cleanup calls Driver.Cleanup when a task is
+// cleanup removes Consul entries and calls Driver.Cleanup when a task is
 // stopping. Errors are logged.
 func (r *TaskRunner) cleanup() {
+	// Remove from Consul
+	r.consul.RemoveTask(r.alloc.ID, r.task)
+
 	drv, err := r.createDriver()
 	if err != nil {
 		r.logger.Printf("[ERR] client: error creating driver to cleanup resources: %v", err)