fix querier serve streaming errors but next() currently takes too long

Author: rubywtl

pkg/api/handlers.go

@@ -169,6 +169,7 @@ func NewQuerierHandler(
 	metadataQuerier querier.MetadataQuerier,
 	reg prometheus.Registerer,
 	logger log.Logger,
+	distributedExecEnabled bool,
 ) http.Handler {
 	// Prometheus histograms for requests to the querier.
 	querierRequestDuration := promauto.With(reg).NewHistogramVec(prometheus.HistogramOpts{
@@ -282,7 +283,7 @@ func NewQuerierHandler(
 	legacyPromRouter := route.New().WithPrefix(path.Join(legacyPrefix, "/api/v1"))
 	api.Register(legacyPromRouter)
 
-	queryAPI := queryapi.NewQueryAPI(engine, queryResultCache, translateSampleAndChunkQueryable, statsRenderer, logger, codecs, corsOrigin)
+	queryAPI := queryapi.NewQueryAPI(engine, queryResultCache, translateSampleAndChunkQueryable, statsRenderer, logger, codecs, corsOrigin, distributedExecEnabled)
 
 	// TODO(gotjosh): This custom handler is temporary until we're able to vendor the changes in:
 	// https://github.com/prometheus/prometheus/pull/7125/files

pkg/api/handlers_test.go

@@ -232,7 +232,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, &FakeLogger{})
+			handler := NewQuerierHandler(cfg, 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/api/queryapi/query_api.go

@@ -26,14 +26,15 @@ import (
 )
 
 type QueryAPI struct {
-	queryable        storage.SampleAndChunkQueryable
-	queryEngine      engine.QueryEngine
-	queryResultCache *distributed_execution.QueryResultCache
-	now              func() time.Time
-	statsRenderer    v1.StatsRenderer
-	logger           log.Logger
-	codecs           []v1.Codec
-	CORSOrigin       *regexp.Regexp
+	queryable              storage.SampleAndChunkQueryable
+	queryEngine            engine.QueryEngine
+	queryResultCache       *distributed_execution.QueryResultCache
+	now                    func() time.Time
+	statsRenderer          v1.StatsRenderer
+	logger                 log.Logger
+	codecs                 []v1.Codec
+	CORSOrigin             *regexp.Regexp
+	distributedExecEnabled bool
 }
 
 func NewQueryAPI(
@@ -44,16 +45,18 @@ func NewQueryAPI(
 	logger log.Logger,
 	codecs []v1.Codec,
 	CORSOrigin *regexp.Regexp,
+	distributedExecEnabled bool,
 ) *QueryAPI {
 	return &QueryAPI{
-		queryEngine:      qe,
-		queryResultCache: queryResultCache,
-		queryable:        q,
-		statsRenderer:    statsRenderer,
-		logger:           logger,
-		codecs:           codecs,
-		CORSOrigin:       CORSOrigin,
-		now:              time.Now,
+		queryEngine:            qe,
+		queryResultCache:       queryResultCache,
+		queryable:              q,
+		statsRenderer:          statsRenderer,
+		logger:                 logger,
+		codecs:                 codecs,
+		CORSOrigin:             CORSOrigin,
+		now:                    time.Now,
+		distributedExecEnabled: distributedExecEnabled,
 	}
 }
 
@@ -139,11 +142,12 @@ func (q *QueryAPI) RangeQueryHandler(r *http.Request) (result apiFuncResult) {
 
 	ctx = httputil.ContextFromRequest(ctx, r)
 
-	// TODO: if distributed exec enabled
-	isRoot, queryID, fragmentID, _, _ := distributed_execution.ExtractFragmentMetaData(ctx)
-	if !isRoot {
-		key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
-		q.queryResultCache.InitWriting(*key)
+	if q.distributedExecEnabled {
+		isRoot, queryID, fragmentID, _, _ := distributed_execution.ExtractFragmentMetaData(ctx)
+		if !isRoot {
+			key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
+			q.queryResultCache.InitWriting(*key)
+		}
 	}
 
 	res := qry.Exec(ctx)
@@ -188,11 +192,14 @@ func (q *QueryAPI) InstantQueryHandler(r *http.Request) (result apiFuncResult) {
 	ctx = engine.AddEngineTypeToContext(ctx, r)
 	ctx = querier.AddBlockStoreTypeToContext(ctx, r.Header.Get(querier.BlockStoreTypeHeader))
 
-	// TODO: if distributed exec enabled
-	isRoot, queryID, fragmentID, _, _ := distributed_execution.ExtractFragmentMetaData(ctx)
-	if !isRoot {
-		key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
-		q.queryResultCache.InitWriting(*key)
+	var isRoot bool
+	var queryID, fragmentID uint64
+	if q.distributedExecEnabled {
+		isRoot, queryID, fragmentID, _, _ = distributed_execution.ExtractFragmentMetaData(ctx)
+		if !isRoot {
+			key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
+			q.queryResultCache.InitWriting(*key)
+		}
 	}
 
 	var qry promql.Query
@@ -202,9 +209,11 @@ func (q *QueryAPI) InstantQueryHandler(r *http.Request) (result apiFuncResult) {
 	if len(byteLP) != 0 {
 		logicalPlan, err := distributed_execution.Unmarshal(byteLP)
 		if err != nil {
-			if !isRoot {
-				key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
-				q.queryResultCache.SetError(*key)
+			if q.distributedExecEnabled {
+				if !isRoot {
+					key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
+					q.queryResultCache.SetError(*key)
+				}
 			}
 			return apiFuncResult{nil, &apiError{errorInternal, fmt.Errorf("invalid logical plan: %v", err)}, nil, nil}
 		}
@@ -264,17 +273,18 @@ func (q *QueryAPI) Wrap(f apiFunc) http.HandlerFunc {
 		}
 
 		if result.data != nil {
-			// TODO: if distributed exec enabled
 			ctx := httputil.ContextFromRequest(r.Context(), r)
-			isRoot, queryID, fragmentID, _, _ := distributed_execution.ExtractFragmentMetaData(ctx)
-			if !isRoot {
-				key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
-				result := distributed_execution.FragmentResult{
-					Data:       result.data,
-					Expiration: time.Now().Add(time.Hour),
+			if q.distributedExecEnabled {
+				isRoot, queryID, fragmentID, _, _ := distributed_execution.ExtractFragmentMetaData(ctx)
+				if !isRoot {
+					key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
+					result := distributed_execution.FragmentResult{
+						Data:       result.data,
+						Expiration: time.Now().Add(time.Hour),
+					}
+					q.queryResultCache.SetComplete(*key, result)
+					return
 				}
-				q.queryResultCache.SetComplete(*key, result)
-				return
 			}
 			q.respond(w, r, result.data, result.warnings, r.FormValue("query"))
 			return

pkg/api/queryapi/query_api_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/cortexproject/cortex/pkg/engine/distributed_execution"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -183,7 +184,7 @@ func Test_CustomAPI(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			c := NewQueryAPI(engine, mockQueryable, querier.StatsRenderer, log.NewNopLogger(), []v1.Codec{v1.JSONCodec{}}, regexp.MustCompile(".*"))
+			c := NewQueryAPI(engine, nil, mockQueryable, querier.StatsRenderer, log.NewNopLogger(), []v1.Codec{v1.JSONCodec{}}, regexp.MustCompile(".*"), false)
 
 			router := mux.NewRouter()
 			router.Path("/api/v1/query").Methods("POST").Handler(c.Wrap(c.InstantQueryHandler))
@@ -244,7 +245,7 @@ func Test_InvalidCodec(t *testing.T) {
 		},
 	}
 
-	queryAPI := NewQueryAPI(engine, mockQueryable, querier.StatsRenderer, log.NewNopLogger(), []v1.Codec{&mockCodec{}}, regexp.MustCompile(".*"))
+	queryAPI := NewQueryAPI(engine, nil, mockQueryable, querier.StatsRenderer, log.NewNopLogger(), []v1.Codec{&mockCodec{}}, regexp.MustCompile(".*"), false)
 	router := mux.NewRouter()
 	router.Path("/api/v1/query").Methods("POST").Handler(queryAPI.Wrap(queryAPI.InstantQueryHandler))
 
@@ -285,7 +286,7 @@ func Test_CustomAPI_StatsRenderer(t *testing.T) {
 		},
 	}
 
-	queryAPI := NewQueryAPI(engine, mockQueryable, querier.StatsRenderer, log.NewNopLogger(), []v1.Codec{v1.JSONCodec{}}, regexp.MustCompile(".*"))
+	queryAPI := NewQueryAPI(engine, nil, mockQueryable, querier.StatsRenderer, log.NewNopLogger(), []v1.Codec{v1.JSONCodec{}}, regexp.MustCompile(".*"), false)
 
 	router := mux.NewRouter()
 	router.Path("/api/v1/query_range").Methods("POST").Handler(queryAPI.Wrap(queryAPI.RangeQueryHandler))
@@ -441,7 +442,7 @@ func Test_Logicalplan_Requests(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			c := NewQueryAPI(engine, mockQueryable, querier.StatsRenderer, log.NewNopLogger(), []v1.Codec{v1.JSONCodec{}}, regexp.MustCompile(".*"))
+			c := NewQueryAPI(engine, &distributed_execution.QueryResultCache{}, mockQueryable, querier.StatsRenderer, log.NewNopLogger(), []v1.Codec{v1.JSONCodec{}}, regexp.MustCompile(".*"), false)
 			router := mux.NewRouter()
 			router.Path("/api/v1/query").Methods("POST").Handler(c.Wrap(c.InstantQueryHandler))
 			router.Path("/api/v1/query_range").Methods("POST").Handler(c.Wrap(c.RangeQueryHandler))

pkg/cortex/modules.go

@@ -407,6 +407,7 @@ func (t *Cortex) initQuerier() (serv services.Service, err error) {
 		t.MetadataQuerier,
 		prometheus.DefaultRegisterer,
 		util_log.Logger,
+		t.Cfg.Querier.DistributedExecEnabled,
 	)
 
 	// If the querier is running standalone without the query-frontend or query-scheduler, we must register it's internal

pkg/engine/distributed_execution/querier_service_server.go

@@ -37,7 +37,7 @@ func (s *QuerierServer) Series(req *querierpb.SeriesRequest, srv querierpb.Queri
 			return fmt.Errorf("fragment not found: %v", key)
 		}
 
-		batchSize := int(req.Batchsize)
+		batchSize := int(req.Batchsize) // TODO: remove this
 		if batchSize <= 0 {
 			batchSize = BATCHSIZE
 		}
@@ -47,7 +47,8 @@ func (s *QuerierServer) Series(req *querierpb.SeriesRequest, srv querierpb.Queri
 			fragmentResult := result.Data.(FragmentResult)
 			v1ResultData := fragmentResult.Data.(*v1.QueryData)
 
-			if v1ResultData.ResultType == parser.ValueTypeMatrix {
+			switch v1ResultData.ResultType {
+			case parser.ValueTypeMatrix:
 				series := v1ResultData.Result.(promql.Matrix)
 				for i := 0; i < len(series); i += batchSize {
 					end := i + batchSize
@@ -72,7 +73,7 @@ func (s *QuerierServer) Series(req *querierpb.SeriesRequest, srv querierpb.Queri
 				}
 				return nil
 
-			} else if v1ResultData.ResultType == parser.ValueTypeVector {
+			case parser.ValueTypeVector:
 				series := v1ResultData.Result.(promql.Vector)
 				for i := 0; i < len(series); i += batchSize {
 					end := i + batchSize
@@ -110,61 +111,46 @@ func (s *QuerierServer) Series(req *querierpb.SeriesRequest, srv querierpb.Queri
 func (s *QuerierServer) Next(req *querierpb.NextRequest, srv querierpb.Querier_NextServer) error {
 	key := MakeFragmentKey(req.QueryID, req.FragmentID)
 
+	batchSize := int(req.Batchsize)
+	if batchSize <= 0 {
+		batchSize = BATCHSIZE
+	}
+
 	for {
-		// TODO: maybe add a timeout thing here too
 		result, ok := s.queryResultCache.Get(*key)
 		if !ok {
 			return fmt.Errorf("fragment not found: %v", key)
 		}
 
 		switch result.Status {
 		case StatusDone:
-			batchSize := int(req.Batchsize)
-			if batchSize <= 0 {
-				batchSize = BATCHSIZE
-			}
-
 			fragmentResult := result.Data.(FragmentResult)
 			v1ResultData := fragmentResult.Data.(*v1.QueryData)
 
-			if v1ResultData.ResultType == parser.ValueTypeMatrix {
+			switch v1ResultData.ResultType {
+			case parser.ValueTypeMatrix:
 				matrix := v1ResultData.Result.(promql.Matrix)
 
-				for i := 0; i < len(matrix); i += batchSize {
-					end := i + batchSize
-					if end > len(matrix) {
-						end = len(matrix)
-					}
-
-					batch := &querierpb.StepVectorBatch{
-						StepVectors: make([]*querierpb.StepVector, 0, end-i),
-					}
+				numTimeSteps := matrix.TotalSamples()
 
-					for _, v := range (matrix)[i:end] {
-						var floats []float64
-						var histograms []*histogram.FloatHistogram
+				for timeStep := 0; timeStep < numTimeSteps; timeStep++ {
 
-						for _, f := range v.Floats {
-							floats = append(floats, f.F)
+					for i, series := range matrix {
+						batch := &querierpb.StepVectorBatch{
+							StepVectors: make([]*querierpb.StepVector, 0, len(matrix)),
 						}
-
-						for _, h := range v.Histograms {
-							histograms = append(histograms, h.H)
+						vector, err := s.createVectorForTimestep(&series, timeStep, uint64(i))
+						if err != nil {
+							return err
 						}
-
-						protoVector := &querierpb.StepVector{
-							T:          0,
-							Samples:    floats,
-							Histograms: FloatHistogramsToFloatHistogramProto(histograms),
+						batch.StepVectors = append(batch.StepVectors, vector)
+						if err := srv.Send(batch); err != nil {
+							return fmt.Errorf("error sending batch: %w", err)
 						}
-						batch.StepVectors = append(batch.StepVectors, protoVector)
 					}
-					if err := srv.Send(batch); err != nil {
-						return err
-					}
-				}
 
-			} else if v1ResultData.ResultType == parser.ValueTypeVector {
+				}
+			case parser.ValueTypeVector:
 				vector := v1ResultData.Result.(promql.Vector)
 
 				for i := 0; i < len(vector); i += batchSize {
@@ -190,23 +176,52 @@ func (s *QuerierServer) Next(req *querierpb.NextRequest, srv querierpb.Querier_N
 						return err
 					}
 				}
-			}
-
-			return nil
 
+			default:
+				return fmt.Errorf("unsupported result type: %v", v1ResultData.ResultType)
+			}
 		case StatusError:
 			return fmt.Errorf("fragment processing failed")
-
 		case StatusWriting:
 			time.Sleep(WritingTimeout)
 			continue
 		}
 	}
 }
 
+func (s *QuerierServer) createVectorForTimestep(series *promql.Series, timeStep int, sampleID uint64) (*querierpb.StepVector, error) {
+	var samples []float64
+	var sampleIDs []uint64
+	var histograms []*histogram.FloatHistogram
+	var histogramIDs []uint64
+	var timestamp int64
+
+	if timeStep < len(series.Floats) {
+		point := series.Floats[timeStep]
+		timestamp = point.T
+		samples = append(samples, point.F)
+		sampleIDs = append(sampleIDs, sampleID)
+	}
+
+	if timeStep < len(series.Histograms) {
+		point := series.Histograms[timeStep]
+		timestamp = point.T
+		histograms = append(histograms, point.H)
+		histogramIDs = append(histogramIDs, uint64(timeStep))
+	}
+
+	return &querierpb.StepVector{
+		T:             timestamp,
+		Sample_IDs:    sampleIDs,
+		Samples:       samples,
+		Histogram_IDs: histogramIDs,
+		Histograms:    FloatHistogramsToFloatHistogramProto(histograms),
+	}, nil
+}
+
 func FloatHistogramsToFloatHistogramProto(histograms []*histogram.FloatHistogram) []querierpb.Histogram {
 	if histograms == nil {
-		return nil
+		return []querierpb.Histogram{}
 	}
 
 	protoHistograms := make([]querierpb.Histogram, 0, len(histograms))