able to pick-up fragments and executed by individual queriers

Author: rubywtl

pkg/api/queryapi/query_api.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"github.com/cortexproject/cortex/pkg/engine/distributed_execution"
-	"github.com/cortexproject/cortex/pkg/scheduler/plan_fragments"
 	"net/http"
 	"strconv"
 	"time"
@@ -141,10 +140,10 @@ 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)
+	isRoot, queryID, fragmentID, _, _ := distributed_execution.ExtractFragmentMetaData(ctx)
 	if !isRoot {
 		key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
-		q.queryResultCache.InitWriting(key)
+		q.queryResultCache.InitWriting(*key)
 	}
 
 	res := qry.Exec(ctx)
@@ -190,10 +189,10 @@ func (q *QueryAPI) InstantQueryHandler(r *http.Request) (result apiFuncResult) {
 	ctx = querier.AddBlockStoreTypeToContext(ctx, r.Header.Get(querier.BlockStoreTypeHeader))
 
 	// TODO: if distributed exec enabled
-	isRoot, queryID, fragmentID, _ := plan_fragments.ExtractFragmentMetaData(ctx)
+	isRoot, queryID, fragmentID, _, _ := distributed_execution.ExtractFragmentMetaData(ctx)
 	if !isRoot {
 		key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
-		q.queryResultCache.InitWriting(key)
+		q.queryResultCache.InitWriting(*key)
 	}
 
 	var qry promql.Query
@@ -205,15 +204,15 @@ func (q *QueryAPI) InstantQueryHandler(r *http.Request) (result apiFuncResult) {
 		if err != nil {
 			if !isRoot {
 				key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
-				q.queryResultCache.SetError(key)
+				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)
+				q.queryResultCache.SetError(*key)
 			}
 			return apiFuncResult{nil, &apiError{errorInternal, fmt.Errorf("failed to create instant query from logical plan: %v", err)}, nil, nil}
 		}
@@ -267,14 +266,14 @@ 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)
+			isRoot, queryID, fragmentID, _, _ := distributed_execution.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)
+				q.queryResultCache.SetComplete(*key, result)
 				return
 			}
 			q.respond(w, r, result.data, result.warnings, r.FormValue("query"))

pkg/cortex/modules.go

@@ -6,13 +6,13 @@ import (
 	"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"
+	"strings"
 	"time"
 
 	"github.com/go-kit/log/level"
@@ -366,6 +366,15 @@ func (t *Cortex) initTenantFederation() (serv services.Service, err error) {
 //	                                          │                  │
 //	                                          └──────────────────┘
 func (t *Cortex) initQuerier() (serv services.Service, err error) {
+
+	t.Cfg.Worker.MaxConcurrentRequests = t.Cfg.Querier.MaxConcurrent
+	t.Cfg.Worker.TargetHeaders = t.Cfg.API.HTTPRequestHeadersToLog
+	ipAddr, err := ring.GetInstanceAddr(t.Cfg.Alertmanager.ShardingRing.InstanceAddr, t.Cfg.Alertmanager.ShardingRing.InstanceInterfaceNames, util_log.Logger)
+	if err != nil {
+		return nil, err
+	}
+	serverAddress := net.JoinHostPort(ipAddr, strconv.Itoa(t.Cfg.Server.GRPCListenPort))
+
 	// Create new map for caching partial results during distributed execution
 	var queryResultCache *distributed_execution.QueryResultCache
 	var queryServer *distributed_execution.QuerierServer
@@ -438,37 +447,22 @@ func (t *Cortex) initQuerier() (serv services.Service, err error) {
 		return nil, nil
 	}
 
-	t.Cfg.Worker.MaxConcurrentRequests = t.Cfg.Querier.MaxConcurrent
-	t.Cfg.Worker.TargetHeaders = t.Cfg.API.HTTPRequestHeadersToLog
-	ipAddr, err := ring.GetInstanceAddr(t.Cfg.Alertmanager.ShardingRing.InstanceAddr, t.Cfg.Alertmanager.ShardingRing.InstanceInterfaceNames, util_log.Logger)
-	if err != nil {
-		return nil, err
-	}
-	serverAddress := net.JoinHostPort(ipAddr, strconv.Itoa(t.Cfg.Server.HTTPListenPort))
-
 	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)
+		querierPool := distributed_execution.NewQuerierPool(t.Cfg.QueryScheduler.GRPCClientConfig, prometheus.DefaultRegisterer, util_log.Logger)
 		internalQuerierRouter = injectPool(internalQuerierRouter, querierPool)
+		//go watchQuerierRingAndUpdatePool(context.Background(), t.Ring, 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) {
+		if r.ProtoMajor == 2 && strings.Contains(r.Header.Get("Content-Type"), "application/grpc") {
+			ctx := distributed_execution.ContextWithPool(r.Context(), pool)
+			next.ServeHTTP(w, r.WithContext(ctx))
+			return
+		}
 		ctx := distributed_execution.ContextWithPool(r.Context(), pool)
 		next.ServeHTTP(w, r.WithContext(ctx))
 	})

