remove unneccessary configs in query frontend + restructure logical-plan-gen middleware helper functions

Signed-off-by: rubywtl <[email protected]>

Author: rubywtl

docs/configuration/config-file-reference.md

@@ -4451,6 +4451,11 @@ grpc_client_config:
 # URL of downstream Prometheus.
 # CLI flag: -frontend.downstream-url
 [downstream_url: <string> | default = ""]
+
+# Experimental: Enables distributed execution of queries by passing logical
+# query plan fragments to downstream components.
+# CLI flag: -frontend.distributed_exec_enabled
+[distributed_exec_enabled: <boolean> | default = false]
 ```
 
 ### `query_range_config`

pkg/cortex/modules.go

@@ -533,16 +533,13 @@ func (t *Cortex) initQueryFrontendTripperware() (serv services.Service, err erro
 		prometheusCodec,
 		shardedPrometheusCodec,
 		t.Cfg.Querier.LookbackDelta,
-		t.Cfg.Querier.EnablePerStepStats,
 		t.Cfg.Frontend.DistributedExecEnabled,
-		t.Cfg.Frontend.DisableDuplicateLabelChecks,
 	)
 	if err != nil {
 		return nil, err
 	}
 
-	instantQueryMiddlewares, err := instantquery.Middlewares(util_log.Logger, t.Overrides, instantQueryCodec, queryAnalyzer, t.Cfg.Querier.LookbackDelta, t.Cfg.Querier.EnablePerStepStats,
-		t.Cfg.Frontend.DistributedExecEnabled, t.Cfg.Frontend.DisableDuplicateLabelChecks)
+	instantQueryMiddlewares, err := instantquery.Middlewares(util_log.Logger, t.Overrides, instantQueryCodec, queryAnalyzer, t.Cfg.Querier.LookbackDelta, t.Cfg.Frontend.DistributedExecEnabled)
 	if err != nil {
 		return nil, err
 	}

pkg/frontend/config.go

@@ -20,9 +20,8 @@ type CombinedFrontendConfig struct {
 	FrontendV1 v1.Config               `yaml:",inline"`
 	FrontendV2 v2.Config               `yaml:",inline"`
 
-	DownstreamURL               string `yaml:"downstream_url"`
-	DistributedExecEnabled      bool   `yaml:"distributed_exec_enabled"`
-	DisableDuplicateLabelChecks bool   `yaml:"disable_duplicate_label_checks"`
+	DownstreamURL          string `yaml:"downstream_url"`
+	DistributedExecEnabled bool   `yaml:"distributed_exec_enabled"`
 }
 
 func (cfg *CombinedFrontendConfig) RegisterFlags(f *flag.FlagSet) {
@@ -32,7 +31,6 @@ func (cfg *CombinedFrontendConfig) RegisterFlags(f *flag.FlagSet) {
 
 	f.StringVar(&cfg.DownstreamURL, "frontend.downstream-url", "", "URL of downstream Prometheus.")
 	f.BoolVar(&cfg.DistributedExecEnabled, "frontend.distributed_exec_enabled", false, "Experimental: Enables distributed execution of queries by passing logical query plan fragments to downstream components.")
-	f.BoolVar(&cfg.DisableDuplicateLabelChecks, "frontend.disable_duplicate_label_checks", false, "Experimental: Disables duplicate label checks during logical plan generation")
 }
 
 // InitFrontend initializes frontend (either V1 -- without scheduler, or V2 -- with scheduler) or no frontend at

pkg/querier/tripperware/instantquery/instant_query.go

@@ -23,6 +23,8 @@ import (
 	"github.com/cortexproject/cortex/pkg/util"
 	"github.com/cortexproject/cortex/pkg/util/limiter"
 	"github.com/cortexproject/cortex/pkg/util/spanlogger"
+
+	"github.com/thanos-io/promql-engine/logicalplan"
 )
 
 var (
@@ -141,6 +143,19 @@ func (c instantQueryCodec) DecodeResponse(ctx context.Context, r *http.Response,
 	return &resp, nil
 }
 
+func (c instantQueryCodec) getSerializedBody(promReq *tripperware.PrometheusRequest) ([]byte, error) {
+	var byteLP []byte
+	var err error
+
+	if promReq.LogicalPlan != nil && *promReq.LogicalPlan != nil {
+		byteLP, err = logicalplan.Marshal((*promReq.LogicalPlan).Root())
+		if err != nil {
+			return nil, err
+		}
+	}
+	return byteLP, nil
+}
+
 func (c instantQueryCodec) EncodeRequest(ctx context.Context, r tripperware.Request) (*http.Request, error) {
 	promReq, ok := r.(*tripperware.PrometheusRequest)
 	if !ok {
@@ -176,11 +191,16 @@ func (c instantQueryCodec) EncodeRequest(ctx context.Context, r tripperware.Requ
 		tripperware.SetRequestHeaders(h, c.defaultCodecType, c.compression)
 	}
 
+	byteBody, err := c.getSerializedBody(promReq)
+	if err != nil {
+		return nil, err
+	}
+
 	req := &http.Request{
 		Method:     "POST",
 		RequestURI: u.String(), // This is what the httpgrpc code looks at.
 		URL:        u,
-		Body:       io.NopCloser(bytes.NewReader(promReq.LogicalPlan)),
+		Body:       io.NopCloser(bytes.NewReader(byteBody)),
 		Header:     h,
 	}
 

pkg/querier/tripperware/instantquery/instant_query_middlewares.go

@@ -15,9 +15,7 @@ func Middlewares(
 	merger tripperware.Merger,
 	queryAnalyzer querysharding.Analyzer,
 	lookbackDelta time.Duration,
-	enablePerStepStats bool,
 	distributedExecEnabled bool,
-	disableDuplicateLabelChecks bool,
 ) ([]tripperware.Middleware, error) {
 	m := []tripperware.Middleware{
 		NewLimitsMiddleware(limits, lookbackDelta),
@@ -26,7 +24,7 @@ func Middlewares(
 
 	if distributedExecEnabled {
 		m = append(m,
-			tripperware.LogicalPlanGenMiddleware(lookbackDelta, enablePerStepStats, disableDuplicateLabelChecks))
+			tripperware.LogicalPlanGenMiddleware())
 	}
 
 	return m, nil

pkg/querier/tripperware/instantquery/instant_query_middlewares_test.go

@@ -2,6 +2,7 @@ package instantquery
 
 import (
 	"context"
+	"github.com/thanos-io/promql-engine/logicalplan"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -60,9 +61,7 @@ func TestRoundTrip(t *testing.T) {
 		nil,
 		qa,
 		5*time.Minute,
-		false,
 		distributedExecEnabled,
-		true,
 	)
 	require.NoError(t, err)
 
@@ -147,8 +146,6 @@ func TestRoundTripWithDistributedExec(t *testing.T) {
 		nil,
 		qa,
 		5*time.Minute,
-		false,
-		true,
 		true,
 	)
 	require.NoError(t, err)
@@ -195,9 +192,14 @@ func TestRoundTripWithDistributedExec(t *testing.T) {
 			req = req.WithContext(ctx)
 			err = user.InjectOrgIDIntoHTTPRequest(ctx, req)
 			require.NoError(t, err)
+
 			body, err := io.ReadAll(req.Body)
 			require.NotEmpty(t, body)
-			require.Equal(t, body, tc.pReq.LogicalPlan)
+			require.NoError(t, err)
+
+			byteLP, err := logicalplan.Marshal((*tc.pReq.LogicalPlan).Root())
+			require.NoError(t, err)
+			require.Equal(t, byteLP, body)
 
 			resp, err := tw(downstream).RoundTrip(req)
 			require.NoError(t, err)
@@ -244,8 +246,6 @@ func TestRoundTripWithoutDistributedExec(t *testing.T) {
 		qa,
 		5*time.Minute,
 		false,
-		false,
-		true,
 	)
 	require.NoError(t, err)
 

pkg/querier/tripperware/instantquery/logical_plan_gen_instant_test.go

@@ -72,11 +72,7 @@ func TestInstantLogicalPlan(t *testing.T) {
 		t.Run(strconv.Itoa(i), func(t *testing.T) {
 			t.Parallel()
 
-			lpm := tripperware.LogicalPlanGenMiddleware(
-				time.Duration(1000),
-				false,
-				true,
-			)
+			lpm := tripperware.LogicalPlanGenMiddleware()
 			handler := lpm.Wrap(tripperware.HandlerFunc(func(_ context.Context, req tripperware.Request) (tripperware.Response, error) {
 				return nil, nil
 			}))
@@ -85,25 +81,34 @@ func TestInstantLogicalPlan(t *testing.T) {
 			_, _ = handler.Do(context.Background(), tc.input)
 			require.NotEmpty(t, tc.input.LogicalPlan, "prom request should not be empty")
 
-			// Test Group 2: Ensure the logical plan can be deserialized back
-			qOpts := query.Options{
-				Start: time.Unix(tc.input.Start, 0),
-				End:   time.Unix(tc.input.End, 0),
-				Step:  time.Duration(1000),
-			}
-			planOpts := logicalplan.PlanOptions{DisableDuplicateLabelCheck: true}
-			_, err := logicalplan.NewFromBytes(tc.input.LogicalPlan, &qOpts, planOpts)
-			require.NoError(t, err)
-
-			// Test 3: Encode the request and validate method and body
+			// Test 2: Encode the request and validate method and body
 			httpReq, err := tripperware.Codec.EncodeRequest(testInstantQueryCodec, context.Background(), tc.input)
 			require.NoError(t, err)
 			require.Equal(t, httpReq.Method, http.MethodPost, "Method should be POST")
 
 			bodyBytes, err := io.ReadAll(httpReq.Body)
 			require.NoError(t, err)
 			require.NotEmpty(t, bodyBytes, "HTTP body should not be empty")
-			require.Equal(t, bodyBytes, tc.input.LogicalPlan, "logical plan in request body does not match expected bytes")
+
+			// Test Group 3: Ensure the logical plan can be deserialized back
+			qOpts := query.Options{
+				Start: time.Unix(tc.input.Start, 0),
+				End:   time.Unix(tc.input.Start, 0),
+				Step:  0,
+			}
+
+			qOpts = query.Options{
+				Start:              time.Unix(tc.input.Start, 0),
+				End:                time.Unix(tc.input.Start, 0),
+				Step:               0,
+				StepsBatch:         10,
+				LookbackDelta:      0,
+				EnablePerStepStats: false,
+			}
+
+			planOpts := logicalplan.PlanOptions{DisableDuplicateLabelCheck: false}
+			_, err = logicalplan.NewFromBytes(bodyBytes, &qOpts, planOpts)
+			require.NoError(t, err)
 		})
 	}
 }

pkg/querier/tripperware/logicalplan_gen.go

@@ -14,33 +14,39 @@ const (
 	stepBatch = 10
 )
 
-func LogicalPlanGenMiddleware(lookbackDelta time.Duration, enablePerStepStats bool, disableDuplicateLabelChecks bool) Middleware {
+func LogicalPlanGenMiddleware() Middleware {
 	return MiddlewareFunc(func(next Handler) Handler {
 		return logicalPlanGen{
-			lookbackDelta:               lookbackDelta,
-			enabledPerStepStats:         enablePerStepStats,
-			next:                        next,
-			disableDuplicateLabelChecks: disableDuplicateLabelChecks,
+			next: next,
 		}
 	})
 }
 
 type logicalPlanGen struct {
-	lookbackDelta               time.Duration
-	enabledPerStepStats         bool
-	next                        Handler
-	disableDuplicateLabelChecks bool
+	next Handler
 }
 
-func (l logicalPlanGen) NewInstantLogicalPlan(qs string, ts time.Time) ([]byte, error) {
-
-	qOpts := query.Options{
-		Start:              ts,
-		End:                ts,
-		Step:               0,
-		StepsBatch:         stepBatch,
-		LookbackDelta:      l.lookbackDelta,
-		EnablePerStepStats: l.disableDuplicateLabelChecks,
+func (l logicalPlanGen) NewLogicalPlan(qs string, start time.Time, end time.Time, step time.Duration) (*logicalplan.Plan, error) {
+
+	qOpts := query.Options{}
+	if step == 0 {
+		qOpts = query.Options{
+			Start:              start,
+			End:                start,
+			Step:               0,
+			StepsBatch:         stepBatch,
+			LookbackDelta:      0,
+			EnablePerStepStats: false,
+		}
+	} else {
+		qOpts = query.Options{
+			Start:              start,
+			End:                end,
+			Step:               step,
+			StepsBatch:         stepBatch,
+			LookbackDelta:      0,
+			EnablePerStepStats: false,
+		}
 	}
 
 	expr, err := parser.NewParser(qs, parser.WithFunctions(parser.Functions)).ParseExpr()
@@ -49,52 +55,15 @@ func (l logicalPlanGen) NewInstantLogicalPlan(qs string, ts time.Time) ([]byte,
 	}
 
 	planOpts := logicalplan.PlanOptions{
-		DisableDuplicateLabelCheck: l.disableDuplicateLabelChecks,
+		DisableDuplicateLabelCheck: false,
 	}
 
 	logicalPlan := logicalplan.NewFromAST(expr, &qOpts, planOpts)
 	optimizedPlan, _ := logicalPlan.Optimize(logicalplan.DefaultOptimizers)
 
-	// TODO: Add distributed optimizer for remote node insertion
-
-	byteLP, err := logicalplan.Marshal(optimizedPlan.Root())
-	if err != nil {
-		return nil, err
-	}
+	//TODO: Add distributed query optimizer
 
-	return byteLP, nil
-}
-
-func (l logicalPlanGen) NewRangeLogicalPlan(qs string, start, end time.Time, interval time.Duration) ([]byte, error) {
-
-	qOpts := query.Options{
-		Start:              start,
-		End:                end,
-		Step:               interval,
-		StepsBatch:         stepBatch,
-		LookbackDelta:      l.lookbackDelta,
-		EnablePerStepStats: l.enabledPerStepStats,
-	}
-
-	expr, err := parser.NewParser(qs, parser.WithFunctions(parser.Functions)).ParseExpr()
-	if err != nil {
-		return nil, err
-	}
-
-	planOpts := logicalplan.PlanOptions{
-		DisableDuplicateLabelCheck: l.disableDuplicateLabelChecks,
-	}
-
-	lPlan := logicalplan.NewFromAST(expr, &qOpts, planOpts)
-	optimizedPlan, _ := lPlan.Optimize(logicalplan.DefaultOptimizers)
-	byteLP, err := logicalplan.Marshal(optimizedPlan.Root())
-	if err != nil {
-		return nil, err
-	}
-
-	// TODO: Add distributed optimizer for remote node insertion
-
-	return byteLP, nil
+	return &optimizedPlan, nil
 }
 
 func (l logicalPlanGen) Do(ctx context.Context, r Request) (Response, error) {
@@ -105,20 +74,16 @@ func (l logicalPlanGen) Do(ctx context.Context, r Request) (Response, error) {
 
 	startTime := time.Unix(0, promReq.Start*int64(time.Millisecond))
 	endTime := time.Unix(0, promReq.End*int64(time.Millisecond))
-	duration := time.Duration(promReq.Step) * time.Millisecond
+	step := time.Duration(promReq.Step) * time.Millisecond
 
-	var byteLP []byte
 	var err error
-	if promReq.Step != 0 {
-		byteLP, err = l.NewRangeLogicalPlan(promReq.Query, startTime, endTime, duration)
-	} else {
-		byteLP, err = l.NewInstantLogicalPlan(promReq.Query, startTime)
-	}
 
+	newLogicalPlan, err := l.NewLogicalPlan(promReq.Query, startTime, endTime, step)
 	if err != nil {
 		return nil, err
 	}
-	promReq.LogicalPlan = byteLP
+
+	promReq.LogicalPlan = newLogicalPlan
 
 	return l.next.Do(ctx, r)
 }