github.com/sequix/cortex@v1.1.6/pkg/chunk/util/util.go (about)

     1  package util
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  
     7  	ot "github.com/opentracing/opentracing-go"
     8  
     9  	"github.com/sequix/cortex/pkg/chunk"
    10  	"github.com/sequix/cortex/pkg/util"
    11  )
    12  
    13  // DoSingleQuery is the interface for indexes that don't support batching yet.
    14  type DoSingleQuery func(
    15  	ctx context.Context, query chunk.IndexQuery,
    16  	callback func(chunk.ReadBatch) bool,
    17  ) error
    18  
    19  // QueryParallelism is the maximum number of subqueries run in
    20  // parallel per higher-level query
    21  var QueryParallelism = 100
    22  
    23  // DoParallelQueries translates between our interface for query batching,
    24  // and indexes that don't yet support batching.
    25  func DoParallelQueries(
    26  	ctx context.Context, doSingleQuery DoSingleQuery, queries []chunk.IndexQuery,
    27  	callback func(chunk.IndexQuery, chunk.ReadBatch) bool,
    28  ) error {
    29  	queue := make(chan chunk.IndexQuery)
    30  	incomingErrors := make(chan error)
    31  	n := util.Min(len(queries), QueryParallelism)
    32  	// Run n parallel goroutines fetching queries from the queue
    33  	for i := 0; i < n; i++ {
    34  		go func() {
    35  			sp, ctx := ot.StartSpanFromContext(ctx, "DoParallelQueries-worker")
    36  			defer sp.Finish()
    37  			for {
    38  				query, ok := <-queue
    39  				if !ok {
    40  					return
    41  				}
    42  				incomingErrors <- doSingleQuery(ctx, query, func(r chunk.ReadBatch) bool {
    43  					return callback(query, r)
    44  				})
    45  			}
    46  		}()
    47  	}
    48  	// Send all the queries into the queue
    49  	go func() {
    50  		for _, query := range queries {
    51  			queue <- query
    52  		}
    53  		close(queue)
    54  	}()
    55  
    56  	// Now receive all the results.
    57  	var lastErr error
    58  	for i := 0; i < len(queries); i++ {
    59  		err := <-incomingErrors
    60  		if err != nil {
    61  
    62  			lastErr = err
    63  		}
    64  	}
    65  	return lastErr
    66  }
    67  
    68  // Callback from an IndexQuery.
    69  type Callback func(chunk.IndexQuery, chunk.ReadBatch) bool
    70  
    71  type filteringBatch struct {
    72  	query chunk.IndexQuery
    73  	chunk.ReadBatch
    74  }
    75  
    76  func (f filteringBatch) Iterator() chunk.ReadBatchIterator {
    77  	return &filteringBatchIter{
    78  		query:             f.query,
    79  		ReadBatchIterator: f.ReadBatch.Iterator(),
    80  	}
    81  }
    82  
    83  type filteringBatchIter struct {
    84  	query chunk.IndexQuery
    85  	chunk.ReadBatchIterator
    86  }
    87  
    88  func (f *filteringBatchIter) Next() bool {
    89  	for f.ReadBatchIterator.Next() {
    90  		rangeValue, value := f.ReadBatchIterator.RangeValue(), f.ReadBatchIterator.Value()
    91  
    92  		if len(f.query.RangeValuePrefix) != 0 && !bytes.HasPrefix(rangeValue, f.query.RangeValuePrefix) {
    93  			continue
    94  		}
    95  		if len(f.query.RangeValueStart) != 0 && bytes.Compare(f.query.RangeValueStart, rangeValue) > 0 {
    96  			continue
    97  		}
    98  		if len(f.query.ValueEqual) != 0 && !bytes.Equal(value, f.query.ValueEqual) {
    99  			continue
   100  		}
   101  
   102  		return true
   103  	}
   104  
   105  	return false
   106  }
   107  
   108  // QueryFilter wraps a callback to ensure the results are filtered correctly;
   109  // useful for the cache and Bigtable backend, which only ever fetches the whole
   110  // row.
   111  func QueryFilter(callback Callback) Callback {
   112  	return func(query chunk.IndexQuery, batch chunk.ReadBatch) bool {
   113  		return callback(query, &filteringBatch{query, batch})
   114  	}
   115  }