Tidy up retry middleward and tests.

Signed-off-by: Tom Wilkie <[email protected]>

Author: tomwilkie

pkg/querier/queryrange/retry.go

@@ -2,7 +2,6 @@ package queryrange
 
 import (
 	"context"
-	"net/http"
 
 	"github.com/go-kit/kit/log"
 	"github.com/go-kit/kit/log/level"
@@ -37,30 +36,25 @@ func NewRetryMiddleware(log log.Logger, maxRetries int) Middleware {
 }
 
 func (r retry) Do(ctx context.Context, req *Request) (*APIResponse, error) {
-	var lastErr error
-	for tries := 0; tries < r.maxRetries; tries++ {
-		retries.Observe(float64(tries))
+	tries := 0
+	defer func() { retries.Observe(float64(tries)) }()
 
+	var lastErr error
+	for ; tries < r.maxRetries; tries++ {
 		resp, err := r.next.Do(ctx, req)
 		if err == nil {
 			return resp, nil
 		}
 
-		// Retry is we get a HTTP 500 or a non-HTTP error.
+		// Retry if we get a HTTP 500 or a non-HTTP error.
 		httpResp, ok := httpgrpc.HTTPResponseFromError(err)
-		if !ok || (ok && httpResp.Code/100 == 5) {
+		if !ok || httpResp.Code/100 == 5 {
 			lastErr = err
 			level.Error(r.log).Log("msg", "error processing request", "try", tries, "err", err)
 			continue
 		}
 
 		return nil, err
 	}
-
-	if lastErr != nil {
-		return nil, lastErr
-	}
-
-	return nil, httpgrpc.Errorf(http.StatusInternalServerError, "Query failed after %d retries.", r.maxRetries)
-
+	return nil, lastErr
 }

pkg/querier/queryrange/retry_test.go

@@ -3,26 +3,65 @@ package queryrange
 import (
 	"context"
 	fmt "fmt"
+	"net/http"
 	"sync/atomic"
 	"testing"
 
 	"github.com/go-kit/kit/log"
 	"github.com/stretchr/testify/require"
+	"github.com/weaveworks/common/httpgrpc"
 )
 
 func TestRetry(t *testing.T) {
-	try := int32(0)
+	var try int32
 
-	h := NewRetryMiddleware(log.NewNopLogger(), 5).Wrap(
-		HandlerFunc(func(_ context.Context, req *Request) (*APIResponse, error) {
-			if atomic.AddInt32(&try, 1) == 5 {
-				return &APIResponse{Status: "Hello World"}, nil
-			}
-			return nil, fmt.Errorf("fail")
-		}),
-	)
-
-	resp, err := h.Do(nil, nil)
-	require.NoError(t, err)
-	require.Equal(t, &APIResponse{Status: "Hello World"}, resp)
+	for _, tc := range []struct {
+		name    string
+		handler Handler
+		resp    *APIResponse
+		err     error
+	}{
+		{
+			name: "retry failures",
+			handler: HandlerFunc(func(_ context.Context, req *Request) (*APIResponse, error) {
+				if atomic.AddInt32(&try, 1) == 5 {
+					return &APIResponse{Status: "Hello World"}, nil
+				}
+				return nil, fmt.Errorf("fail")
+			}),
+			resp: &APIResponse{Status: "Hello World"},
+		},
+		{
+			name: "don't retry 400s",
+			handler: HandlerFunc(func(_ context.Context, req *Request) (*APIResponse, error) {
+				return nil, httpgrpc.Errorf(http.StatusBadRequest, "Bad Request")
+			}),
+			err: httpgrpc.Errorf(http.StatusBadRequest, "Bad Request"),
+		},
+		{
+			name: "retry 500s",
+			handler: HandlerFunc(func(_ context.Context, req *Request) (*APIResponse, error) {
+				return nil, httpgrpc.Errorf(http.StatusInternalServerError, "Internal Server Error")
+			}),
+			err: httpgrpc.Errorf(http.StatusInternalServerError, "Internal Server Error"),
+		},
+		{
+			name: "last error",
+			handler: HandlerFunc(func(_ context.Context, req *Request) (*APIResponse, error) {
+				if atomic.AddInt32(&try, 1) == 5 {
+					return nil, httpgrpc.Errorf(http.StatusBadRequest, "Bad Request")
+				}
+				return nil, httpgrpc.Errorf(http.StatusInternalServerError, "Internal Server Error")
+			}),
+			err: httpgrpc.Errorf(http.StatusBadRequest, "Bad Request"),
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			try = 0
+			h := NewRetryMiddleware(log.NewNopLogger(), 5).Wrap(tc.handler)
+			resp, err := h.Do(nil, nil)
+			require.Equal(t, tc.err, err)
+			require.Equal(t, tc.resp, resp)
+		})
+	}
 }