docker/swarmkit

package scheduler

import (
    "container/heap"
    "errors"
    "strings"

    "github.com/moby/swarmkit/v2/api"
    "github.com/moby/swarmkit/v2/manager/constraint"
)

var errNodeNotFound = errors.New("node not found in scheduler dataset")

type nodeSet struct {
    nodes map[string]NodeInfo // map from node id to node info
}

func (ns *nodeSet) alloc(n int) {
    ns.nodes = make(map[string]NodeInfo, n)
}

// nodeInfo returns the NodeInfo struct for a given node identified by its ID.
func (ns *nodeSet) nodeInfo(nodeID string) (NodeInfo, error) {
    node, ok := ns.nodes[nodeID]
    if ok {
        return node, nil
    }
    return NodeInfo{}, errNodeNotFound
}

// addOrUpdateNode sets the number of tasks for a given node. It adds the node
// to the set if it wasn't already tracked.
func (ns *nodeSet) addOrUpdateNode(n NodeInfo) {
    ns.nodes[n.ID] = n
}

// updateNode sets the number of tasks for a given node. It ignores the update
// if the node isn't already tracked in the set.
func (ns *nodeSet) updateNode(n NodeInfo) {
    _, ok := ns.nodes[n.ID]
    if ok {
        ns.nodes[n.ID] = n
    }
}

func (ns *nodeSet) remove(nodeID string) {
    delete(ns.nodes, nodeID)
}

func (ns *nodeSet) tree(serviceID string, preferences []*api.PlacementPreference, maxAssignments int, meetsConstraints func(*NodeInfo) bool, nodeLess func(*NodeInfo, *NodeInfo) bool) decisionTree {
    var root decisionTree

    if maxAssignments == 0 {
        return root
    }

    for _, node := range ns.nodes {
        tree := &root
        for _, pref := range preferences {
            // Only spread is supported so far
            spread := pref.GetSpread()
            if spread == nil {
                continue
            }

            descriptor := spread.SpreadDescriptor
            var value string
            switch {
            case len(descriptor) > len(constraint.NodeLabelPrefix) && strings.EqualFold(descriptor[:len(constraint.NodeLabelPrefix)], constraint.NodeLabelPrefix):
                if node.Spec.Annotations.Labels != nil {
                    value = node.Spec.Annotations.Labels[descriptor[len(constraint.NodeLabelPrefix):]]
                }
            case len(descriptor) > len(constraint.EngineLabelPrefix) && strings.EqualFold(descriptor[:len(constraint.EngineLabelPrefix)], constraint.EngineLabelPrefix):
                if node.Description != nil && node.Description.Engine != nil && node.Description.Engine.Labels != nil {
                    value = node.Description.Engine.Labels[descriptor[len(constraint.EngineLabelPrefix):]]
                }
            // TODO(aaronl): Support other items from constraint
            // syntax like node ID, hostname, os/arch, etc?
            default:
                continue
            }

            // If value is still uninitialized, the value used for
            // the node at this level of the tree is "". This makes
            // sure that the tree structure is not affected by
            // which properties nodes have and don't have.

            if node.ActiveTasksCountByService != nil {
                tree.tasks += node.ActiveTasksCountByService[serviceID]
            }

            if tree.next == nil {
                tree.next = make(map[string]*decisionTree)
            }
            next := tree.next[value]
            if next == nil {
                next = &decisionTree{}
                tree.next[value] = next
            }
            tree = next
        }

        if node.ActiveTasksCountByService != nil {
            tree.tasks += node.ActiveTasksCountByService[serviceID]
        }

        if tree.nodeHeap.lessFunc == nil {
            tree.nodeHeap.lessFunc = nodeLess
        }

        if tree.nodeHeap.Len() < maxAssignments {
            if meetsConstraints(&node) {
                heap.Push(&tree.nodeHeap, node)
            }
        } else if nodeLess(&node, &tree.nodeHeap.nodes[0]) {
            if meetsConstraints(&node) {
                tree.nodeHeap.nodes[0] = node
                heap.Fix(&tree.nodeHeap, 0)
            }
        }
    }

    return root
}