Merge pull request #3795 from hashicorp/f-next-alloc

Track next allocation

api/allocations.go

@@ -91,6 +91,7 @@ type Allocation struct {
 	DeploymentID       string
 	DeploymentStatus   *AllocDeploymentStatus
 	PreviousAllocation string
+	NextAllocation     string
 	CreateIndex        uint64
 	ModifyIndex        uint64
 	AllocModifyIndex   uint64

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() {

nomad/state/state_store_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func testStateStore(t *testing.T) *StateStore {
@@ -4344,6 +4345,49 @@ func TestStateStore_Allocs(t *testing.T) {
 	}
 }
 
+func TestStateStore_Allocs_PrevAlloc(t *testing.T) {
+	state := testStateStore(t)
+	var allocs []*structs.Allocation
+
+	require := require.New(t)
+	for i := 0; i < 5; i++ {
+		alloc := mock.Alloc()
+		allocs = append(allocs, alloc)
+	}
+	for i, alloc := range allocs {
+		state.UpsertJobSummary(uint64(900+i), mock.JobSummary(alloc.JobID))
+	}
+	// Set some previous alloc ids
+	allocs[1].PreviousAllocation = allocs[0].ID
+	allocs[2].PreviousAllocation = allocs[1].ID
+
+	err := state.UpsertAllocs(1000, allocs)
+	require.Nil(err)
+
+	ws := memdb.NewWatchSet()
+	iter, err := state.Allocs(ws)
+	require.Nil(err)
+
+	var out []*structs.Allocation
+	for {
+		raw := iter.Next()
+		if raw == nil {
+			break
+		}
+		out = append(out, raw.(*structs.Allocation))
+	}
+
+	// Set expected NextAllocation fields
+	allocs[0].NextAllocation = allocs[1].ID
+	allocs[1].NextAllocation = allocs[2].ID
+
+	sort.Sort(AllocIDSort(allocs))
+	sort.Sort(AllocIDSort(out))
+
+	require.Equal(allocs, out)
+	require.False(watchFired(ws))
+}
+
 func TestStateStore_RestoreAlloc(t *testing.T) {
 	state := testStateStore(t)
 	alloc := mock.Alloc()

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

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))