decision_tree.go 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455
  1. package scheduler
  2. import (
  3. "container/heap"
  4. )
  5. type decisionTree struct {
  6. // Count of tasks for the service scheduled to this subtree
  7. tasks int
  8. // Non-leaf point to the next level of the tree. The key is the
  9. // value that the subtree covers.
  10. next map[string]*decisionTree
  11. // Leaf nodes contain a list of nodes
  12. nodeHeap nodeMaxHeap
  13. }
  14. // orderedNodes returns the nodes in this decision tree entry, sorted best
  15. // (lowest) first according to the sorting function. Must be called on a leaf
  16. // of the decision tree.
  17. //
  18. // The caller may modify the nodes in the returned slice. This has the effect
  19. // of changing the nodes in the decision tree entry. The next node to
  20. // findBestNodes on this decisionTree entry will take into account the changes
  21. // that were made to the nodes.
  22. func (dt *decisionTree) orderedNodes(meetsConstraints func(*NodeInfo) bool, nodeLess func(*NodeInfo, *NodeInfo) bool) []NodeInfo {
  23. if dt.nodeHeap.length != len(dt.nodeHeap.nodes) {
  24. // We already collapsed the heap into a sorted slice, so
  25. // re-heapify. There may have been modifications to the nodes
  26. // so we can't return dt.nodeHeap.nodes as-is. We also need to
  27. // reevaluate constraints because of the possible modifications.
  28. for i := 0; i < len(dt.nodeHeap.nodes); {
  29. if meetsConstraints(&dt.nodeHeap.nodes[i]) {
  30. i++
  31. } else {
  32. last := len(dt.nodeHeap.nodes) - 1
  33. dt.nodeHeap.nodes[i] = dt.nodeHeap.nodes[last]
  34. dt.nodeHeap.nodes = dt.nodeHeap.nodes[:last]
  35. }
  36. }
  37. dt.nodeHeap.length = len(dt.nodeHeap.nodes)
  38. heap.Init(&dt.nodeHeap)
  39. }
  40. // Popping every element orders the nodes from best to worst. The
  41. // first pop gets the worst node (since this a max-heap), and puts it
  42. // at position n-1. Then the next pop puts the next-worst at n-2, and
  43. // so on.
  44. for dt.nodeHeap.Len() > 0 {
  45. heap.Pop(&dt.nodeHeap)
  46. }
  47. return dt.nodeHeap.nodes
  48. }