Fail fast if context errors out

Signed-off-by: 🌲 Harry 🌊 John 🏔 <[email protected]>
cortexproject · Sep 14, 2023 · 26c024e · 26c024e
1 parent b514ae6
commit 26c024e
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 10 deletions.
diff --git a/pkg/frontend/transport/retry.go b/pkg/frontend/transport/retry.go
@@ -25,7 +25,7 @@ func NewRetry(maxRetries int, reg prometheus.Registerer) *Retry {
 	}
 }
 
-func (r *Retry) Do(f func() (*httpgrpc.HTTPResponse, error)) (*httpgrpc.HTTPResponse, 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()
@@ -39,8 +39,11 @@ func (r *Retry) Do(f func() (*httpgrpc.HTTPResponse, error)) (*httpgrpc.HTTPResp
 		err  error
 	)
 	for ; tries < r.maxRetries; tries++ {
-		resp, err = f()
+		if ctx.Err() != nil {
+			return nil, ctx.Err()
+		}
 
+		resp, err = f()
 		if err != nil && err != context.Canceled {
 			continue // Retryable
 		} else if resp != nil && resp.Code/100 == 5 {

diff --git a/pkg/frontend/transport/retry_test.go b/pkg/frontend/transport/retry_test.go
@@ -1,6 +1,7 @@
 package transport
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -11,18 +12,18 @@ import (
 func TestRetry(t *testing.T) {
 	tries := atomic.NewInt64(3)
 	r := NewRetry(3, nil)
-
-	res, err := r.Do(func() (*httpgrpc.HTTPResponse, error) {
+	ctx := context.Background()
+	res, err := r.Do(ctx, func() (*httpgrpc.HTTPResponse, error) {
 		try := tries.Dec()
 		if try > 1 {
 			return &httpgrpc.HTTPResponse{
 				Code: 500,
 			}, nil
-		} else {
-			return &httpgrpc.HTTPResponse{
-				Code: 200,
-			}, nil
 		}
+		return &httpgrpc.HTTPResponse{
+			Code: 200,
+		}, nil
+
 	})
 
 	require.NoError(t, err)

diff --git a/pkg/frontend/v1/frontend.go b/pkg/frontend/v1/frontend.go
@@ -176,7 +176,7 @@ func (f *Frontend) RoundTripGRPC(ctx context.Context, req *httpgrpc.HTTPRequest)
 		}
 	}
 
-	return f.retry.Do(func() (*httpgrpc.HTTPResponse, error) {
+	return f.retry.Do(ctx, func() (*httpgrpc.HTTPResponse, error) {
 		request := request{
 			request:     req,
 			originalCtx: ctx,

diff --git a/pkg/frontend/v2/frontend.go b/pkg/frontend/v2/frontend.go
@@ -188,7 +188,7 @@ func (f *Frontend) RoundTripGRPC(ctx context.Context, req *httpgrpc.HTTPRequest)
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	return f.retry.Do(func() (*httpgrpc.HTTPResponse, error) {
+	return f.retry.Do(ctx, func() (*httpgrpc.HTTPResponse, error) {
 		freq := &frontendRequest{
 			queryID:      f.lastQueryID.Inc(),
 			request:      req,