pkg/engine/distributed_execution/id.go

@@ -0,0 +1,36 @@
+package distributed_execution
+
+import (
+	"context"
+)
+
+type fragmentMetadataKey struct{}
+
+type fragmentMetadata struct {
+	queryID       uint64
+	fragmentID    uint64
+	childIDToAddr map[uint64]string
+	isRoot        bool
+}
+
+func InjectFragmentMetaData(ctx context.Context, fragmentID uint64, queryID uint64, isRoot bool, childIDs []uint64, childAddr []string) context.Context {
+	childIDToAddr := make(map[uint64]string)
+	for i, childID := range childIDs {
+		childIDToAddr[childID] = childAddr[i]
+	}
+
+	return context.WithValue(ctx, fragmentMetadataKey{}, fragmentMetadata{
+		queryID:       queryID,
+		fragmentID:    fragmentID,
+		childIDToAddr: childIDToAddr,
+		isRoot:        isRoot,
+	})
+}
+
+func ExtractFragmentMetaData(ctx context.Context) (isRoot bool, queryID uint64, fragmentID uint64, childAddrs map[uint64]string, ok bool) {
+	metadata, ok := ctx.Value(fragmentMetadataKey{}).(fragmentMetadata)
+	if !ok {
+		return false, 0, 0, nil, false
+	}
+	return metadata.isRoot, metadata.queryID, metadata.fragmentID, metadata.childIDToAddr, true
+}

pkg/scheduler/plan_fragments/id_test.go …/engine/distributed_execution/id_test.gopkg/scheduler/plan_fragments/id_test.go renamed to pkg/engine/distributed_execution/id_test.go pkg/scheduler/plan_fragments/id_test.go renamed to pkg/engine/distributed_execution/id_test.go

@@ -1,7 +1,8 @@
-package plan_fragments
+package distributed_execution
 
 import (
 	"context"
+	"github.com/cortexproject/cortex/pkg/scheduler/plan_fragments"
 	"github.com/stretchr/testify/require"
 	"testing"
 )
