github.com/thanos-io/thanos@v0.32.5/pkg/queryfrontend/shard_query.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  // This is a modified copy from
     5  // https://github.com/cortexproject/cortex/blob/master/pkg/querier/queryrange/split_by_interval.go.
     6  
     7  package queryfrontend
     8  
     9  import (
    10  	"context"
    11  
    12  	"github.com/prometheus/client_golang/prometheus"
    13  	"github.com/prometheus/client_golang/prometheus/promauto"
    14  
    15  	"github.com/thanos-io/thanos/internal/cortex/querier/queryrange"
    16  
    17  	"github.com/thanos-io/thanos/pkg/querysharding"
    18  	"github.com/thanos-io/thanos/pkg/store/storepb"
    19  )
    20  
    21  // PromQLShardingMiddleware creates a new Middleware that shards PromQL aggregations using grouping labels.
    22  func PromQLShardingMiddleware(queryAnalyzer querysharding.Analyzer, numShards int, limits queryrange.Limits, merger queryrange.Merger, registerer prometheus.Registerer) queryrange.Middleware {
    23  	return queryrange.MiddlewareFunc(func(next queryrange.Handler) queryrange.Handler {
    24  		queriesTotal := promauto.With(registerer).NewCounterVec(prometheus.CounterOpts{
    25  			Namespace: "thanos",
    26  			Name:      "frontend_sharding_middleware_queries_total",
    27  			Help:      "Total number of queries analyzed by the sharding middleware",
    28  		}, []string{"shardable"})
    29  
    30  		queriesTotal.WithLabelValues("true")
    31  		queriesTotal.WithLabelValues("false")
    32  
    33  		return querySharder{
    34  			next:          next,
    35  			limits:        limits,
    36  			queryAnalyzer: queryAnalyzer,
    37  			numShards:     numShards,
    38  			merger:        merger,
    39  			queriesTotal:  queriesTotal,
    40  		}
    41  	})
    42  }
    43  
    44  type querySharder struct {
    45  	next   queryrange.Handler
    46  	limits queryrange.Limits
    47  
    48  	queryAnalyzer querysharding.Analyzer
    49  	numShards     int
    50  	merger        queryrange.Merger
    51  
    52  	// Metrics
    53  	queriesTotal *prometheus.CounterVec
    54  }
    55  
    56  func (s querySharder) Do(ctx context.Context, r queryrange.Request) (queryrange.Response, error) {
    57  	analysis, err := s.queryAnalyzer.Analyze(r.GetQuery())
    58  
    59  	if err != nil || !analysis.IsShardable() {
    60  		s.queriesTotal.WithLabelValues("false").Inc()
    61  		return s.next.Do(ctx, r)
    62  	}
    63  
    64  	s.queriesTotal.WithLabelValues("true").Inc()
    65  	reqs := s.shardQuery(r, analysis)
    66  
    67  	reqResps, err := queryrange.DoRequests(ctx, s.next, reqs, s.limits)
    68  	if err != nil {
    69  		return nil, err
    70  	}
    71  
    72  	resps := make([]queryrange.Response, 0, len(reqResps))
    73  	for _, reqResp := range reqResps {
    74  		resps = append(resps, reqResp.Response)
    75  	}
    76  
    77  	response, err := s.merger.MergeResponse(r, resps...)
    78  	if err != nil {
    79  		return nil, err
    80  	}
    81  	return response, nil
    82  }
    83  
    84  func (s querySharder) shardQuery(r queryrange.Request, analysis querysharding.QueryAnalysis) []queryrange.Request {
    85  	tr, ok := r.(ShardedRequest)
    86  	if !ok {
    87  		return []queryrange.Request{r}
    88  	}
    89  
    90  	reqs := make([]queryrange.Request, s.numShards)
    91  	for i := 0; i < s.numShards; i++ {
    92  		reqs[i] = tr.WithShardInfo(&storepb.ShardInfo{
    93  			TotalShards: int64(s.numShards),
    94  			ShardIndex:  int64(i),
    95  			By:          analysis.ShardBy(),
    96  			Labels:      analysis.ShardingLabels(),
    97  		})
    98  	}
    99  
   100  	return reqs
   101  }