scheduler: determine if any allocations need to be migrated

scheduler/util.go

@@ -36,14 +36,15 @@ func indexAllocs(allocs []*structs.Allocation) map[string][]*structs.Allocation
 }
 
 // diffAllocs is used to do a set difference between the target allocations
-// and the existing allocations. This returns 4 sets of results, the list of
+// and the existing allocations. This returns 5 sets of results, the list of
 // named task groups that need to be placed (no existing allocation), the
-// allocations that need to be updated (job definition is newer), the allocs
-// that need to be evicted (no longer required), and those that should be
-// ignored.
+// allocations that need to be updated (job definition is newer), allocs that
+// need to be migrated (node is draining), the allocs that need to be evicted
+// (no longer required), and those that should be ignored.
 func diffAllocs(job *structs.Job,
+	taintedNodes map[string]bool,
 	required map[string]*structs.TaskGroup,
-	existing map[string][]*structs.Allocation) (place, update, evict, ignore []allocNameID) {
+	existing map[string][]*structs.Allocation) (place, update, migrate, evict, ignore []allocNameID) {
 	// Scan the existing updates
 	for name, existList := range existing {
 		for _, exist := range existList {
@@ -56,6 +57,12 @@ func diffAllocs(job *structs.Job,
 				continue
 			}
 
+			// If we are on a tainted node, we must migrate
+			if taintedNodes[exist.NodeID] {
+				migrate = append(migrate, allocNameID{name, exist.ID})
+				continue
+			}
+
 			// If the definition is updated we need to update
 			// XXX: This is an extremely conservative approach. We can check
 			// if the job definition has changed in a way that affects

scheduler/util_test.go

@@ -54,6 +54,11 @@ func TestDiffAllocs(t *testing.T) {
 	*oldJob = *job
 	oldJob.ModifyIndex -= 1
 
+	tainted := map[string]bool{
+		"dead": true,
+		"zip":  false,
+	}
+
 	allocs := []*structs.Allocation{
 		// Update the 1st
 		&structs.Allocation{
@@ -77,10 +82,17 @@ func TestDiffAllocs(t *testing.T) {
 			NodeID: "zip",
 			Name:   "my-job.web[10]",
 		},
+
+		// Migrate the 3rd
+		&structs.Allocation{
+			ID:     mock.GenerateUUID(),
+			NodeID: "dead",
+			Name:   "my-job.web[2]",
+		},
 	}
 	existing := indexAllocs(allocs)
 
-	place, update, evict, ignore := diffAllocs(job, required, existing)
+	place, update, migrate, evict, ignore := diffAllocs(job, tainted, required, existing)
 
 	// We should update the first alloc
 	if len(update) != 1 || update[0].ID != allocs[0].ID {
@@ -97,8 +109,13 @@ func TestDiffAllocs(t *testing.T) {
 		t.Fatalf("bad: %#v", evict)
 	}
 
-	// We should place 8
-	if len(place) != 8 {
+	// We should migrate the 4rd alloc
+	if len(migrate) != 1 || migrate[0].ID != allocs[3].ID {
+		t.Fatalf("bad: %#v", migrate)
+	}
+
+	// We should place 7
+	if len(place) != 7 {
 		t.Fatalf("bad: %#v", place)
 	}
 }