initial implementation

Author: rubywtl

pkg/api/api.go

@@ -3,6 +3,8 @@ package api
 import (
 	"context"
 	"flag"
+	"github.com/cortexproject/cortex/pkg/engine/distributed_execution"
+	"github.com/cortexproject/cortex/pkg/engine/distributed_execution/querierpb"
 	"net/http"
 	"path"
 	"strings"
@@ -480,6 +482,10 @@ func (a *API) RegisterQueryScheduler(f *scheduler.Scheduler) {
 	schedulerpb.RegisterSchedulerForQuerierServer(a.server.GRPC, f)
 }
 
+func (a *API) RegisterQuerierServer(q *distributed_execution.QuerierServer) {
+	querierpb.RegisterQuerierServer(a.server.GRPC, q)
+}
+
 // RegisterServiceMapHandler registers the Cortex structs service handler
 // TODO: Refactor this code to be accomplished using the services.ServiceManager
 // or a future module manager #2291

pkg/api/handlers.go

@@ -3,6 +3,7 @@ package api
 import (
 	"context"
 	"encoding/json"
+	"github.com/cortexproject/cortex/pkg/engine/distributed_execution"
 	"html/template"
 	"net/http"
 	"path"
@@ -164,6 +165,7 @@ func NewQuerierHandler(
 	queryable storage.SampleAndChunkQueryable,
 	exemplarQueryable storage.ExemplarQueryable,
 	engine engine.QueryEngine,
+	queryResultCache *distributed_execution.QueryResultCache,
 	metadataQuerier querier.MetadataQuerier,
 	reg prometheus.Registerer,
 	logger log.Logger,
@@ -280,7 +282,7 @@ func NewQuerierHandler(
 	legacyPromRouter := route.New().WithPrefix(path.Join(legacyPrefix, "/api/v1"))
 	api.Register(legacyPromRouter)
 
-	queryAPI := queryapi.NewQueryAPI(engine, translateSampleAndChunkQueryable, statsRenderer, logger, codecs, corsOrigin)
+	queryAPI := queryapi.NewQueryAPI(engine, queryResultCache, translateSampleAndChunkQueryable, statsRenderer, logger, codecs, corsOrigin)
 
 	// TODO(gotjosh): This custom handler is temporary until we're able to vendor the changes in:
 	// https://github.com/prometheus/prometheus/pull/7125/files

pkg/api/queryapi/query_api.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"github.com/cortexproject/cortex/pkg/engine/distributed_execution"
+	"github.com/cortexproject/cortex/pkg/scheduler/plan_fragments"
 	"net/http"
 	"strconv"
 	"time"
@@ -26,31 +27,34 @@ import (
 )
 
 type QueryAPI struct {
-	queryable     storage.SampleAndChunkQueryable
-	queryEngine   engine.QueryEngine
-	now           func() time.Time
-	statsRenderer v1.StatsRenderer
-	logger        log.Logger
-	codecs        []v1.Codec
-	CORSOrigin    *regexp.Regexp
+	queryable        storage.SampleAndChunkQueryable
+	queryEngine      engine.QueryEngine
+	queryResultCache *distributed_execution.QueryResultCache
+	now              func() time.Time
+	statsRenderer    v1.StatsRenderer
+	logger           log.Logger
+	codecs           []v1.Codec
+	CORSOrigin       *regexp.Regexp
 }
 
 func NewQueryAPI(
 	qe engine.QueryEngine,
+	queryResultCache *distributed_execution.QueryResultCache,
 	q storage.SampleAndChunkQueryable,
 	statsRenderer v1.StatsRenderer,
 	logger log.Logger,
 	codecs []v1.Codec,
 	CORSOrigin *regexp.Regexp,
 ) *QueryAPI {
 	return &QueryAPI{
-		queryEngine:   qe,
-		queryable:     q,
-		statsRenderer: statsRenderer,
-		logger:        logger,
-		codecs:        codecs,
-		CORSOrigin:    CORSOrigin,
-		now:           time.Now,
+		queryEngine:      qe,
+		queryResultCache: queryResultCache,
+		queryable:        q,
+		statsRenderer:    statsRenderer,
+		logger:           logger,
+		codecs:           codecs,
+		CORSOrigin:       CORSOrigin,
+		now:              time.Now,
 	}
 }
 
@@ -136,6 +140,13 @@ func (q *QueryAPI) RangeQueryHandler(r *http.Request) (result apiFuncResult) {
 
 	ctx = httputil.ContextFromRequest(ctx, r)
 
+	// TODO: if distributed exec enabled
+	isRoot, queryID, fragmentID, _ := plan_fragments.ExtractFragmentMetaData(ctx)
+	if !isRoot {
+		key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
+		q.queryResultCache.InitWriting(key)
+	}
+
 	res := qry.Exec(ctx)
 	if res.Err != nil {
 		return apiFuncResult{nil, returnAPIError(res.Err), res.Warnings, qry.Close}
@@ -178,17 +189,32 @@ func (q *QueryAPI) InstantQueryHandler(r *http.Request) (result apiFuncResult) {
 	ctx = engine.AddEngineTypeToContext(ctx, r)
 	ctx = querier.AddBlockStoreTypeToContext(ctx, r.Header.Get(querier.BlockStoreTypeHeader))
 
+	// TODO: if distributed exec enabled
+	isRoot, queryID, fragmentID, _ := plan_fragments.ExtractFragmentMetaData(ctx)
+	if !isRoot {
+		key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
+		q.queryResultCache.InitWriting(key)
+	}
+
 	var qry promql.Query
 	tsTime := convertMsToTime(ts)
 
 	byteLP := []byte(r.PostFormValue("plan"))
 	if len(byteLP) != 0 {
 		logicalPlan, err := distributed_execution.Unmarshal(byteLP)
 		if err != nil {
+			if !isRoot {
+				key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
+				q.queryResultCache.SetError(key)
+			}
 			return apiFuncResult{nil, &apiError{errorInternal, fmt.Errorf("invalid logical plan: %v", err)}, nil, nil}
 		}
 		qry, err = q.queryEngine.MakeInstantQueryFromPlan(ctx, q.queryable, opts, logicalPlan, tsTime, r.FormValue("query"))
 		if err != nil {
+			if !isRoot {
+				key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
+				q.queryResultCache.SetError(key)
+			}
 			return apiFuncResult{nil, &apiError{errorInternal, fmt.Errorf("failed to create instant query from logical plan: %v", err)}, nil, nil}
 		}
 	} else { // if there is logical plan field is empty, fall back
@@ -239,6 +265,18 @@ func (q *QueryAPI) Wrap(f apiFunc) http.HandlerFunc {
 		}
 
 		if result.data != nil {
+			// TODO: if distributed exec enabled
+			ctx := httputil.ContextFromRequest(r.Context(), r)
+			isRoot, queryID, fragmentID, _ := plan_fragments.ExtractFragmentMetaData(ctx)
+			if !isRoot {
+				key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
+				result := distributed_execution.FragmentResult{
+					Data:       result.data,
+					Expiration: time.Now().Add(time.Hour),
+				}
+				q.queryResultCache.SetComplete(key, result)
+				return
+			}
 			q.respond(w, r, result.data, result.warnings, r.FormValue("query"))
 			return
 		}

pkg/cortex/modules.go

@@ -4,13 +4,16 @@ import (
 	"context"
 	"flag"
 	"fmt"
-
+	"github.com/cortexproject/cortex/pkg/engine/distributed_execution"
+	"github.com/cortexproject/cortex/pkg/ring/client"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 	"log/slog"
 	"net"
 	"net/http"
 	"runtime"
 	"runtime/debug"
 	"strconv"
+	"time"
 
 	"github.com/go-kit/log/level"
 	"github.com/opentracing-contrib/go-stdlib/nethttp"
@@ -363,13 +366,35 @@ func (t *Cortex) initTenantFederation() (serv services.Service, err error) {
 //	                                          │                  │
 //	                                          └──────────────────┘
 func (t *Cortex) initQuerier() (serv services.Service, err error) {
+	// Create new map for caching partial results during distributed execution
+	var queryResultCache *distributed_execution.QueryResultCache
+	var queryServer *distributed_execution.QuerierServer
+
+	if t.Cfg.Querier.DistributedExecEnabled {
+		// set up querier server service and register it
+		queryResultCache = distributed_execution.NewQueryResultCache()
+		queryServer = distributed_execution.NewQuerierServer(queryResultCache)
+
+		go func() {
+			// TODO: make expire time a config var
+			ticker := time.NewTicker(5 * time.Minute)
+			defer ticker.Stop()
+
+			for range ticker.C {
+				queryResultCache.CleanExpired()
+			}
+		}()
+		t.API.RegisterQuerierServer(queryServer)
+	}
+
 	// Create a internal HTTP handler that is configured with the Prometheus API routes and points
 	// to a Prometheus API struct instantiated with the Cortex Queryable.
 	internalQuerierRouter := api.NewQuerierHandler(
 		t.Cfg.API,
 		t.QuerierQueryable,
 		t.ExemplarQueryable,
 		t.QuerierEngine,
+		queryResultCache,
 		t.MetadataQuerier,
 		prometheus.DefaultRegisterer,
 		util_log.Logger,
@@ -420,7 +445,33 @@ func (t *Cortex) initQuerier() (serv services.Service, err error) {
 		return nil, err
 	}
 	serverAddress := net.JoinHostPort(ipAddr, strconv.Itoa(t.Cfg.Server.HTTPListenPort))
-	return querier_worker.NewQuerierWorker(t.Cfg.Worker, httpgrpc_server.NewServer(internalQuerierRouter), util_log.Logger, prometheus.DefaultRegisterer, serverAddress)
+
+	if t.Cfg.Querier.DistributedExecEnabled {
+		poolConfig := client.PoolConfig{
+			CheckInterval:      5 * time.Second,
+			HealthCheckEnabled: true,
+			HealthCheckTimeout: 1 * time.Second,
+		}
+
+		querierClientsGauge := promauto.With(prometheus.DefaultRegisterer).NewGauge(prometheus.GaugeOpts{
+			Name: "cortex_querier_service_clients",
+			Help: "The current number of clients connected to this querier",
+		})
+		factory := func(addr string) (client.PoolClient, error) {
+			return distributed_execution.CreateQuerierClient(t.Cfg.QueryScheduler.GRPCClientConfig, addr)
+		}
+		querierPool := client.NewPool("querier", poolConfig, nil, factory, querierClientsGauge, util_log.Logger)
+		internalQuerierRouter = injectPool(internalQuerierRouter, querierPool)
+	}
+
+	return querier_worker.NewQuerierWorker(t.Cfg.Worker, httpgrpc_server.NewServer(internalQuerierRouter), util_log.Logger, prometheus.DefaultRegisterer, serverAddress, t.Cfg.Querier.DistributedExecEnabled, queryResultCache)
+}
+
+func injectPool(next http.Handler, pool *client.Pool) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ctx := distributed_execution.ContextWithPool(r.Context(), pool)
+		next.ServeHTTP(w, r.WithContext(ctx))
+	})
 }
 
 func (t *Cortex) initStoreQueryables() (services.Service, error) {

pkg/engine/distributed_execution/querier_service_client.go

@@ -0,0 +1,82 @@
+package distributed_execution
+
+import (
+	"github.com/cortexproject/cortex/pkg/engine/distributed_execution/querierpb"
+	"github.com/cortexproject/cortex/pkg/ring/client"
+	"github.com/cortexproject/cortex/pkg/util/grpcclient"
+	otgrpc "github.com/opentracing-contrib/go-grpc"
+	"github.com/opentracing/opentracing-go"
+	"github.com/prometheus/prometheus/model/histogram"
+	"github.com/weaveworks/common/middleware"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/health/grpc_health_v1"
+)
+
+func CreateQuerierClient(grpcConfig grpcclient.Config, addr string) (client.PoolClient, error) {
+	opts, err := grpcConfig.DialOption([]grpc.UnaryClientInterceptor{
+		otgrpc.OpenTracingClientInterceptor(opentracing.GlobalTracer()),
+		middleware.ClientUserHeaderInterceptor,
+	}, nil)
+
+	if err != nil {
+		return nil, err
+	}
+
+	conn, err := grpc.NewClient(addr, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	return &querierClient{
+		QuerierClient: querierpb.NewQuerierClient(conn),
+		HealthClient:  grpc_health_v1.NewHealthClient(conn),
+		conn:          conn,
+	}, nil
+}
+
+type querierClient struct {
+	querierpb.QuerierClient
+	grpc_health_v1.HealthClient
+	conn *grpc.ClientConn
+}
+
+func (qc *querierClient) Close() error {
+	return qc.conn.Close()
+}
+
+func FloatHistogramProtoToFloatHistograms(hps []querierpb.Histogram) []*histogram.FloatHistogram {
+	floatHistograms := make([]*histogram.FloatHistogram, len(hps))
+	for _, hp := range hps {
+		newHist := FloatHistogramProtoToFloatHistogram(hp)
+		floatHistograms = append(floatHistograms, newHist)
+	}
+	return floatHistograms
+}
+
+func FloatHistogramProtoToFloatHistogram(hp querierpb.Histogram) *histogram.FloatHistogram {
+	_, IsFloatHist := hp.GetCount().(*querierpb.Histogram_CountFloat)
+	if !IsFloatHist {
+		panic("FloatHistogramProtoToFloatHistogram called with an integer histogram")
+	}
+	return &histogram.FloatHistogram{
+		CounterResetHint: histogram.CounterResetHint(hp.ResetHint),
+		Schema:           hp.Schema,
+		ZeroThreshold:    hp.ZeroThreshold,
+		ZeroCount:        hp.GetZeroCountFloat(),
+		Count:            hp.GetCountFloat(),
+		Sum:              hp.Sum,
+		PositiveSpans:    spansProtoToSpans(hp.GetPositiveSpans()),
+		PositiveBuckets:  hp.GetPositiveCounts(),
+		NegativeSpans:    spansProtoToSpans(hp.GetNegativeSpans()),
+		NegativeBuckets:  hp.GetNegativeCounts(),
+	}
+}
+
+func spansProtoToSpans(s []querierpb.BucketSpan) []histogram.Span {
+	spans := make([]histogram.Span, len(s))
+	for i := 0; i < len(s); i++ {
+		spans[i] = histogram.Span{Offset: s[i].Offset, Length: s[i].Length}
+	}
+
+	return spans
+}