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