Fix client backoff mechanism

axiom/client.go

@@ -304,8 +304,10 @@ func (c *Client) do(req *http.Request, v interface{}) (*response, error) {
 		}, err
 	}
 
-	var resp *response
-	var httpResp *http.Response
+	var (
+		resp     *response
+		httpResp *http.Response
+	)
 
 	bck := backoff.NewExponentialBackOff()
 	bck.InitialInterval = 200 * time.Millisecond
@@ -314,16 +316,15 @@ func (c *Client) do(req *http.Request, v interface{}) (*response, error) {
 
 	err := backoff.Retry(func() error {
 		var err error
-		httpResp, err = c.httpClient.Do(req) //nolint:bodyclose // We close the body in the defer func below.
-		if err != nil {
+		if httpResp, err = c.httpClient.Do(req); err != nil { //nolint:bodyclose // We close the body in the defer func below.
 			return err
 		}
 
 		resp = newResponse(httpResp)
 
 		// We should only retry in the case the status code is >= 500, anything below isn't worth retrying.
-		if resp.StatusCode >= 500 {
-			return fmt.Errorf("got status code %d", resp.StatusCode)
+		if code := resp.StatusCode; code >= 500 {
+			return fmt.Errorf("got status code %d", code)
 		}
 
 		return nil
@@ -337,7 +338,7 @@ func (c *Client) do(req *http.Request, v interface{}) (*response, error) {
 	}()
 
 	if err != nil {
-		return nil, fmt.Errorf("error making request: %v", err)
+		return resp, err
 	}
 
 	key := limitKey(resp.Limit.limitType, resp.Limit.Scope)

axiom/client_test.go

@@ -674,6 +674,55 @@ func TestClient_do_ValidOnlyAPITokenPaths(t *testing.T) {
 	}
 }
 
+func TestClient_do_Backoff(t *testing.T) {
+	var currentCalls int
+	hf := func(w http.ResponseWriter, r *http.Request) {
+		currentCalls++
+		switch currentCalls {
+		case 1:
+			w.WriteHeader(http.StatusInternalServerError)
+		case 2:
+			w.WriteHeader(http.StatusBadGateway)
+		case 3:
+			w.WriteHeader(http.StatusGatewayTimeout)
+		default:
+			w.WriteHeader(http.StatusOK)
+		}
+	}
+
+	client, teardown := setup(t, "/", hf)
+	defer teardown()
+
+	req, err := client.newRequest(context.Background(), http.MethodGet, "/", nil)
+	require.NoError(t, err)
+
+	resp, err := client.do(req, nil)
+	require.NoError(t, err)
+
+	assert.Equal(t, 4, currentCalls)
+	assert.Equal(t, http.StatusOK, resp.StatusCode)
+}
+
+func TestClient_do_Backoff_NoRetryOn400(t *testing.T) {
+	var currentCalls int
+	hf := func(w http.ResponseWriter, r *http.Request) {
+		currentCalls++
+		w.WriteHeader(http.StatusBadRequest)
+	}
+
+	client, teardown := setup(t, "/", hf)
+	defer teardown()
+
+	req, err := client.newRequest(context.Background(), http.MethodGet, "/", nil)
+	require.NoError(t, err)
+
+	resp, err := client.do(req, nil)
+	require.Error(t, err, "got status code 400")
+
+	assert.Equal(t, 1, currentCalls)
+	assert.Equal(t, http.StatusBadRequest, resp.StatusCode)
+}
+
 func TestAPITokenPathRegex(t *testing.T) {
 	tests := []struct {
 		input string
@@ -723,49 +772,6 @@ func TestAPITokenPathRegex(t *testing.T) {
 	}
 }
 
-func TestClient_do_Backoff(t *testing.T) {
-	var currentCalls int
-	r := http.NewServeMux()
-	r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
-		currentCalls++
-		switch currentCalls {
-		case 1:
-			w.WriteHeader(http.StatusInternalServerError)
-		case 2:
-			w.WriteHeader(http.StatusBadGateway)
-		case 3:
-			w.WriteHeader(http.StatusGatewayTimeout)
-		default:
-			w.WriteHeader(http.StatusOK)
-		}
-	})
-	srv := httptest.NewServer(r)
-	defer srv.Close()
-
-	clientOptions := []Option{
-		SetURL(srv.URL),
-		SetAccessToken(personalToken),
-		SetOrgID(orgID),
-		SetClient(srv.Client()),
-		SetStrictDecoding(true),
-		SetNoEnv(),
-	}
-
-	client, err := NewClient(clientOptions...)
-	require.NoError(t, err)
-
-	req, err := client.newRequest(context.Background(), http.MethodGet, "/", nil)
-	require.NoError(t, err)
-
-	resp, err := client.do(req, nil)
-	require.NoError(t, err)
-	if currentCalls > 4 {
-		t.Fatal("expected to attempt 4 times")
-	}
-	assert.Equal(t, 4, currentCalls)
-	assert.Equal(t, http.StatusOK, resp.StatusCode)
-}
-
 // setup sets up a test HTTP server along with a client that is configured to
 // talk to that test server. Tests should pass a handler function which provides
 // the response for the API method being tested.