From 06f3aac9de9d5263271d6900266beecc5da607ad Mon Sep 17 00:00:00 2001
From: Pete Wildsmith <pete@weargoggles.co.uk>
Date: Thu, 27 Apr 2017 23:27:58 +0100
Subject: [PATCH] clean up consul earlier when destroying a task

---
 client/task_runner.go | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/client/task_runner.go b/client/task_runner.go
index c578ca5ea..be32215ec 100644
--- a/client/task_runner.go
+++ b/client/task_runner.go
@@ -918,6 +918,7 @@ func (r *TaskRunner) run() {
 			select {
 			case success := <-prestartResultCh:
 				if !success {
+					r.consulCleanup()
 					r.cleanup()
 					r.setState(structs.TaskStateDead, nil)
 					return
@@ -1025,6 +1026,11 @@ 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)
@@ -1059,6 +1065,7 @@ 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
@@ -1073,12 +1080,14 @@ func (r *TaskRunner) run() {
 	}
 }
 
-// cleanup removes Consul entries and calls Driver.Cleanup when a task is
-// stopping. Errors are logged.
-func (r *TaskRunner) cleanup() {
+func (r *TaskRunner) consulCleanup() {
 	// Remove from Consul
 	r.consul.RemoveTask(r.alloc.ID, r.task)
+}
 
+// cleanup calls Driver.Cleanup when a task is
+// stopping. Errors are logged.
+func (r *TaskRunner) cleanup() {
 	drv, err := r.createDriver()
 	if err != nil {
 		r.logger.Printf("[ERR] client: error creating driver to cleanup resources: %v", err)