@@ -13,13 +14,13 @@ func TestFragmentMetadata(t *testing.T) {
 	t.Run("basic injection and extraction", func(t *testing.T) {
 
 		ctx := context.Background()
-		fragment := Fragment{
+		fragment := plan_fragments.Fragment{
 			FragmentID: 123,
 			IsRoot:     true,
 		}
 		queryID := uint64(456)
 
-		ctx = InjectFragmentMetaData(ctx, fragment, queryID)
+		ctx = InjectFragmentMetaData(ctx, fragment.FragmentID, queryID, fragment.IsRoot, fragment.ChildIDs)
 
 		isRoot, qID, fID, ok := ExtractFragmentMetaData(ctx)
 		require.True(t, ok)
@@ -40,8 +41,8 @@ func TestFragmentMetadata(t *testing.T) {
 
 	t.Run("zero values", func(t *testing.T) {
 		ctx := context.Background()
-		fragment := Fragment{}
-		ctx = InjectFragmentMetaData(ctx, fragment, 0)
+		fragment := plan_fragments.Fragment{}
+		ctx = InjectFragmentMetaData(ctx, fragment.FragmentID, 0, fragment.IsRoot, fragment.ChildIDs)
 
 		isRoot, queryID, fragmentID, ok := ExtractFragmentMetaData(ctx)
 		require.True(t, ok)

pkg/engine/distributed_execution/logicalplan.go

@@ -3,10 +3,9 @@ package distributed_execution
 import (
 	"bytes"
 	"encoding/json"
-	"math"
-
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/thanos-io/promql-engine/logicalplan"
+	"math"
 )
 
 type jsonNode struct {
@@ -21,6 +20,46 @@ const (
 	negInfVal = `"-Inf"`
 )
 
+func Marshal(node Node) ([]byte, error) {
+	clone := node.Clone()
+	return marshalNode(clone)
+}
+
+func marshalNode(node Node) ([]byte, error) {
+	children := make([]json.RawMessage, 0, len(node.Children()))
+	for _, c := range node.Children() {
+		childData, err := marshalNode(*c)
+		if err != nil {
+			return nil, err
+		}
+		children = append(children, childData)
+	}
+	var data json.RawMessage = nil
+	if n, ok := node.(*logicalplan.NumberLiteral); ok {
+		if math.IsInf(n.Val, 1) {
+			data = json.RawMessage(infVal)
+		}
+		if math.IsInf(n.Val, -1) {
+			data = json.RawMessage(negInfVal)
+		}
+		if math.IsNaN(n.Val) {
+			data = json.RawMessage(nanVal)
+		}
+	}
+	if data == nil {
+		var err error
+		data, err = json.Marshal(node)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return json.Marshal(jsonNode{
+		Type:     node.Type(),
+		Data:     data,
+		Children: children,
+	})
+}
+
 func Unmarshal(data []byte) (logicalplan.Node, error) {
 	return unmarshalNode(data)
 }

pkg/engine/distributed_execution/logicalplan_test.go

@@ -0,0 +1 @@
+package distributed_execution

pkg/engine/distributed_execution/querier_service_client.go

@@ -4,12 +4,17 @@ import (
 	"github.com/cortexproject/cortex/pkg/engine/distributed_execution/querierpb"
 	"github.com/cortexproject/cortex/pkg/ring/client"
 	"github.com/cortexproject/cortex/pkg/util/grpcclient"
+	cortexmiddleware "github.com/cortexproject/cortex/pkg/util/middleware"
+	"github.com/go-kit/log"
 	otgrpc "github.com/opentracing-contrib/go-grpc"
 	"github.com/opentracing/opentracing-go"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/weaveworks/common/middleware"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/health/grpc_health_v1"
+	"time"
 )
 
 func CreateQuerierClient(grpcConfig grpcclient.Config, addr string) (client.PoolClient, error) {
@@ -44,6 +49,60 @@ func (qc *querierClient) Close() error {
 	return qc.conn.Close()
 }
 
+func NewQuerierPool(cfg grpcclient.Config, reg prometheus.Registerer, log log.Logger) *client.Pool {
+	requestDuration := promauto.With(reg).NewHistogramVec(prometheus.HistogramOpts{
+		Name:    "cortex_querier_query_request_duration_seconds",
+		Help:    "Time spent doing requests to querier.",
+		Buckets: prometheus.ExponentialBuckets(0.001, 4, 6),
+	}, []string{"operation", "status_code"})
+
+	clientsGauge := promauto.With(reg).NewGauge(prometheus.GaugeOpts{
+		Name: "cortex_querier_query_clients",
+		Help: "The current number of clients connected to querier.",
+	})
+
+	poolConfig := client.PoolConfig{
+		CheckInterval:      time.Minute,
+		HealthCheckEnabled: true,
+		HealthCheckTimeout: 10 * time.Second,
+	}
+
+	q := &querierPool{
+		grpcConfig:      cfg,
+		requestDuration: requestDuration,
+	}
+
+	return client.NewPool("querier", poolConfig, nil, q.createQuerierClient, clientsGauge, log)
+}
+
+type querierPool struct {
+	grpcConfig      grpcclient.Config
+	requestDuration *prometheus.HistogramVec
+}
+
+func (q *querierPool) createQuerierClient(addr string) (client.PoolClient, error) {
+	opts, err := q.grpcConfig.DialOption([]grpc.UnaryClientInterceptor{
+		otgrpc.OpenTracingClientInterceptor(opentracing.GlobalTracer()),
+		middleware.ClientUserHeaderInterceptor,
+		cortexmiddleware.PrometheusGRPCUnaryInstrumentation(q.requestDuration),
+	}, 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
+}
+
 func FloatHistogramProtoToFloatHistograms(hps []querierpb.Histogram) []*histogram.FloatHistogram {
 	floatHistograms := make([]*histogram.FloatHistogram, len(hps))
 	for _, hp := range hps {