extend retry client to understand a Retry-After header

Author: ptzianos

retry/go.mod

@@ -1,10 +1,12 @@
 module github.com/smithy-security/pkg/retry
 
-go 1.23
+go 1.23.2
 
 require (
 	github.com/cenkalti/backoff/v5 v5.0.2
-	github.com/stretchr/testify v1.9.0
+	github.com/go-errors/errors v1.5.1
+	github.com/smithy-security/pkg/utils v0.0.2
+	github.com/stretchr/testify v1.10.0
 )
 
 require (

retry/go.sum

@@ -3,6 +3,8 @@ github.com/cenkalti/backoff/v5 v5.0.2/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F9
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/go-errors/errors v1.5.1 h1:ZwEMSLRCapFLflTpT7NKaAc7ukJ8ZPEjzlxt8rPN8bk=
+github.com/go-errors/errors v1.5.1/go.mod h1:sIVyrIiJhuEF+Pj9Ebtd6P/rEYROXFi3BopGUQ5a5Og=
 github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
@@ -16,8 +18,10 @@ github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZN
 github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
 github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII=
 github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o=
-github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
-github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+github.com/smithy-security/pkg/utils v0.0.2 h1:r1Gz5eki8xUJXShw4i5ZaizkiKgZlYNYtKE2PDwpoHQ=
+github.com/smithy-security/pkg/utils v0.0.2/go.mod h1:bzCtRv/q9BdCrALRkcWWW3y8DzugbZrEQPwgZ/iepig=
+github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
+github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=

retry/retry.go

@@ -1,13 +1,16 @@
 package retry
 
 import (
-	"errors"
-	"fmt"
 	"log/slog"
 	"net/http"
 	"net/http/httputil"
+	"strconv"
+	"strings"
+	"time"
 
 	"github.com/cenkalti/backoff/v5"
+	"github.com/go-errors/errors"
+	"github.com/smithy-security/pkg/utils"
 )
 
 const defaultMaxRetries uint = 5
@@ -29,12 +32,17 @@ var (
 )
 
 type (
+	// ResponseInfoLogger is allowed to inspect the retryable response and
+	// print more info about it that might be useful to the caller for
+	// debugging
+	ResponseInfoLogger func(res *http.Response, logger Logger)
+
 	// Logger allows to inject a custom logger in the client.
 	Logger interface {
-		Error(msg string, keysAndValues ...interface{})
-		Info(msg string, keysAndValues ...interface{})
-		Debug(msg string, keysAndValues ...interface{})
-		Warn(msg string, keysAndValues ...interface{})
+		Error(msg string, keysAndValues ...any)
+		Info(msg string, keysAndValues ...any)
+		Debug(msg string, keysAndValues ...any)
+		Warn(msg string, keysAndValues ...any)
 	}
 
 	// NextRetryInSeconds allows customising the behaviour for the calculating the next retry.
@@ -59,6 +67,9 @@ type (
 		// AcceptedStatusCodes allows to specify the non-retryable status codes.
 		// defaultAcceptedStatusCodes are the default.
 		AcceptedStatusCodes map[int]struct{}
+		// ResponseInfoLoggerFunc when set will be used to check the response
+		// returned by the API
+		ResponseInfoLoggerFunc ResponseInfoLogger
 	}
 
 	retry struct {
@@ -124,7 +135,7 @@ func applyConfig(cfg Config) (Config, error) {
 func NewClient(config Config) (*http.Client, error) {
 	config, err := applyConfig(config)
 	if err != nil {
-		return nil, fmt.Errorf("failed to apply config: %w", err)
+		return nil, errors.Errorf("failed to apply config: %w", err)
 	}
 
 	config.BaseClient.Transport = &retry{
@@ -139,7 +150,7 @@ func NewClient(config Config) (*http.Client, error) {
 func NewRoundTripper(config Config) (http.RoundTripper, error) {
 	config, err := applyConfig(config)
 	if err != nil {
-		return nil, fmt.Errorf("failed to apply config: %w", err)
+		return nil, errors.Errorf("failed to apply config: %w", err)
 	}
 
 	return &retry{
@@ -148,6 +159,41 @@ func NewRoundTripper(config Config) (http.RoundTripper, error) {
 	}, nil
 }
 
+const noRetryHeader = -1
+
+// parseRetryHeader does a best effort parsing of the retry header
+func parseRetryHeader(logger Logger, resp *http.Response) int {
+	vals, ok := resp.Header["Retry-After"]
+	if !ok {
+		return noRetryHeader
+	}
+
+	logger.Debug("response contains retry after header", slog.String("vals", strings.Join(vals, ",")))
+	if len(vals) > 1 {
+		logger.Error("retry header has multiple values")
+		return noRetryHeader
+	}
+
+	retrySeconds, err := strconv.ParseInt(vals[0], 10, 32)
+	if err == nil {
+		return int(retrySeconds)
+	}
+
+	logger.Debug(
+		"could not parse `retry after` value into seconds, trying as a date",
+		slog.String("err", err.Error()),
+	)
+
+	retryAfterTime, err := time.Parse(http.TimeFormat, vals[0])
+	if err == nil {
+		logger.Debug("parsed successfully time from retry after header")
+		return int(time.Until(retryAfterTime).Seconds()) + 1
+	}
+
+	logger.Error("could not parse http time in retry header", slog.String("err", err.Error()))
+	return noRetryHeader
+}
+
 // RoundTrip implements a http transport RoundTripper with retry capabilities.
 func (re *retry) RoundTrip(req *http.Request) (*http.Response, error) {
 	var (
@@ -168,21 +214,30 @@ func (re *retry) RoundTrip(req *http.Request) (*http.Response, error) {
 			switch {
 			case !isAcceptedStatus && currAttempt >= re.config.MaxRetries:
 				return resp, backoff.Permanent(
-					fmt.Errorf(
+					errors.Errorf(
 						"maximum number of retries exceeded: %d",
 						currAttempt,
 					),
 				)
 			case !isAcceptedStatus && isRetryableStatus:
-				nextRetryInSeconds := re.config.NextRetryInSecondsFunc(currAttempt)
+				nextRetryInSeconds := parseRetryHeader(logger, resp)
+				if nextRetryInSeconds == noRetryHeader {
+					nextRetryInSeconds = re.config.NextRetryInSecondsFunc(currAttempt)
+				}
 
 				logger.Debug(
 					"retryable status code, retrying",
 					slog.Int("retry_in_seconds", nextRetryInSeconds),
 					slog.Int("curr_attempt", int(currAttempt)),
 					slog.Int("status_code", resp.StatusCode),
 				)
+
+				if !utils.IsNil(re.config.ResponseInfoLoggerFunc) {
+					re.config.ResponseInfoLoggerFunc(resp, logger)
+				}
+
 				currAttempt++
+
 				return resp, backoff.RetryAfter(nextRetryInSeconds)
 			case !isAcceptedStatus && !isRetryableStatus:
 				bb, err := httputil.DumpResponse(resp, true)
@@ -197,7 +252,8 @@ func (re *retry) RoundTrip(req *http.Request) (*http.Response, error) {
 					slog.Int("status_code", resp.StatusCode),
 					slog.String("raw_body", string(bb)),
 				)
-				return resp, backoff.Permanent(fmt.Errorf("invalid status code: %d", resp.StatusCode))
+
+				return resp, backoff.Permanent(errors.Errorf("invalid status code: %d", resp.StatusCode))
 			}
 
 			return resp, nil
@@ -209,7 +265,7 @@ func (re *retry) RoundTrip(req *http.Request) (*http.Response, error) {
 		retryableOp,
 	)
 	if err != nil {
-		return result, fmt.Errorf("could not process backoff result: %w", err)
+		return result, errors.Errorf("could not process backoff result: %w", err)
 	}
 
 	return result, nil

retry/retry_test.go

@@ -46,6 +46,15 @@ func TestRoundTripper(t *testing.T) {
 				Body:       io.NopCloser(bytes.NewBufferString(body)),
 			}
 		}
+		makeResponseWithHeader = func(statusCode int, headerName, headerVal, body string) *http.Response {
+			resp := &http.Response{
+				Header:     http.Header{},
+				StatusCode: statusCode,
+				Body:       io.NopCloser(bytes.NewBufferString(body)),
+			}
+			resp.Header.Add(headerName, headerVal)
+			return resp
+		}
 		zeroDelayRetry = func(uint) int { return 0 }
 	)
 
@@ -172,6 +181,30 @@ func TestRoundTripper(t *testing.T) {
 			expectedBody:   "success",
 			expectError:    false,
 		},
+		{
+			name: "response has retry after header",
+			responses: []*http.Response{
+				makeResponseWithHeader(http.StatusTooManyRequests, "Retry-After", "1", "rate limited 1"),
+				makeResponse(http.StatusOK, "success"),
+			},
+			errors:         []error{nil, nil, nil},
+			maxRetries:     3,
+			expectedStatus: http.StatusOK,
+			expectedBody:   "success",
+			expectError:    false,
+		},
+		{
+			name: "response has wrong retry after header",
+			responses: []*http.Response{
+				makeResponseWithHeader(http.StatusTooManyRequests, "Retry-After", "bla", "rate limited 1"),
+				makeResponse(http.StatusOK, "success"),
+			},
+			errors:         []error{nil, nil, nil},
+			maxRetries:     3,
+			expectedStatus: http.StatusOK,
+			expectedBody:   "success",
+			expectError:    false,
+		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			retryFunc := tc.retryFunc