refactor query frontend to return prometheus error response

Signed-off-by: Ben Ye <[email protected]>

Author: yeya24

integration/query_frontend_test.go

@@ -530,3 +530,91 @@ func TestQueryFrontendNoRetryChunkPool(t *testing.T) {
 	// We shouldn't be able to see any retries.
 	require.NoError(t, queryFrontend.WaitSumMetricsWithOptions(e2e.Equals(0), []string{"cortex_query_frontend_retries"}, e2e.WaitMissingMetrics))
 }
+
+func TestQueryFrontendErrorResponseFormat(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(),
+		"-blocks-storage.bucket-store.max-chunk-pool-bytes": "1",
+	})
+
+	// Start dependencies.
+	consul := e2edb.NewConsul()
+	minio := e2edb.NewMinio(9000, flags["-blocks-storage.s3.bucket-name"])
+	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))
+
+	// Wait until the distributor has updated the ring.
+	require.NoError(t, distributor.WaitSumMetrics(e2e.Equals(512), "cortex_ring_tokens_total"))
+
+	// Push some series to Cortex.
+	c, err := e2ecortex.NewClient(distributor.HTTPEndpoint(), "", "", "", "user-1")
+	require.NoError(t, err)
+
+	seriesTimestamp := time.Now()
+	series2Timestamp := seriesTimestamp.Add(blockRangePeriod * 2)
+	series1, _ := generateSeries("series_1", seriesTimestamp, prompb.Label{Name: "job", Value: "test"})
+	series2, _ := generateSeries("series_2", series2Timestamp, 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)
+
+	// Wait until the TSDB head is compacted and shipped to the storage.
+	// The shipped block contains the 1st series, while the 2ns series is in the head.
+	require.NoError(t, ingester.WaitSumMetrics(e2e.Equals(1), "cortex_ingester_shipper_uploads_total"))
+	require.NoError(t, ingester.WaitSumMetrics(e2e.Equals(2), "cortex_ingester_memory_series_created_total"))
+	require.NoError(t, ingester.WaitSumMetrics(e2e.Equals(1), "cortex_ingester_memory_series_removed_total"))
+	require.NoError(t, ingester.WaitSumMetrics(e2e.Equals(1), "cortex_ingester_memory_series"))
+
+	queryFrontend := e2ecortex.NewQueryFrontendWithConfigFile("query-frontend", "", flags, "")
+	require.NoError(t, s.Start(queryFrontend))
+
+	// Start the querier and store-gateway, and configure them to frequently sync blocks fast enough to trigger consistency check.
+	storeGateway := e2ecortex.NewStoreGateway("store-gateway", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), mergeFlags(flags, map[string]string{
+		"-blocks-storage.bucket-store.sync-interval": "5s",
+	}), "")
+	querier := e2ecortex.NewQuerier("querier", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), mergeFlags(flags, map[string]string{
+		"-blocks-storage.bucket-store.sync-interval": "1s",
+		"-querier.frontend-address":                  queryFrontend.NetworkGRPCEndpoint(),
+	}), "")
+	require.NoError(t, s.StartAndWaitReady(querier, storeGateway))
+
+	// Wait until the querier and store-gateway have updated the ring, and wait until the blocks are old enough for consistency check
+	require.NoError(t, querier.WaitSumMetrics(e2e.Equals(512*2), "cortex_ring_tokens_total"))
+	require.NoError(t, storeGateway.WaitSumMetrics(e2e.Equals(512), "cortex_ring_tokens_total"))
+	require.NoError(t, querier.WaitSumMetricsWithOptions(e2e.GreaterOrEqual(4), []string{"cortex_querier_blocks_scan_duration_seconds"}, e2e.WithMetricCount))
+
+	// Sleep 3 * bucket sync interval to make sure consistency checker
+	// doesn't consider block is uploaded recently.
+	time.Sleep(3 * time.Second)
+
+	// Query back the series.
+	c, err = e2ecortex.NewClient("", queryFrontend.HTTPEndpoint(), "", "", "user-1")
+	require.NoError(t, err)
+
+	// We expect request to hit chunk pool exhaustion.
+	resp, body, err := c.QueryRaw(`{job="test"}`, series2Timestamp)
+	require.NoError(t, err)
+	require.Equal(t, http.StatusInternalServerError, resp.StatusCode)
+	require.Contains(t, string(body), pool.ErrPoolExhausted.Error())
+	// We shouldn't be able to see any retries.
+	require.NoError(t, queryFrontend.WaitSumMetricsWithOptions(e2e.Equals(0), []string{"cortex_query_frontend_retries"}, e2e.WaitMissingMetrics))
+}

