github.com/kardianos/nomad@v0.1.3-0.20151022182107-b13df73ee850/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 artifically 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  }