add engine unit test and distributed exec unit test

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

Author: rubywtl

integration/querier_test.go

@@ -1310,3 +1310,109 @@ func TestQuerierMaxSamplesLimit(t *testing.T) {
 		Error:     "query processing would load too many samples into memory in query execution",
 	})
 }
+
+func TestQuerierDistributedExecution(t *testing.T) {
+	// e2e test setup
+	s, err := e2e.NewScenario(networkName)
+	require.NoError(t, err)
+	defer s.Close()
+
+	consul := e2edb.NewConsulWithName("consul")
+	memcached := e2ecache.NewMemcached()
+	require.NoError(t, s.StartAndWaitReady(consul, memcached))
+
+	// initialize the flags
+	baseFlags := mergeFlags(AlertmanagerLocalFlags(), BlocksStorageFlags())
+	flags := mergeFlags(
+		baseFlags,
+		map[string]string{
+			"-blocks-storage.tsdb.head-compaction-interval":    "4m",
+			"-blocks-storage.tsdb.block-ranges-period":         "2h",
+			"-blocks-storage.tsdb.ship-interval":               "1h",
+			"-blocks-storage.bucket-store.sync-interval":       "1s",
+			"-blocks-storage.tsdb.retention-period":            "24h",
+			"-blocks-storage.bucket-store.index-cache.backend": tsdb.IndexCacheBackendInMemory,
+			"-querier.query-store-for-labels-enabled":          "true",
+			// Ingester.
+			"-ring.store":      "consul",
+			"-consul.hostname": consul.NetworkHTTPEndpoint(),
+			// Distributor.
+			"-distributor.replication-factor": "1",
+			// Store-gateway.
+			"-store-gateway.sharding-enabled": "false",
+			// Alert manager
+			"-alertmanager.web.external-url":      "http://localhost/alertmanager",
+			"-frontend.query-vertical-shard-size": "1",
+			"-frontend.max-cache-freshness":       "1m",
+			// enable experimental promQL funcs
+			"-querier.enable-promql-experimental-functions": "true",
+			// enable distributed execution (logical plan execution)
+			"-querier.distributed-exec-enabled": "true",
+		},
+	)
+
+	minio := e2edb.NewMinio(9000, flags["-blocks-storage.s3.bucket-name"])
+	require.NoError(t, s.StartAndWaitReady(minio))
+
+	// start services
+	var queryScheduler *e2ecortex.CortexService
+	queryScheduler = e2ecortex.NewQueryScheduler("query-scheduler", flags, "")
+	require.NoError(t, s.StartAndWaitReady(queryScheduler))
+	flags["-frontend.scheduler-address"] = queryScheduler.NetworkGRPCEndpoint()
+	flags["-querier.scheduler-address"] = queryScheduler.NetworkGRPCEndpoint()
+
+	queryFrontend := e2ecortex.NewQueryFrontend("query-frontend", flags, "")
+	require.NoError(t, s.Start(queryFrontend))
+
+	ingester := e2ecortex.NewIngester("ingester", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), flags, "")
+	distributor := e2ecortex.NewDistributor("distributor", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), flags, "")
+	querier1 := e2ecortex.NewQuerier("querier-1", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), flags, "")
+	querier2 := e2ecortex.NewQuerier("querier-2", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), flags, "")
+
+	require.NoError(t, s.StartAndWaitReady(querier1, querier2, ingester, distributor))
+	require.NoError(t, s.WaitReady(queryFrontend))
+
+	// wait until distributor and queriers have updated the ring.
+	require.NoError(t, distributor.WaitSumMetrics(e2e.Equals(512), "cortex_ring_tokens_total"))
+	require.NoError(t, querier1.WaitSumMetrics(e2e.Equals(512), "cortex_ring_tokens_total"))
+	require.NoError(t, querier2.WaitSumMetrics(e2e.Equals(512), "cortex_ring_tokens_total"))
+
+	// push some series to Cortex.
+	distClient, err := e2ecortex.NewClient(distributor.HTTPEndpoint(), "", "", "", userID)
+	require.NoError(t, err)
+
+	series1Timestamp := time.Now()
+	series2Timestamp := series1Timestamp.Add(blockRangePeriod * 2)
+	series1, expectedVector1 := generateSeries("series_1", series1Timestamp, prompb.Label{Name: "series_1", Value: "series_1"})
+	series2, expectedVector2 := generateSeries("series_2", series2Timestamp, prompb.Label{Name: "series_2", Value: "series_2"})
+
+	res, err := distClient.Push(series1)
+	require.NoError(t, err)
+	require.Equal(t, 200, res.StatusCode)
+
+	res, err = distClient.Push(series2)
+	require.NoError(t, err)
+	require.Equal(t, 200, res.StatusCode)
+
+	for _, q := range []*e2ecortex.CortexService{querier1, querier2} {
+		c, err := e2ecortex.NewClient("", q.HTTPEndpoint(), "", "", userID)
+		require.NoError(t, err)
+
+		_, err = c.Query("series_1", now)
+		require.NoError(t, err)
+	}
+
+	require.NoError(t, queryScheduler.WaitSumMetrics(e2e.Equals(2), "cortex_query_scheduler_connected_querier_clients"))
+
+	// main tests
+	// - make sure queries are still executable with distributed execution enabled
+	res, body, err = c.QueryRaw(`sum({job="test"})`, series1Timestamp, map[string]string{})
+	require.NoError(t, err)
+	require.Equal(t, 200, res.StatusCode)
+	require.Equal(t, expectedVector1, string(body))
+
+	res, body, err = c.QueryRaw(`sum({job="test"})`, series2Timestamp, map[string]string{})
+	require.NoError(t, err)
+	require.Equal(t, 200, res.StatusCode)
+	require.Equal(t, expectedVector2, string(body))
+}

