From 95c3c637ba7f41be01755911bd7dd2f34826ff41 Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Mon, 26 Mar 2018 17:54:46 -0700
Subject: [PATCH] Correct status desc on draining system allocs

---
 scheduler/generic_sched.go | 4 ++++
 scheduler/system_sched.go  | 9 +++++----
 scheduler/util.go          | 5 -----
 scheduler/util_test.go     | 4 ++--
 4 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/scheduler/generic_sched.go b/scheduler/generic_sched.go
index 501052979..a81fbb62b 100644
--- a/scheduler/generic_sched.go
+++ b/scheduler/generic_sched.go
@@ -35,6 +35,10 @@ const (
 	// allocInPlace is the status used when speculating on an in-place update
 	allocInPlace = "alloc updating in-place"
 
+	// allocNodeTainted is the status used when stopping an alloc because it's
+	// node is tainted.
+	allocNodeTainted = "alloc not needed as node is tainted"
+
 	// blockedEvalMaxPlanDesc is the description used for blocked evals that are
 	// a result of hitting the max number of plan attempts
 	blockedEvalMaxPlanDesc = "created due to placement conflicts"
diff --git a/scheduler/system_sched.go b/scheduler/system_sched.go
index 4fa2d20f6..4a10d0731 100644
--- a/scheduler/system_sched.go
+++ b/scheduler/system_sched.go
@@ -14,10 +14,6 @@ const (
 	// we will attempt to schedule if we continue to hit conflicts for system
 	// jobs.
 	maxSystemScheduleAttempts = 5
-
-	// allocNodeTainted is the status used when stopping an alloc because it's
-	// node is tainted.
-	allocNodeTainted = "alloc not needed as node is tainted"
 )
 
 // SystemScheduler is used for 'system' jobs. This scheduler is
@@ -212,6 +208,11 @@ func (s *SystemScheduler) computeJobAllocs() error {
 		s.plan.AppendUpdate(e.Alloc, structs.AllocDesiredStatusStop, allocNotNeeded, "")
 	}
 
+	// Add all the allocs to migrate
+	for _, e := range diff.migrate {
+		s.plan.AppendUpdate(e.Alloc, structs.AllocDesiredStatusStop, allocNodeTainted, "")
+	}
+
 	// Lost allocations should be transitioned to desired status stop and client
 	// status lost.
 	for _, e := range diff.lost {
diff --git a/scheduler/util.go b/scheduler/util.go
index c0943e126..6b589fa50 100644
--- a/scheduler/util.go
+++ b/scheduler/util.go
@@ -214,11 +214,6 @@ func diffSystemAllocs(job *structs.Job, nodes []*structs.Node, taintedNodes map[
 			}
 		}
 
-		// Migrate does not apply to system jobs and instead should be marked as
-		// stop because if a node is tainted, the job is invalid on that node.
-		diff.stop = append(diff.stop, diff.migrate...)
-		diff.migrate = nil
-
 		result.Append(diff)
 	}
 
diff --git a/scheduler/util_test.go b/scheduler/util_test.go
index 7fde4fa65..8912e9c8d 100644
--- a/scheduler/util_test.go
+++ b/scheduler/util_test.go
@@ -265,12 +265,12 @@ func TestDiffSystemAllocs(t *testing.T) {
 	}
 
 	// We should stop the third alloc
-	if len(stop) != 1 || stop[0].Alloc != allocs[2] {
+	if len(stop) != 0 {
 		t.Fatalf("bad: %#v", stop)
 	}
 
 	// There should be no migrates.
-	if len(migrate) != 0 {
+	if len(migrate) != 1 || migrate[0].Alloc != allocs[2] {
 		t.Fatalf("bad: %#v", migrate)
 	}