Fix a bug in node tree when all nodes in a zone are removed

pkg/scheduler/internal/cache/node_tree.go

@@ -21,7 +21,6 @@ import (
 	"sync"
 
 	"k8s.io/api/core/v1"
-	"k8s.io/apimachinery/pkg/util/sets"
 	utilnode "k8s.io/kubernetes/pkg/util/node"
 
 	"github.com/golang/glog"
@@ -30,12 +29,11 @@ import (
 // NodeTree is a tree-like data structure that holds node names in each zone. Zone names are
 // keys to "NodeTree.tree" and values of "NodeTree.tree" are arrays of node names.
 type NodeTree struct {
-	tree           map[string]*nodeArray // 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)
-	zoneIndex      int
-	exhaustedZones sets.String // set of zones that all of their nodes are returned by next()
-	NumNodes       int
-	mu             sync.RWMutex
+	tree      map[string]*nodeArray // 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)
+	zoneIndex int
+	NumNodes  int
+	mu        sync.RWMutex
 }
 
 // nodeArray is a struct that has nodes that are in a zone.
@@ -62,8 +60,7 @@ func (na *nodeArray) next() (nodeName string, exhausted bool) {
 // newNodeTree creates a NodeTree from nodes.
 func newNodeTree(nodes []*v1.Node) *NodeTree {
 	nt := &NodeTree{
-		tree:           make(map[string]*nodeArray),
-		exhaustedZones: sets.NewString(),
+		tree: make(map[string]*nodeArray),
 	}
 	for _, n := range nodes {
 		nt.AddNode(n)
@@ -156,7 +153,7 @@ func (nt *NodeTree) resetExhausted() {
 	for _, na := range nt.tree {
 		na.lastIndex = 0
 	}
-	nt.exhaustedZones = sets.NewString()
+	nt.zoneIndex = 0
 }
 
 // Next returns the name of the next node. NodeTree iterates over zones and in each zone iterates
@@ -167,18 +164,19 @@ func (nt *NodeTree) Next() string {
 	if len(nt.zones) == 0 {
 		return ""
 	}
+	numExhaustedZones := 0
 	for {
 		if nt.zoneIndex >= len(nt.zones) {
 			nt.zoneIndex = 0
 		}
 		zone := nt.zones[nt.zoneIndex]
 		nt.zoneIndex++
-		// We do not check the set of exhausted zones before calling next() on the zone. This ensures
+		// We do not check the exhausted zones before calling next() on the zone. This ensures
 		// that if more nodes are added to a zone after it is exhausted, we iterate over the new nodes.
 		nodeName, exhausted := nt.tree[zone].next()
 		if exhausted {
-			nt.exhaustedZones.Insert(zone)
-			if len(nt.exhaustedZones) == len(nt.zones) { // all zones are exhausted. we should reset.
+			numExhaustedZones++
+			if numExhaustedZones >= len(nt.zones) { // all zones are exhausted. we should reset.
 				nt.resetExhausted()
 			}
 		} else {

pkg/scheduler/internal/cache/node_tree_test.go

@@ -403,6 +403,13 @@ func TestNodeTreeMultiOperations(t *testing.T) {
 			operations:     []string{"add", "add", "add", "add", "add", "next", "next", "next", "next", "add", "next", "next", "next"},
 			expectedOutput: []string{"node-4", "node-5", "node-6", "node-7", "node-3", "node-8", "node-4"},
 		},
+		{
+			name:           "remove zone and add new to ensure exhausted is reset correctly",
+			nodesToAdd:     append(allNodes[3:5], allNodes[6:8]...),
+			nodesToRemove:  allNodes[3:5],
+			operations:     []string{"add", "add", "next", "next", "remove", "add", "add", "next", "next", "remove", "next", "next"},
+			expectedOutput: []string{"node-3", "node-4", "node-6", "node-7", "node-6", "node-7"},
+		},
 	}
 
 	for _, test := range tests {