From cde908e35d8b3dedc2ed68d2d7b2eb832aaf06ad Mon Sep 17 00:00:00 2001
From: Michael Schurter <schmichael@hashicorp.com>
Date: Fri, 15 Sep 2017 14:54:37 -0700
Subject: [PATCH] Cleanup and test restart failure code

---
 client/restarts.go      | 62 ++++++++++++++++++++---------------------
 client/restarts_test.go | 13 +++++++++
 2 files changed, 44 insertions(+), 31 deletions(-)

diff --git a/client/restarts.go b/client/restarts.go
index ca0456f32d21..c403b6f05d86 100644
--- a/client/restarts.go
+++ b/client/restarts.go
@@ -145,42 +145,42 @@ func (r *RestartTracker) GetState() (string, time.Duration) {
 	}
 
 	// Handle restarts due to failures
-	if r.failure {
-		if r.startErr != nil {
-			// If the error is not recoverable, do not restart.
-			if !structs.IsRecoverable(r.startErr) {
-				r.reason = ReasonUnrecoverableErrror
-				return structs.TaskNotRestarting, 0
-			}
-		} else if r.waitRes != nil {
-			// If the task started successfully and restart on success isn't specified,
-			// don't restart but don't mark as failed.
-			if r.waitRes.Successful() && !r.onSuccess {
-				r.reason = "Restart unnecessary as task terminated successfully"
-				return structs.TaskTerminated, 0
-			}
-		}
+	if !r.failure {
+		return "", 0
+	}
 
-		// If this task has been restarted due to failures more times
-		// than the restart policy allows within an interval fail
-		// according to the restart policy's mode.
-		if r.count > r.policy.Attempts {
-			if r.policy.Mode == structs.RestartPolicyModeFail {
-				r.reason = fmt.Sprintf(
-					`Exceeded allowed attempts %d in interval %v and mode is "fail"`,
-					r.policy.Attempts, r.policy.Interval)
-				return structs.TaskNotRestarting, 0
-			} else {
-				r.reason = ReasonDelay
-				return structs.TaskRestarting, r.getDelay()
-			}
+	if r.startErr != nil {
+		// If the error is not recoverable, do not restart.
+		if !structs.IsRecoverable(r.startErr) {
+			r.reason = ReasonUnrecoverableErrror
+			return structs.TaskNotRestarting, 0
+		}
+	} else if r.waitRes != nil {
+		// If the task started successfully and restart on success isn't specified,
+		// don't restart but don't mark as failed.
+		if r.waitRes.Successful() && !r.onSuccess {
+			r.reason = "Restart unnecessary as task terminated successfully"
+			return structs.TaskTerminated, 0
 		}
+	}
 
-		r.reason = ReasonWithinPolicy
-		return structs.TaskRestarting, r.jitter()
+	// If this task has been restarted due to failures more times
+	// than the restart policy allows within an interval fail
+	// according to the restart policy's mode.
+	if r.count > r.policy.Attempts {
+		if r.policy.Mode == structs.RestartPolicyModeFail {
+			r.reason = fmt.Sprintf(
+				`Exceeded allowed attempts %d in interval %v and mode is "fail"`,
+				r.policy.Attempts, r.policy.Interval)
+			return structs.TaskNotRestarting, 0
+		} else {
+			r.reason = ReasonDelay
+			return structs.TaskRestarting, r.getDelay()
+		}
 	}
 
-	return "", 0
+	r.reason = ReasonWithinPolicy
+	return structs.TaskRestarting, r.jitter()
 }
 
 // getDelay returns the delay time to enter the next interval.
diff --git a/client/restarts_test.go b/client/restarts_test.go
index 0d4c4e2068bb..b0cad5b1a3c1 100644
--- a/client/restarts_test.go
+++ b/client/restarts_test.go
@@ -104,6 +104,19 @@ func TestClient_RestartTracker_RestartTriggered(t *testing.T) {
 	}
 }
 
+func TestClient_RestartTracker_RestartTriggered_Failure(t *testing.T) {
+	t.Parallel()
+	p := testPolicy(true, structs.RestartPolicyModeFail)
+	p.Attempts = 1
+	rt := newRestartTracker(p, structs.JobTypeService)
+	if state, when := rt.SetRestartTriggered(true).GetState(); state != structs.TaskRestarting || when == 0 {
+		t.Fatalf("expect restart got %v %v", state, when)
+	}
+	if state, when := rt.SetRestartTriggered(true).GetState(); state != structs.TaskNotRestarting || when != 0 {
+		t.Fatalf("expect failed got %v %v", state, when)
+	}
+}
+
 func TestClient_RestartTracker_StartError_Recoverable_Fail(t *testing.T) {
 	t.Parallel()
 	p := testPolicy(true, structs.RestartPolicyModeFail)