Skip to content

Commit

Permalink
Add custom retry decider for fail http requests
Browse files Browse the repository at this point in the history
  • Loading branch information
vural committed May 27, 2020
1 parent 2f92c68 commit 77dcda3
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 1 deletion.
22 changes: 21 additions & 1 deletion client.go
Original file line number Diff line number Diff line change
Expand Up @@ -289,6 +289,11 @@ type Client struct {
// By default will not waiting, return ErrNoFreeConns immediately
MaxConnWaitTimeout time.Duration

// RetryIf controls whether a retry should be attempted after an error.
//
// By default will use isIdempotent function
RetryIf RetryIfFunc

mLock sync.Mutex
m map[string]*HostClient
ms map[string]*HostClient
Expand Down Expand Up @@ -493,6 +498,7 @@ func (c *Client) Do(req *Request, resp *Response) error {
DisableHeaderNamesNormalizing: c.DisableHeaderNamesNormalizing,
DisablePathNormalizing: c.DisablePathNormalizing,
MaxConnWaitTimeout: c.MaxConnWaitTimeout,
RetryIf: c.RetryIf,
}
m[string(host)] = hc
if len(m) == 1 {
Expand Down Expand Up @@ -560,6 +566,11 @@ const DefaultMaxIdemponentCallAttempts = 5
// - foobar.com:8080
type DialFunc func(addr string) (net.Conn, error)

// RetryIfFunc signature of retry if function
//
// Request argument passed to RetryIfFunc, if there are any request errors.
type RetryIfFunc func(request *Request) bool

// HostClient balances http requests among hosts listed in Addr.
//
// HostClient may be used for balancing load among multiple upstream hosts.
Expand Down Expand Up @@ -698,6 +709,11 @@ type HostClient struct {
// By default will not waiting, return ErrNoFreeConns immediately
MaxConnWaitTimeout time.Duration

// RetryIf controls whether a retry should be attempted after an error.
//
// By default will use isIdempotent function
RetryIf RetryIfFunc

clientName atomic.Value
lastUseTime uint32

Expand Down Expand Up @@ -1183,6 +1199,10 @@ func (c *HostClient) Do(req *Request, resp *Response) error {
if maxAttempts <= 0 {
maxAttempts = DefaultMaxIdemponentCallAttempts
}
isRequestRetryable := isIdempotent
if c.RetryIf != nil {
isRequestRetryable = c.RetryIf
}
attempts := 0
hasBodyStream := req.IsBodyStream()

Expand All @@ -1196,7 +1216,7 @@ func (c *HostClient) Do(req *Request, resp *Response) error {
if hasBodyStream {
break
}
if !isIdempotent(req) {
if !isRequestRetryable(req) {
// Retry non-idempotent requests if the server closes
// the connection before sending the response.
//
Expand Down
52 changes: 52 additions & 0 deletions client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -1648,6 +1648,58 @@ func TestClientIdempotentRequest(t *testing.T) {
}
}

func TestClientRetryRequestWithCustomDecider(t *testing.T) {
t.Parallel()

dialsCount := 0
c := &Client{
Dial: func(addr string) (net.Conn, error) {
dialsCount++
switch dialsCount {
case 1:
return &singleReadConn{
s: "invalid response",
}, nil
case 2:
return &writeErrorConn{}, nil
case 3:
return &readErrorConn{}, nil
case 4:
return &singleReadConn{
s: "HTTP/1.1 345 OK\r\nContent-Type: foobar\r\nContent-Length: 7\r\n\r\n0123456",
}, nil
default:
t.Fatalf("unexpected number of dials: %d", dialsCount)
}
panic("unreachable")
},
RetryIf: func(req *Request) bool {
return req.URI().String() == "http://foobar/a/b"
},
}

var args Args

// Post must succeed for http://foobar/a/b uri.
statusCode, body, err := c.Post(nil, "http://foobar/a/b", &args)
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
if statusCode != 345 {
t.Fatalf("unexpected status code: %d. Expecting 345", statusCode)
}
if string(body) != "0123456" {
t.Fatalf("unexpected body: %q. Expecting %q", body, "0123456")
}

// POST must fail for http://foobar/a/b/c uri.
dialsCount = 0
_, _, err = c.Post(nil, "http://foobar/a/b/c", &args)
if err == nil {
t.Fatalf("expecting error")
}
}

type writeErrorConn struct {
net.Conn
}
Expand Down

0 comments on commit 77dcda3

Please sign in to comment.