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/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) {