refactor

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

fix e2e test

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

Author: yeya24

integration/query_frontend_test.go

@@ -7,10 +7,8 @@ import (
 	"crypto/x509"
 	"crypto/x509/pkix"
 	"fmt"
-	"github.com/thanos-io/thanos/pkg/pool"
 	"net/http"
 	"os"
-	"path"
 	"path/filepath"
 	"strconv"
 	"sync"
@@ -23,6 +21,7 @@ import (
 	"github.com/prometheus/prometheus/prompb"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/thanos-io/thanos/pkg/pool"
 
 	"github.com/cortexproject/cortex/integration/ca"
 	"github.com/cortexproject/cortex/integration/e2e"
@@ -452,7 +451,7 @@ func TestQueryFrontendNoRetryChunkPool(t *testing.T) {
 		"-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",
+		"-blocks-storage.bucket-store.max-chunk-pool-bytes": "1",
 	})
 
 	// Start dependencies.
@@ -492,24 +491,26 @@ func TestQueryFrontendNoRetryChunkPool(t *testing.T) {
 	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",
 	}), "")
-	queryFrontend := e2ecortex.NewQueryFrontendWithConfigFile("query-frontend", "", flags, "")
 	querier := e2ecortex.NewQuerier("querier", e2ecortex.RingStoreConsul, consul.NetworkHTTPEndpoint(), mergeFlags(flags, map[string]string{
 		"-blocks-storage.bucket-store.sync-interval": "5s",
 		"-querier.frontend-address":                  queryFrontend.NetworkGRPCEndpoint(),
 	}), "")
-	require.NoError(t, s.StartAndWaitReady(queryFrontend, querier, storeGateway))
+	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))
 
 	// Query back the series.
-	c, err = e2ecortex.NewClient("", path.Join(queryFrontend.HTTPEndpoint(), "/prometheus"), "", "", "user-1")
+	c, err = e2ecortex.NewClient("", queryFrontend.HTTPEndpoint(), "", "", "user-1")
 	require.NoError(t, err)
 
 	// We expect request to hit chunk pool exhaustion.

pkg/frontend/config.go

@@ -59,15 +59,15 @@ func InitFrontend(cfg CombinedFrontendConfig, limits v1.Limits, grpcListenPort i
 			cfg.FrontendV2.Port = grpcListenPort
 		}
 
-		fr, err := v2.NewFrontend(cfg.FrontendV2, log, reg)
-		return transport.AdaptGrpcRoundTripperToHTTPRoundTripper(fr, retry), nil, fr, err
+		fr, err := v2.NewFrontend(cfg.FrontendV2, log, reg, retry)
+		return transport.AdaptGrpcRoundTripperToHTTPRoundTripper(fr), nil, fr, err
 
 	default:
 		// No scheduler = use original frontend.
-		fr, err := v1.New(cfg.FrontendV1, limits, log, reg)
+		fr, err := v1.New(cfg.FrontendV1, limits, log, reg, retry)
 		if err != nil {
 			return nil, nil, nil, err
 		}
-		return transport.AdaptGrpcRoundTripperToHTTPRoundTripper(fr, retry), fr, nil, nil
+		return transport.AdaptGrpcRoundTripperToHTTPRoundTripper(fr), fr, nil, nil
 	}
 }

pkg/frontend/transport/retry.go

@@ -1,17 +1,15 @@
 package transport
 
 import (
-	"bytes"
 	"context"
 	"errors"
-	"io"
-	"net/http"
 	"strings"
 	"unsafe"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/thanos-io/thanos/pkg/pool"
+	"github.com/weaveworks/common/httpgrpc"
 
 	"github.com/cortexproject/cortex/pkg/querier/tripperware"
 )
@@ -33,7 +31,7 @@ func NewRetry(maxRetries int, reg prometheus.Registerer) *Retry {
 	}
 }
 
