-
Notifications
You must be signed in to change notification settings - Fork 6
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
79858d7
commit 5d4cef0
Showing
3 changed files
with
100 additions
and
16 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,87 @@ | ||
| package utils | ||
|
|
||
| import ( | ||
| "bytes" | ||
| "io" | ||
| "math" | ||
| "net/http" | ||
| "time" | ||
| ) | ||
|
|
||
| const RetryCount = 3 | ||
|
|
||
| type retryableTransport struct { | ||
| transport http.RoundTripper | ||
| } | ||
|
|
||
| func backoff(retries int) time.Duration { | ||
| return time.Duration(math.Pow(2, float64(retries))) * time.Second | ||
| } | ||
|
|
||
| func shouldRetry(err error, resp *http.Response) bool { | ||
| if err != nil { | ||
| return true | ||
| } | ||
|
|
||
| if resp.StatusCode == 401 { | ||
| return true | ||
| } | ||
|
|
||
| return false | ||
| } | ||
|
|
||
| func drainBody(resp *http.Response) { | ||
| if resp.Body != nil { | ||
| io.Copy(io.Discard, resp.Body) | ||
| resp.Body.Close() | ||
| } | ||
| } | ||
|
|
||
| func (t *retryableTransport) RoundTrip(req *http.Request) (*http.Response, error) { | ||
| // Clone the request body | ||
| var bodyBytes []byte | ||
| if req.Body != nil { | ||
| bodyBytes, _ = io.ReadAll(req.Body) | ||
| req.Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) | ||
| } | ||
|
|
||
| // Send the request | ||
| resp, err := t.transport.RoundTrip(req) | ||
|
|
||
| // Retry logic | ||
| retries := 0 | ||
| for shouldRetry(err, resp) && retries < RetryCount { | ||
| // Wait for the specified backoff period | ||
| time.Sleep(backoff(retries)) | ||
|
|
||
| // We're going to retry, consume any response to reuse the connection. | ||
| drainBody(resp) | ||
|
|
||
| // Clone the request body again | ||
| if req.Body != nil { | ||
| req.Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) | ||
| } | ||
|
|
||
| // Retry the request | ||
| resp, err = t.transport.RoundTrip(req) | ||
|
|
||
| retries++ | ||
| } | ||
|
|
||
| // Return the response | ||
| return resp, err | ||
| } | ||
|
|
||
| func NewRetryableClient() *http.Client { | ||
| transport := &retryableTransport{ | ||
| transport: &http.Transport{ | ||
| MaxIdleConns: 20, | ||
| MaxIdleConnsPerHost: 20, | ||
| DisableKeepAlives: true}, | ||
| } | ||
|
|
||
| return &http.Client{ | ||
| Transport: transport, | ||
| Timeout: 2 * time.Minute, | ||
| } | ||
| } |