Skip to content
This repository has been archived by the owner on Jan 15, 2024. It is now read-only.

Commit

Permalink
Allow setting status codes to retry (#152)
Browse files Browse the repository at this point in the history
For grafana/terraform-provider-grafana#893
We currently retry 429 and 5xx. This is retained as a default in this PR but there's a new `RetryStatusCodes` that, if set, overrides these statuses
That way, the user that opened the issue above will be able to also retry 407 on their Azure instance

**Had to do some fixes in dashboard tests, the httptest servers were not properly used, not closed and leaked out of the tests**
  • Loading branch information
julienduchesne authored Jun 6, 2023
1 parent 655166f commit 023cfdf
Show file tree
Hide file tree
Showing 4 changed files with 153 additions and 57 deletions.
39 changes: 37 additions & 2 deletions client.go
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,8 @@ type Config struct {
NumRetries int
// RetryTimeout says how long to wait before retrying a request
RetryTimeout time.Duration
// RetryStatusCodes contains the list of status codes to retry, use "x" as a wildcard for a single digit (default: [429, 5xx])
RetryStatusCodes []string
}

// New creates a new Grafana client.
Expand Down Expand Up @@ -80,6 +82,10 @@ func (c *Client) request(method, requestPath string, query url.Values, body []by
err error
bodyContents []byte
)
retryStatusCodes := c.config.RetryStatusCodes
if len(retryStatusCodes) == 0 {
retryStatusCodes = []string{"429", "5xx"}
}

// retry logic
for n := 0; n <= c.config.NumRetries; n++ {
Expand Down Expand Up @@ -115,8 +121,11 @@ func (c *Client) request(method, requestPath string, query url.Values, body []by
continue
}

// Exit the loop if we have something final to return. This is anything < 500, if it's not a 429.
if resp.StatusCode < http.StatusInternalServerError && resp.StatusCode != http.StatusTooManyRequests {
shouldRetry, err := matchRetryCode(resp.StatusCode, retryStatusCodes)
if err != nil {
return err
}
if !shouldRetry {
break
}
}
Expand Down Expand Up @@ -179,3 +188,29 @@ func (c *Client) newRequest(method, requestPath string, query url.Values, body i
req.Header.Add("Content-Type", "application/json")
return req, err
}

// matchRetryCode checks if the status code matches any of the configured retry status codes.
func matchRetryCode(gottenCode int, retryCodes []string) (bool, error) {
gottenCodeStr := strconv.Itoa(gottenCode)
for _, retryCode := range retryCodes {
if len(retryCode) != 3 {
return false, fmt.Errorf("invalid retry status code: %s", retryCode)
}
matched := true
for i := range retryCode {
c := retryCode[i]
if c == 'x' {
continue
}
if gottenCodeStr[i] != c {
matched = false
break
}
}
if matched {
return true, nil
}
}

return false, nil
}
50 changes: 49 additions & 1 deletion client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ import (
"net/http"
"net/http/httptest"
"net/url"
"strconv"
"strings"
"testing"
"time"
)
Expand Down Expand Up @@ -209,8 +211,11 @@ func TestClient_requestWithRetries(t *testing.T) {
case 2:
http.Error(w, `{"error":"calm down"}`, http.StatusTooManyRequests)

default:
case 3:
w.Write([]byte(`{"foo":"bar"}`)) //nolint:errcheck

default:
t.Errorf("unexpected retry %d", try)
}
}))
defer ts.Close()
Expand Down Expand Up @@ -255,6 +260,49 @@ func TestClient_requestWithRetries(t *testing.T) {
t.Logf("request successful after %d retries", try)
}

func TestClient_CustomRetryStatusCode(t *testing.T) {
body := []byte(`lorem ipsum dolor sit amet`)
var try int
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
defer r.Body.Close()

try++

switch try {
case 1, 2:
http.Error(w, `{"error":"weird error"}`, http.StatusUpgradeRequired)
default:
http.Error(w, `{"error":"failed"}`, http.StatusInternalServerError)
}
}))
defer ts.Close()

httpClient := &http.Client{
Transport: &customRoundTripper{},
}

c, err := New(ts.URL, Config{
NumRetries: 5,
Client: httpClient,
RetryTimeout: 50 * time.Millisecond,
RetryStatusCodes: []string{strconv.Itoa(http.StatusUpgradeRequired)},
})
if err != nil {
t.Fatalf("unexpected error creating client: %v", err)
}

var got interface{}
err = c.request(http.MethodPost, "/", nil, body, &got)
expectedErr := "status: 500, body: {\"error\":\"failed\"}" // The 500 is not retried because it's not in RetryStatusCodes
if strings.TrimSpace(err.Error()) != expectedErr {
t.Fatalf("expected err: %s, got err: %v", expectedErr, err)
}

if try != 3 {
t.Fatalf("unexpected number of tries: %d", try)
}
}

type customRoundTripper struct {
try int
}
Expand Down
111 changes: 62 additions & 49 deletions dashboard_test.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package gapi

