diff --git a/scheduler/util.go b/scheduler/util.go index 81ccfbd73..d45c7906a 100644 --- a/scheduler/util.go +++ b/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 diff --git a/scheduler/util_test.go b/scheduler/util_test.go index 851fcf096..5b67d2fad 100644 --- a/scheduler/util_test.go +++ b/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) } }