From 7fa7655ebebb84a1b7fa47de6b4c376d9f08159e Mon Sep 17 00:00:00 2001
From: Preetha Appan <preetha@hashicorp.com>
Date: Wed, 4 Apr 2018 14:38:15 -0500
Subject: [PATCH] Moves setting finishedAt to the right place and adds two unit
 tests.

---
 client/alloc_runner.go      | 49 +++++++++++++++++++++++--------------
 client/alloc_runner_test.go | 47 +++++++++++++++++++++++++++++++++++
 nomad/structs/structs.go    |  5 ++++
 3 files changed, 83 insertions(+), 18 deletions(-)

diff --git a/client/alloc_runner.go b/client/alloc_runner.go
index 5bf30dfde..d640b5d84 100644
--- a/client/alloc_runner.go
+++ b/client/alloc_runner.go
@@ -523,29 +523,41 @@ func copyTaskStates(states map[string]*structs.TaskState) map[string]*structs.Ta
 	return copy
 }
 
+// finalizeTerminalAlloc sets any missing required fields like
+// finishedAt in the alloc runner's task States. finishedAt is used
+// to calculate reschedule time for failed allocs, so we make sure that
+// it is set
+func (r *AllocRunner) finalizeTerminalAlloc(alloc *structs.Allocation) {
+	if !alloc.ClientTerminalStatus() {
+		return
+	}
+	r.taskStatusLock.Lock()
+	defer r.taskStatusLock.Unlock()
+
+	group := alloc.Job.LookupTaskGroup(alloc.TaskGroup)
+	if r.taskStates == nil {
+		r.taskStates = make(map[string]*structs.TaskState)
+	}
+	now := time.Now()
+	for _, task := range group.Tasks {
+		ts, ok := r.taskStates[task.Name]
+		if !ok {
+			ts = &structs.TaskState{}
+			r.taskStates[task.Name] = ts
+		}
+		if ts.FinishedAt.IsZero() {
+			ts.FinishedAt = now
+		}
+	}
+	alloc.TaskStates = copyTaskStates(r.taskStates)
+}
+
 // Alloc returns the associated allocation
 func (r *AllocRunner) Alloc() *structs.Allocation {
 	r.allocLock.Lock()
 	// We rely upon FinishedAt to determine rescheduling eligibility so
 	// this makes sure that it is set for every task group.
 	// If the alloc never started FinishedAt may not be set
-	if r.alloc.TerminalStatus() {
-		group := r.alloc.Job.LookupTaskGroup(r.alloc.TaskGroup)
-		if r.alloc.TaskStates == nil {
-			r.alloc.TaskStates = make(map[string]*structs.TaskState)
-		}
-		now := time.Now()
-		for _, task := range group.Tasks {
-			ts, ok := r.alloc.TaskStates[task.Name]
-			if !ok {
-				ts = &structs.TaskState{}
-				r.alloc.TaskStates[task.Name] = ts
-			}
-			if ts.FinishedAt.IsZero() {
-				ts.FinishedAt = now
-			}
-		}
-	}
 
 	// Don't do a deep copy of the job
 	alloc := r.alloc.CopySkipJob()
@@ -561,6 +573,7 @@ func (r *AllocRunner) Alloc() *structs.Allocation {
 		r.taskStatusLock.RUnlock()
 
 		r.allocLock.Unlock()
+		r.finalizeTerminalAlloc(alloc)
 		return alloc
 	}
 
@@ -589,7 +602,7 @@ func (r *AllocRunner) Alloc() *structs.Allocation {
 		}
 	}
 	r.allocLock.Unlock()
-
+	r.finalizeTerminalAlloc(alloc)
 	return alloc
 }
 
diff --git a/client/alloc_runner_test.go b/client/alloc_runner_test.go
index 4952e551b..6b42b0f31 100644
--- a/client/alloc_runner_test.go
+++ b/client/alloc_runner_test.go
@@ -26,6 +26,7 @@ import (
 
 	"github.com/hashicorp/nomad/client/config"
 	"github.com/hashicorp/nomad/client/vaultclient"
+	"github.com/stretchr/testify/require"
 )
 
 type MockAllocStateUpdater struct {
@@ -109,6 +110,52 @@ func TestAllocRunner_SimpleRun(t *testing.T) {
 	})
 }
 
+// Test that FinisheAt is set when the alloc is in a terminal state
+func TestAllocRunner_FinishedAtSet(t *testing.T) {
+	t.Parallel()
+	require := require.New(t)
+	_, ar := testAllocRunner(t, false)
+	//ar.alloc.ClientStatus = structs.AllocClientStatusFailed
+	ar.allocClientStatus = structs.AllocClientStatusFailed
+	alloc := ar.Alloc()
+	taskFinishedAt := make(map[string]time.Time)
+	require.NotEmpty(alloc.TaskStates)
+	for name, s := range alloc.TaskStates {
+		require.False(s.FinishedAt.IsZero())
+		taskFinishedAt[name] = s.FinishedAt
+	}
+
+	// Verify that calling again should not mutate finishedAt
+	alloc2 := ar.Alloc()
+	for name, s := range alloc2.TaskStates {
+		require.Equal(taskFinishedAt[name], s.FinishedAt)
+	}
+
+}
+
+// Test that FinisheAt is set when the alloc is in a terminal state
+func TestAllocRunner_FinishedAtSet_TaskEvents(t *testing.T) {
+	t.Parallel()
+	require := require.New(t)
+	_, ar := testAllocRunner(t, false)
+	ar.taskStates[ar.alloc.Job.TaskGroups[0].Tasks[0].Name] = &structs.TaskState{State: structs.TaskStateDead, Failed: true}
+
+	alloc := ar.Alloc()
+	taskFinishedAt := make(map[string]time.Time)
+	require.NotEmpty(alloc.TaskStates)
+	for name, s := range alloc.TaskStates {
+		require.False(s.FinishedAt.IsZero())
+		taskFinishedAt[name] = s.FinishedAt
+	}
+
+	// Verify that calling again should not mutate finishedAt
+	alloc2 := ar.Alloc()
+	for name, s := range alloc2.TaskStates {
+		require.Equal(taskFinishedAt[name], s.FinishedAt)
+	}
+
+}
+
 // Test that the watcher will mark the allocation as unhealthy.
 func TestAllocRunner_DeploymentHealth_Unhealthy_BadStart(t *testing.T) {
 	t.Parallel()
diff --git a/nomad/structs/structs.go b/nomad/structs/structs.go
index b73869fbf..2ad736525 100644
--- a/nomad/structs/structs.go
+++ b/nomad/structs/structs.go
@@ -5696,6 +5696,11 @@ func (a *Allocation) TerminalStatus() bool {
 	default:
 	}
 
+	return a.ClientTerminalStatus()
+}
+
+// ClientTerminalStatus returns if the client status is terminal and will no longer transition
+func (a *Allocation) ClientTerminalStatus() bool {
 	switch a.ClientStatus {
 	case AllocClientStatusComplete, AllocClientStatusFailed, AllocClientStatusLost:
 		return true