new option to set status codes to be retried

hystrix/hystrix_client.go

@@ -32,6 +32,7 @@ type Client struct {
 	retrier                heimdall.Retriable
 	fallbackFunc           func(err error) error
 	statsD                 *plugins.StatsdCollectorConfig
+	statusCodeToRetry      map[int]struct{}
 }
 
 const (
@@ -49,6 +50,7 @@ const (
 
 var _ heimdall.Client = (*Client)(nil)
 var err5xx = errors.New("server returned 5xx status code")
+var errCodeToRetry = errors.New("server returned status code to retry")
 
 // NewClient returns a new instance of hystrix Client
 func NewClient(opts ...Option) *Client {
@@ -62,6 +64,7 @@ func NewClient(opts ...Option) *Client {
 		requestVolumeThreshold: defaultRequestVolumeThreshold,
 		retryCount:             defaultHystrixRetryCount,
 		retrier:                heimdall.NewNoRetrier(),
+		statusCodeToRetry:      make(map[int]struct{}),
 	}
 
 	for _, opt := range opts {
@@ -198,6 +201,14 @@ func (hhc *Client) Do(request *http.Request) (*http.Response, error) {
 				return err
 			}
 
+			if len(hhc.statusCodeToRetry) > 0 {
+				_, ok := hhc.statusCodeToRetry[response.StatusCode]
+				if ok {
+					return errCodeToRetry
+				}
+				return nil
+			}
+
 			if response.StatusCode >= http.StatusInternalServerError {
 				return err5xx
 			}

hystrix/hystrix_client_test.go

@@ -545,3 +545,44 @@ func TestDurationToInt(t *testing.T) {
 		assert.Equal(t, 30000, timeoutInMs)
 	})
 }
+
+func TestHystrixHTTPClientRetriesOnSpecificStatusCodes(t *testing.T) {
+	client := NewClient(
+		WithHTTPTimeout(10*time.Millisecond),
+		WithCommandName("some_new_command_name"),
+		WithHystrixTimeout(10*time.Millisecond),
+		WithMaxConcurrentRequests(100),
+		WithErrorPercentThreshold(10),
+		WithSleepWindow(100),
+		WithRequestVolumeThreshold(10),
+		WithHTTPClient(&myHTTPClient{
+			client: http.Client{Timeout: 25 * time.Millisecond}}),
+		WithStatusCodeToRetry(400),
+		WithRetryCount(3),
+	)
+
+	counter := 0
+	dummyHandler := func(w http.ResponseWriter, r *http.Request) {
+		counter++
+		if counter < 2 {
+			w.WriteHeader(http.StatusBadRequest)
+		} else {
+			w.WriteHeader(http.StatusOK)
+		}
+		assert.Equal(t, r.Header.Get("foo"), "bar")
+		assert.NotEqual(t, r.Header.Get("foo"), "baz")
+		_, _ = w.Write([]byte(`{ "response": "ok" }`))
+	}
+
+	server := httptest.NewServer(http.HandlerFunc(dummyHandler))
+	defer server.Close()
+
+	req, err := http.NewRequest(http.MethodGet, server.URL, nil)
+	require.NoError(t, err)
+	response, err := client.Do(req)
+	assert.Equal(t, http.StatusOK, response.StatusCode)
+
+	body, err := ioutil.ReadAll(response.Body)
+	require.NoError(t, err)
+	assert.Equal(t, "{ \"response\": \"ok\" }", string(body))
+}

hystrix/options.go

@@ -95,3 +95,12 @@ func WithStatsDCollector(addr, prefix string) Option {
 		c.statsD = &plugins.StatsdCollectorConfig{StatsdAddr: addr, Prefix: prefix}
 	}
 }
+
+// WithStatusCodeToRetry sets status codes to be retried
+func WithStatusCodeToRetry(statusCodes ...int) Option {
+	return func(c *Client) {
+		for _, code := range statusCodes {
+			c.statusCodeToRetry[code] = struct{}{}
+		}
+	}
+}

hystrix/options_test.go

@@ -19,6 +19,7 @@ func TestOptionsAreSet(t *testing.T) {
 		WithSleepWindow(5),
 		WithRequestVolumeThreshold(5),
 		WithStatsDCollector("localhost:8125", "myapp.hystrix"),
+		WithStatusCodeToRetry(200, 400),
 	)
 
 	assert.Equal(t, 10*time.Second, c.timeout)
@@ -30,6 +31,7 @@ func TestOptionsAreSet(t *testing.T) {
 	assert.Equal(t, 5, c.requestVolumeThreshold)
 	assert.Equal(t, "localhost:8125", c.statsD.StatsdAddr)
 	assert.Equal(t, "myapp.hystrix", c.statsD.Prefix)
+	assert.Equal(t, map[int]struct{}{200: {}, 400: {}}, c.statusCodeToRetry)
 }
 
 func TestOptionsHaveDefaults(t *testing.T) {