pkg/engine/engine.go

@@ -137,6 +137,7 @@ prom:
 func (qf *Engine) MakeInstantQueryFromPlan(ctx context.Context, q storage.Queryable, opts promql.QueryOpts, root logicalplan.Node, ts time.Time, qs string) (promql.Query, error) {
 	if engineType := GetEngineType(ctx); engineType == Prometheus {
 		qf.engineSwitchQueriesTotal.WithLabelValues(string(Prometheus)).Inc()
+		goto prom
 	} else if engineType == Thanos {
 		qf.engineSwitchQueriesTotal.WithLabelValues(string(Thanos)).Inc()
 	}
@@ -161,6 +162,7 @@ prom:
 func (qf *Engine) MakeRangeQueryFromPlan(ctx context.Context, q storage.Queryable, opts promql.QueryOpts, root logicalplan.Node, start time.Time, end time.Time, interval time.Duration, qs string) (promql.Query, error) {
 	if engineType := GetEngineType(ctx); engineType == Prometheus {
 		qf.engineSwitchQueriesTotal.WithLabelValues(string(Prometheus)).Inc()
+		goto prom
 	} else if engineType == Thanos {
 		qf.engineSwitchQueriesTotal.WithLabelValues(string(Thanos)).Inc()
 	}

pkg/engine/engine_test.go

@@ -14,6 +14,8 @@ import (
 	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/prometheus/prometheus/promql/promqltest"
 	"github.com/stretchr/testify/require"
+	"github.com/thanos-io/promql-engine/logicalplan"
+	"github.com/thanos-io/promql-engine/query"
 
 	utillog "github.com/cortexproject/cortex/pkg/util/log"
 )
@@ -96,3 +98,97 @@ func TestEngine_Switch(t *testing.T) {
 		cortex_engine_switch_queries_total{engine_type="thanos"} 2
 	`), "cortex_engine_switch_queries_total"))
 }
+
+func TestEngine_With_Logical_Plan(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, true, reg)
+
+	range_lp := createTestLogicalPlan(t, start, now, step, "up")
+	instant_lp := createTestLogicalPlan(t, now, now, 0, "up")
+
+	r := &http.Request{Header: http.Header{}}
+	r.Header.Set(TypeHeader, string(Thanos))
+	ctx = AddEngineTypeToContext(ctx, r)
+
+	// Case 1: Executing logical plan with thanos engine
+	_, _ = queryEngine.MakeInstantQueryFromPlan(ctx, queryable, nil, instant_lp.Root(), now, "up")
+	require.NoError(t, testutil.GatherAndCompare(reg, bytes.NewBufferString(`
+		# HELP cortex_thanos_engine_fallback_queries_total Total number of fallback queries due to not implementation in thanos engine
+		# TYPE cortex_thanos_engine_fallback_queries_total counter
+		cortex_thanos_engine_fallback_queries_total 0
+	`), "cortex_thanos_engine_fallback_queries_total"))
+
+	_, _ = queryEngine.MakeRangeQueryFromPlan(ctx, queryable, nil, range_lp.Root(), start, now, step, "up")
+	require.NoError(t, testutil.GatherAndCompare(reg, bytes.NewBufferString(`
+		# HELP cortex_thanos_engine_fallback_queries_total Total number of fallback queries due to not implementation in thanos engine
+		# TYPE cortex_thanos_engine_fallback_queries_total counter
+		cortex_thanos_engine_fallback_queries_total 0
+	`), "cortex_thanos_engine_fallback_queries_total"))
+
+	// Case 2: Logical plan that thanos engine cannot execute (so it will fall back to prometheus engine)
+	err_range_lp := createTestLogicalPlan(t, start, now, step, "up[10]")
+	_, _ = queryEngine.MakeRangeQueryFromPlan(ctx, queryable, nil, err_range_lp.Root(), start, now, step, "up")
+	require.NoError(t, testutil.GatherAndCompare(reg, bytes.NewBufferString(`
+		# HELP cortex_thanos_engine_fallback_queries_total Total number of fallback queries due to not implementation in thanos engine
+		# TYPE cortex_thanos_engine_fallback_queries_total counter
+		cortex_thanos_engine_fallback_queries_total 1
+	`), "cortex_thanos_engine_fallback_queries_total"))
+
+	// Case 3: executing with prometheus engine
+	r.Header.Set(TypeHeader, string(Prometheus))
+	ctx = AddEngineTypeToContext(ctx, r)
+
+	_, _ = queryEngine.MakeInstantQueryFromPlan(ctx, queryable, nil, instant_lp.Root(), now, "up")
+	require.NoError(t, testutil.GatherAndCompare(reg, bytes.NewBufferString(`
+		# HELP cortex_engine_switch_queries_total Total number of queries where engine_type is set explicitly
+		# TYPE cortex_engine_switch_queries_total counter
+		cortex_engine_switch_queries_total{engine_type="prometheus"} 1
+		cortex_engine_switch_queries_total{engine_type="thanos"} 3
+	`), "cortex_engine_switch_queries_total"))
+
+	_, _ = queryEngine.MakeRangeQueryFromPlan(ctx, queryable, nil, range_lp.Root(), start, now, step, "up")
+	require.NoError(t, testutil.GatherAndCompare(reg, bytes.NewBufferString(`
+		# HELP cortex_engine_switch_queries_total Total number of queries where engine_type is set explicitly
+		# TYPE cortex_engine_switch_queries_total counter
+		cortex_engine_switch_queries_total{engine_type="prometheus"} 2
+		cortex_engine_switch_queries_total{engine_type="thanos"} 3
+	`), "cortex_engine_switch_queries_total"))
+}
+
+func createTestLogicalPlan(t *testing.T, startTime time.Time, endTime time.Time, step time.Duration, q string) logicalplan.Plan {
+
+	qOpts := query.Options{
+		Start:              startTime,
+		End:                startTime,
+		Step:               0,
+		StepsBatch:         10,
+		LookbackDelta:      0,
+		EnablePerStepStats: false,
+	}
+
+	if step != 0 {
+		qOpts.End = endTime
+		qOpts.Step = step
+	}
+
+	expr, err := parser.NewParser(q, parser.WithFunctions(parser.Functions)).ParseExpr()
+	require.NoError(t, err)
+
+	planOpts := logicalplan.PlanOptions{
+		DisableDuplicateLabelCheck: false,
+	}
+
+	logicalPlan := logicalplan.NewFromAST(expr, &qOpts, planOpts)
+
+	return logicalPlan
+}