pkg/frontend/transport/handler.go

@@ -16,16 +16,18 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
+	v1 "github.com/prometheus/client_golang/api/prometheus/v1"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/weaveworks/common/httpgrpc"
-	"github.com/weaveworks/common/httpgrpc/server"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
 	querier_stats "github.com/cortexproject/cortex/pkg/querier/stats"
 	"github.com/cortexproject/cortex/pkg/querier/tripperware"
 	"github.com/cortexproject/cortex/pkg/tenant"
 	"github.com/cortexproject/cortex/pkg/util"
+	util_api "github.com/cortexproject/cortex/pkg/util/api"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
 )
 
@@ -239,8 +241,9 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		writeServiceTimingHeader(queryResponseTime, hs, stats)
 	}
 
+	logger := util_log.WithContext(r.Context(), f.log)
 	if err != nil {
-		writeError(w, err, hs)
+		writeError(logger, w, err, hs)
 		return
 	}
 
@@ -252,7 +255,7 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	// log copy response body error so that we will know even though success response code returned
 	bytesCopied, err := io.Copy(w, resp.Body)
 	if err != nil && !errors.Is(err, syscall.EPIPE) {
-		level.Error(util_log.WithContext(r.Context(), f.log)).Log("msg", "write response body error", "bytesCopied", bytesCopied, "err", err)
+		level.Error(logger).Log("msg", "write response body error", "bytesCopied", bytesCopied, "err", err)
 	}
 }
 
@@ -441,7 +444,7 @@ func formatQueryString(queryString url.Values) (fields []interface{}) {
 	return fields
 }
 
