more testing and test fixes

Author: rubywtl

integration/querier_test.go

@@ -1450,3 +1450,124 @@ func TestQuerierDistributedExecution(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, expectedVector2, val.(model.Vector))
 }
+
+func TestDistributedExecResults(t *testing.T) {
+	// e2e test setup
+	s, err := e2e.NewScenario(networkName)
+	require.NoError(t, err)
+	defer s.Close()
+
+	// initialize the flags
+	flags := mergeFlags(
+		BlocksStorageFlags(),
+		map[string]string{
+			"-blocks-storage.tsdb.block-ranges-period": (5 * time.Second).String(),
+			"-blocks-storage.tsdb.ship-interval":       "1s",
+			"-blocks-storage.tsdb.retention-period":    ((5 * time.Second * 2) - 1).String(),
+			"-querier.thanos-engine":                   "true",
+			"-querier.distributed-exec-enabled":        "true",
+		},
+	)
+
+	flags2 := mergeFlags(
+		BlocksStorageFlags(),
+		map[string]string{
+			"-blocks-storage.tsdb.block-ranges-period": (5 * time.Second).String(),
+			"-blocks-storage.tsdb.ship-interval":       "1s",
+			"-blocks-storage.tsdb.retention-period":    ((5 * time.Second * 2) - 1).String(),
+			"-querier.thanos-engine":                   "true",
+			"-querier.distributed-exec-enabled":        "false",
+		},
+	)
+
+	minio := e2edb.NewMinio(9000, flags["-blocks-storage.s3.bucket-name"])
+	consul := e2edb.NewConsul()
+	require.NoError(t, s.StartAndWaitReady(consul, minio))
+
+	// START SERVICE 1
+	distributor := e2ecortex.NewDistributor("distributor", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), flags, "")
+	ingester := e2ecortex.NewIngester("ingester", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), flags, "")
+	queryScheduler := e2ecortex.NewQueryScheduler("query-scheduler", flags, "")
+	storeGateway := e2ecortex.NewStoreGateway("store-gateway", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), flags, "")
+	require.NoError(t, s.StartAndWaitReady(queryScheduler, distributor, ingester, storeGateway))
+	flags = mergeFlags(flags, map[string]string{
+		"-querier.store-gateway-addresses": strings.Join([]string{storeGateway.NetworkGRPCEndpoint()}, ","),
+	})
+	queryFrontend := e2ecortex.NewQueryFrontend("query-frontend", mergeFlags(flags, map[string]string{
+		"-frontend.scheduler-address": queryScheduler.NetworkGRPCEndpoint(),
+	}), "")
+	require.NoError(t, s.Start(queryFrontend))
+	querier := e2ecortex.NewQuerier("querier", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), mergeFlags(flags, map[string]string{
+		"-querier.scheduler-address": queryScheduler.NetworkGRPCEndpoint(),
+	}), "")
+	require.NoError(t, s.StartAndWaitReady(querier))
+
+	require.NoError(t, distributor.WaitSumMetrics(e2e.Equals(512), "cortex_ring_tokens_total"))
+	require.NoError(t, querier.WaitSumMetrics(e2e.Equals(2*512), "cortex_ring_tokens_total"))
+
+	c, err := e2ecortex.NewClient(distributor.HTTPEndpoint(), queryFrontend.HTTPEndpoint(), "", "", "user-1")
+	require.NoError(t, err)
+
+	// START SERVICE 2
+	distributor2 := e2ecortex.NewDistributor("distributor", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), flags2, "")
+	ingester2 := e2ecortex.NewIngester("ingester", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), flags2, "")
+	queryScheduler2 := e2ecortex.NewQueryScheduler("query-scheduler", flags2, "")
+	storeGateway2 := e2ecortex.NewStoreGateway("store-gateway", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), flags, "")
+	require.NoError(t, s.StartAndWaitReady(queryScheduler2, distributor2, ingester2, storeGateway2))
+	flags2 = mergeFlags(flags2, map[string]string{
+		"-querier.store-gateway-addresses": strings.Join([]string{storeGateway2.NetworkGRPCEndpoint()}, ","),
+	})
+	queryFrontend2 := e2ecortex.NewQueryFrontend("query-frontend", mergeFlags(flags2, map[string]string{
+		"-frontend.scheduler-address": queryScheduler2.NetworkGRPCEndpoint(),
+	}), "")
+	require.NoError(t, s.Start(queryFrontend2))
+	querier2 := e2ecortex.NewQuerier("querier", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), mergeFlags(flags2, map[string]string{
+		"-querier.scheduler-address": queryScheduler2.NetworkGRPCEndpoint(),
+	}), "")
+	require.NoError(t, s.StartAndWaitReady(querier2))
+	require.NoError(t, distributor2.WaitSumMetrics(e2e.Equals(512), "cortex_ring_tokens_total"))
+	require.NoError(t, querier2.WaitSumMetrics(e2e.Equals(2*512), "cortex_ring_tokens_total"))
+	c2, err := e2ecortex.NewClient(distributor.HTTPEndpoint(), queryFrontend2.HTTPEndpoint(), "", "", "user-2")
+	require.NoError(t, err)
+
+	// INJECT DATA
+	series1Timestamp := time.Now()
+	series2Timestamp := series1Timestamp
+	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 := c.Push(series1)
+	require.NoError(t, err)
+	require.Equal(t, 200, res.StatusCode)
+	res, err = c2.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)
+	res, err = c2.Push(series2)
+	require.NoError(t, err)
+	require.Equal(t, 200, res.StatusCode)
+
+	var val model.Value
+	val, err = c.Query("series_1", series1Timestamp)
+	require.NoError(t, err)
+	require.Equal(t, expectedVector1, val.(model.Vector))
+	val, err = c.Query("series_2", series2Timestamp)
+	require.NoError(t, err)
+	require.Equal(t, expectedVector2, val.(model.Vector))
+
+	val, err = c2.Query("series_1", series1Timestamp)
+	require.NoError(t, err)
+	require.Equal(t, expectedVector1, val.(model.Vector))
+	val, err = c2.Query("series_2", series2Timestamp)
+	require.NoError(t, err)
+	require.Equal(t, expectedVector2, val.(model.Vector))
+
+	val, err = c.Query("sum(series_1)+sum(series_2)", series1Timestamp)
+	require.NoError(t, err)
+	val_distributed, err := c2.Query("sum(series_1)+sum(series_2)", series1Timestamp)
+	require.NoError(t, err)
+	require.Equal(t, val, val_distributed)
+}

