github.com/jrxfive/nomad@v0.6.1-0.20170802162750-1fef470e89bf/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 seen int 10 } 11 12 // NewLimitIterator is returns a LimitIterator with a fixed limit of returned options 13 func NewLimitIterator(ctx Context, source RankIterator, limit int) *LimitIterator { 14 iter := &LimitIterator{ 15 ctx: ctx, 16 source: source, 17 limit: limit, 18 } 19 return iter 20 } 21 22 func (iter *LimitIterator) SetLimit(limit int) { 23 iter.limit = limit 24 } 25 26 func (iter *LimitIterator) Next() *RankedNode { 27 if iter.seen == iter.limit { 28 return nil 29 } 30 31 option := iter.source.Next() 32 if option == nil { 33 return nil 34 } 35 36 iter.seen += 1 37 return option 38 } 39 40 func (iter *LimitIterator) Reset() { 41 iter.source.Reset() 42 iter.seen = 0 43 } 44 45 // MaxScoreIterator is a RankIterator used to return only a single result 46 // of the item with the highest score. This iterator will consume all of the 47 // possible inputs and only returns the highest ranking result. 48 type MaxScoreIterator struct { 49 ctx Context 50 source RankIterator 51 max *RankedNode 52 } 53 54 // MaxScoreIterator returns a MaxScoreIterator over the given source 55 func NewMaxScoreIterator(ctx Context, source RankIterator) *MaxScoreIterator { 56 iter := &MaxScoreIterator{ 57 ctx: ctx, 58 source: source, 59 } 60 return iter 61 } 62 63 func (iter *MaxScoreIterator) Next() *RankedNode { 64 // Check if we've found the max, return nil 65 if iter.max != nil { 66 return nil 67 } 68 69 // Consume and determine the max 70 for { 71 option := iter.source.Next() 72 if option == nil { 73 return iter.max 74 } 75 76 if iter.max == nil || option.Score > iter.max.Score { 77 iter.max = option 78 } 79 } 80 } 81 82 func (iter *MaxScoreIterator) Reset() { 83 iter.source.Reset() 84 iter.max = nil 85 }