From a508bb970935043259e4285c2a4844883ab5c2be Mon Sep 17 00:00:00 2001
From: Michael Schurter <schmichael@hashicorp.com>
Date: Thu, 14 Sep 2017 15:27:39 -0700
Subject: [PATCH] Fold SetFailure into SetRestartTriggered

---
 client/restarts.go      | 17 ++++++-----------
 client/restarts_test.go |  2 +-
 client/task_runner.go   |  8 +-------
 3 files changed, 8 insertions(+), 19 deletions(-)

diff --git a/client/restarts.go b/client/restarts.go
index 022c1cac3..ccf344fd1 100644
--- a/client/restarts.go
+++ b/client/restarts.go
@@ -75,19 +75,14 @@ func (r *RestartTracker) SetWaitResult(res *dstructs.WaitResult) *RestartTracker
 
 // SetRestartTriggered is used to mark that the task has been signalled to be
 // restarted
-func (r *RestartTracker) SetRestartTriggered() *RestartTracker {
+func (r *RestartTracker) SetRestartTriggered(failure bool) *RestartTracker {
 	r.lock.Lock()
 	defer r.lock.Unlock()
-	r.restartTriggered = true
-	return r
-}
-
-// SetFailure is used to mark that a task should be restarted due to failure
-// such as a failed Consul healthcheck.
-func (r *RestartTracker) SetFailure() *RestartTracker {
-	r.lock.Lock()
-	defer r.lock.Unlock()
-	r.failure = true
+	if failure {
+		r.failure = true
+	} else {
+		r.restartTriggered = true
+	}
 	return r
 }
 
diff --git a/client/restarts_test.go b/client/restarts_test.go
index 851052576..0d4c4e206 100644
--- a/client/restarts_test.go
+++ b/client/restarts_test.go
@@ -99,7 +99,7 @@ func TestClient_RestartTracker_RestartTriggered(t *testing.T) {
 	p := testPolicy(true, structs.RestartPolicyModeFail)
 	p.Attempts = 0
 	rt := newRestartTracker(p, structs.JobTypeService)
-	if state, when := rt.SetRestartTriggered().GetState(); state != structs.TaskRestarting && when != 0 {
+	if state, when := rt.SetRestartTriggered(false).GetState(); state != structs.TaskRestarting && when != 0 {
 		t.Fatalf("expect restart immediately, got %v %v", state, when)
 	}
 }
diff --git a/client/task_runner.go b/client/task_runner.go
index a33df2c39..06acb6dc6 100644
--- a/client/task_runner.go
+++ b/client/task_runner.go
@@ -1175,13 +1175,7 @@ func (r *TaskRunner) run() {
 					<-handleWaitCh
 				}
 
-				if restartEvent.failure {
-					r.restartTracker.SetFailure()
-				} else {
-					// Since the restart isn't from a failure, restart immediately
-					// and don't count against the restart policy
-					r.restartTracker.SetRestartTriggered()
-				}
+				r.restartTracker.SetRestartTriggered(restartEvent.failure)
 				break WAIT
 
 			case <-r.destroyCh: