Merge pull request #491 from hashicorp/b-batch-task-restart

Not restarting if a task exited properly

CHANGELOG.md

@@ -12,6 +12,7 @@ BUG FIXES:
   * driver/docker: Support `port_map` for static ports [GH-476]
   * driver/docker: Pass 0.2.0-style port environment variables to the docker container [GH-476]
   * client/service discovery: Make Service IDs unique [GH-479]
+  * client/restart policy: Not restarting Batch Jobs if the exit code is 0 [GH-491]
 
 ## 0.2.0 (November 18, 2015)
 

client/restarts.go

@@ -11,7 +11,7 @@ import (
 // For Batch jobs, the interval is set to zero value since the takss
 // will be restarted only upto maxAttempts times
 type restartTracker interface {
-	nextRestart() (bool, time.Duration)
+	nextRestart(exitCode int) (bool, time.Duration)
 }
 
 func newRestartTracker(jobType string, restartPolicy *structs.RestartPolicy) restartTracker {
@@ -47,8 +47,8 @@ func (b *batchRestartTracker) increment() {
 	b.count += 1
 }
 
-func (b *batchRestartTracker) nextRestart() (bool, time.Duration) {
-	if b.count < b.maxAttempts {
+func (b *batchRestartTracker) nextRestart(exitCode int) (bool, time.Duration) {
+	if b.count < b.maxAttempts && exitCode > 0 {
 		b.increment()
 		return true, b.delay
 	}
@@ -68,7 +68,7 @@ func (s *serviceRestartTracker) increment() {
 	s.count += 1
 }
 
-func (s *serviceRestartTracker) nextRestart() (bool, time.Duration) {
+func (s *serviceRestartTracker) nextRestart(exitCode int) (bool, time.Duration) {
 	defer s.increment()
 	windowEndTime := s.startTime.Add(s.interval)
 	now := time.Now()

client/restarts_test.go

@@ -13,7 +13,7 @@ func TestTaskRunner_ServiceRestartCounter(t *testing.T) {
 	rt := newRestartTracker(structs.JobTypeService, &structs.RestartPolicy{Attempts: attempts, Interval: interval, Delay: delay})
 
 	for i := 0; i < attempts; i++ {
-		actual, when := rt.nextRestart()
+		actual, when := rt.nextRestart(127)
 		if !actual {
 			t.Fatalf("should restart returned %v, actual %v", actual, true)
 		}
@@ -24,7 +24,7 @@ func TestTaskRunner_ServiceRestartCounter(t *testing.T) {
 
 	time.Sleep(1 * time.Second)
 	for i := 0; i < 3; i++ {
-		actual, when := rt.nextRestart()
+		actual, when := rt.nextRestart(127)
 		if !actual {
 			t.Fail()
 		}
@@ -46,16 +46,33 @@ func TestTaskRunner_BatchRestartCounter(t *testing.T) {
 		},
 	)
 	for i := 0; i < attempts; i++ {
-		shouldRestart, when := rt.nextRestart()
+		shouldRestart, when := rt.nextRestart(127)
 		if !shouldRestart {
 			t.Fatalf("should restart returned %v, actual %v", shouldRestart, true)
 		}
 		if when != delay {
 			t.Fatalf("Delay should be %v, actual: %v", delay, when)
 		}
 	}
-	actual, _ := rt.nextRestart()
+	actual, _ := rt.nextRestart(1)
 	if actual {
 		t.Fatalf("Expect %v, Actual: %v", false, actual)
 	}
 }
+
+func TestTaskRunner_BatchRestartOnSuccess(t *testing.T) {
+	attempts := 2
+	interval := 1 * time.Second
+	delay := 1 * time.Second
+	rt := newRestartTracker(structs.JobTypeBatch,
+		&structs.RestartPolicy{Attempts: attempts,
+			Interval: interval,
+			Delay:    delay,
+		},
+	)
+	shouldRestart, _ := rt.nextRestart(0)
+	if shouldRestart {
+		t.Fatalf("should restart returned %v, expected: %v", shouldRestart, false)
+	}
+
+}

client/task_runner.go

@@ -284,8 +284,8 @@ func (r *TaskRunner) run() {
 		}
 
 		// Check if we should restart. If not mark task as dead and exit.
+		shouldRestart, when := r.restartTracker.nextRestart(waitRes.ExitCode)
 		waitEvent := r.waitErrorToEvent(waitRes)
-		shouldRestart, when := r.restartTracker.nextRestart()
 		if !shouldRestart {
 			r.logger.Printf("[INFO] client: Not restarting task: %v for alloc: %v ", r.task.Name, r.allocID)
 			r.setState(structs.TaskStateDead, waitEvent)