Skip to content

Commit

Permalink
ensure to drain
Browse files Browse the repository at this point in the history
  • Loading branch information
Dean Karn committed Mar 25, 2024
1 parent d296ddd commit e6f2e43
Showing 1 changed file with 16 additions and 5 deletions.
21 changes: 16 additions & 5 deletions net/http/retrier.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ package httpext
import (
"context"
"errors"
"io"
"net/http"
"strconv"
"time"
Expand Down Expand Up @@ -107,11 +108,17 @@ func NewRetryer() Retryer {
wait := time.Millisecond * 200

var sce ErrStatusCode
if errors.As(err, &sce) && (sce.Response.StatusCode == http.StatusTooManyRequests || sce.Response.StatusCode == http.StatusServiceUnavailable) && sce.Response.Header != nil {
defer sce.Response.Body.Close()

if ra := HasRetryAfter(sce.Response.Header); ra.IsSome() {
wait = ra.Unwrap()
if errors.As(err, &sce) {
// can close in backoff because if this function is hit we will be retrying again.
defer func() {
_, _ = io.Copy(io.Discard, sce.Response.Body)
_ = sce.Response.Body.Close()
}()

if sce.Response.Header != nil && (sce.Response.StatusCode == http.StatusTooManyRequests || sce.Response.StatusCode == http.StatusServiceUnavailable) {
if ra := HasRetryAfter(sce.Response.Header); ra.IsSome() {
wait = ra.Unwrap()
}
}
}

Expand Down Expand Up @@ -176,6 +183,9 @@ func (r Retryer) Backoff(fn errorsext.BackoffFn[error]) Retryer {
}

// MaxMemory sets the maximum memory to use when decoding the response body.
//
// This max memory will NOT be enforced when *http.Response is returned through an error and is up to the caller to
// handle/limit it if desired.
func (r Retryer) MaxMemory(maxMemory bytesext.Bytes) Retryer {
r.maxMemory = maxMemory
return r
Expand Down Expand Up @@ -254,6 +264,7 @@ func (r Retryer) Do(ctx context.Context, fn BuildRequestFn2, v any, expectedResp
}
defer func() {
if closeBody {
_, _ = io.Copy(io.Discard, resp.Body)
_ = resp.Body.Close()
}
}()
Expand Down

0 comments on commit e6f2e43

Please sign in to comment.