Retry sidecar API 429s honouring Retry-After (30s budget)

internal/circleci/client.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"time"
 
 	hc "github.com/CircleCI-Public/chunk-cli/internal/httpcl"
 	"github.com/CircleCI-Public/chunk-cli/internal/version"
@@ -33,10 +34,11 @@ func NewClient(cfg Config) (*Client, error) {
 		return nil, ErrTokenNotFound
 	}
 	cl := hc.New(hc.Config{
-		BaseURL:    cfg.BaseURL,
-		AuthToken:  cfg.Token,
-		AuthHeader: "Circle-Token",
-		UserAgent:  version.UserAgent(),
+		BaseURL:          cfg.BaseURL,
+		AuthToken:        cfg.Token,
+		AuthHeader:       "Circle-Token",
+		UserAgent:        version.UserAgent(),
+		RetryOn429Budget: 30 * time.Second,
 	})
 	return &Client{cl: cl}, nil
 }

internal/httpcl/client.go

@@ -11,11 +11,22 @@ import (
 	"io"
 	"net/http"
 	"net/url"
+	"strconv"
 	"time"
 
 	"github.com/hashicorp/go-retryablehttp"
 )
 
+// retryCtxKey is the context key for the per-call retry state.
+type retryCtxKey struct{}
+
+// retryState tracks per-call retry counters stored in the request context.
+// Using a pointer allows mutation across CheckRetry invocations for the same call.
+type retryState struct {
+	start                time.Time
+	nonRateLimitAttempts int
+}
+
 const jsonContentType = "application/json; charset=utf-8"
 
 // Config configures a Client.
@@ -34,18 +45,24 @@ type Config struct {
 	// DisableRetries disables automatic retries. By default requests are
 	// retried up to 3 times with exponential backoff.
 	DisableRetries bool
+	// RetryOn429Budget, when non-zero, enables retrying HTTP 429 responses by
+	// honouring the Retry-After response header. Retries stop when the
+	// cumulative wait time would exceed this budget, or when a single
+	// Retry-After value exceeds it, and a RateLimitError is returned.
+	RetryOn429Budget time.Duration
 	// Transport overrides the HTTP transport (useful for testing).
 	Transport http.RoundTripper
 }
 
 // Client is a simple HTTP client with JSON defaults and automatic retries.
 type Client struct {
-	baseURL    string
-	authToken  string
-	authHeader string
-	userAgent  string
-	timeout    time.Duration
-	http       *retryablehttp.Client
+	baseURL          string
+	authToken        string
+	authHeader       string
+	userAgent        string
+	timeout          time.Duration
+	retryOn429Budget time.Duration
+	http             *retryablehttp.Client
 }
 
 // New creates a Client from the given config.
@@ -64,20 +81,75 @@ func New(cfg Config) *Client {
 	rc.RetryWaitMax = 2 * time.Second
 	rc.Logger = nil // suppress default log output
 
+	if cfg.RetryOn429Budget > 0 {
+		budget := cfg.RetryOn429Budget
+		origMax := 3
+		if cfg.DisableRetries {
+			origMax = 0
+		}
+		// Raise RetryMax so it never binds before the budget does.
+		// Each 429 retry consumes ≥1s (Retry-After floor), so budget/s + origMax is sufficient.
+		rc.RetryMax = int(budget/time.Second) + origMax
+		rc.CheckRetry = func(ctx context.Context, resp *http.Response, err error) (bool, error) {
+			state, _ := ctx.Value(retryCtxKey{}).(*retryState)
+			if resp != nil && resp.StatusCode == http.StatusTooManyRequests {
+				retryAfter := parseRetryAfter(resp)
+				elapsed := time.Duration(0)
+				if state != nil {
+					elapsed = time.Since(state.start)
+				}
+				if elapsed+retryAfter > budget {
+					return false, &RateLimitError{RetryAfter: retryAfter, Budget: budget}
+				}
+				return true, nil
+			}
+			// Cap non-429 retries at the original limit.
+			if state != nil {
+				state.nonRateLimitAttempts++
+				if state.nonRateLimitAttempts > origMax {
+					return false, nil
+				}
+			}
+			return retryablehttp.DefaultRetryPolicy(ctx, resp, err)
+		}
+		// DefaultBackoff already honours Retry-After; keep it.
+	}
+
 	if cfg.Transport != nil {
 		rc.HTTPClient.Transport = cfg.Transport
 	}
 
 	return &Client{
-		baseURL:    cfg.BaseURL,
-		authToken:  cfg.AuthToken,
-		authHeader: cfg.AuthHeader,
-		userAgent:  cfg.UserAgent,
-		timeout:    timeout,
-		http:       rc,
+		baseURL:          cfg.BaseURL,
+		authToken:        cfg.AuthToken,
+		authHeader:       cfg.AuthHeader,
+		userAgent:        cfg.UserAgent,
+		timeout:          timeout,
+		retryOn429Budget: cfg.RetryOn429Budget,
+		http:             rc,
 	}
 }
 
