240 lines
4.9 KiB
Go
240 lines
4.9 KiB
Go
package zarinhub
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"net/http"
|
|
"time"
|
|
)
|
|
|
|
type RetryTransport struct {
|
|
base http.RoundTripper
|
|
timeout time.Duration
|
|
maxRetries int
|
|
retryWait func(ctx context.Context, delay time.Duration) bool
|
|
logger *slog.Logger
|
|
}
|
|
|
|
func NewRetryTransport(
|
|
base http.RoundTripper,
|
|
timeout time.Duration,
|
|
maxRetries int,
|
|
loggers ...*slog.Logger,
|
|
) *RetryTransport {
|
|
if base == nil {
|
|
base = http.DefaultTransport
|
|
}
|
|
if maxRetries < 0 {
|
|
maxRetries = 0
|
|
}
|
|
|
|
var logger *slog.Logger
|
|
if len(loggers) > 0 && loggers[0] != nil {
|
|
logger = loggers[0]
|
|
}
|
|
|
|
return &RetryTransport{
|
|
base: base,
|
|
timeout: timeout,
|
|
maxRetries: maxRetries,
|
|
retryWait: waitForRetry,
|
|
logger: logger,
|
|
}
|
|
}
|
|
|
|
func NewHTTPClient(
|
|
timeout time.Duration,
|
|
maxRetries int,
|
|
loggers ...*slog.Logger,
|
|
) *http.Client {
|
|
transport := NewRetryTransport(
|
|
http.DefaultTransport,
|
|
timeout,
|
|
maxRetries,
|
|
loggers...,
|
|
)
|
|
|
|
return &http.Client{
|
|
Transport: transport,
|
|
}
|
|
}
|
|
|
|
func (t *RetryTransport) log() *slog.Logger {
|
|
if t.logger != nil {
|
|
return t.logger
|
|
}
|
|
return slog.New(slog.NewTextHandler(io.Discard, nil))
|
|
}
|
|
|
|
func (t *RetryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
|
waitFn := t.retryWait
|
|
if waitFn == nil {
|
|
waitFn = waitForRetry
|
|
}
|
|
|
|
if req.Body != nil && req.GetBody == nil {
|
|
bodyBytes, err := io.ReadAll(req.Body)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
_ = req.Body.Close()
|
|
req.GetBody = func() (io.ReadCloser, error) {
|
|
return io.NopCloser(bytes.NewReader(bodyBytes)), nil
|
|
}
|
|
}
|
|
|
|
var (
|
|
resp *http.Response
|
|
err error
|
|
)
|
|
|
|
for attempt := 0; attempt <= t.maxRetries; attempt++ {
|
|
if attempt > 0 {
|
|
t.log().InfoContext(
|
|
req.Context(),
|
|
fmt.Sprintf("[RETRY] zarinhub request attempt %d", attempt),
|
|
"is_retry", true,
|
|
"retry_count", attempt,
|
|
)
|
|
} else {
|
|
t.log().InfoContext(
|
|
req.Context(),
|
|
"[INITIAL] zarinhub request",
|
|
"is_retry", false,
|
|
"retry_count", 0,
|
|
)
|
|
}
|
|
|
|
var (
|
|
attemptCtx context.Context
|
|
cancelAttempt context.CancelFunc
|
|
)
|
|
if t.timeout > 0 {
|
|
attemptCtx, cancelAttempt = context.WithTimeout(req.Context(), t.timeout)
|
|
} else {
|
|
attemptCtx = req.Context()
|
|
cancelAttempt = func() {}
|
|
}
|
|
|
|
attemptReq := req.Clone(attemptCtx)
|
|
if req.GetBody != nil {
|
|
body, bodyErr := req.GetBody()
|
|
if bodyErr != nil {
|
|
cancelAttempt()
|
|
return nil, bodyErr
|
|
}
|
|
attemptReq.Body = body
|
|
}
|
|
|
|
resp, err = t.base.RoundTrip(attemptReq)
|
|
|
|
if req.Context().Err() != nil {
|
|
cancelAttempt()
|
|
if resp != nil && resp.Body != nil {
|
|
_ = resp.Body.Close()
|
|
}
|
|
if errors.Is(req.Context().Err(), context.Canceled) {
|
|
return nil, fmt.Errorf("zarinhub request canceled: %w", req.Context().Err())
|
|
}
|
|
return nil, req.Context().Err()
|
|
}
|
|
|
|
isErrRetryable := false
|
|
var retryReason error
|
|
|
|
if err != nil {
|
|
isErrRetryable = isTransportErrorRetryable(err)
|
|
retryReason = err
|
|
} else if isStatusCodeRetryable(resp.StatusCode) {
|
|
isErrRetryable = true
|
|
retryReason = fmt.Errorf("upstream returned status %d", resp.StatusCode)
|
|
}
|
|
|
|
if !isErrRetryable {
|
|
if resp != nil && resp.Body != nil {
|
|
resp.Body = &bodyWithCancel{
|
|
ReadCloser: resp.Body,
|
|
cancel: cancelAttempt,
|
|
}
|
|
} else {
|
|
cancelAttempt()
|
|
}
|
|
return resp, err
|
|
}
|
|
|
|
if attempt == t.maxRetries {
|
|
if attempt > 0 {
|
|
t.log().ErrorContext(
|
|
req.Context(),
|
|
"[RETRY] zarinhub request failed: all retries exhausted",
|
|
"is_retry", true,
|
|
"retry_count", attempt,
|
|
"error", retryReason,
|
|
)
|
|
}
|
|
if resp != nil && resp.Body != nil {
|
|
resp.Body = &bodyWithCancel{
|
|
ReadCloser: resp.Body,
|
|
cancel: cancelAttempt,
|
|
}
|
|
} else {
|
|
cancelAttempt()
|
|
}
|
|
return resp, err
|
|
}
|
|
|
|
if resp != nil && resp.Body != nil {
|
|
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4096))
|
|
_ = resp.Body.Close()
|
|
}
|
|
cancelAttempt()
|
|
|
|
delay := retryDelay(attempt)
|
|
t.log().WarnContext(
|
|
req.Context(),
|
|
fmt.Sprintf("[RETRY] zarinhub request failed, retrying in %v (retry %d)", delay, attempt+1),
|
|
"is_retry", true,
|
|
"next_retry", attempt+1,
|
|
"delay_ms", delay.Milliseconds(),
|
|
"error", retryReason,
|
|
)
|
|
|
|
if !waitFn(req.Context(), delay) {
|
|
if req.Context().Err() != nil {
|
|
if errors.Is(req.Context().Err(), context.Canceled) {
|
|
return nil, fmt.Errorf("zarinhub request canceled: %w", req.Context().Err())
|
|
}
|
|
return nil, req.Context().Err()
|
|
}
|
|
return nil, context.Canceled
|
|
}
|
|
}
|
|
|
|
return resp, err
|
|
}
|
|
|
|
func isTransportErrorRetryable(err error) bool {
|
|
return !errors.Is(err, context.Canceled)
|
|
}
|
|
|
|
func isStatusCodeRetryable(statusCode int) bool {
|
|
return statusCode == http.StatusRequestTimeout ||
|
|
statusCode == http.StatusTooManyRequests ||
|
|
statusCode >= http.StatusInternalServerError
|
|
}
|
|
|
|
type bodyWithCancel struct {
|
|
io.ReadCloser
|
|
cancel context.CancelFunc
|
|
}
|
|
|
|
func (b *bodyWithCancel) Close() error {
|
|
err := b.ReadCloser.Close()
|
|
b.cancel()
|
|
return err
|
|
}
|