2
0

httpclient.go 1.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  1. package httpclient
  2. import (
  3. "io"
  4. "net/http"
  5. "strconv"
  6. "time"
  7. "github.com/opencost/opencost/core/pkg/log"
  8. )
  9. const (
  10. defaultMaxRetries = 5
  11. defaultRetryBaseWait = 2 * time.Second
  12. defaultRetryMaxWait = 60 * time.Second
  13. )
  14. // retryTransport is an http.RoundTripper that retries requests on 429 and 503s
  15. type retryTransport struct {
  16. wrapped http.RoundTripper
  17. maxRetries int
  18. baseWait time.Duration
  19. maxWait time.Duration
  20. }
  21. func (t *retryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
  22. wait := t.baseWait
  23. for attempt := 0; attempt <= t.maxRetries; attempt++ {
  24. resp, err := t.wrapped.RoundTrip(req)
  25. if err != nil {
  26. return nil, err
  27. }
  28. if resp.StatusCode != http.StatusTooManyRequests && resp.StatusCode != http.StatusServiceUnavailable {
  29. return resp, nil
  30. }
  31. if attempt == t.maxRetries {
  32. // Return the final error response untouched so the caller can
  33. // read the body and status code.
  34. return resp, nil
  35. }
  36. // Consume and discard the error body so the connection can be reused,
  37. // then close it before sleeping.
  38. _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4096))
  39. _ = resp.Body.Close()
  40. delay := wait
  41. if ra := resp.Header.Get("Retry-After"); ra != "" {
  42. if secs, err := strconv.Atoi(ra); err == nil {
  43. delay = time.Duration(secs) * time.Second
  44. }
  45. }
  46. if delay > t.maxWait {
  47. delay = t.maxWait
  48. }
  49. log.Warnf("pricing httpclient: HTTP %d, retrying in %s (attempt %d/%d)",
  50. resp.StatusCode, delay, attempt+1, t.maxRetries)
  51. time.Sleep(delay)
  52. wait *= 2
  53. }
  54. return nil, nil
  55. }
  56. // NewClient returns an *http.Client whose transport automatically retries
  57. // on HTTP 429 / 503 with exponential backoff
  58. func NewClient(timeout time.Duration) *http.Client {
  59. return &http.Client{
  60. Timeout: timeout,
  61. Transport: &retryTransport{
  62. wrapped: http.DefaultTransport,
  63. maxRetries: defaultMaxRetries,
  64. baseWait: defaultRetryBaseWait,
  65. maxWait: defaultRetryMaxWait,
  66. },
  67. }
  68. }