-func (r *Retry) Do(ctx context.Context, f func() (*http.Response, error)) (*http.Response, error) {
+func (r *Retry) Do(ctx context.Context, f func() (*httpgrpc.HTTPResponse, error)) (*httpgrpc.HTTPResponse, error) {
 	if r.maxRetries == 0 {
 		// Retries are disabled. Try only once.
 		return f()
@@ -43,7 +41,7 @@ func (r *Retry) Do(ctx context.Context, f func() (*http.Response, error)) (*http
 	defer func() { r.retriesCount.Observe(float64(tries)) }()
 
 	var (
-		resp *http.Response
+		resp *httpgrpc.HTTPResponse
 		err  error
 	)
 	for ; tries < r.maxRetries; tries++ {
@@ -54,8 +52,11 @@ func (r *Retry) Do(ctx context.Context, f func() (*http.Response, error)) (*http
 		resp, err = f()
 		if err != nil && !errors.Is(err, context.Canceled) {
 			continue // Retryable
-		} else if resp != nil && resp.StatusCode/100 == 5 {
-			body, err := tripperware.BodyBuffer(resp, nil)
+		} else if resp != nil && resp.Code/100 == 5 {
+			// This is not that efficient as we might decode the body multiple
+			// times. But error response should be too large so we should be fine.
+			// TODO: investigate ways to decode only once.
+			body, err := tripperware.BodyBufferFromHTTPGRPCResponse(resp, nil)
 			if err != nil {
 				return nil, err
 			}
@@ -64,22 +65,14 @@ func (r *Retry) Do(ctx context.Context, f func() (*http.Response, error)) (*http
 				continue
 			}
 
-			resp.Body = &buffer{buff: body, ReadCloser: io.NopCloser(bytes.NewReader(body))}
-			resp.ContentLength = int64(len(body))
 			return resp, nil
 		}
 		break
 	}
 	if err != nil {
 		return nil, err
 	}
-	// We always want to return decoded response body if possible.
-	body, err := tripperware.BodyBuffer(resp, nil)
-	if err != nil {
-		return nil, err
-	}
-	resp.Body = &buffer{buff: body, ReadCloser: io.NopCloser(bytes.NewReader(body))}
-	resp.ContentLength = int64(len(body))
+
 	return resp, err
 }
 

pkg/frontend/transport/retry_test.go

@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"github.com/thanos-io/thanos/pkg/pool"
 	"github.com/weaveworks/common/httpgrpc"
 	"go.uber.org/atomic"
 )
@@ -39,6 +40,7 @@ func TestNoRetryOnChunkPoolExhaustion(t *testing.T) {
 		if try > 1 {
 			return &httpgrpc.HTTPResponse{
 				Code: 500,
+				Body: []byte(pool.ErrPoolExhausted.Error()),
 			}, nil
 		}
 		return &httpgrpc.HTTPResponse{

pkg/frontend/transport/roundtripper.go

@@ -15,14 +15,13 @@ type GrpcRoundTripper interface {
 	RoundTripGRPC(context.Context, *httpgrpc.HTTPRequest) (*httpgrpc.HTTPResponse, error)
 }
 
-func AdaptGrpcRoundTripperToHTTPRoundTripper(r GrpcRoundTripper, retry *Retry) http.RoundTripper {
-	return &grpcRoundTripperAdapter{roundTripper: r, retry: retry}
+func AdaptGrpcRoundTripperToHTTPRoundTripper(r GrpcRoundTripper) http.RoundTripper {
+	return &grpcRoundTripperAdapter{roundTripper: r}
 }
 
 // This adapter wraps GrpcRoundTripper and converted it into http.RoundTripper
 type grpcRoundTripperAdapter struct {
 	roundTripper GrpcRoundTripper
-	retry        *Retry
 }
 
 type buffer struct {
@@ -40,16 +39,11 @@ func (a *grpcRoundTripperAdapter) RoundTrip(r *http.Request) (*http.Response, er
 		return nil, err
 	}
 
-	return a.retry.Do(r.Context(), func() (*http.Response, error) {
-		resp, err := a.roundTripper.RoundTripGRPC(r.Context(), req)
-		if err != nil {
-			return nil, err
-		}
-		return httpGRPCRespToHTTPResp(resp), nil
-	})
-}
+	resp, err := a.roundTripper.RoundTripGRPC(r.Context(), req)
+	if err != nil {
+		return nil, err
+	}
 
-func httpGRPCRespToHTTPResp(resp *httpgrpc.HTTPResponse) *http.Response {
 	httpResp := &http.Response{
 		StatusCode:    int(resp.Code),
 		Body:          &buffer{buff: resp.Body, ReadCloser: io.NopCloser(bytes.NewReader(resp.Body))},
@@ -59,5 +53,5 @@ func httpGRPCRespToHTTPResp(resp *httpgrpc.HTTPResponse) *http.Response {
 	for _, h := range resp.Headers {
 		httpResp.Header[h.Key] = h.Values
 	}
-	return httpResp
+	return httpResp, nil
 }

pkg/frontend/v1/frontend.go

@@ -94,11 +94,12 @@ type request struct {
 }
 
 // New creates a new frontend. Frontend implements service, and must be started and stopped.
-func New(cfg Config, limits Limits, log log.Logger, registerer prometheus.Registerer) (*Frontend, error) {
+func New(cfg Config, limits Limits, log log.Logger, registerer prometheus.Registerer, retry *transport.Retry) (*Frontend, error) {
 	f := &Frontend{
 		cfg:    cfg,
 		log:    log,
 		limits: limits,
+		retry:  retry,
 		queueLength: promauto.With(registerer).NewGaugeVec(prometheus.GaugeOpts{
 			Name: "cortex_query_frontend_queue_length",
 			Help: "Number of queries in the queue.",
@@ -175,31 +176,33 @@ func (f *Frontend) RoundTripGRPC(ctx context.Context, req *httpgrpc.HTTPRequest)
 		}
 	}
 
-	request := request{
-		request:     req,
-		originalCtx: ctx,
+	return f.retry.Do(ctx, func() (*httpgrpc.HTTPResponse, error) {
+		request := request{
+			request:     req,
+			originalCtx: ctx,
 
-		// Buffer of 1 to ensure response can be written by the server side
-		// of the Process stream, even if this goroutine goes away due to
-		// client context cancellation.
-		err:      make(chan error, 1),
-		response: make(chan *httpgrpc.HTTPResponse, 1),
-	}
+			// Buffer of 1 to ensure response can be written by the server side
+			// of the Process stream, even if this goroutine goes away due to
+			// client context cancellation.
+			err:      make(chan error, 1),
+			response: make(chan *httpgrpc.HTTPResponse, 1),
+		}
 
-	if err := f.queueRequest(ctx, &request); err != nil {
-		return nil, err
-	}
+		if err := f.queueRequest(ctx, &request); err != nil {
+			return nil, err
+		}
 
-	select {
-	case <-ctx.Done():
-		return nil, ctx.Err()
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
 
-	case resp := <-request.response:
-		return resp, nil
+		case resp := <-request.response:
+			return resp, nil
 
-	case err := <-request.err:
-		return nil, err
-	}
+		case err := <-request.err:
+			return nil, err
+		}
+	})
 }
 
 // Process allows backends to pull requests from the frontend.