github.com/djenriquez/nomad-1@v0.8.1/scheduler/select.go (about) 1 package scheduler 2 3 // LimitIterator is a RankIterator used to limit the number of options 4 // that are returned before we artificially end the stream. 5 type LimitIterator struct { 6 ctx Context 7 source RankIterator 8 limit int 9 maxSkip int 10 scoreThreshold float64 11 seen int 12 skippedNodes []*RankedNode 13 skippedNodeIndex int 14 } 15 16 // NewLimitIterator returns a LimitIterator with a fixed limit of returned options. 17 // Up to maxSkip options whose score is below scoreThreshold are skipped 18 // if there are additional options available in the source iterator 19 func NewLimitIterator(ctx Context, source RankIterator, limit int, scoreThreshold float64, maxSkip int) *LimitIterator { 20 iter := &LimitIterator{ 21 ctx: ctx, 22 source: source, 23 limit: limit, 24 maxSkip: maxSkip, 25 scoreThreshold: scoreThreshold, 26 skippedNodes: make([]*RankedNode, 0, maxSkip), 27 } 28 return iter 29 } 30 31 func (iter *LimitIterator) SetLimit(limit int) { 32 iter.limit = limit 33 } 34 35 func (iter *LimitIterator) Next() *RankedNode { 36 if iter.seen == iter.limit { 37 return nil 38 } 39 option := iter.nextOption() 40 if option == nil { 41 return nil 42 } 43 44 if len(iter.skippedNodes) < iter.maxSkip { 45 // Try skipping ahead up to maxSkip to find an option with score lesser than the threshold 46 for option != nil && option.Score <= iter.scoreThreshold && len(iter.skippedNodes) < iter.maxSkip { 47 iter.skippedNodes = append(iter.skippedNodes, option) 48 option = iter.source.Next() 49 } 50 } 51 iter.seen += 1 52 if option == nil { // Didn't find anything, so use the skipped nodes instead 53 return iter.nextOption() 54 } 55 return option 56 } 57 58 // nextOption uses the iterator's list of skipped nodes if the source iterator is exhausted 59 func (iter *LimitIterator) nextOption() *RankedNode { 60 sourceOption := iter.source.Next() 61 if sourceOption == nil && iter.skippedNodeIndex < len(iter.skippedNodes) { 62 skippedOption := iter.skippedNodes[iter.skippedNodeIndex] 63 iter.skippedNodeIndex += 1 64 return skippedOption 65 } 66 return sourceOption 67 } 68 69 func (iter *LimitIterator) Reset() { 70 iter.source.Reset() 71 iter.seen = 0 72 iter.skippedNodes = make([]*RankedNode, 0, iter.maxSkip) 73 iter.skippedNodeIndex = 0 74 } 75 76 // MaxScoreIterator is a RankIterator used to return only a single result 77 // of the item with the highest score. This iterator will consume all of the 78 // possible inputs and only returns the highest ranking result. 79 type MaxScoreIterator struct { 80 ctx Context 81 source RankIterator 82 max *RankedNode 83 } 84 85 // MaxScoreIterator returns a MaxScoreIterator over the given source 86 func NewMaxScoreIterator(ctx Context, source RankIterator) *MaxScoreIterator { 87 iter := &MaxScoreIterator{ 88 ctx: ctx, 89 source: source, 90 } 91 return iter 92 } 93 94 func (iter *MaxScoreIterator) Next() *RankedNode { 95 // Check if we've found the max, return nil 96 if iter.max != nil { 97 return nil 98 } 99 100 // Consume and determine the max 101 for { 102 option := iter.source.Next() 103 if option == nil { 104 return iter.max 105 } 106 107 if iter.max == nil || option.Score > iter.max.Score { 108 iter.max = option 109 } 110 } 111 } 112 113 func (iter *MaxScoreIterator) Reset() { 114 iter.source.Reset() 115 iter.max = nil 116 }