github.com/kaisenlinux/docker.io@v0.0.0-20230510090727-ea55db55fac7/swarmkit/manager/scheduler/nodeset.go (about)

     1  package scheduler
     2  
     3  import (
     4  	"container/heap"
     5  	"errors"
     6  	"strings"
     7  
     8  	"github.com/docker/swarmkit/api"
     9  	"github.com/docker/swarmkit/manager/constraint"
    10  )
    11  
    12  var errNodeNotFound = errors.New("node not found in scheduler dataset")
    13  
    14  type nodeSet struct {
    15  	nodes map[string]NodeInfo // map from node id to node info
    16  }
    17  
    18  func (ns *nodeSet) alloc(n int) {
    19  	ns.nodes = make(map[string]NodeInfo, n)
    20  }
    21  
    22  // nodeInfo returns the NodeInfo struct for a given node identified by its ID.
    23  func (ns *nodeSet) nodeInfo(nodeID string) (NodeInfo, error) {
    24  	node, ok := ns.nodes[nodeID]
    25  	if ok {
    26  		return node, nil
    27  	}
    28  	return NodeInfo{}, errNodeNotFound
    29  }
    30  
    31  // addOrUpdateNode sets the number of tasks for a given node. It adds the node
    32  // to the set if it wasn't already tracked.
    33  func (ns *nodeSet) addOrUpdateNode(n NodeInfo) {
    34  	ns.nodes[n.ID] = n
    35  }
    36  
    37  // updateNode sets the number of tasks for a given node. It ignores the update
    38  // if the node isn't already tracked in the set.
    39  func (ns *nodeSet) updateNode(n NodeInfo) {
    40  	_, ok := ns.nodes[n.ID]
    41  	if ok {
    42  		ns.nodes[n.ID] = n
    43  	}
    44  }
    45  
    46  func (ns *nodeSet) remove(nodeID string) {
    47  	delete(ns.nodes, nodeID)
    48  }
    49  
    50  func (ns *nodeSet) tree(serviceID string, preferences []*api.PlacementPreference, maxAssignments int, meetsConstraints func(*NodeInfo) bool, nodeLess func(*NodeInfo, *NodeInfo) bool) decisionTree {
    51  	var root decisionTree
    52  
    53  	if maxAssignments == 0 {
    54  		return root
    55  	}
    56  
    57  	for _, node := range ns.nodes {
    58  		tree := &root
    59  		for _, pref := range preferences {
    60  			// Only spread is supported so far
    61  			spread := pref.GetSpread()
    62  			if spread == nil {
    63  				continue
    64  			}
    65  
    66  			descriptor := spread.SpreadDescriptor
    67  			var value string
    68  			switch {
    69  			case len(descriptor) > len(constraint.NodeLabelPrefix) && strings.EqualFold(descriptor[:len(constraint.NodeLabelPrefix)], constraint.NodeLabelPrefix):
    70  				if node.Spec.Annotations.Labels != nil {
    71  					value = node.Spec.Annotations.Labels[descriptor[len(constraint.NodeLabelPrefix):]]
    72  				}
    73  			case len(descriptor) > len(constraint.EngineLabelPrefix) && strings.EqualFold(descriptor[:len(constraint.EngineLabelPrefix)], constraint.EngineLabelPrefix):
    74  				if node.Description != nil && node.Description.Engine != nil && node.Description.Engine.Labels != nil {
    75  					value = node.Description.Engine.Labels[descriptor[len(constraint.EngineLabelPrefix):]]
    76  				}
    77  			// TODO(aaronl): Support other items from constraint
    78  			// syntax like node ID, hostname, os/arch, etc?
    79  			default:
    80  				continue
    81  			}
    82  
    83  			// If value is still uninitialized, the value used for
    84  			// the node at this level of the tree is "". This makes
    85  			// sure that the tree structure is not affected by
    86  			// which properties nodes have and don't have.
    87  
    88  			if node.ActiveTasksCountByService != nil {
    89  				tree.tasks += node.ActiveTasksCountByService[serviceID]
    90  			}
    91  
    92  			if tree.next == nil {
    93  				tree.next = make(map[string]*decisionTree)
    94  			}
    95  			next := tree.next[value]
    96  			if next == nil {
    97  				next = &decisionTree{}
    98  				tree.next[value] = next
    99  			}
   100  			tree = next
   101  		}
   102  
   103  		if node.ActiveTasksCountByService != nil {
   104  			tree.tasks += node.ActiveTasksCountByService[serviceID]
   105  		}
   106  
   107  		if tree.nodeHeap.lessFunc == nil {
   108  			tree.nodeHeap.lessFunc = nodeLess
   109  		}
   110  
   111  		if tree.nodeHeap.Len() < maxAssignments {
   112  			if meetsConstraints(&node) {
   113  				heap.Push(&tree.nodeHeap, node)
   114  			}
   115  		} else if nodeLess(&node, &tree.nodeHeap.nodes[0]) {
   116  			if meetsConstraints(&node) {
   117  				tree.nodeHeap.nodes[0] = node
   118  				heap.Fix(&tree.nodeHeap, 0)
   119  			}
   120  		}
   121  	}
   122  
   123  	return root
   124  }