1/*
2Copyright 2018 The Kubernetes Authors.
3
4Licensed under the Apache License, Version 2.0 (the "License");
5you may not use this file except in compliance with the License.
6You may obtain a copy of the License at
7
8    http://www.apache.org/licenses/LICENSE-2.0
9
10Unless required by applicable law or agreed to in writing, software
11distributed under the License is distributed on an "AS IS" BASIS,
12WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13See the License for the specific language governing permissions and
14limitations under the License.
15*/
16
17package cache
18
19import (
20	"errors"
21	"fmt"
22
23	v1 "k8s.io/api/core/v1"
24	utilnode "k8s.io/component-helpers/node/topology"
25	"k8s.io/klog/v2"
26)
27
28// nodeTree is a tree-like data structure that holds node names in each zone. Zone names are
29// keys to "NodeTree.tree" and values of "NodeTree.tree" are arrays of node names.
30// NodeTree is NOT thread-safe, any concurrent updates/reads from it must be synchronized by the caller.
31// It is used only by schedulerCache, and should stay as such.
32type nodeTree struct {
33	tree     map[string][]string // a map from zone (region-zone) to an array of nodes in the zone.
34	zones    []string            // a list of all the zones in the tree (keys)
35	numNodes int
36}
37
38// newNodeTree creates a NodeTree from nodes.
39func newNodeTree(nodes []*v1.Node) *nodeTree {
40	nt := &nodeTree{
41		tree: make(map[string][]string),
42	}
43	for _, n := range nodes {
44		nt.addNode(n)
45	}
46	return nt
47}
48
49// addNode adds a node and its corresponding zone to the tree. If the zone already exists, the node
50// is added to the array of nodes in that zone.
51func (nt *nodeTree) addNode(n *v1.Node) {
52	zone := utilnode.GetZoneKey(n)
53	if na, ok := nt.tree[zone]; ok {
54		for _, nodeName := range na {
55			if nodeName == n.Name {
56				klog.Warningf("node %q already exist in the NodeTree", n.Name)
57				return
58			}
59		}
60		nt.tree[zone] = append(na, n.Name)
61	} else {
62		nt.zones = append(nt.zones, zone)
63		nt.tree[zone] = []string{n.Name}
64	}
65	klog.V(2).Infof("Added node %q in group %q to NodeTree", n.Name, zone)
66	nt.numNodes++
67}
68
69// removeNode removes a node from the NodeTree.
70func (nt *nodeTree) removeNode(n *v1.Node) error {
71	zone := utilnode.GetZoneKey(n)
72	if na, ok := nt.tree[zone]; ok {
73		for i, nodeName := range na {
74			if nodeName == n.Name {
75				nt.tree[zone] = append(na[:i], na[i+1:]...)
76				if len(nt.tree[zone]) == 0 {
77					nt.removeZone(zone)
78				}
79				klog.V(2).Infof("Removed node %q in group %q from NodeTree", n.Name, zone)
80				nt.numNodes--
81				return nil
82			}
83		}
84	}
85	klog.Errorf("Node %q in group %q was not found", n.Name, zone)
86	return fmt.Errorf("node %q in group %q was not found", n.Name, zone)
87}
88
89// removeZone removes a zone from tree.
90// This function must be called while writer locks are hold.
91func (nt *nodeTree) removeZone(zone string) {
92	delete(nt.tree, zone)
93	for i, z := range nt.zones {
94		if z == zone {
95			nt.zones = append(nt.zones[:i], nt.zones[i+1:]...)
96			return
97		}
98	}
99}
100
101// updateNode updates a node in the NodeTree.
102func (nt *nodeTree) updateNode(old, new *v1.Node) {
103	var oldZone string
104	if old != nil {
105		oldZone = utilnode.GetZoneKey(old)
106	}
107	newZone := utilnode.GetZoneKey(new)
108	// If the zone ID of the node has not changed, we don't need to do anything. Name of the node
109	// cannot be changed in an update.
110	if oldZone == newZone {
111		return
112	}
113	nt.removeNode(old) // No error checking. We ignore whether the old node exists or not.
114	nt.addNode(new)
115}
116
117// list returns the list of names of the node. NodeTree iterates over zones and in each zone iterates
118// over nodes in a round robin fashion.
119func (nt *nodeTree) list() ([]string, error) {
120	if len(nt.zones) == 0 {
121		return nil, nil
122	}
123	nodesList := make([]string, 0, nt.numNodes)
124	numExhaustedZones := 0
125	nodeIndex := 0
126	for len(nodesList) < nt.numNodes {
127		if numExhaustedZones >= len(nt.zones) { // all zones are exhausted.
128			return nodesList, errors.New("all zones exhausted before reaching count of nodes expected")
129		}
130		for zoneIndex := 0; zoneIndex < len(nt.zones); zoneIndex++ {
131			na := nt.tree[nt.zones[zoneIndex]]
132			if nodeIndex >= len(na) { // If the zone is exhausted, continue
133				if nodeIndex == len(na) { // If it is the first time the zone is exhausted
134					numExhaustedZones++
135				}
136				continue
137			}
138			nodesList = append(nodesList, na[nodeIndex])
139		}
140		nodeIndex++
141	}
142	return nodesList, nil
143}
144