pkg/api/handlers_test.go

@@ -2,6 +2,7 @@ package api
 
 import (
 	"encoding/json"
+	"github.com/cortexproject/cortex/pkg/querier"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -232,7 +233,7 @@ func TestBuildInfoAPI(t *testing.T) {
 			version.Version = tc.version
 			version.Branch = tc.branch
 			version.Revision = tc.revision
-			handler := NewQuerierHandler(cfg, nil, nil, nil, nil, nil, nil, &FakeLogger{}, false)
+			handler := NewQuerierHandler(cfg, querier.Config{}, nil, nil, nil, nil, nil, nil, &FakeLogger{}, false)
 			writer := httptest.NewRecorder()
 			req := httptest.NewRequest("GET", "/api/v1/status/buildinfo", nil)
 			req = req.WithContext(user.InjectOrgID(req.Context(), "test"))

pkg/cortex/modules.go

@@ -7,11 +7,9 @@ import (
 	"github.com/cortexproject/cortex/pkg/distributed_execution"
 	"github.com/cortexproject/cortex/pkg/ring/client"
 	"log/slog"
-	"net"
 	"net/http"
 	"runtime"
 	"runtime/debug"
-	"strconv"
 	"strings"
 	"time"
 
@@ -370,12 +368,6 @@ 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
@@ -451,12 +443,12 @@ func (t *Cortex) initQuerier() (serv services.Service, err error) {
 	}
 
 	if t.Cfg.Querier.DistributedExecEnabled {
-		querierPool := distributed_execution.NewQuerierPool(t.Cfg.QueryScheduler.GRPCClientConfig, prometheus.DefaultRegisterer, util_log.Logger)
+		querierPool := distributed_execution.NewQuerierPool(t.Cfg.Worker.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)
+	return querier_worker.NewQuerierWorker(t.Cfg.Worker, httpgrpc_server.NewServer(internalQuerierRouter), util_log.Logger, prometheus.DefaultRegisterer, t.Cfg.Querier.DistributedExecEnabled, queryResultCache)
 }
 
 func injectPool(next http.Handler, pool *client.Pool) http.Handler {

pkg/distributed_execution/distributed_exec_test.go

@@ -0,0 +1,29 @@
+package distributed_execution
+
+//func TestDistributedExec(t *testing.T) {
+//
+//	// non-fragmented query
+//	now := time.Now()
+//	plan := createTestLogicalPlan(t, now, now, 0, "sum() + sum()")
+//	oneFrag := []plan_fragments.Fragment{
+//		{
+//			Node:       plan.Root(),
+//			FragmentID: 0,
+//			ChildIDs:   []uint64{},
+//			IsRoot:     true,
+//		},
+//	}
+//
+//	// fragmented query
+//	plan2 := createTestLogicalPlan(t, now, now, 0, "sum() + sum()")
+//	dOptimizer := DistributedOptimizer{}
+//	node2, _, err := dOptimizer.Optimize(plan2.Root())
+//	require.NoError(t, err)
+//
+//	fragments, err := FragmentLogicalPlanNode(uint64(2), node2)
+//	require.NoError(t, err)
+//
+//	// get results
+//
+//	// compare if they are the same
+//}

pkg/distributed_execution/id.go

@@ -13,12 +13,7 @@ type fragmentMetadata struct {
 	isRoot        bool
 }
 
-func InjectFragmentMetaData(ctx context.Context, fragmentID uint64, queryID uint64, isRoot bool, childIDs []uint64, childAddr []string) context.Context {
-
-	childIDToAddr := make(map[uint64]string, len(childIDs))
-	for i, childID := range childIDs {
-		childIDToAddr[childID] = childAddr[i]
-	}
+func InjectFragmentMetaData(ctx context.Context, fragmentID uint64, queryID uint64, isRoot bool, childIDToAddr map[uint64]string) context.Context {
 
 	return context.WithValue(ctx, fragmentMetadataKey{}, fragmentMetadata{
 		queryID:       queryID,

pkg/distributed_execution/id_test.go

@@ -45,7 +45,12 @@ func TestFragmentMetadata(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			// injection
 			ctx := context.Background()
-			newCtx := InjectFragmentMetaData(ctx, tt.fragID, tt.queryID, tt.isRoot, tt.childIDs, tt.childAddr)
+
+			childIDToAddr := make(map[uint64]string)
+			for i, childID := range tt.childIDs {
+				childIDToAddr[childID] = tt.childAddr[i]
+			}
+			newCtx := InjectFragmentMetaData(ctx, tt.fragID, tt.queryID, tt.isRoot, childIDToAddr)
 
 			// extraction
 			isRoot, queryID, fragmentID, childAddrs, ok := ExtractFragmentMetaData(newCtx)

pkg/distributed_execution/querier_service_client.go

@@ -35,8 +35,10 @@ func NewQuerierPool(cfg grpcclient.Config, reg prometheus.Registerer, log log.Lo
 	}, []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.",
+		Namespace:   "cortex",
+		Name:        "cortex_querier_query_clients",
+		Help:        "TThe current number of clients connected to querier.",
+		ConstLabels: map[string]string{"client": "querier"},
 	})
 
 	poolConfig := client.PoolConfig{