+// parseRetryAfter parses the Retry-After header as seconds or an HTTP date.
+func parseRetryAfter(resp *http.Response) time.Duration {
+	ra := resp.Header.Get("Retry-After")
+	if ra == "" {
+		return 0
+	}
+	if secs, err := strconv.ParseInt(ra, 10, 64); err == nil {
+		if secs > 0 {
+			return time.Duration(secs) * time.Second
+		}
+		return 0
+	}
+	if t, err := http.ParseTime(ra); err == nil {
+		if d := time.Until(t); d > 0 {
+			return d
+		}
+	}
+	return 0
+}
+
 // Call executes the request and returns the HTTP status code.
 // Non-2xx responses return an *HTTPError. If a decoder is set and the
 // response is 2xx, the response body is decoded.
@@ -99,7 +171,12 @@ func (c *Client) Call(ctx context.Context, r Request) (int, error) {
 		bodyReader = bytes.NewReader(b)
 	}
 
-	ctx, cancel := context.WithTimeout(ctx, c.timeout)
+	ctxTimeout := c.timeout
+	if c.retryOn429Budget > 0 {
+		ctx = context.WithValue(ctx, retryCtxKey{}, &retryState{start: time.Now()})
+		ctxTimeout = c.retryOn429Budget + c.timeout // extend deadline to cover retry waits
+	}
+	ctx, cancel := context.WithTimeout(ctx, ctxTimeout)
 	defer cancel()
 
 	req, err := retryablehttp.NewRequestWithContext(ctx, r.method, u.String(), bodyReader)
@@ -131,13 +208,15 @@ func (c *Client) Call(ctx context.Context, r Request) (int, error) {
 	}
 
 	resp, err := c.http.Do(req)
+	if resp != nil {
+		defer func() {
+			_, _ = io.Copy(io.Discard, resp.Body)
+			_ = resp.Body.Close()
+		}()
+	}
 	if err != nil {
 		return 0, err
 	}
-	defer func() {
-		_, _ = io.Copy(io.Discard, resp.Body)
-		_ = resp.Body.Close()
-	}()
 
 	status := resp.StatusCode
 

internal/httpcl/client_internal_test.go

@@ -0,0 +1,56 @@
+package httpcl
+
+import (
+	"net/http"
+	"testing"
+	"time"
+)
+
+func TestParseRetryAfter_Seconds(t *testing.T) {
+	resp := &http.Response{Header: http.Header{"Retry-After": []string{"30"}}}
+	if d := parseRetryAfter(resp); d != 30*time.Second {
+		t.Fatalf("expected 30s, got %v", d)
+	}
+}
+
+func TestParseRetryAfter_ZeroSeconds(t *testing.T) {
+	resp := &http.Response{Header: http.Header{"Retry-After": []string{"0"}}}
+	if d := parseRetryAfter(resp); d != 0 {
+		t.Fatalf("expected 0, got %v", d)
+	}
+}
+
+func TestParseRetryAfter_HTTPDate_Future(t *testing.T) {
+	future := time.Now().Add(5 * time.Second)
+	resp := &http.Response{
+		Header: http.Header{"Retry-After": []string{future.UTC().Format(http.TimeFormat)}},
+	}
+	d := parseRetryAfter(resp)
+	if d < 3*time.Second || d > 6*time.Second {
+		t.Fatalf("expected ~5s from HTTP-date, got %v", d)
+	}
+}
+
+func TestParseRetryAfter_HTTPDate_Past(t *testing.T) {
+	past := time.Now().Add(-5 * time.Second)
+	resp := &http.Response{
+		Header: http.Header{"Retry-After": []string{past.UTC().Format(http.TimeFormat)}},
+	}
+	if d := parseRetryAfter(resp); d != 0 {
+		t.Fatalf("expected 0 for past date, got %v", d)
+	}
+}
+
+func TestParseRetryAfter_Missing(t *testing.T) {
+	resp := &http.Response{Header: http.Header{}}
+	if d := parseRetryAfter(resp); d != 0 {
+		t.Fatalf("expected 0 for missing header, got %v", d)
+	}
+}
+
+func TestParseRetryAfter_Invalid(t *testing.T) {
+	resp := &http.Response{Header: http.Header{"Retry-After": []string{"garbage"}}}
+	if d := parseRetryAfter(resp); d != 0 {
+		t.Fatalf("expected 0 for invalid value, got %v", d)
+	}
+}