import (
"fmt"
"strings"
"testing"

Expand Down Expand Up @@ -78,67 +79,79 @@ func TestDashboardCreateAndUpdate(t *testing.T) {
}

func TestDashboardGet(t *testing.T) {
client := gapiTestTools(t, 200, getDashboardResponse)
t.Run("By slug", func(t *testing.T) {
client := gapiTestTools(t, 200, getDashboardResponse)

resp, err := client.Dashboard("test")
if err != nil {
t.Error(err)
}
uid, ok := resp.Model["uid"]
if !ok || uid != "cIBgcSjkk" {
t.Errorf("Invalid uid - %s, Expected %s", uid, "cIBgcSjkk")
}

client = gapiTestTools(t, 200, getDashboardResponse)
resp, err := client.Dashboard("test")
if err != nil {
t.Error(err)
}
uid, ok := resp.Model["uid"]
if !ok || uid != "cIBgcSjkk" {
t.Errorf("Invalid uid - %s, Expected %s", uid, "cIBgcSjkk")
}
})

resp, err = client.DashboardByUID("cIBgcSjkk")
if err != nil {
t.Fatal(err)
}
uid, ok = resp.Model["uid"]
if !ok || uid != "cIBgcSjkk" {
t.Fatalf("Invalid UID - %s, Expected %s", uid, "cIBgcSjkk")
}
t.Run("By UID", func(t *testing.T) {
client := gapiTestTools(t, 200, getDashboardResponse)

for _, code := range []int{401, 403, 404} {
client = gapiTestTools(t, code, "error")
_, err = client.Dashboard("test")
if err == nil {
t.Errorf("%d not detected", code)
resp, err := client.DashboardByUID("cIBgcSjkk")
if err != nil {
t.Fatal(err)
}

_, err = client.DashboardByUID("cIBgcSjkk")
if err == nil {
t.Errorf("%d not detected", code)
uid, ok := resp.Model["uid"]
if !ok || uid != "cIBgcSjkk" {
t.Fatalf("Invalid UID - %s, Expected %s", uid, "cIBgcSjkk")
}
})

for _, code := range []int{401, 403, 404} {
t.Run(fmt.Sprintf("Dashboard error: %d", code), func(t *testing.T) {
client := gapiTestToolsFromCalls(t, []mockServerCall{{code, "error"}, {code, "error"}})
_, err := client.Dashboard("test")
if err == nil {
t.Errorf("%d not detected", code)
}

_, err = client.DashboardByUID("cIBgcSjkk")
if err == nil {
t.Errorf("%d not detected", code)
}
})
}
}

func TestDashboardDelete(t *testing.T) {
client := gapiTestTools(t, 200, "")
err := client.DeleteDashboard("test")
if err != nil {
t.Error(err)
}

client = gapiTestTools(t, 200, "")
err = client.DeleteDashboardByUID("cIBgcSjkk")
if err != nil {
t.Fatal(err)
}

for _, code := range []int{401, 403, 404, 412} {
client = gapiTestTools(t, code, "error")

err = client.DeleteDashboard("test")
if err == nil {
t.Errorf("%d not detected", code)
t.Run("By slug", func(t *testing.T) {
client := gapiTestTools(t, 200, "")
err := client.DeleteDashboard("test")
if err != nil {
t.Error(err)
}
})

err = client.DeleteDashboardByUID("cIBgcSjkk")
if err == nil {
t.Errorf("%d not detected", code)
t.Run("By UID", func(t *testing.T) {
client := gapiTestTools(t, 200, "")
err := client.DeleteDashboardByUID("cIBgcSjkk")
if err != nil {
t.Fatal(err)
}
})

for _, code := range []int{401, 403, 404, 412} {
t.Run(fmt.Sprintf("Dashboard error: %d", code), func(t *testing.T) {
client := gapiTestToolsFromCalls(t, []mockServerCall{{code, "error"}, {code, "error"}})

err := client.DeleteDashboard("test")
if err == nil {
t.Errorf("%d not detected", code)
}

err = client.DeleteDashboardByUID("cIBgcSjkk")
if err == nil {
t.Errorf("%d not detected", code)
}
})
}
}

Expand Down
10 changes: 5 additions & 5 deletions mock.go → mock_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,8 @@ type mockServer struct {
server *httptest.Server
}

func (m *mockServer) Close() {
m.server.Close()
}

func gapiTestTools(t *testing.T, code int, body string) *Client {
t.Helper()
return gapiTestToolsFromCalls(t, []mockServerCall{{code, body}})
}

Expand All @@ -35,6 +32,9 @@ func gapiTestToolsFromCalls(t *testing.T, calls []mockServerCall) *Client {
}

mock.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if len(mock.upcomingCalls) == 0 {
t.Fatalf("unexpected call to %s %s", r.Method, r.URL)
}
call := mock.upcomingCalls[0]
if len(calls) > 1 {
mock.upcomingCalls = mock.upcomingCalls[1:]
Expand All @@ -61,7 +61,7 @@ func gapiTestToolsFromCalls(t *testing.T, calls []mockServerCall) *Client {
}

t.Cleanup(func() {
mock.Close()
mock.server.Close()
})

return client
Expand Down

0 comments on commit 023cfdf

Please sign in to comment.