Refactor engine config

Signed-off-by: 🌲 Harry 🌊 John 🏔 <[email protected]>

Author: harry671003

integration/querier_test.go

@@ -416,6 +416,7 @@ func TestQuerierWithBlocksStorageRunningInSingleBinaryMode(t *testing.T) {
 						"-blocks-storage.bucket-store.bucket-index.enabled": strconv.FormatBool(testCfg.bucketIndexEnabled),
 						"-querier.query-store-for-labels-enabled":           "true",
 						"-querier.thanos-engine":                            strconv.FormatBool(thanosEngine),
+						"-querier.enable-x-functions":                       strconv.FormatBool(thanosEngine),
 						// Ingester.
 						"-ring.store":      "consul",
 						"-consul.hostname": consul.NetworkHTTPEndpoint(),
@@ -1310,3 +1311,63 @@ func TestQuerierMaxSamplesLimit(t *testing.T) {
 		Error:     "query processing would load too many samples into memory in query execution",
 	})
 }
+
+func TestQuerierXFunctions(t *testing.T) {
+	const blockRangePeriod = 5 * time.Second
+
+	s, err := e2e.NewScenario(networkName)
+	require.NoError(t, err)
+	defer s.Close()
+
+	// Configure the blocks storage to frequently compact TSDB head
+	// and ship blocks to the storage.
+	flags := mergeFlags(BlocksStorageFlags(), map[string]string{
+		"-blocks-storage.tsdb.block-ranges-period": blockRangePeriod.String(),
+		"-blocks-storage.tsdb.ship-interval":       "1s",
+		"-blocks-storage.tsdb.retention-period":    ((blockRangePeriod * 2) - 1).String(),
+		"-querier.thanos_engine":                   "true",
+		"-querier.enable-x-functions":              "true",
+	})
+
+	// Start dependencies.
+	minio := e2edb.NewMinio(9000, flags["-blocks-storage.s3.bucket-name"])
+	consul := e2edb.NewConsul()
+	require.NoError(t, s.StartAndWaitReady(consul, minio))
+
+	// Start Cortex components for the write path.
+	distributor := e2ecortex.NewDistributor("distributor", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), flags, "")
+	ingester := e2ecortex.NewIngester("ingester", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), flags, "")
+	require.NoError(t, s.StartAndWaitReady(distributor, ingester))
+
+	queryFrontend := e2ecortex.NewQueryFrontendWithConfigFile("query-frontend", "", flags, "")
+	require.NoError(t, s.Start(queryFrontend))
+
+	querier := e2ecortex.NewQuerier("querier", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), mergeFlags(flags, map[string]string{
+		"-querier.frontend-address": queryFrontend.NetworkGRPCEndpoint(),
+	}), "")
+	require.NoError(t, s.StartAndWaitReady(querier))
+
+	// Wait until the distributor and querier has updated the ring.
+	require.NoError(t, distributor.WaitSumMetrics(e2e.Equals(512), "cortex_ring_tokens_total"))
+	require.NoError(t, querier.WaitSumMetrics(e2e.Equals(512), "cortex_ring_tokens_total"))
+
+	c, err := e2ecortex.NewClient(distributor.HTTPEndpoint(), queryFrontend.HTTPEndpoint(), "", "", "user-1")
+	require.NoError(t, err)
+
+	// Push some series to Cortex.
+	series1Timestamp := time.Now()
+	series1, _ := generateSeries("series_1", series1Timestamp, prompb.Label{Name: "job", Value: "test"})
+	series2, _ := generateSeries("series_2", series1Timestamp, prompb.Label{Name: "job", Value: "test"})
+
+	res, err := c.Push(series1)
+	require.NoError(t, err)
+	require.Equal(t, 200, res.StatusCode)
+	res, err = c.Push(series2)
+	require.NoError(t, err)
+	require.Equal(t, 200, res.StatusCode)
+
+	// Query back the series (1 only in the storage, 1 only in the ingesters, 1 on both).
+	result, err := c.Query(`xrate({job="test"}[1m])`, series1Timestamp)
+	require.NoError(t, err)
+	require.Equal(t, model.ValVector, result.Type())
+}

pkg/cortex/modules.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"flag"
 	"fmt"
+
 	"log/slog"
 	"net/http"
 	"runtime"
