mirror of
https://github.com/moby/moby.git
synced 2022-11-09 12:21:53 -05:00
128 lines
3.3 KiB
Go
128 lines
3.3 KiB
Go
package scheduler
|
|
|
|
import (
|
|
"container/heap"
|
|
"errors"
|
|
"time"
|
|
|
|
"github.com/docker/swarmkit/api"
|
|
)
|
|
|
|
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) {
|
|
if n.Tasks == nil {
|
|
n.Tasks = make(map[string]*api.Task)
|
|
}
|
|
if n.DesiredRunningTasksCountByService == nil {
|
|
n.DesiredRunningTasksCountByService = make(map[string]int)
|
|
}
|
|
if n.recentFailures == nil {
|
|
n.recentFailures = make(map[string][]time.Time)
|
|
}
|
|
|
|
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)
|
|
}
|
|
|
|
type nodeMaxHeap struct {
|
|
nodes []NodeInfo
|
|
lessFunc func(*NodeInfo, *NodeInfo) bool
|
|
length int
|
|
}
|
|
|
|
func (h nodeMaxHeap) Len() int {
|
|
return h.length
|
|
}
|
|
|
|
func (h nodeMaxHeap) Swap(i, j int) {
|
|
h.nodes[i], h.nodes[j] = h.nodes[j], h.nodes[i]
|
|
}
|
|
|
|
func (h nodeMaxHeap) Less(i, j int) bool {
|
|
// reversed to make a max-heap
|
|
return h.lessFunc(&h.nodes[j], &h.nodes[i])
|
|
}
|
|
|
|
func (h *nodeMaxHeap) Push(x interface{}) {
|
|
h.nodes = append(h.nodes, x.(NodeInfo))
|
|
h.length++
|
|
}
|
|
|
|
func (h *nodeMaxHeap) Pop() interface{} {
|
|
h.length--
|
|
// return value is never used
|
|
return nil
|
|
}
|
|
|
|
// findBestNodes returns n nodes (or < n if fewer nodes are available) that
|
|
// rank best (lowest) according to the sorting function.
|
|
func (ns *nodeSet) findBestNodes(n int, meetsConstraints func(*NodeInfo) bool, nodeLess func(*NodeInfo, *NodeInfo) bool) []NodeInfo {
|
|
if n == 0 {
|
|
return []NodeInfo{}
|
|
}
|
|
|
|
nodeHeap := nodeMaxHeap{lessFunc: nodeLess}
|
|
|
|
// TODO(aaronl): Is is possible to avoid checking constraints on every
|
|
// node? Perhaps we should try to schedule with n*2 nodes that weren't
|
|
// prescreened, and repeat the selection if there weren't enough nodes
|
|
// meeting the constraints.
|
|
for _, node := range ns.nodes {
|
|
// If there are fewer then n nodes in the heap, we add this
|
|
// node if it meets the constraints. Otherwise, the heap has
|
|
// n nodes, and if this node is better than the worst node in
|
|
// the heap, we replace the worst node and then fix the heap.
|
|
if nodeHeap.Len() < n {
|
|
if meetsConstraints(&node) {
|
|
heap.Push(&nodeHeap, node)
|
|
}
|
|
} else if nodeLess(&node, &nodeHeap.nodes[0]) {
|
|
if meetsConstraints(&node) {
|
|
nodeHeap.nodes[0] = node
|
|
heap.Fix(&nodeHeap, 0)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Popping every element orders the nodes from best to worst. The
|
|
// first pop gets the worst node (since this a max-heap), and puts it
|
|
// at position n-1. Then the next pop puts the next-worst at n-2, and
|
|
// so on.
|
|
for nodeHeap.Len() > 0 {
|
|
heap.Pop(&nodeHeap)
|
|
}
|
|
|
|
return nodeHeap.nodes
|
|
}
|