frontend: Implement logical query plan middleware under new distributed execution feature flag

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

Author: rubywtl

pkg/cortex/modules.go

@@ -533,12 +533,16 @@ 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)
+	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)
 	if err != nil {
 		return nil, err
 	}

pkg/frontend/config.go

@@ -20,7 +20,9 @@ type CombinedFrontendConfig struct {
 	FrontendV1 v1.Config               `yaml:",inline"`
 	FrontendV2 v2.Config               `yaml:",inline"`
 
-	DownstreamURL string `yaml:"downstream_url"`
+	DownstreamURL               string `yaml:"downstream_url"`
+	DistributedExecEnabled      bool   `yaml:"distributed_exec_enabled"`
+	DisableDuplicateLabelChecks bool   `yaml:"disable_duplicate_label_checks"`
 }
 
 func (cfg *CombinedFrontendConfig) RegisterFlags(f *flag.FlagSet) {
@@ -29,6 +31,8 @@ func (cfg *CombinedFrontendConfig) RegisterFlags(f *flag.FlagSet) {
 	cfg.FrontendV2.RegisterFlags(f)
 
 	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_logicalplan_gen.go

@@ -0,0 +1,84 @@
+package instantquery
+
+import (
+	"context"
+	"github.com/cortexproject/cortex/pkg/querier/tripperware"
+	"github.com/prometheus/prometheus/promql/parser"
+	"github.com/thanos-io/promql-engine/logicalplan"
+	"github.com/thanos-io/promql-engine/query"
+	"github.com/weaveworks/common/httpgrpc"
+	"net/http"
+	"time"
+)
+
+const (
+	stepBatch = 10
+)
+
+func InstantLogicalPlanGenMiddleware(lookbackDelta time.Duration, enablePerStepStats bool, disableDuplicateLabelChecks bool) tripperware.Middleware {
+	return tripperware.MiddlewareFunc(func(next tripperware.Handler) tripperware.Handler {
+		return instantLogicalPlanGen{
+			lookbackDelta:               lookbackDelta,
+			enabledPerStepStats:         enablePerStepStats,
+			disableDuplicateLabelChecks: disableDuplicateLabelChecks,
+			next:                        next,
+		}
+	})
+}
+
+type instantLogicalPlanGen struct {
+	next                        tripperware.Handler
+	lookbackDelta               time.Duration
+	enabledPerStepStats         bool
+	disableDuplicateLabelChecks bool
+}
+
+func (l instantLogicalPlanGen) 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,
+	}
+
+	expr, err := parser.NewParser(qs, parser.WithFunctions(parser.Functions)).ParseExpr()
+	if err != nil {
+		return nil, err
+	}
+
+	planOpts := logicalplan.PlanOptions{
+		DisableDuplicateLabelCheck: l.disableDuplicateLabelChecks,
+	}
+
+	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
+	}
+
+	return byteLP, nil
+}
+
+func (l instantLogicalPlanGen) Do(ctx context.Context, r tripperware.Request) (tripperware.Response, error) {
+	promReq, ok := r.(*tripperware.PrometheusRequest)
+	if !ok {
+		return nil, httpgrpc.Errorf(http.StatusBadRequest, "invalid request format")
+	}
+
+	instantTime := time.Unix(0, promReq.Start*int64(time.Millisecond))
+
+	byteLP, err := l.NewInstantLogicalPlan(promReq.Query, instantTime)
+	if err != nil {
+		return nil, err
+	}
+	promReq.LogicalPlan = byteLP
+
+	return l.next.Do(ctx, r)
+}