@@ -18,8 +19,6 @@ import (
 	"github.com/prometheus/prometheus/rules"
 	prom_storage "github.com/prometheus/prometheus/storage"
 	"github.com/thanos-io/objstore"
-	"github.com/thanos-io/promql-engine/engine"
-	"github.com/thanos-io/promql-engine/logicalplan"
 	"github.com/thanos-io/thanos/pkg/discovery/dns"
 	"github.com/thanos-io/thanos/pkg/querysharding"
 	httpgrpc_server "github.com/weaveworks/common/httpgrpc/server"
@@ -32,6 +31,7 @@ import (
 	configAPI "github.com/cortexproject/cortex/pkg/configs/api"
 	"github.com/cortexproject/cortex/pkg/configs/db"
 	"github.com/cortexproject/cortex/pkg/distributor"
+	"github.com/cortexproject/cortex/pkg/engine"
 	"github.com/cortexproject/cortex/pkg/flusher"
 	"github.com/cortexproject/cortex/pkg/frontend"
 	"github.com/cortexproject/cortex/pkg/frontend/transport"
@@ -634,7 +634,6 @@ func (t *Cortex) initRuler() (serv services.Service, err error) {
 	if t.Cfg.ExternalPusher != nil && t.Cfg.ExternalQueryable != nil {
 		rulerRegisterer := prometheus.WrapRegistererWith(prometheus.Labels{"engine": "ruler"}, prometheus.DefaultRegisterer)
 
-		var queryEngine promql.QueryEngine
 		opts := promql.EngineOpts{
 			Logger:               util_log.SLogger,
 			Reg:                  rulerRegisterer,
@@ -649,15 +648,7 @@ func (t *Cortex) initRuler() (serv services.Service, err error) {
 				return t.Cfg.Querier.DefaultEvaluationInterval.Milliseconds()
 			},
 		}
-		if t.Cfg.Querier.ThanosEngine {
-			queryEngine = engine.New(engine.Opts{
-				EngineOpts:        opts,
-				LogicalOptimizers: logicalplan.AllOptimizers,
-				EnableAnalysis:    true,
-			})
-		} else {
-			queryEngine = promql.NewEngine(opts)
-		}
+		queryEngine := engine.New(opts, t.Cfg.Querier.Engine, rulerRegisterer)
 
 		managerFactory := ruler.DefaultTenantManagerFactory(t.Cfg.Ruler, t.Cfg.ExternalPusher, t.Cfg.ExternalQueryable, queryEngine, t.Overrides, metrics, prometheus.DefaultRegisterer)
 		manager, err = ruler.NewDefaultMultiTenantManager(t.Cfg.Ruler, t.Overrides, managerFactory, metrics, prometheus.DefaultRegisterer, util_log.Logger)

pkg/engine/config.go

@@ -0,0 +1,63 @@
+package engine
+
+import (
+	"flag"
+	"fmt"
+	"github.com/thanos-io/promql-engine/logicalplan"
+	"strings"
+)
+
+var supportedOptimizers = []string{"default", "all", "propagate-matchers", "sort-matchers", "merge-selects", "detect-histogram-stats"}
+
+// Config contains the configuration to create engine
+type Config struct {
+	EnableThanosEngine bool                    `yaml:"enable_thanos_engine"`
+	EnableXFunctions   bool                    `yaml:"enable_x_functions"`
+	Optimizers         string                  `yaml:"optimizers"`
+	LogicalOptimizers  []logicalplan.Optimizer `yaml:"-"`
+}
+
+// RegisterFlags adds the flags required to config this to the given FlagSet.
+func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
+	f.BoolVar(&cfg.EnableThanosEngine, "querier.thanos-engine", false, "Experimental. Use Thanos promql engine https://github.com/thanos-io/promql-engine rather than the Prometheus promql engine.")
+	f.BoolVar(&cfg.EnableXFunctions, "querier.enable-x-functions", false, "Enable xincrease, xdelta, xrate etc from Thanos engine.")
+	f.StringVar(&cfg.Optimizers, "querier.optimizers", "default", "Logical plan optimizers. Multiple optimizers can be provided as a comma-separated list. Supported values: "+strings.Join(supportedOptimizers, ", "))
+}
+
+func (cfg *Config) Validate() error {
+	splitOptimizers := strings.Split(cfg.Optimizers, ",")
+
+	for _, optimizer := range splitOptimizers {
+		if optimizer == "all" || optimizer == "default" {
+			if len(splitOptimizers) > 1 {
+				return fmt.Errorf("special optimizer %s cannot be combined with other optimizers", optimizer)
+			}
+		}
+		optimizers, err := getOptimizer(optimizer)
+		if err != nil {
+			return err
+		}
+		cfg.LogicalOptimizers = append(cfg.LogicalOptimizers, optimizers...)
+	}
+
+	return nil
+}
+
+func getOptimizer(name string) ([]logicalplan.Optimizer, error) {
+	switch name {
+	case "default":
+		return logicalplan.DefaultOptimizers, nil
+	case "all":
+		return logicalplan.AllOptimizers, nil
+	case "propagate-matchers":
+		return []logicalplan.Optimizer{logicalplan.PropagateMatchersOptimizer{}}, nil
+	case "sort-matchers":
+		return []logicalplan.Optimizer{logicalplan.SortMatchers{}}, nil
+	case "merge-selects":
+		return []logicalplan.Optimizer{logicalplan.MergeSelectsOptimizer{}}, nil
+	case "detect-histogram-stats":
+		return []logicalplan.Optimizer{logicalplan.DetectHistogramStatsOptimizer{}}, nil
+	default:
+		return nil, fmt.Errorf("unknown optimizer %s", name)
+	}
+}

pkg/engine/engine.go

@@ -10,7 +10,6 @@ import (
 	"github.com/prometheus/prometheus/promql"
 	"github.com/prometheus/prometheus/storage"
 	thanosengine "github.com/thanos-io/promql-engine/engine"
-	"github.com/thanos-io/promql-engine/logicalplan"
 )
 
 type engineKeyType struct{}
@@ -52,15 +51,16 @@ type Engine struct {
 	engineSwitchQueriesTotal *prometheus.CounterVec
 }
 
-func New(opts promql.EngineOpts, enableThanosEngine bool, reg prometheus.Registerer) *Engine {
+func New(opts promql.EngineOpts, engineCfg *Config, reg prometheus.Registerer) *Engine {
 	prometheusEngine := promql.NewEngine(opts)
 
 	var thanosEngine *thanosengine.Engine
-	if enableThanosEngine {
+	if engineCfg.EnableThanosEngine {
 		thanosEngine = thanosengine.New(thanosengine.Opts{
 			EngineOpts:        opts,
-			LogicalOptimizers: logicalplan.DefaultOptimizers,
+			LogicalOptimizers: engineCfg.LogicalOptimizers,
 			EnableAnalysis:    true,
+			EnableXFunctions:  engineCfg.EnableXFunctions,
 		})
 	}
 

pkg/engine/engine_test.go

@@ -3,6 +3,8 @@ package engine
 import (
 	"bytes"
 	"context"
+	"fmt"
+	"github.com/thanos-io/promql-engine/execution/parse"
 	"net/http"
 	"testing"
 	"time"
@@ -37,7 +39,7 @@ func TestEngine_Fallback(t *testing.T) {
 		Logger: utillog.GoKitLogToSlog(log.NewNopLogger()),
 		Reg:    reg,
 	}
-	queryEngine := New(opts, true, reg)
+	queryEngine := New(opts, &Config{EnableThanosEngine: true}, reg)
 
 	// instant query, should go to fallback
 	_, _ = queryEngine.NewInstantQuery(ctx, queryable, nil, "unimplemented(foo)", now)
@@ -68,7 +70,7 @@ func TestEngine_Switch(t *testing.T) {
 		Logger: utillog.GoKitLogToSlog(log.NewNopLogger()),
 		Reg:    reg,
 	}
-	queryEngine := New(opts, true, reg)
+	queryEngine := New(opts, &Config{EnableThanosEngine: true}, reg)
 
 	// Query Prometheus engine
 	r := &http.Request{Header: http.Header{}}
@@ -96,3 +98,28 @@ func TestEngine_Switch(t *testing.T) {
 		cortex_engine_switch_queries_total{engine_type="thanos"} 2
 	`), "cortex_engine_switch_queries_total"))
 }
+
+func TestEngine_XFunctions(t *testing.T) {
+	ctx := context.Background()
+	reg := prometheus.NewRegistry()
+
+	now := time.Now()
+	start := time.Now().Add(-time.Minute * 5)
+	step := time.Minute
+	queryable := promqltest.LoadedStorage(t, "")
+	opts := promql.EngineOpts{
+		Logger: utillog.GoKitLogToSlog(log.NewNopLogger()),
+		Reg:    reg,
+	}
+	queryEngine := New(opts, &Config{EnableThanosEngine: true, EnableXFunctions: true}, reg)
+
+	for name, _ := range parse.XFunctions {
+		t.Run(name, func(t *testing.T) {
+			_, err := queryEngine.NewInstantQuery(ctx, queryable, nil, fmt.Sprintf("%s(foo[1m])", name), now)
+			require.NoError(t, err)
+
+			_, err = queryEngine.NewRangeQuery(ctx, queryable, nil, fmt.Sprintf("%s(foo[1m])", name), start, now, step)
+			require.NoError(t, err)
+		})
+	}
+}

pkg/querier/querier.go

@@ -85,9 +85,7 @@ type Config struct {
 
 	ShuffleShardingIngestersLookbackPeriod time.Duration `yaml:"shuffle_sharding_ingesters_lookback_period"`
 
-	// Experimental. Use https://github.com/thanos-io/promql-engine rather than
-	// the Prometheus query engine.
-	ThanosEngine bool `yaml:"thanos_engine"`
+	Engine *engine.Config `yaml:"thanos_engine"`
 
 	// Ignore max query length check at Querier.
 	IgnoreMaxQueryLength              bool `yaml:"ignore_max_query_length"`
@@ -139,13 +137,14 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.IntVar(&cfg.IngesterQueryMaxAttempts, "querier.ingester-query-max-attempts", 1, "The maximum number of times we attempt fetching data from ingesters for retryable errors (ex. partial data returned).")
 	f.DurationVar(&cfg.LookbackDelta, "querier.lookback-delta", 5*time.Minute, "Time since the last sample after which a time series is considered stale and ignored by expression evaluations.")
 	f.DurationVar(&cfg.ShuffleShardingIngestersLookbackPeriod, "querier.shuffle-sharding-ingesters-lookback-period", 0, "When distributor's sharding strategy is shuffle-sharding and this setting is > 0, queriers fetch in-memory series from the minimum set of required ingesters, selecting only ingesters which may have received series since 'now - lookback period'. The lookback period should be greater or equal than the configured 'query store after' and 'query ingesters within'. If this setting is 0, queriers always query all ingesters (ingesters shuffle sharding on read path is disabled).")
-	f.BoolVar(&cfg.ThanosEngine, "querier.thanos-engine", false, "Experimental. Use Thanos promql engine https://github.com/thanos-io/promql-engine rather than the Prometheus promql engine.")
 	f.Int64Var(&cfg.MaxSubQuerySteps, "querier.max-subquery-steps", 0, "Max number of steps allowed for every subquery expression in query. Number of steps is calculated using subquery range / step. A value > 0 enables it.")
 	f.BoolVar(&cfg.IgnoreMaxQueryLength, "querier.ignore-max-query-length", false, "If enabled, ignore max query length check at Querier select method. Users can choose to ignore it since the validation can be done before Querier evaluation like at Query Frontend or Ruler.")
 	f.BoolVar(&cfg.EnablePromQLExperimentalFunctions, "querier.enable-promql-experimental-functions", false, "[Experimental] If true, experimental promQL functions are enabled.")
 	f.BoolVar(&cfg.EnableParquetQueryable, "querier.enable-parquet-queryable", false, "[Experimental] If true, querier will try to query the parquet files if available.")
 	f.IntVar(&cfg.ParquetQueryableShardCacheSize, "querier.parquet-queryable-shard-cache-size", 512, "[Experimental] [Experimental] Maximum size of the Parquet queryable shard cache. 0 to disable.")
 	f.StringVar(&cfg.ParquetQueryableDefaultBlockStore, "querier.parquet-queryable-default-block-store", string(parquetBlockStore), "Parquet queryable's default block store to query. Valid options are tsdb and parquet. If it is set to tsdb, parquet queryable always fallback to store gateway.")
+
+	cfg.Engine.RegisterFlags(f)
 }
 
 // Validate the config
@@ -181,6 +180,10 @@ func (cfg *Config) Validate() error {
 		}
 	}
 
+	if err := cfg.Engine.Validate(); err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -246,7 +249,7 @@ func New(cfg Config, limits *validation.Overrides, distributor Distributor, stor
 			return cfg.DefaultEvaluationInterval.Milliseconds()
 		},
 	}
-	queryEngine := engine.New(opts, cfg.ThanosEngine, reg)
+	queryEngine := engine.New(opts, cfg.Engine, reg)
 	return NewSampleAndChunkQueryable(lazyQueryable), exemplarQueryable, queryEngine
 }