Allow the HTTP fetcher to retry failed requests (#388)

Author: AlCutter

go.mod

@@ -13,6 +13,8 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.81.0
 	github.com/aws/aws-sdk-go-v2/service/secretsmanager v1.35.7
 	github.com/aws/smithy-go v1.22.4
+	github.com/cenkalti/backoff/v5 v5.0.2
+	github.com/dustin/go-humanize v1.0.1
 	github.com/gdamore/tcell/v2 v2.8.1
 	github.com/go-sql-driver/mysql v1.9.3
 	github.com/google/go-cmp v0.7.0
@@ -68,7 +70,6 @@ require (
 	github.com/cncf/xds/go v0.0.0-20250326154945-ae57f3c0d45f // indirect
 	github.com/dgraph-io/badger/v4 v4.7.0 // indirect
 	github.com/dgraph-io/ristretto/v2 v2.2.0 // indirect
-	github.com/dustin/go-humanize v1.0.1 // indirect
 	github.com/envoyproxy/go-control-plane/envoy v1.32.4 // indirect
 	github.com/envoyproxy/protoc-gen-validate v1.2.1 // indirect
 	github.com/felixge/httpsnoop v1.0.4 // indirect

go.sum

@@ -692,6 +692,8 @@ github.com/aws/smithy-go v1.22.4 h1:uqXzVZNuNexwc/xrh6Tb56u89WDlJY6HS+KC0S4QSjw=
 github.com/aws/smithy-go v1.22.4/go.mod h1:t1ufH5HMublsJYulve2RKmHDC15xu1f26kHCp/HgceI=
 github.com/boombuler/barcode v1.0.0/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8=
 github.com/boombuler/barcode v1.0.1/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8=
+github.com/cenkalti/backoff/v5 v5.0.2 h1:rIfFVxEf1QsI7E1ZHfp/B4DF/6QBAUhmgkxc0H7Zss8=
+github.com/cenkalti/backoff/v5 v5.0.2/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw=
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
 github.com/census-instrumentation/opencensus-proto v0.3.0/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
 github.com/census-instrumentation/opencensus-proto v0.4.1/go.mod h1:4T9NM4+4Vw91VeyqjLS6ao50K5bOcLKN6Q42XnYaRYw=

internal/client/fetcher.go

@@ -22,8 +22,10 @@ import (
 	"net/url"
 	"os"
 	"path"
+	"strconv"
 	"strings"
 
+	"github.com/cenkalti/backoff/v5"
 	"github.com/transparency-dev/tessera/api/layout"
 	"k8s.io/klog/v2"
 )
@@ -43,6 +45,7 @@ func NewHTTPFetcher(rootURL *url.URL, c *http.Client) (*HTTPFetcher, error) {
 	return &HTTPFetcher{
 		c:       c,
 		rootURL: rootURL,
+		backOff: []backoff.RetryOption{backoff.WithMaxTries(1)},
 	}, nil
 }
 
@@ -51,6 +54,7 @@ type HTTPFetcher struct {
 	c          *http.Client
 	rootURL    *url.URL
 	authHeader string
+	backOff    []backoff.RetryOption
 }
 
 // SetAuthorizationHeader sets the value to be used with an Authorization: header
@@ -59,39 +63,57 @@ func (h *HTTPFetcher) SetAuthorizationHeader(v string) {
 	h.authHeader = v
 }
 
-func (h HTTPFetcher) fetch(ctx context.Context, p string) ([]byte, error) {
-	u, err := h.rootURL.Parse(p)
-	if err != nil {
-		return nil, fmt.Errorf("invalid URL: %v", err)
-	}
-	req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
-	if err != nil {
-		return nil, fmt.Errorf("NewRequestWithContext(%q): %v", u.String(), err)
-	}
-	if h.authHeader != "" {
-		req.Header.Add("Authorization", h.authHeader)
-	}
-	r, err := h.c.Do(req)
-	if err != nil {
-		return nil, fmt.Errorf("get(%q): %v", u.String(), err)
-	}
-	switch r.StatusCode {
-	case http.StatusOK:
-		// All good, continue below
-		break
-	case http.StatusNotFound:
-		// Need to return ErrNotExist here, by contract.
-		return nil, fmt.Errorf("get(%q): %w", u.String(), os.ErrNotExist)
-	default:
-		return nil, fmt.Errorf("get(%q): %v", u.String(), r.StatusCode)
-	}
+// EnableRetries causes requests which result in a non-permanent error to be retried with up to maxRetries attempts.
+func (h *HTTPFetcher) EnableRetries(maxRetries uint) {
+	h.backOff = []backoff.RetryOption{backoff.WithBackOff(backoff.NewExponentialBackOff()), backoff.WithMaxTries(10)}
+}
 
-	defer func() {
-		if err := r.Body.Close(); err != nil {
-			klog.Errorf("resp.Body.Close(): %v", err)
+func (h HTTPFetcher) fetch(ctx context.Context, p string) ([]byte, error) {
+	return backoff.Retry(ctx, func() ([]byte, error) {
+		u, err := h.rootURL.Parse(p)
+		if err != nil {
+			return nil, fmt.Errorf("invalid URL: %v", err)
+		}
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
+		if err != nil {
+			return nil, fmt.Errorf("NewRequestWithContext(%q): %v", u.String(), err)
+		}
+		if h.authHeader != "" {
+			req.Header.Add("Authorization", h.authHeader)
 		}
-	}()
-	return io.ReadAll(r.Body)
+		r, err := h.c.Do(req)
+		if err != nil {
+			return nil, fmt.Errorf("get(%q): %v", u.String(), err)
+		}
+		switch r.StatusCode {
+		case http.StatusOK:
+			// All good, continue below
+			break
+		case http.StatusTooManyRequests:
+			seconds, err := strconv.ParseInt(r.Header.Get("Retry-After"), 10, 32)
+			if err != nil {
+				// The server didn't say how long to wait, so we'll wait an arbitrary amount of time.
+				seconds = 10
+			}
+			return nil, backoff.RetryAfter(int(seconds))
+		case http.StatusNotFound:
+			// Need to return ErrNotExist here, by contract, and also let the backoff know not to retry).
+			return nil, backoff.Permanent(fmt.Errorf("get(%q): %w", u.String(), os.ErrNotExist))
+		case http.StatusBadRequest, http.StatusUnauthorized, http.StatusForbidden, http.StatusMethodNotAllowed, http.StatusConflict, http.StatusUnprocessableEntity:
+			// Should not retry for any of these status codes.
+			return nil, backoff.Permanent(fmt.Errorf("get(%q): %v", u.String(), r.StatusCode))
+		default:
+			// Everything else will be retried
+			return nil, fmt.Errorf("get(%q): %v", u.String(), r.StatusCode)
+		}
+
+		defer func() {
+			if err := r.Body.Close(); err != nil {
+				klog.Errorf("resp.Body.Close(): %v", err)
+			}
+		}()
+		return io.ReadAll(r.Body)
+	}, h.backOff...)
 }
 
 func (h HTTPFetcher) ReadCheckpoint(ctx context.Context) ([]byte, error) {

internal/hammer/hammer.go

@@ -426,6 +426,7 @@ func mustCreateReaders(ctx context.Context, us []string) loadtest.LogReader {
 			if err != nil {
 				klog.Exitf("Failed to create HTTP fetcher for %q: %v", u, err)
 			}
+			c.EnableRetries(5)
 			if *bearerToken != "" {
 				c.SetAuthorizationHeader(fmt.Sprintf("Bearer %s", *bearerToken))
 			}