Merge pull request #93473 from Nordix/sched-fix-mael

Change nodeInfolist building logic in scheduler

pkg/scheduler/internal/cache/cache.go

@@ -279,9 +279,11 @@ func (cache *schedulerCache) updateNodeInfoSnapshotList(snapshot *Snapshot, upda
 	if updateAll {
 		// Take a snapshot of the nodes order in the tree
 		snapshot.nodeInfoList = make([]*framework.NodeInfo, 0, cache.nodeTree.numNodes)
-		cache.nodeTree.resetExhausted()
+		nodesList, err := cache.nodeTree.list()
-		for i := 0; i < cache.nodeTree.numNodes; i++ {
+		if err != nil {
-			nodeName := cache.nodeTree.next()
+			klog.Error(err)
+		}
+		for _, nodeName := range nodesList {
 			if n := snapshot.nodeInfoMap[nodeName]; n != nil {
 				snapshot.nodeInfoList = append(snapshot.nodeInfoList, n)
 				if len(n.PodsWithAffinity) > 0 {

pkg/scheduler/internal/cache/cache_test.go

@@ -1090,7 +1090,11 @@ func TestNodeOperators(t *testing.T) {
 			if !found {
 				t.Errorf("Failed to find node %v in internalcache.", node.Name)
 			}
-			if cache.nodeTree.numNodes != 1 || cache.nodeTree.next() != node.Name {
+			nodesList, err := cache.nodeTree.list()
+			if err != nil {
+				t.Fatal(err)
+			}
+			if cache.nodeTree.numNodes != 1 || nodesList[len(nodesList)-1] != node.Name {
 				t.Errorf("cache.nodeTree is not updated correctly after adding node: %v", node.Name)
 			}
 
@@ -1134,7 +1138,11 @@ func TestNodeOperators(t *testing.T) {
 				t.Errorf("Failed to update node in schedulertypes:\n got: %+v \nexpected: %+v", got, expected)
 			}
 			// Check nodeTree after update
-			if cache.nodeTree.numNodes != 1 || cache.nodeTree.next() != node.Name {
+			nodesList, err = cache.nodeTree.list()
+			if err != nil {
+				t.Fatal(err)
+			}
+			if cache.nodeTree.numNodes != 1 || nodesList[len(nodesList)-1] != node.Name {
 				t.Errorf("unexpected cache.nodeTree after updating node: %v", node.Name)
 			}
 
@@ -1147,8 +1155,12 @@ func TestNodeOperators(t *testing.T) {
 			} else if n != nil {
 				t.Errorf("The node object for %v should be nil", node.Name)
 			}
-			// Check node is removed from nodeTree.
+			// Check node is removed from nodeTree as well.
-			if cache.nodeTree.numNodes != 0 || cache.nodeTree.next() != "" {
+			nodesList, err = cache.nodeTree.list()
+			if err != nil {
+				t.Fatal(err)
+			}
+			if cache.nodeTree.numNodes != 0 || len(nodesList) != 0 {
 				t.Errorf("unexpected cache.nodeTree after removing node: %v", node.Name)
 			}
 			// Pods are still in the pods cache.
@@ -1306,7 +1318,7 @@ func TestSchedulerCache_UpdateSnapshot(t *testing.T) {
 	updateSnapshot := func() operation {
 		return func() {
 			cache.UpdateSnapshot(snapshot)
-			if err := compareCacheWithNodeInfoSnapshot(cache, snapshot); err != nil {
+			if err := compareCacheWithNodeInfoSnapshot(t, cache, snapshot); err != nil {
 				t.Error(err)
 			}
 		}
@@ -1487,14 +1499,14 @@ func TestSchedulerCache_UpdateSnapshot(t *testing.T) {
 			if err := cache.UpdateSnapshot(snapshot); err != nil {
 				t.Error(err)
 			}
-			if err := compareCacheWithNodeInfoSnapshot(cache, snapshot); err != nil {
+			if err := compareCacheWithNodeInfoSnapshot(t, cache, snapshot); err != nil {
 				t.Error(err)
 			}
 		})
 	}
 }
 
-func compareCacheWithNodeInfoSnapshot(cache *schedulerCache, snapshot *Snapshot) error {
+func compareCacheWithNodeInfoSnapshot(t *testing.T, cache *schedulerCache, snapshot *Snapshot) error {
 	// Compare the map.
 	if len(snapshot.nodeInfoMap) != len(cache.nodes) {
 		return fmt.Errorf("unexpected number of nodes in the snapshot. Expected: %v, got: %v", len(cache.nodes), len(snapshot.nodeInfoMap))
@@ -1512,8 +1524,11 @@ func compareCacheWithNodeInfoSnapshot(cache *schedulerCache, snapshot *Snapshot)
 
 	expectedNodeInfoList := make([]*framework.NodeInfo, 0, cache.nodeTree.numNodes)
 	expectedHavePodsWithAffinityNodeInfoList := make([]*framework.NodeInfo, 0, cache.nodeTree.numNodes)
-	for i := 0; i < cache.nodeTree.numNodes; i++ {
+	nodesList, err := cache.nodeTree.list()
-		nodeName := cache.nodeTree.next()
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, nodeName := range nodesList {
 		if n := snapshot.nodeInfoMap[nodeName]; n != nil {
 			expectedNodeInfoList = append(expectedNodeInfoList, n)
 			if len(n.PodsWithAffinity) > 0 {
@@ -1576,7 +1591,7 @@ func TestSchedulerCache_updateNodeInfoSnapshotList(t *testing.T) {
 
 	updateSnapshot := func(t *testing.T) {
 		cache.updateNodeInfoSnapshotList(snapshot, true)
-		if err := compareCacheWithNodeInfoSnapshot(cache, snapshot); err != nil {
+		if err := compareCacheWithNodeInfoSnapshot(t, cache, snapshot); err != nil {
 			t.Error(err)
 		}
 	}
@@ -1672,7 +1687,7 @@ func TestSchedulerCache_updateNodeInfoSnapshotList(t *testing.T) {
 
 			// Always update the snapshot at the end of operations and compare it.
 			cache.updateNodeInfoSnapshotList(snapshot, true)
-			if err := compareCacheWithNodeInfoSnapshot(cache, snapshot); err != nil {
+			if err := compareCacheWithNodeInfoSnapshot(t, cache, snapshot); err != nil {
 				t.Error(err)
 			}
 			nodeNames := make([]string, len(snapshot.nodeInfoList))

pkg/scheduler/internal/cache/node_tree.go

@@ -17,9 +17,10 @@ limitations under the License.
 package cache
 
 import (
+	"errors"
 	"fmt"
 
-	"k8s.io/api/core/v1"
+	v1 "k8s.io/api/core/v1"
 	"k8s.io/klog/v2"
 	utilnode "k8s.io/kubernetes/pkg/util/node"
 )
@@ -29,37 +30,15 @@ import (
 // NodeTree is NOT thread-safe, any concurrent updates/reads from it must be synchronized by the caller.
 // It is used only by schedulerCache, and should stay as such.
 type nodeTree struct {
-	tree      map[string]*nodeArray // a map from zone (region-zone) to an array of nodes in the zone.
+	tree     map[string][]string // a map from zone (region-zone) to an array of nodes in the zone.
-	zones     []string              // a list of all the zones in the tree (keys)
+	zones    []string            // a list of all the zones in the tree (keys)
-	zoneIndex int
+	numNodes int
-	numNodes  int
-}
-
-// nodeArray is a struct that has nodes that are in a zone.
-// We use a slice (as opposed to a set/map) to store the nodes because iterating over the nodes is
-// a lot more frequent than searching them by name.
-type nodeArray struct {
-	nodes     []string
-	lastIndex int
-}
-
-func (na *nodeArray) next() (nodeName string, exhausted bool) {
-	if len(na.nodes) == 0 {
-		klog.Error("The nodeArray is empty. It should have been deleted from NodeTree.")
-		return "", false
-	}
-	if na.lastIndex >= len(na.nodes) {
-		return "", true
-	}
-	nodeName = na.nodes[na.lastIndex]
-	na.lastIndex++
-	return nodeName, false
 }
 
 // newNodeTree creates a NodeTree from nodes.
 func newNodeTree(nodes []*v1.Node) *nodeTree {
 	nt := &nodeTree{
-		tree: make(map[string]*nodeArray),
+		tree: make(map[string][]string),
 	}
 	for _, n := range nodes {
 		nt.addNode(n)
@@ -72,16 +51,16 @@ func newNodeTree(nodes []*v1.Node) *nodeTree {
 func (nt *nodeTree) addNode(n *v1.Node) {
 	zone := utilnode.GetZoneKey(n)
 	if na, ok := nt.tree[zone]; ok {
-		for _, nodeName := range na.nodes {
+		for _, nodeName := range na {
 			if nodeName == n.Name {
 				klog.Warningf("node %q already exist in the NodeTree", n.Name)
 				return
 			}
 		}
-		na.nodes = append(na.nodes, n.Name)
+		nt.tree[zone] = append(na, n.Name)
 	} else {
 		nt.zones = append(nt.zones, zone)
-		nt.tree[zone] = &nodeArray{nodes: []string{n.Name}, lastIndex: 0}
+		nt.tree[zone] = []string{n.Name}
 	}
 	klog.V(2).Infof("Added node %q in group %q to NodeTree", n.Name, zone)
 	nt.numNodes++
@@ -91,10 +70,10 @@ func (nt *nodeTree) addNode(n *v1.Node) {
 func (nt *nodeTree) removeNode(n *v1.Node) error {
 	zone := utilnode.GetZoneKey(n)
 	if na, ok := nt.tree[zone]; ok {
-		for i, nodeName := range na.nodes {
+		for i, nodeName := range na {
 			if nodeName == n.Name {
-				na.nodes = append(na.nodes[:i], na.nodes[i+1:]...)
+				nt.tree[zone] = append(na[:i], na[i+1:]...)
-				if len(na.nodes) == 0 {
+				if len(nt.tree[zone]) == 0 {
 					nt.removeZone(zone)
 				}
 				klog.V(2).Infof("Removed node %q in group %q from NodeTree", n.Name, zone)
@@ -135,36 +114,30 @@ func (nt *nodeTree) updateNode(old, new *v1.Node) {
 	nt.addNode(new)
 }
 
-func (nt *nodeTree) resetExhausted() {
+// list returns the list of names of the node. NodeTree iterates over zones and in each zone iterates
-	for _, na := range nt.tree {
-		na.lastIndex = 0
-	}
-	nt.zoneIndex = 0
-}
-
-// next returns the name of the next node. NodeTree iterates over zones and in each zone iterates
 // over nodes in a round robin fashion.
-func (nt *nodeTree) next() string {
+func (nt *nodeTree) list() ([]string, error) {
 	if len(nt.zones) == 0 {
-		return ""
+		return nil, nil
 	}
+	nodesList := make([]string, 0, nt.numNodes)
 	numExhaustedZones := 0
-	for {
+	nodeIndex := 0
-		if nt.zoneIndex >= len(nt.zones) {
+	for len(nodesList) < nt.numNodes {
-			nt.zoneIndex = 0
+		if numExhaustedZones >= len(nt.zones) { // all zones are exhausted.
+			return nodesList, errors.New("all zones exhausted before reaching count of nodes expected")
 		}
-		zone := nt.zones[nt.zoneIndex]
+		for zoneIndex := 0; zoneIndex < len(nt.zones); zoneIndex++ {
-		nt.zoneIndex++
+			na := nt.tree[nt.zones[zoneIndex]]
-		// We do not check the exhausted zones before calling next() on the zone. This ensures
+			if nodeIndex >= len(na) { // If the zone is exhausted, continue
-		// that if more nodes are added to a zone after it is exhausted, we iterate over the new nodes.
+				if nodeIndex == len(na) { // If it is the first time the zone is exhausted
-		nodeName, exhausted := nt.tree[zone].next()
+					numExhaustedZones++
-		if exhausted {
+				}
-			numExhaustedZones++
+				continue
-			if numExhaustedZones >= len(nt.zones) { // all zones are exhausted. we should reset.
-				nt.resetExhausted()
 			}
-		} else {
+			nodesList = append(nodesList, na[nodeIndex])
-			return nodeName
 		}
+		nodeIndex++
 	}
+	return nodesList, nil
 }

pkg/scheduler/internal/cache/node_tree_test.go

@@ -20,7 +20,7 @@ import (
 	"reflect"
 	"testing"
 
-	"k8s.io/api/core/v1"
+	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
@@ -133,10 +133,10 @@ var allNodes = []*v1.Node{
 	},
 }
 
-func verifyNodeTree(t *testing.T, nt *nodeTree, expectedTree map[string]*nodeArray) {
+func verifyNodeTree(t *testing.T, nt *nodeTree, expectedTree map[string][]string) {
 	expectedNumNodes := int(0)
 	for _, na := range expectedTree {
-		expectedNumNodes += len(na.nodes)
+		expectedNumNodes += len(na)
 	}
 	if numNodes := nt.numNodes; numNodes != expectedNumNodes {
 		t.Errorf("unexpected nodeTree.numNodes. Expected: %v, Got: %v", expectedNumNodes, numNodes)
@@ -158,41 +158,41 @@ func TestNodeTree_AddNode(t *testing.T) {
 	tests := []struct {
 		name         string
 		nodesToAdd   []*v1.Node
-		expectedTree map[string]*nodeArray
+		expectedTree map[string][]string
 	}{
 		{
 			name:         "single node no labels",
 			nodesToAdd:   allNodes[:1],
-			expectedTree: map[string]*nodeArray{"": {[]string{"node-0"}, 0}},
+			expectedTree: map[string][]string{"": {"node-0"}},
 		},
 		{
 			name:       "mix of nodes with and without proper labels",
 			nodesToAdd: allNodes[:4],
-			expectedTree: map[string]*nodeArray{
+			expectedTree: map[string][]string{
-				"":                     {[]string{"node-0"}, 0},
+				"":                     {"node-0"},
-				"region-1:\x00:":       {[]string{"node-1"}, 0},
+				"region-1:\x00:":       {"node-1"},
-				":\x00:zone-2":         {[]string{"node-2"}, 0},
+				":\x00:zone-2":         {"node-2"},
-				"region-1:\x00:zone-2": {[]string{"node-3"}, 0},
+				"region-1:\x00:zone-2": {"node-3"},
 			},
 		},
 		{
 			name:       "mix of nodes with and without proper labels and some zones with multiple nodes",
 			nodesToAdd: allNodes[:7],
-			expectedTree: map[string]*nodeArray{
+			expectedTree: map[string][]string{
-				"":                     {[]string{"node-0"}, 0},
+				"":                     {"node-0"},
-				"region-1:\x00:":       {[]string{"node-1"}, 0},
+				"region-1:\x00:":       {"node-1"},
-				":\x00:zone-2":         {[]string{"node-2"}, 0},
+				":\x00:zone-2":         {"node-2"},
-				"region-1:\x00:zone-2": {[]string{"node-3", "node-4"}, 0},
+				"region-1:\x00:zone-2": {"node-3", "node-4"},
-				"region-1:\x00:zone-3": {[]string{"node-5"}, 0},
+				"region-1:\x00:zone-3": {"node-5"},
-				"region-2:\x00:zone-2": {[]string{"node-6"}, 0},
+				"region-2:\x00:zone-2": {"node-6"},
 			},
 		},
 		{
 			name:       "nodes also using deprecated zone/region label",
 			nodesToAdd: allNodes[9:],
-			expectedTree: map[string]*nodeArray{
+			expectedTree: map[string][]string{
-				"region-2:\x00:zone-2": {[]string{"node-9"}, 0},
+				"region-2:\x00:zone-2": {"node-9"},
-				"region-2:\x00:zone-3": {[]string{"node-10"}, 0},
+				"region-2:\x00:zone-3": {"node-10"},
 			},
 		},
 	}
@@ -213,43 +213,43 @@ func TestNodeTree_RemoveNode(t *testing.T) {
 		name          string
 		existingNodes []*v1.Node
 		nodesToRemove []*v1.Node
-		expectedTree  map[string]*nodeArray
+		expectedTree  map[string][]string
 		expectError   bool
 	}{
 		{
 			name:          "remove a single node with no labels",
 			existingNodes: allNodes[:7],
 			nodesToRemove: allNodes[:1],
-			expectedTree: map[string]*nodeArray{
+			expectedTree: map[string][]string{
-				"region-1:\x00:":       {[]string{"node-1"}, 0},
+				"region-1:\x00:":       {"node-1"},
-				":\x00:zone-2":         {[]string{"node-2"}, 0},
+				":\x00:zone-2":         {"node-2"},
-				"region-1:\x00:zone-2": {[]string{"node-3", "node-4"}, 0},
+				"region-1:\x00:zone-2": {"node-3", "node-4"},
-				"region-1:\x00:zone-3": {[]string{"node-5"}, 0},
+				"region-1:\x00:zone-3": {"node-5"},
-				"region-2:\x00:zone-2": {[]string{"node-6"}, 0},
+				"region-2:\x00:zone-2": {"node-6"},
 			},
 		},
 		{
 			name:          "remove a few nodes including one from a zone with multiple nodes",
 			existingNodes: allNodes[:7],
 			nodesToRemove: allNodes[1:4],
-			expectedTree: map[string]*nodeArray{
+			expectedTree: map[string][]string{
-				"":                     {[]string{"node-0"}, 0},
+				"":                     {"node-0"},
-				"region-1:\x00:zone-2": {[]string{"node-4"}, 0},
+				"region-1:\x00:zone-2": {"node-4"},
-				"region-1:\x00:zone-3": {[]string{"node-5"}, 0},
+				"region-1:\x00:zone-3": {"node-5"},
-				"region-2:\x00:zone-2": {[]string{"node-6"}, 0},
+				"region-2:\x00:zone-2": {"node-6"},
 			},
 		},
 		{
 			name:          "remove all nodes",
 			existingNodes: allNodes[:7],
 			nodesToRemove: allNodes[:7],
-			expectedTree:  map[string]*nodeArray{},
+			expectedTree:  map[string][]string{},
 		},
 		{
 			name:          "remove non-existing node",
 			existingNodes: nil,
 			nodesToRemove: allNodes[:5],
-			expectedTree:  map[string]*nodeArray{},
+			expectedTree:  map[string][]string{},
 			expectError:   true,
 		},
 	}
@@ -273,7 +273,7 @@ func TestNodeTree_UpdateNode(t *testing.T) {
 		name          string
 		existingNodes []*v1.Node
 		nodeToUpdate  *v1.Node
-		expectedTree  map[string]*nodeArray
+		expectedTree  map[string][]string
 	}{
 		{
 			name:          "update a node without label",
@@ -287,12 +287,12 @@ func TestNodeTree_UpdateNode(t *testing.T) {
 					},
 				},
 			},
-			expectedTree: map[string]*nodeArray{
+			expectedTree: map[string][]string{
-				"region-1:\x00:":       {[]string{"node-1"}, 0},
+				"region-1:\x00:":       {"node-1"},
-				":\x00:zone-2":         {[]string{"node-2"}, 0},
+				":\x00:zone-2":         {"node-2"},
-				"region-1:\x00:zone-2": {[]string{"node-3", "node-4", "node-0"}, 0},
+				"region-1:\x00:zone-2": {"node-3", "node-4", "node-0"},
-				"region-1:\x00:zone-3": {[]string{"node-5"}, 0},
+				"region-1:\x00:zone-3": {"node-5"},
-				"region-2:\x00:zone-2": {[]string{"node-6"}, 0},
+				"region-2:\x00:zone-2": {"node-6"},
 			},
 		},
 		{
@@ -307,8 +307,8 @@ func TestNodeTree_UpdateNode(t *testing.T) {
 					},
 				},
 			},
-			expectedTree: map[string]*nodeArray{
+			expectedTree: map[string][]string{
-				"region-1:\x00:zone-2": {[]string{"node-0"}, 0},
+				"region-1:\x00:zone-2": {"node-0"},
 			},
 		},
 		{
@@ -323,9 +323,9 @@ func TestNodeTree_UpdateNode(t *testing.T) {
 					},
 				},
 			},
-			expectedTree: map[string]*nodeArray{
+			expectedTree: map[string][]string{
-				"":                     {[]string{"node-0"}, 0},
+				"":                     {"node-0"},
-				"region-1:\x00:zone-2": {[]string{"node-new"}, 0},
+				"region-1:\x00:zone-2": {"node-new"},
 			},
 		},
 	}
@@ -349,36 +349,31 @@ func TestNodeTree_UpdateNode(t *testing.T) {
 	}
 }
 
-func TestNodeTree_Next(t *testing.T) {
+func TestNodeTree_List(t *testing.T) {
 	tests := []struct {
 		name           string
 		nodesToAdd     []*v1.Node
-		numRuns        int // number of times to run Next()
 		expectedOutput []string
 	}{
 		{
 			name:           "empty tree",
 			nodesToAdd:     nil,
-			numRuns:        2,
+			expectedOutput: nil,
-			expectedOutput: []string{"", ""},
 		},
 		{
-			name:           "should go back to the first node after finishing a round",
+			name:           "one node",
 			nodesToAdd:     allNodes[:1],
-			numRuns:        2,
+			expectedOutput: []string{"node-0"},
-			expectedOutput: []string{"node-0", "node-0"},
 		},
 		{
-			name:           "should go back to the first node after going over all nodes",
+			name:           "four nodes",
 			nodesToAdd:     allNodes[:4],
-			numRuns:        5,
+			expectedOutput: []string{"node-0", "node-1", "node-2", "node-3"},
-			expectedOutput: []string{"node-0", "node-1", "node-2", "node-3", "node-0"},
 		},
 		{
-			name:           "should go to all zones before going to the second nodes in the same zone",
+			name:           "all nodes",
 			nodesToAdd:     allNodes[:9],
-			numRuns:        11,
+			expectedOutput: []string{"node-0", "node-1", "node-2", "node-3", "node-5", "node-6", "node-4", "node-7", "node-8"},
-			expectedOutput: []string{"node-0", "node-1", "node-2", "node-3", "node-5", "node-6", "node-4", "node-7", "node-8", "node-0", "node-1"},
 		},
 	}
 
@@ -386,9 +381,9 @@ func TestNodeTree_Next(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			nt := newNodeTree(test.nodesToAdd)
 
-			var output []string
+			output, err := nt.list()
-			for i := 0; i < test.numRuns; i++ {
+			if err != nil {
-				output = append(output, nt.next())
+				t.Fatal(err)
 			}
 			if !reflect.DeepEqual(output, test.expectedOutput) {
 				t.Errorf("unexpected output. Expected: %v, Got: %v", test.expectedOutput, output)
@@ -397,6 +392,15 @@ func TestNodeTree_Next(t *testing.T) {
 	}
 }
 
+func TestNodeTree_List_Exhausted(t *testing.T) {
+	nt := newNodeTree(allNodes[:9])
+	nt.numNodes++
+	_, err := nt.list()
+	if err == nil {
+		t.Fatal("Expected an error from zone exhaustion")
+	}
+}
+
 func TestNodeTreeMultiOperations(t *testing.T) {
 	tests := []struct {
 		name           string
@@ -406,39 +410,39 @@ func TestNodeTreeMultiOperations(t *testing.T) {
 		expectedOutput []string
 	}{
 		{
-			name:           "add and remove all nodes between two Next operations",
+			name:           "add and remove all nodes",
 			nodesToAdd:     allNodes[2:9],
 			nodesToRemove:  allNodes[2:9],
-			operations:     []string{"add", "add", "next", "add", "remove", "remove", "remove", "next"},
+			operations:     []string{"add", "add", "add", "remove", "remove", "remove"},
-			expectedOutput: []string{"node-2", ""},
+			expectedOutput: nil,
 		},
 		{
-			name:           "add and remove some nodes between two Next operations",
+			name:           "add and remove some nodes",
 			nodesToAdd:     allNodes[2:9],
 			nodesToRemove:  allNodes[2:9],
-			operations:     []string{"add", "add", "next", "add", "remove", "remove", "next"},
+			operations:     []string{"add", "add", "add", "remove"},
-			expectedOutput: []string{"node-2", "node-4"},
+			expectedOutput: []string{"node-3", "node-4"},
 		},
 		{
-			name:           "remove nodes already iterated on and add new nodes",
+			name:           "remove three nodes",
 			nodesToAdd:     allNodes[2:9],
 			nodesToRemove:  allNodes[2:9],
-			operations:     []string{"add", "add", "next", "next", "add", "remove", "remove", "next"},
+			operations:     []string{"add", "add", "add", "remove", "remove", "remove", "add"},
-			expectedOutput: []string{"node-2", "node-3", "node-4"},
+			expectedOutput: []string{"node-5"},
 		},
 		{
 			name:           "add more nodes to an exhausted zone",
 			nodesToAdd:     append(allNodes[4:9:9], allNodes[3]),
 			nodesToRemove:  nil,
-			operations:     []string{"add", "add", "add", "add", "add", "next", "next", "next", "next", "add", "next", "next", "next"},
+			operations:     []string{"add", "add", "add", "add", "add", "add"},
-			expectedOutput: []string{"node-4", "node-5", "node-6", "node-7", "node-3", "node-8", "node-4"},
+			expectedOutput: []string{"node-4", "node-5", "node-6", "node-3", "node-7", "node-8"},
 		},
 		{
-			name:           "remove zone and add new to ensure exhausted is reset correctly",
+			name:           "remove zone and add new",
 			nodesToAdd:     append(allNodes[3:5:5], allNodes[6:8]...),
 			nodesToRemove:  allNodes[3:5],
-			operations:     []string{"add", "add", "next", "next", "remove", "add", "add", "next", "next", "remove", "next", "next"},
+			operations:     []string{"add", "add", "remove", "add", "add", "remove"},
-			expectedOutput: []string{"node-3", "node-4", "node-6", "node-7", "node-6", "node-7"},
+			expectedOutput: []string{"node-6", "node-7"},
 		},
 	}
 
@@ -447,7 +451,6 @@ func TestNodeTreeMultiOperations(t *testing.T) {
 			nt := newNodeTree(nil)
 			addIndex := 0
 			removeIndex := 0
-			var output []string
 			for _, op := range test.operations {
 				switch op {
 				case "add":
@@ -464,12 +467,14 @@ func TestNodeTreeMultiOperations(t *testing.T) {
 						nt.removeNode(test.nodesToRemove[removeIndex])
 						removeIndex++
 					}
-				case "next":
-					output = append(output, nt.next())
 				default:
 					t.Errorf("unknow operation: %v", op)
 				}
 			}
+			output, err := nt.list()
+			if err != nil {
+				t.Fatal(err)
+			}
 			if !reflect.DeepEqual(output, test.expectedOutput) {
 				t.Errorf("unexpected output. Expected: %v, Got: %v", test.expectedOutput, output)
 			}