From 64e9cf9733b10d08ac42f1e14f4ab36a94fdffa2 Mon Sep 17 00:00:00 2001
From: Preetha Appan <preetha@hashicorp.com>
Date: Wed, 24 Jan 2018 14:56:57 -0600
Subject: [PATCH] Add a field to track the next allocation during a replacement

---
 api/allocations.go              |  1 +
 nomad/state/state_store.go      | 15 +++++++++++++++
 nomad/structs/structs.go        |  3 +++
 scheduler/generic_sched_test.go |  2 +-
 4 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/api/allocations.go b/api/allocations.go
index f23b83551..0b2823bd2 100644
--- a/api/allocations.go
+++ b/api/allocations.go
@@ -91,6 +91,7 @@ type Allocation struct {
 	DeploymentID       string
 	DeploymentStatus   *AllocDeploymentStatus
 	PreviousAllocation string
+	NextAllocation     string
 	CreateIndex        uint64
 	ModifyIndex        uint64
 	AllocModifyIndex   uint64
diff --git a/nomad/state/state_store.go b/nomad/state/state_store.go
index f837a6106..de0fc7cc1 100644
--- a/nomad/state/state_store.go
+++ b/nomad/state/state_store.go
@@ -1889,6 +1889,21 @@ func (s *StateStore) upsertAllocsImpl(index uint64, allocs []*structs.Allocation
 			return fmt.Errorf("alloc insert failed: %v", err)
 		}
 
+		if alloc.PreviousAllocation != "" {
+			prevAlloc, err := txn.First("allocs", "id", alloc.PreviousAllocation)
+			if err != nil {
+				return fmt.Errorf("alloc lookup failed: %v", err)
+			}
+			existingPrevAlloc, _ := prevAlloc.(*structs.Allocation)
+			if existingPrevAlloc != nil {
+				prevAllocCopy := existingPrevAlloc.Copy()
+				prevAllocCopy.NextAllocation = alloc.ID
+				if err := txn.Insert("allocs", prevAllocCopy); err != nil {
+					return fmt.Errorf("alloc insert failed: %v", err)
+				}
+			}
+		}
+
 		// If the allocation is running, force the job to running status.
 		forceStatus := ""
 		if !alloc.TerminalStatus() {
diff --git a/nomad/structs/structs.go b/nomad/structs/structs.go
index c820caa6d..465df281b 100644
--- a/nomad/structs/structs.go
+++ b/nomad/structs/structs.go
@@ -4910,6 +4910,9 @@ type Allocation struct {
 	// PreviousAllocation is the allocation that this allocation is replacing
 	PreviousAllocation string
 
+	// NextAllocation is the allocation that this allocation is being replaced by
+	NextAllocation string
+
 	// DeploymentID identifies an allocation as being created from a
 	// particular deployment
 	DeploymentID string
diff --git a/scheduler/generic_sched_test.go b/scheduler/generic_sched_test.go
index 39eff3773..e7649a238 100644
--- a/scheduler/generic_sched_test.go
+++ b/scheduler/generic_sched_test.go
@@ -3443,7 +3443,7 @@ func TestGenericSched_ChainedAlloc(t *testing.T) {
 	}
 	sort.Strings(prevAllocs)
 
-	// Ensure that the new allocations has their corresponging original
+	// Ensure that the new allocations has their corresponding original
 	// allocation ids
 	if !reflect.DeepEqual(prevAllocs, allocIDs) {
 		t.Fatalf("expected: %v, actual: %v", len(allocIDs), len(prevAllocs))