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: