github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/querier/batch/batch.go (about)

     1  package batch
     2  
     3  import (
     4  	"github.com/prometheus/common/model"
     5  	"github.com/prometheus/prometheus/tsdb/chunkenc"
     6  
     7  	"github.com/cortexproject/cortex/pkg/chunk"
     8  	"github.com/cortexproject/cortex/pkg/chunk/encoding"
     9  	promchunk "github.com/cortexproject/cortex/pkg/chunk/encoding"
    10  )
    11  
    12  // GenericChunk is a generic chunk used by the batch iterator, in order to make the batch
    13  // iterator general purpose.
    14  type GenericChunk struct {
    15  	MinTime int64
    16  	MaxTime int64
    17  
    18  	iterator func(reuse encoding.Iterator) encoding.Iterator
    19  }
    20  
    21  func NewGenericChunk(minTime, maxTime int64, iterator func(reuse encoding.Iterator) encoding.Iterator) GenericChunk {
    22  	return GenericChunk{
    23  		MinTime:  minTime,
    24  		MaxTime:  maxTime,
    25  		iterator: iterator,
    26  	}
    27  }
    28  
    29  func (c GenericChunk) Iterator(reuse encoding.Iterator) encoding.Iterator {
    30  	return c.iterator(reuse)
    31  }
    32  
    33  // iterator iterates over batches.
    34  type iterator interface {
    35  	// Seek to the batch at (or after) time t.
    36  	Seek(t int64, size int) bool
    37  
    38  	// Next moves to the next batch.
    39  	Next(size int) bool
    40  
    41  	// AtTime returns the start time of the next batch.  Must only be called after
    42  	// Seek or Next have returned true.
    43  	AtTime() int64
    44  
    45  	// Batch returns the current batch.  Must only be called after Seek or Next
    46  	// have returned true.
    47  	Batch() promchunk.Batch
    48  
    49  	Err() error
    50  }
    51  
    52  // NewChunkMergeIterator returns a chunkenc.Iterator that merges Cortex chunks together.
    53  func NewChunkMergeIterator(chunks []chunk.Chunk, _, _ model.Time) chunkenc.Iterator {
    54  	converted := make([]GenericChunk, len(chunks))
    55  	for i, c := range chunks {
    56  		converted[i] = NewGenericChunk(int64(c.From), int64(c.Through), c.Data.NewIterator)
    57  	}
    58  
    59  	return NewGenericChunkMergeIterator(converted)
    60  }
    61  
    62  // NewGenericChunkMergeIterator returns a chunkenc.Iterator that merges generic chunks together.
    63  func NewGenericChunkMergeIterator(chunks []GenericChunk) chunkenc.Iterator {
    64  	iter := newMergeIterator(chunks)
    65  	return newIteratorAdapter(iter)
    66  }
    67  
    68  // iteratorAdapter turns a batchIterator into a chunkenc.Iterator.
    69  // It fetches ever increasing batchSizes (up to promchunk.BatchSize) on each
    70  // call to Next; on calls to Seek, resets batch size to 1.
    71  type iteratorAdapter struct {
    72  	batchSize  int
    73  	curr       promchunk.Batch
    74  	underlying iterator
    75  }
    76  
    77  func newIteratorAdapter(underlying iterator) chunkenc.Iterator {
    78  	return &iteratorAdapter{
    79  		batchSize:  1,
    80  		underlying: underlying,
    81  	}
    82  }
    83  
    84  // Seek implements chunkenc.Iterator.
    85  func (a *iteratorAdapter) Seek(t int64) bool {
    86  
    87  	// Optimisation: fulfill the seek using current batch if possible.
    88  	if a.curr.Length > 0 && a.curr.Index < a.curr.Length {
    89  		if t <= a.curr.Timestamps[a.curr.Index] {
    90  			//In this case, the interface's requirement is met, so state of this
    91  			//iterator does not need any change.
    92  			return true
    93  		} else if t <= a.curr.Timestamps[a.curr.Length-1] {
    94  			//In this case, some timestamp between current sample and end of batch can fulfill
    95  			//the seek. Let's find it.
    96  			for a.curr.Index < a.curr.Length && t > a.curr.Timestamps[a.curr.Index] {
    97  				a.curr.Index++
    98  			}
    99  			return true
   100  		}
   101  	}
   102  
   103  	a.curr.Length = -1
   104  	a.batchSize = 1
   105  	if a.underlying.Seek(t, a.batchSize) {
   106  		a.curr = a.underlying.Batch()
   107  		return a.curr.Index < a.curr.Length
   108  	}
   109  	return false
   110  }
   111  
   112  // Next implements chunkenc.Iterator.
   113  func (a *iteratorAdapter) Next() bool {
   114  	a.curr.Index++
   115  	for a.curr.Index >= a.curr.Length && a.underlying.Next(a.batchSize) {
   116  		a.curr = a.underlying.Batch()
   117  		a.batchSize = a.batchSize * 2
   118  		if a.batchSize > promchunk.BatchSize {
   119  			a.batchSize = promchunk.BatchSize
   120  		}
   121  	}
   122  	return a.curr.Index < a.curr.Length
   123  }
   124  
   125  // At implements chunkenc.Iterator.
   126  func (a *iteratorAdapter) At() (int64, float64) {
   127  	return a.curr.Timestamps[a.curr.Index], a.curr.Values[a.curr.Index]
   128  }
   129  
   130  // Err implements chunkenc.Iterator.
   131  func (a *iteratorAdapter) Err() error {
   132  	return nil
   133  }