From 831848bf5e1049fab2dd1b0cccd7fc707a8f2235 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Mon, 7 Sep 2015 11:26:16 -0700
Subject: [PATCH] scheduler: allow StaticIterator to update base set

---
 scheduler/feasible.go      |  6 ++++++
 scheduler/feasible_test.go | 17 +++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/scheduler/feasible.go b/scheduler/feasible.go
index 8912c6552..d11a2034e 100644
--- a/scheduler/feasible.go
+++ b/scheduler/feasible.go
@@ -62,6 +62,12 @@ func (iter *StaticIterator) Reset() {
 	iter.seen = 0
 }
 
+func (iter *StaticIterator) SetNodes(nodes []*structs.Node) {
+	iter.nodes = nodes
+	iter.offset = 0
+	iter.seen = 0
+}
+
 // NewRandomIterator constructs a static iterator from a list of nodes
 // after applying the Fisher-Yates algorithm for a random shuffle. This
 // is applied in-place
diff --git a/scheduler/feasible_test.go b/scheduler/feasible_test.go
index 6f9ee0ca6..e651dec6f 100644
--- a/scheduler/feasible_test.go
+++ b/scheduler/feasible_test.go
@@ -39,6 +39,23 @@ func TestStaticIterator_Reset(t *testing.T) {
 	}
 }
 
+func TestStaticIterator_SetNodes(t *testing.T) {
+	_, ctx := testContext(t)
+	var nodes []*structs.Node
+	for i := 0; i < 3; i++ {
+		nodes = append(nodes, mock.Node())
+	}
+	static := NewStaticIterator(ctx, nodes)
+
+	newNodes := []*structs.Node{mock.Node()}
+	static.SetNodes(newNodes)
+
+	out := collectFeasible(static)
+	if !reflect.DeepEqual(out, newNodes) {
+		t.Fatalf("bad: %#v", out)
+	}
+}
+
 func TestRandomIterator(t *testing.T) {
 	_, ctx := testContext(t)
 	var nodes []*structs.Node