Fold SetFailure into SetRestartTriggered

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
 }
 

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)
 	}
 }

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: