scheduler: allow StaticIterator to update base set

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

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