From efdf717991478fedb5df8b04a4a6d4e3c42e2916 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Mon, 7 Sep 2015 11:30:13 -0700
Subject: [PATCH] scheduler: allow updating the base nodes

---
 scheduler/stack.go      | 16 ++++++++++++++--
 scheduler/stack_test.go | 15 +++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/scheduler/stack.go b/scheduler/stack.go
index 859a0c5b0..14fdb08a0 100644
--- a/scheduler/stack.go
+++ b/scheduler/stack.go
@@ -22,6 +22,9 @@ const (
 // make placement decisions. Different schedulers may customize the
 // stack they use to vary the way placements are made.
 type Stack interface {
+	// SetNodes is used to set the base set of potential nodes
+	SetNodes([]*structs.Node)
+
 	// SetTaskGroup is used to set the job for selection
 	SetJob(job *structs.Job)
 
@@ -34,6 +37,7 @@ type Stack interface {
 type GenericStack struct {
 	batch               bool
 	ctx                 Context
+	source              *StaticIterator
 	jobConstraint       *ConstraintIterator
 	taskGroupDrivers    *DriverIterator
 	taskGroupConstraint *ConstraintIterator
@@ -53,10 +57,10 @@ func NewGenericStack(batch bool, ctx Context, baseNodes []*structs.Node) *Generi
 	// Create the source iterator. We randomize the order we visit nodes
 	// to reduce collisions between schedulers and to do a basic load
 	// balancing across eligible nodes.
-	source := NewRandomIterator(ctx, baseNodes)
+	stack.source = NewRandomIterator(ctx, baseNodes)
 
 	// Attach the job constraints. The job is filled in later.
-	stack.jobConstraint = NewConstraintIterator(ctx, source, nil)
+	stack.jobConstraint = NewConstraintIterator(ctx, stack.source, nil)
 
 	// Filter on task group drivers first as they are faster
 	stack.taskGroupDrivers = NewDriverIterator(ctx, stack.jobConstraint, nil)
@@ -101,6 +105,14 @@ func NewGenericStack(batch bool, ctx Context, baseNodes []*structs.Node) *Generi
 	return stack
 }
 
+func (s *GenericStack) SetNodes(baseNodes []*structs.Node) {
+	// Shuffle base nodes
+	shuffleNodes(baseNodes)
+
+	// Update the set of base nodes
+	s.source.SetNodes(baseNodes)
+}
+
 func (s *GenericStack) SetJob(job *structs.Job) {
 	s.jobConstraint.SetConstraints(job.Constraints)
 	s.binPack.SetPriority(job.Priority)
diff --git a/scheduler/stack_test.go b/scheduler/stack_test.go
index d52c9e310..c39246d3c 100644
--- a/scheduler/stack_test.go
+++ b/scheduler/stack_test.go
@@ -8,6 +8,21 @@ import (
 	"github.com/hashicorp/nomad/nomad/structs"
 )
 
+func TestServiceStack_SetNodes(t *testing.T) {
+	_, ctx := testContext(t)
+	stack := NewGenericStack(false, ctx, nil)
+
+	nodes := []*structs.Node{
+		mock.Node(),
+	}
+	stack.SetNodes(nodes)
+
+	out := collectFeasible(stack.source)
+	if !reflect.DeepEqual(out, nodes) {
+		t.Fatalf("bad: %#v", out)
+	}
+}
+
 func TestServiceStack_SetJob(t *testing.T) {
 	_, ctx := testContext(t)
 	stack := NewGenericStack(false, ctx, nil)