From 117b926e2b2d8307d88d659b043ae266dc6fbecf Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Tue, 17 May 2016 15:37:37 -0700
Subject: [PATCH] inplaceUpdate returns the allocs that were updated in-place

---
 scheduler/generic_sched.go |  3 +--
 scheduler/system_sched.go  |  3 +--
 scheduler/util.go          | 15 ++++++++-------
 scheduler/util_test.go     | 16 ++++++++++------
 4 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/scheduler/generic_sched.go b/scheduler/generic_sched.go
index dd50cc558..2cce204a1 100644
--- a/scheduler/generic_sched.go
+++ b/scheduler/generic_sched.go
@@ -324,8 +324,7 @@ func (s *GenericScheduler) computeJobAllocs() error {
 	}
 
 	// Attempt to do the upgrades in place
-	destructiveUpdates := inplaceUpdate(s.ctx, s.eval, s.job, s.stack, diff.update)
-	inplaceUpdates := diff.update[len(destructiveUpdates):]
+	destructiveUpdates, inplaceUpdates := inplaceUpdate(s.ctx, s.eval, s.job, s.stack, diff.update)
 	diff.update = destructiveUpdates
 
 	if s.eval.AnnotatePlan {
diff --git a/scheduler/system_sched.go b/scheduler/system_sched.go
index 6454c16c2..3499059fa 100644
--- a/scheduler/system_sched.go
+++ b/scheduler/system_sched.go
@@ -186,8 +186,7 @@ func (s *SystemScheduler) computeJobAllocs() error {
 	}
 
 	// Attempt to do the upgrades in place
-	destructiveUpdates := inplaceUpdate(s.ctx, s.eval, s.job, s.stack, diff.update)
-	inplaceUpdates := diff.update[len(destructiveUpdates):]
+	destructiveUpdates, inplaceUpdates := inplaceUpdate(s.ctx, s.eval, s.job, s.stack, diff.update)
 	diff.update = destructiveUpdates
 
 	if s.eval.AnnotatePlan {
diff --git a/scheduler/util.go b/scheduler/util.go
index e6ece1417..4dc411dc9 100644
--- a/scheduler/util.go
+++ b/scheduler/util.go
@@ -367,12 +367,13 @@ func setStatus(logger *log.Logger, planner Planner, eval, nextEval *structs.Eval
 	return planner.UpdateEval(newEval)
 }
 
-// inplaceUpdate attempts to update allocations in-place where possible.
+// inplaceUpdate attempts to update allocations in-place where possible. It
+// returns the allocs that couldn't be done inplace and then those that could.
 func inplaceUpdate(ctx Context, eval *structs.Evaluation, job *structs.Job,
-	stack Stack, updates []allocTuple) []allocTuple {
+	stack Stack, updates []allocTuple) (destructive, inplace []allocTuple) {
 
 	n := len(updates)
-	inplace := 0
+	inplaceCount := 0
 	for i := 0; i < n; i++ {
 		// Get the update
 		update := updates[i]
@@ -441,15 +442,15 @@ func inplaceUpdate(ctx Context, eval *structs.Evaluation, job *structs.Job,
 		ctx.Plan().AppendAlloc(newAlloc)
 
 		// Remove this allocation from the slice
-		updates[i] = updates[n-1]
+		updates[i], updates[n-1] = updates[n-1], updates[i]
 		i--
 		n--
-		inplace++
+		inplaceCount++
 	}
 	if len(updates) > 0 {
-		ctx.Logger().Printf("[DEBUG] sched: %#v: %d in-place updates of %d", eval, inplace, len(updates))
+		ctx.Logger().Printf("[DEBUG] sched: %#v: %d in-place updates of %d", eval, inplaceCount, len(updates))
 	}
-	return updates[:n]
+	return updates[:n], updates[n:]
 }
 
 // evictAndPlace is used to mark allocations for evicts and add them to the
diff --git a/scheduler/util_test.go b/scheduler/util_test.go
index 0d5cc915b..1d689fe2b 100644
--- a/scheduler/util_test.go
+++ b/scheduler/util_test.go
@@ -549,9 +549,9 @@ func TestInplaceUpdate_ChangedTaskGroup(t *testing.T) {
 	stack := NewGenericStack(false, ctx)
 
 	// Do the inplace update.
-	unplaced := inplaceUpdate(ctx, eval, job, stack, updates)
+	unplaced, inplace := inplaceUpdate(ctx, eval, job, stack, updates)
 
-	if len(unplaced) != 1 {
+	if len(unplaced) != 1 || len(inplace) != 0 {
 		t.Fatal("inplaceUpdate incorrectly did an inplace update")
 	}
 
@@ -594,9 +594,9 @@ func TestInplaceUpdate_NoMatch(t *testing.T) {
 	stack := NewGenericStack(false, ctx)
 
 	// Do the inplace update.
-	unplaced := inplaceUpdate(ctx, eval, job, stack, updates)
+	unplaced, inplace := inplaceUpdate(ctx, eval, job, stack, updates)
 
-	if len(unplaced) != 1 {
+	if len(unplaced) != 1 || len(inplace) != 0 {
 		t.Fatal("inplaceUpdate incorrectly did an inplace update")
 	}
 
@@ -665,9 +665,9 @@ func TestInplaceUpdate_Success(t *testing.T) {
 	stack.SetJob(job)
 
 	// Do the inplace update.
-	unplaced := inplaceUpdate(ctx, eval, job, stack, updates)
+	unplaced, inplace := inplaceUpdate(ctx, eval, job, stack, updates)
 
-	if len(unplaced) != 0 {
+	if len(unplaced) != 0 || len(inplace) != 1 {
 		t.Fatal("inplaceUpdate did not do an inplace update")
 	}
 
@@ -675,6 +675,10 @@ func TestInplaceUpdate_Success(t *testing.T) {
 		t.Fatal("inplaceUpdate did not do an inplace update")
 	}
 
+	if inplace[0].Alloc.ID != alloc.ID {
+		t.Fatalf("inplaceUpdate returned the wrong, inplace updated alloc: %#v", inplace)
+	}
+
 	// Get the alloc we inserted.
 	a := ctx.plan.NodeAllocation[alloc.NodeID][0]
 	if len(a.Services) != 3 {