github.com/uchennaokeke444/nomad@v0.11.8/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.FinalScore <= 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.FinalScore > iter.max.FinalScore {
   108  			iter.max = option
   109  		}
   110  	}
   111  }
   112  
   113  func (iter *MaxScoreIterator) Reset() {
   114  	iter.source.Reset()
   115  	iter.max = nil
   116  }