-func writeError(w http.ResponseWriter, err error, additionalHeaders http.Header) {
+func writeError(logger log.Logger, w http.ResponseWriter, err error, additionalHeaders http.Header) {
 	switch err {
 	case context.Canceled:
 		err = errCanceled
@@ -453,20 +456,35 @@ func writeError(w http.ResponseWriter, err error, additionalHeaders http.Header)
 		}
 	}
 
+	headers := w.Header()
+	for k, values := range additionalHeaders {
+		for _, value := range values {
+			headers.Set(k, value)
+		}
+	}
 	resp, ok := httpgrpc.HTTPResponseFromError(err)
 	if ok {
-		for k, values := range additionalHeaders {
-			resp.Headers = append(resp.Headers, &httpgrpc.Header{Key: k, Values: values})
+		code := int(resp.Code)
+		var errTyp v1.ErrorType
+		switch resp.Code {
+		case http.StatusBadRequest, http.StatusRequestEntityTooLarge:
+			errTyp = v1.ErrBadData
+		case StatusClientClosedRequest:
+			errTyp = v1.ErrCanceled
+		case http.StatusGatewayTimeout:
+			errTyp = v1.ErrTimeout
+		case http.StatusUnprocessableEntity:
+			errTyp = v1.ErrExec
+		case int32(codes.PermissionDenied):
+			// Convert gRPC status code to HTTP status code.
+			code = http.StatusUnprocessableEntity
+			errTyp = v1.ErrBadData
+		default:
+			errTyp = v1.ErrServer
 		}
-		_ = server.WriteResponse(w, resp)
+		util_api.RespondError(logger, w, errTyp, string(resp.Body), code)
 	} else {
-		headers := w.Header()
-		for k, values := range additionalHeaders {
-			for _, value := range values {
-				headers.Set(k, value)
-			}
-		}
-		http.Error(w, err.Error(), http.StatusInternalServerError)
+		util_api.RespondError(logger, w, v1.ErrServer, err.Error(), http.StatusInternalServerError)
 	}
 }
 

pkg/frontend/transport/handler_test.go

@@ -21,6 +21,7 @@ import (
 	"github.com/weaveworks/common/user"
 
 	querier_stats "github.com/cortexproject/cortex/pkg/querier/stats"
+	util_log "github.com/cortexproject/cortex/pkg/util/log"
 )
 
 type roundTripperFunc func(*http.Request) (*http.Response, error)
@@ -46,7 +47,7 @@ func TestWriteError(t *testing.T) {
 	} {
 		t.Run(test.err.Error(), func(t *testing.T) {
 			w := httptest.NewRecorder()
-			writeError(w, test.err, test.additionalHeaders)
+			writeError(util_log.Logger, w, test.err, test.additionalHeaders)
 			require.Equal(t, test.status, w.Result().StatusCode)
 			expectedAdditionalHeaders := test.additionalHeaders
 			if expectedAdditionalHeaders != nil {

pkg/querier/tripperware/instantquery/limits.go

@@ -48,8 +48,7 @@ func (l limitsMiddleware) Do(ctx context.Context, r tripperware.Request) (trippe
 	if maxQueryLength := validation.SmallestPositiveNonZeroDurationPerTenant(tenantIDs, l.MaxQueryLength); maxQueryLength > 0 {
 		expr, err := parser.ParseExpr(r.GetQuery())
 		if err != nil {
-			// Let Querier propagates the parsing error.
-			return l.next.Do(ctx, r)
+			return nil, httpgrpc.Errorf(http.StatusBadRequest, err.Error())
 		}
 
 		// Enforce query length across all selectors in the query.

pkg/querier/tripperware/instantquery/limits_test.go

@@ -2,12 +2,15 @@ package instantquery
 
 import (
 	"context"
+	"net/http"
 	"testing"
 	"time"
 
+	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
+	"github.com/weaveworks/common/httpgrpc"
 	"github.com/weaveworks/common/user"
 
 	"github.com/cortexproject/cortex/pkg/querier/tripperware"
@@ -20,6 +23,9 @@ func TestLimitsMiddleware_MaxQueryLength(t *testing.T) {
 		thirtyDays = 30 * 24 * time.Hour
 	)
 
+	wrongQuery := `up[`
+	_, parserErr := parser.ParseExpr(wrongQuery)
+
 	tests := map[string]struct {
 		maxQueryLength time.Duration
 		query          string
@@ -31,6 +37,7 @@ func TestLimitsMiddleware_MaxQueryLength(t *testing.T) {
 		"even though failed to parse expression, should return no error since request will pass to next middleware": {
 			query:          `up[`,
 			maxQueryLength: thirtyDays,
+			expectedErr:    httpgrpc.Errorf(http.StatusBadRequest, parserErr.Error()).Error(),
 		},
 		"should succeed on a query not exceeding time range": {
 			query:          `up`,

pkg/querier/tripperware/queryrange/limits.go

@@ -84,8 +84,7 @@ func (l limitsMiddleware) Do(ctx context.Context, r tripperware.Request) (trippe
 
 		expr, err := parser.ParseExpr(r.GetQuery())
 		if err != nil {
-			// Let Querier propagates the parsing error.
-			return l.next.Do(ctx, r)
+			return nil, httpgrpc.Errorf(http.StatusBadRequest, err.Error())
 		}
 
 		// Enforce query length across all selectors in the query.

pkg/querier/tripperware/queryrange/limits_test.go

@@ -2,6 +2,9 @@ package queryrange
 
 import (
 	"context"
+	"github.com/prometheus/prometheus/promql/parser"
+	"github.com/weaveworks/common/httpgrpc"
+	"net/http"
 	"testing"
 	"time"
 
@@ -115,6 +118,9 @@ func TestLimitsMiddleware_MaxQueryLength(t *testing.T) {
 
 	now := time.Now()
 
+	wrongQuery := `up[`
+	_, parserErr := parser.ParseExpr(wrongQuery)
+
 	tests := map[string]struct {
 		maxQueryLength time.Duration
 		query          string
@@ -132,6 +138,7 @@ func TestLimitsMiddleware_MaxQueryLength(t *testing.T) {
 			reqStartTime:   now.Add(-time.Hour),
 			reqEndTime:     now,
 			maxQueryLength: thirtyDays,
+			expectedErr:    httpgrpc.Errorf(http.StatusBadRequest, parserErr.Error()).Error(),
 		},
 		"should succeed on a query on short time range, ending now": {
 			maxQueryLength: thirtyDays,

pkg/querier/tripperware/queryrange/split_by_interval.go

@@ -47,12 +47,7 @@ func (s splitByInterval) Do(ctx context.Context, r tripperware.Request) (tripper
 	// to line up the boundaries with step.
 	reqs, err := splitQuery(r, s.interval(r))
 	if err != nil {
-		// If the query itself is bad, we don't return error but send the query
-		// to querier to return the expected error message. This is not very efficient
-		// but should be okay for now.
-		// TODO(yeya24): query frontend can reuse the Prometheus API handler and return
-		// expected error message locally without passing it to the querier through network.
-		return s.next.Do(ctx, r)
+		return nil, httpgrpc.Errorf(http.StatusBadRequest, err.Error())
 	}
 	s.splitByCounter.Add(float64(len(reqs)))
 

pkg/querier/tripperware/roundtrip.go

@@ -143,7 +143,7 @@ func NewQueryTripperware(
 				tenantIDs, err := tenant.TenantIDs(r.Context())
 				// This should never happen anyways because we have auth middleware before this.
 				if err != nil {
-					return nil, err
+					return nil, httpgrpc.Errorf(http.StatusBadRequest, err.Error())
 				}
 				now := time.Now()
 				userStr := tenant.JoinTenantIDs(tenantIDs)
@@ -162,8 +162,7 @@ func NewQueryTripperware(
 
 					expr, err := parser.ParseExpr(query)
 					if err != nil {
-						// If query is invalid, no need to go through tripperwares for further splitting.
-						return next.RoundTrip(r)
+						return nil, httpgrpc.Errorf(http.StatusBadRequest, err.Error())
 					}
 
 					reqStats := stats.FromContext(r.Context())
@@ -189,7 +188,10 @@ func NewQueryTripperware(
 						return next.RoundTrip(r)
 					}
 					analysis, err := queryAnalyzer.Analyze(query)
-					if err != nil || !analysis.IsShardable() {
+					if err != nil {
+						return nil, httpgrpc.Errorf(http.StatusBadRequest, err.Error())
+					}
+					if !analysis.IsShardable() {
 						return next.RoundTrip(r)
 					}
 					return instantQuery.RoundTrip(r)

pkg/querier/tripperware/shard_by.go

@@ -55,6 +55,7 @@ func (s shardBy) Do(ctx context.Context, r Request) (Response, error) {
 	analysis, err := s.analyzer.Analyze(r.GetQuery())
 	if err != nil {
 		level.Warn(logger).Log("msg", "error analyzing query", "q", r.GetQuery(), "err", err)
+		return nil, httpgrpc.Errorf(http.StatusBadRequest, err.Error())
 	}
 
 	stats.AddExtraFields(
@@ -63,7 +64,7 @@ func (s shardBy) Do(ctx context.Context, r Request) (Response, error) {
 		"shard_by.sharding_labels", analysis.ShardingLabels(),
 	)
 
-	if err != nil || !analysis.IsShardable() {
+	if !analysis.IsShardable() {
 		return s.next.Do(ctx, r)
 	}