From 80a636bc4217d8c201adaa65b6b1c6f01b669d5b Mon Sep 17 00:00:00 2001 From: Elizabeth Worstell Date: Mon, 9 Nov 2026 16:00:53 -0800 Subject: [PATCH 1/3] fix: keep mirrors available during upstream cooldowns Amp-Thread-ID: https://ampcode.com/threads/T-01a1c2a3-3b00-7292-bb00-b4a066f5cd6b Co-authored-by: Amp --- README.md | 4 +- internal/httputil/backoff.go | 69 ++++++++++++++--------- internal/httputil/backoff_test.go | 57 +++++++++++++++++-- internal/strategy/github_releases_test.go | 10 +++- internal/strategy/handler/handler.go | 10 +++- internal/strategy/handler/handler_test.go | 16 ++++-- internal/strategy/hermit.go | 2 +- internal/strategy/hermit_test.go | 44 +++++++++++++++ internal/strategy/host_test.go | 38 ++++++++++--- internal/strategy/oci.go | 2 +- internal/strategy/oci_auth.go | 54 +++++++++++++++--- internal/strategy/oci_test.go | 60 ++++++++++++++++++++ 12 files changed, 304 insertions(+), 62 deletions(-) diff --git a/README.md b/README.md index a723f2d0..f05de460 100644 --- a/README.md +++ b/README.md @@ -219,7 +219,9 @@ spool { ## Upstream Rate Limits -After an upstream returns 429, cachewd stops sending upstream requests for that strategy until a cooldown ends. This includes its fallbacks and redirect targets. During the cooldown, misses get 429 with `Retry-After` and `Cache-Control: no-store` without contacting any upstream. Cache hits and other strategies are unaffected. Requests made outside the strategy handlers, such as Git, pause per origin (scheme, host and port). Nothing is retried automatically. +After an upstream returns 429, cachewd pauses requests to that origin (scheme, host and port) within the strategy until a cooldown ends. A rate-limited redirect target also pauses the origins that redirected to it. Cache hits, other strategies, and unrelated origins remain available. OCI blob, manifest, authentication, and pass-through requests share the registry strategy's cooldowns; Hermit's direct and release handlers share its cooldowns. + +Configured fallbacks are tried on the first 429 and on subsequent cache misses while the primary is cooling. A mirror's 404 or 5xx does not freeze it; a mirror's 429 starts its own cooldown. When every applicable upstream is cooling, misses get a local 429 with `Retry-After` and `Cache-Control: no-store` without contacting upstreams. Requests outside strategy handlers, such as Git, share per-origin cooldowns without a strategy scope. Cooldowns do not schedule retries. The cooldown uses the upstream's `Retry-After` (seconds or HTTP date) when it is valid, and `upstream-backoff` otherwise. Cooldowns are held in memory per process. diff --git a/internal/httputil/backoff.go b/internal/httputil/backoff.go index dd40cc5d..d8251a54 100644 --- a/internal/httputil/backoff.go +++ b/internal/httputil/backoff.go @@ -15,36 +15,27 @@ import ( ) // BackoffTransport shares upstream cooldowns across requests and redirect hops. -// Requests in a [BackoffScope] share one cooldown; other requests share one -// per origin. +// Cooldowns are per origin and isolated by [BackoffScope]. type BackoffTransport struct { base http.RoundTripper delay time.Duration mu sync.Mutex - until map[any]time.Time + until map[backoffKey]time.Time +} + +type backoffKey struct { + scope *BackoffScope + origin string } // NewBackoffTransport uses delay when a 429 has no valid Retry-After header. func NewBackoffTransport(base http.RoundTripper, delay time.Duration) *BackoffTransport { - return &BackoffTransport{base: base, delay: delay, until: map[any]time.Time{}} + return &BackoffTransport{base: base, delay: delay, until: map[backoffKey]time.Time{}} } -// RoundTrip rejects requests in a cooling scope or origin without retrying them. +// RoundTrip rejects requests to a cooling origin within their scope without retrying them. func (t *BackoffTransport) RoundTrip(r *http.Request) (*http.Response, error) { - var cooldown any - if scope, ok := r.Context().Value(backoffScopeKey{}).(*BackoffScope); ok { - cooldown = scope - } else { - port := r.URL.Port() - if port == "" { - if r.URL.Scheme == "https" { - port = "443" - } else { - port = "80" - } - } - cooldown = strings.ToLower(r.URL.Scheme + "://" + net.JoinHostPort(r.URL.Hostname(), port)) - } + cooldown := backoffKeyFor(r) now := time.Now() t.mu.Lock() for k, until := range t.until { @@ -55,6 +46,7 @@ func (t *BackoffTransport) RoundTrip(r *http.Request) (*http.Response, error) { until := t.until[cooldown] t.mu.Unlock() if until.After(now) { + t.recordCooldown(r, until) if r.Body != nil { _ = r.Body.Close() } @@ -78,11 +70,7 @@ func (t *BackoffTransport) RoundTrip(r *http.Request) (*http.Response, error) { now = time.Now() var valid bool until, valid = retryAfter(resp.Header.Get("Retry-After"), now, t.delay) - t.mu.Lock() - if until.After(t.until[cooldown]) { - t.until[cooldown] = until - } - t.mu.Unlock() + t.recordCooldown(r, until) // The cooldown ends the 429; a downstream cache must not keep it longer. resp.Header.Set("Cache-Control", "no-store") if !valid { @@ -92,10 +80,39 @@ func (t *BackoffTransport) RoundTrip(r *http.Request) (*http.Response, error) { return resp, nil } +func (t *BackoffTransport) recordCooldown(r *http.Request, until time.Time) { + scope := backoffKeyFor(r).scope + t.mu.Lock() + defer t.mu.Unlock() + for hop := r; hop != nil; { + key := backoffKeyFor(hop) + key.scope = scope + if until.After(t.until[key]) { + t.until[key] = until + } + if hop.Response == nil { + break + } + hop = hop.Response.Request + } +} + +func backoffKeyFor(r *http.Request) backoffKey { + port := r.URL.Port() + if port == "" { + if r.URL.Scheme == "https" { + port = "443" + } else { + port = "80" + } + } + scope, _ := r.Context().Value(backoffScopeKey{}).(*BackoffScope) + return backoffKey{scope: scope, origin: strings.ToLower(r.URL.Scheme + "://" + net.JoinHostPort(r.URL.Hostname(), port))} +} + type backoffScopeKey struct{} -// BackoffScope sends every request through it in one [BackoffTransport] -// cooldown, whatever origin the request reaches. +// BackoffScope isolates a strategy's per-origin cooldowns from other strategies. type BackoffScope struct { base http.RoundTripper } diff --git a/internal/httputil/backoff_test.go b/internal/httputil/backoff_test.go index 61ebcb6f..8cb80548 100644 --- a/internal/httputil/backoff_test.go +++ b/internal/httputil/backoff_test.go @@ -99,16 +99,16 @@ func TestBackoffRedirectDestination(t *testing.T) { base := roundTripFunc(func(r *http.Request) (*http.Response, error) { calls[r.URL.Host]++ resp := &http.Response{StatusCode: http.StatusTooManyRequests, Header: make(http.Header), Body: http.NoBody, Request: r} - if r.URL.Host == "redirect.example" { + if r.URL.Host != "limited.example" { resp.StatusCode = http.StatusFound resp.Header.Set("Location", "https://limited.example/file") } return resp, nil }) transport := httputil.NewBackoffTransport(base, time.Minute) - for range 2 { - client := &http.Client{Transport: transport} - req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, "https://redirect.example/file", nil) + client := &http.Client{Transport: httputil.NewBackoffScope(transport)} + for _, host := range []string{"redirect.example", "redirect.example", "another.example", "another.example"} { + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, "https://"+host+"/file", nil) assert.NoError(t, err) resp, err := client.Do(req) assert.NoError(t, err) @@ -116,7 +116,52 @@ func TestBackoffRedirectDestination(t *testing.T) { assert.Equal(t, http.StatusTooManyRequests, resp.StatusCode) assert.Equal(t, "60", resp.Header.Get("Retry-After")) } - assert.Equal(t, 2, calls["redirect.example"]) - assert.Equal(t, 1, calls["limited.example"]) + assert.Equal(t, map[string]int{"redirect.example": 1, "another.example": 1, "limited.example": 1}, calls) + time.Sleep(time.Minute) + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, "https://redirect.example/file", nil) + assert.NoError(t, err) + resp, err := client.Do(req) + assert.NoError(t, err) + assert.NoError(t, resp.Body.Close()) + assert.Equal(t, map[string]int{"redirect.example": 2, "another.example": 1, "limited.example": 2}, calls) + }) +} + +func TestBackoffScopes(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + calls := map[string]int{} + base := roundTripFunc(func(r *http.Request) (*http.Response, error) { + calls[r.URL.Host]++ + status := http.StatusOK + if r.URL.Path == "/limited" { + status = http.StatusTooManyRequests + } + return &http.Response{StatusCode: status, Header: make(http.Header), Body: http.NoBody, Request: r}, nil + }) + transport := httputil.NewBackoffTransport(base, time.Minute) + first := &http.Client{Transport: httputil.NewBackoffScope(transport)} + second := &http.Client{Transport: httputil.NewBackoffScope(transport)} + unscoped := &http.Client{Transport: transport} + for _, tt := range []struct { + client *http.Client + url string + status int + }{ + {first, "https://primary.example/limited", http.StatusTooManyRequests}, + {first, "https://PRIMARY.example:443/other", http.StatusTooManyRequests}, + {first, "https://mirror.example/other", http.StatusOK}, + {second, "https://primary.example/other", http.StatusOK}, + {unscoped, "https://primary.example/other", http.StatusOK}, + {unscoped, "https://mirror.example/limited", http.StatusTooManyRequests}, + {first, "https://mirror.example/another", http.StatusOK}, + } { + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, tt.url, nil) + assert.NoError(t, err) + resp, err := tt.client.Do(req) + assert.NoError(t, err) + assert.NoError(t, resp.Body.Close()) + assert.Equal(t, tt.status, resp.StatusCode) + } + assert.Equal(t, map[string]int{"primary.example": 3, "mirror.example": 3}, calls) }) } diff --git a/internal/strategy/github_releases_test.go b/internal/strategy/github_releases_test.go index bc9bd69d..2fd5090b 100644 --- a/internal/strategy/github_releases_test.go +++ b/internal/strategy/github_releases_test.go @@ -273,20 +273,24 @@ func TestGitHubReleasesPrivateRepoAPIError(t *testing.T) { assert.Equal(t, 0, mock.downloadCallCount) } -func TestGitHubReleasesRateLimitPausesStrategy(t *testing.T) { +func TestGitHubReleasesRateLimitPausesAPI(t *testing.T) { mock, mux, ctx := setupTest(t, strategy.GitHubReleasesConfig{ Token: "test-token", PrivateOrgs: []string{"privateorg"}, }) - for _, path := range []string{"privateorg/repo/releases/download/limited/binary.tar.gz", "publicorg/repo/releases/download/v1.0.0/binary.tar.gz"} { + for _, path := range []string{"privateorg/repo/releases/download/limited/binary.tar.gz", "privateorg/repo/releases/download/v1.0.0/binary.tar.gz"} { w := httptest.NewRecorder() mux.ServeHTTP(w, httptest.NewRequestWithContext(ctx, http.MethodGet, "/github.com/"+path, nil)) assert.Equal(t, http.StatusTooManyRequests, w.Code) assert.Equal(t, "no-store", w.Header().Get("Cache-Control")) assert.Equal(t, "30", w.Header().Get("Retry-After")) } + w := httptest.NewRecorder() + mux.ServeHTTP(w, httptest.NewRequestWithContext(ctx, http.MethodGet, + "/github.com/publicorg/repo/releases/download/v1.0.0/binary.tar.gz", nil)) + assert.Equal(t, http.StatusOK, w.Code) assert.Equal(t, 1, mock.apiCallCount) - assert.Equal(t, 0, mock.publicCallCount) + assert.Equal(t, 1, mock.publicCallCount) } func TestGitHubReleasesPublicRepoNotFound(t *testing.T) { diff --git a/internal/strategy/handler/handler.go b/internal/strategy/handler/handler.go index ea18bd52..05517ad9 100644 --- a/internal/strategy/handler/handler.go +++ b/internal/strategy/handler/handler.go @@ -100,14 +100,18 @@ const checkedAtHeader = "X-Cachew-Checked-At" // New creates a new Handler with the given HTTP client and cache. Concurrent // cache misses for the same key share one upstream fetch through a spool group. -// All upstream requests of the handler, including requests its transform makes, -// fallbacks, redirects and parallel chunks, share one 429 cooldown. +// The client's BackoffScope is reused when present, otherwise a new scope is +// created. Upstream requests, including transforms and parallel chunks, share +// per-origin cooldowns within that scope; fallback origins remain independent. // By default: // - Cache key is derived from the request URL // - No request transformation is performed // - Standard error handling is used. func New(client *http.Client, c cache.Cache, spools *spool.Registry) *Handler { - scope := httputil.NewBackoffScope(client.Transport) + scope, ok := client.Transport.(*httputil.BackoffScope) + if !ok { + scope = httputil.NewBackoffScope(client.Transport) + } scoped := *client scoped.Transport = scope return &Handler{ diff --git a/internal/strategy/handler/handler_test.go b/internal/strategy/handler/handler_test.go index e6e4c9f9..1a16565f 100644 --- a/internal/strategy/handler/handler_test.go +++ b/internal/strategy/handler/handler_test.go @@ -1033,7 +1033,7 @@ func TestBackoffPausesHandler(t *testing.T) { {limited.URL, "/cached"}, {limited.URL, "/limited"}, {limited.URL, "/cached"}, {limited.URL, "/miss"}, {other.URL, "/miss"}, } { w := get(h, req.upstream, req.path) - if req.path == "/cached" { + if req.path == "/cached" || req.upstream == other.URL { assert.Equal(t, http.StatusOK, w.Code) assert.Equal(t, "body", w.Body.String()) } else { @@ -1041,10 +1041,18 @@ func TestBackoffPausesHandler(t *testing.T) { assert.Equal(t, "60", w.Header().Get("Retry-After")) } } - assert.Equal(t, int32(2), calls.Load()) + assert.Equal(t, int32(3), calls.Load()) - w := get(newHandler(), other.URL, "/miss") + w := get(newHandler(), limited.URL, "/independent") assert.Equal(t, http.StatusOK, w.Code) assert.Equal(t, "body", w.Body.String()) - assert.Equal(t, int32(3), calls.Load()) + assert.Equal(t, int32(4), calls.Load()) + + client.Transport = httputil.NewBackoffScope(client.Transport) + first, second := newHandler(), newHandler() + assert.Equal(t, http.StatusTooManyRequests, get(first, limited.URL, "/limited").Code) + w = get(second, limited.URL, "/shared") + assert.Equal(t, http.StatusTooManyRequests, w.Code) + assert.Equal(t, "60", w.Header().Get("Retry-After")) + assert.Equal(t, int32(5), calls.Load()) } diff --git a/internal/strategy/hermit.go b/internal/strategy/hermit.go index bcc524af..4c49b402 100644 --- a/internal/strategy/hermit.go +++ b/internal/strategy/hermit.go @@ -62,7 +62,7 @@ func NewHermit(ctx context.Context, config HermitConfig, _ jobscheduler.Schedule s := &Hermit{ config: config, cache: c, - client: http.DefaultClient, + client: &http.Client{Transport: httputil.NewBackoffScope(http.DefaultTransport)}, logger: logger, mux: mux, } diff --git a/internal/strategy/hermit_test.go b/internal/strategy/hermit_test.go index 598d3db3..1bbd1705 100644 --- a/internal/strategy/hermit_test.go +++ b/internal/strategy/hermit_test.go @@ -13,6 +13,7 @@ import ( "github.com/block/cachew/internal/cache" "github.com/block/cachew/internal/githubapp" + "github.com/block/cachew/internal/httputil" "github.com/block/cachew/internal/logging" "github.com/block/cachew/internal/strategy" ) @@ -36,6 +37,49 @@ func setupHermitTest(t *testing.T) (*http.ServeMux, context.Context, cache.Cache return mux, ctx, memCache } +func TestHermitBackoff(t *testing.T) { + calls := 0 + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + if r.URL.Path == "/org/repo/releases/download/v1/file" { + w.WriteHeader(http.StatusTooManyRequests) + return + } + _, _ = w.Write([]byte("package")) + })) + defer backend.Close() + original := http.DefaultTransport + http.DefaultTransport = httputil.NewBackoffTransport( //nolint:reassign + &mockTransport{backend: backend, originalTransport: original}, time.Minute) + t.Cleanup(func() { http.DefaultTransport = original }) //nolint:reassign + _, ctx := logging.Configure(t.Context(), logging.Config{Level: slog.LevelError}) + c, err := cache.NewMemory(ctx, cache.MemoryConfig{MaxTTL: time.Hour}) + assert.NoError(t, err) + defer c.Close() + mux := http.NewServeMux() + _, err = strategy.NewHermit(ctx, strategy.HermitConfig{GitHubBaseURL: "https://github.com"}, + nil, c, mux, newTestSpools(t)) + assert.NoError(t, err) + for _, tt := range []struct { + path string + status int + }{ + {"/hermit/github.com/cached", http.StatusOK}, + {"/hermit/github.com/org/repo/releases/download/v1/file", http.StatusTooManyRequests}, + {"/hermit/github.com/another", http.StatusTooManyRequests}, + {"/hermit/example.com/another", http.StatusOK}, + {"/hermit/github.com/cached", http.StatusOK}, + } { + w := httptest.NewRecorder() + mux.ServeHTTP(w, httptest.NewRequestWithContext(ctx, http.MethodGet, tt.path, nil)) + assert.Equal(t, tt.status, w.Code, tt.path) + if tt.status == http.StatusTooManyRequests { + assert.Equal(t, "60", w.Header().Get("Retry-After")) + } + } + assert.Equal(t, 3, calls) +} + func TestHermitBinaryShortTTL(t *testing.T) { httpTransportMutexHermit.Lock() defer httpTransportMutexHermit.Unlock() diff --git a/internal/strategy/host_test.go b/internal/strategy/host_test.go index 52d08b62..6ce3fc37 100644 --- a/internal/strategy/host_test.go +++ b/internal/strategy/host_test.go @@ -304,7 +304,7 @@ func TestHostFallbackOrder(t *testing.T) { assert.Equal(t, []int{0, 1, 2}, order) } -func TestHostBackoffPausesFallbacks(t *testing.T) { +func TestHostBackoffKeepsFallbacksAvailable(t *testing.T) { defaultTransport := http.DefaultTransport http.DefaultTransport = httputil.NewBackoffTransport(defaultTransport, time.Minute) //nolint:reassign t.Cleanup(func() { http.DefaultTransport = defaultTransport }) //nolint:reassign @@ -314,9 +314,16 @@ func TestHostBackoffPausesFallbacks(t *testing.T) { defer c.Close() calls := map[string]int{} newServer := func(name string, status int) *httptest.Server { - return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls[name]++ - w.WriteHeader(status) + responseStatus := status + if r.URL.Path == "/missing.jar" { + responseStatus = http.StatusNotFound + } else if r.URL.Path == "/limited.jar" { + responseStatus = http.StatusTooManyRequests + } + w.WriteHeader(responseStatus) + _, _ = io.WriteString(w, name+r.URL.Path) })) } primary := newServer("primary", http.StatusTooManyRequests) @@ -331,13 +338,28 @@ func TestHostBackoffPausesFallbacks(t *testing.T) { assert.NoError(t, err) u, err := url.Parse(primary.URL) assert.NoError(t, err) - for _, path := range []string{"/a.jar", "/b.jar"} { + for _, tt := range []struct { + path string + status int + }{ + {"/a.jar", http.StatusOK}, + {"/missing.jar", http.StatusNotFound}, + {"/b.jar", http.StatusOK}, + {"/limited.jar", http.StatusTooManyRequests}, + {"/c.jar", http.StatusTooManyRequests}, + {"/a.jar", http.StatusOK}, + } { w := httptest.NewRecorder() - mux.ServeHTTP(w, httptest.NewRequestWithContext(ctx, http.MethodGet, "/"+u.Host+path, nil)) - assert.Equal(t, http.StatusTooManyRequests, w.Code) - assert.Equal(t, "60", w.Header().Get("Retry-After")) + mux.ServeHTTP(w, httptest.NewRequestWithContext(ctx, http.MethodGet, "/"+u.Host+tt.path, nil)) + assert.Equal(t, tt.status, w.Code, tt.path) + if tt.status == http.StatusTooManyRequests { + assert.Equal(t, "60", w.Header().Get("Retry-After")) + assert.Equal(t, "no-store", w.Header().Get("Cache-Control")) + } else { + assert.Equal(t, "mirror"+tt.path, w.Body.String()) + } } - assert.Equal(t, map[string]int{"primary": 1}, calls) + assert.Equal(t, map[string]int{"primary": 1, "mirror": 4}, calls) } func TestHostFallbackIncompleteResponse(t *testing.T) { diff --git a/internal/strategy/oci.go b/internal/strategy/oci.go index 0e8ca08e..cc58c87f 100644 --- a/internal/strategy/oci.go +++ b/internal/strategy/oci.go @@ -113,7 +113,7 @@ func NewOCI(ctx context.Context, config OCIConfig, c cache.Cache, mux Mux, spool upstream: upstream, rewrites: rewrites, tagTTL: config.TagTTL, - client: &http.Client{Transport: transport}, + client: &http.Client{Transport: httputil.NewBackoffScope(transport)}, logger: logging.FromContext(ctx), } o.blobs = handler.New(o.client, c, spools).Metrics("oci", upstream.Hostname()). diff --git a/internal/strategy/oci_auth.go b/internal/strategy/oci_auth.go index bf57e341..8107bb68 100644 --- a/internal/strategy/oci_auth.go +++ b/internal/strategy/oci_auth.go @@ -1,6 +1,8 @@ package strategy import ( + "bytes" + "io" "net/http" "net/url" "sync" @@ -45,22 +47,26 @@ func newOCITransport(upstream *url.URL, username, password string, base http.Rou return &ociTransport{ registry: registry, auth: auth, - base: base, + base: ociRateLimitTransport{base: base}, repositories: map[string]*ociRepositoryTransport{}, }, nil } func (t *ociTransport) RoundTrip(req *http.Request) (*http.Response, error) { + var resp *http.Response rt, err := t.forRepository(req) - if err != nil { - return nil, err + if err == nil { + // Client credentials belong to cachew, not to the upstream registry. + req = req.Clone(req.Context()) + req.Header.Del("Authorization") + req.Header.Del("Cookie") + resp, err = rt.RoundTrip(req) + } + if limited, ok := errors.AsType[*ociRateLimitError](err); ok { + limited.response.Request = req + return limited.response, nil } - // Client credentials are for cachew, and the upstream must only see the - // strategy's own. RoundTrippers must not modify the request they receive. - req = req.Clone(req.Context()) - req.Header.Del("Authorization") - req.Header.Del("Cookie") - return errors.WithStack2(rt.RoundTrip(req)) + return resp, errors.WithStack(err) } func (t *ociTransport) forRepository(req *http.Request) (http.RoundTripper, error) { @@ -88,3 +94,33 @@ func (t *ociTransport) forRepository(req *http.Request) (http.RoundTripper, erro entry.transport = rt return rt, nil } + +// ociRateLimitTransport preserves 429 headers that the auth library discards +// when turning registry ping and token responses into errors. +type ociRateLimitTransport struct { + base http.RoundTripper +} + +func (t ociRateLimitTransport) RoundTrip(req *http.Request) (*http.Response, error) { + resp, err := t.base.RoundTrip(req) + if err != nil || resp.StatusCode != http.StatusTooManyRequests { + return resp, errors.WithStack(err) + } + defer resp.Body.Close() + body, err := io.ReadAll(io.LimitReader(resp.Body, 64<<10)) + if err != nil { + return nil, errors.Wrap(err, "read registry rate limit response") + } + limited := *resp + limited.Body = io.NopCloser(bytes.NewReader(body)) + limited.ContentLength = int64(len(body)) + limited.Header = resp.Header.Clone() + limited.Header.Del("Content-Length") + return nil, &ociRateLimitError{response: &limited} +} + +type ociRateLimitError struct { + response *http.Response +} + +func (e *ociRateLimitError) Error() string { return "registry rate limited" } diff --git a/internal/strategy/oci_test.go b/internal/strategy/oci_test.go index c9771053..cda0e7b9 100644 --- a/internal/strategy/oci_test.go +++ b/internal/strategy/oci_test.go @@ -16,6 +16,7 @@ import ( "github.com/opencontainers/go-digest" "github.com/block/cachew/internal/cache" + "github.com/block/cachew/internal/httputil" "github.com/block/cachew/internal/logging" "github.com/block/cachew/internal/strategy" ) @@ -48,6 +49,8 @@ type fakeRegistry struct { corrupt bool headStatus int requireToken bool + limitedPath string + exchanges int requests []fakeRequest cdnRequests []fakeRequest authorizations []string @@ -87,6 +90,13 @@ func (f *fakeRegistry) redirectBlobs(t *testing.T) { func (f *fakeRegistry) serve(w http.ResponseWriter, r *http.Request) { f.mu.Lock() defer f.mu.Unlock() + f.exchanges++ + if r.URL.Path == f.limitedPath { + w.Header().Set("Retry-After", "120") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte("registry limit")) + return + } if r.URL.Path == "/token" { username, _, _ := r.BasicAuth() f.tokenRequests = append(f.tokenRequests, fakeTokenRequest{Scope: r.URL.Query().Get("scope"), Username: username}) @@ -501,3 +511,53 @@ func TestOCITokenAuth(t *testing.T) { }) } } + +func TestOCIBackoff(t *testing.T) { + blob := []byte("blob") + blobPath := "/v2/library/alpine/blobs/" + digest.FromBytes(blob).String() + for _, tt := range []struct { + name, limitedPath, trigger string + }{ + {"Blob", blobPath, "/v2/alpine/blobs/" + digest.FromBytes(blob).String()}, + {"Ping", "/v2/", "/v2/other/manifests/latest"}, + {"Token", "/token", "/v2/other/tags/list"}, + } { + t.Run(tt.name, func(t *testing.T) { + f := newFakeRegistry(t, blob) + f.requireToken = true + defaultTransport := http.DefaultTransport + http.DefaultTransport = httputil.NewBackoffTransport(defaultTransport, time.Minute) //nolint:reassign + t.Cleanup(func() { http.DefaultTransport = defaultTransport }) //nolint:reassign + mux := newOCITestMux(t, f, strategy.OCIConfig{TagTTL: time.Hour}) + cached := "/oci/example.io/v2/alpine/manifests/cached" + assert.Equal(t, http.StatusOK, serveOCI(t, mux, http.MethodGet, cached, nil).Code) + f.set(func(f *fakeRegistry) { f.limitedPath = tt.limitedPath }) + first := serveOCI(t, mux, http.MethodGet, "/oci/example.io"+tt.trigger, nil) + assert.Equal(t, http.StatusTooManyRequests, first.Code) + assert.Equal(t, "registry limit", first.Body.String()) + assert.Equal(t, "120", first.Header().Get("Retry-After")) + assert.Equal(t, "no-store", first.Header().Get("Cache-Control")) + before := read(f, &f.exchanges) + for _, path := range []string{ + "/v2/alpine/blobs/" + digest.FromBytes(blob).String(), + "/v2/alpine/manifests/latest", "/v2/alpine/tags/list", + "/v2/alpine/referrers/" + digest.FromBytes(blob).String(), + "/v2/cold/manifests/latest", "/v2/cold/tags/list", + } { + for _, method := range []string{http.MethodGet, http.MethodHead} { + w := serveOCI(t, mux, method, "/oci/example.io"+path, nil) + assert.Equal(t, http.StatusTooManyRequests, w.Code, path) + assert.Equal(t, "120", w.Header().Get("Retry-After")) + assert.Equal(t, "no-store", w.Header().Get("Cache-Control")) + } + } + assert.Equal(t, http.StatusOK, serveOCI(t, mux, http.MethodGet, cached, nil).Code) + assert.Equal(t, before, read(f, &f.exchanges)) + f.set(func(f *fakeRegistry) { f.limitedPath = "" }) + independent := newOCITestMux(t, f, strategy.OCIConfig{}) + assert.Equal(t, http.StatusOK, serveOCI(t, independent, http.MethodGet, + "/oci/example.io/v2/alpine/manifests/latest", nil).Code) + assert.True(t, read(f, &f.exchanges) > before) + }) + } +} From 55bb6d165659bc39fe63f7f7fbbbc8a1c2237aad Mon Sep 17 00:00:00 2001 From: Elizabeth Worstell Date: Mon, 9 Nov 2026 16:21:03 -0800 Subject: [PATCH 2/3] fix: share upstream cooldowns by outbound hostname Amp-Thread-ID: https://ampcode.com/threads/T-01a1c2a3-3b00-7292-bb00-b4a066f5cd6b Co-authored-by: Amp --- README.md | 6 +- cmd/cachewd/main.go | 2 +- internal/httputil/backoff.go | 85 +++-------------------- internal/httputil/backoff_test.go | 27 ++++--- internal/strategy/handler/handler.go | 15 +--- internal/strategy/handler/handler_test.go | 11 +-- internal/strategy/hermit.go | 2 +- internal/strategy/host_test.go | 1 + internal/strategy/oci.go | 2 +- internal/strategy/oci_test.go | 4 +- 10 files changed, 36 insertions(+), 119 deletions(-) diff --git a/README.md b/README.md index f05de460..e242a3d7 100644 --- a/README.md +++ b/README.md @@ -219,11 +219,11 @@ spool { ## Upstream Rate Limits -After an upstream returns 429, cachewd pauses requests to that origin (scheme, host and port) within the strategy until a cooldown ends. A rate-limited redirect target also pauses the origins that redirected to it. Cache hits, other strategies, and unrelated origins remain available. OCI blob, manifest, authentication, and pass-through requests share the registry strategy's cooldowns; Hermit's direct and release handlers share its cooldowns. +After an upstream returns 429, cachewd pauses HTTP requests to that outbound hostname across strategies, schemes, and ports until a cooldown ends. Each redirect destination is checked independently. Cache hits and other hostnames remain available. OCI blob, manifest, authentication, and pass-through requests use the same shared cooldowns. -Configured fallbacks are tried on the first 429 and on subsequent cache misses while the primary is cooling. A mirror's 404 or 5xx does not freeze it; a mirror's 429 starts its own cooldown. When every applicable upstream is cooling, misses get a local 429 with `Retry-After` and `Cache-Control: no-store` without contacting upstreams. Requests outside strategy handlers, such as Git, share per-origin cooldowns without a strategy scope. Cooldowns do not schedule retries. +Configured fallbacks are tried on the first 429 and on subsequent cache misses while the primary is cooling. A mirror on another hostname remains available until it returns its own 429; its 404 or 5xx does not start a cooldown. Requests to cooling hosts get a local 429 with `Retry-After` and `Cache-Control: no-store`. Cooldowns do not schedule retries. -The cooldown uses the upstream's `Retry-After` (seconds or HTTP date) when it is valid, and `upstream-backoff` otherwise. Cooldowns are held in memory per process. +The cooldown uses the upstream's `Retry-After` (seconds or HTTP date) when it is valid, and `upstream-backoff` otherwise. Cooldowns are held in memory per process and apply to clients that use the shared HTTP transport, not subprocess traffic. Other replicas and requests already in flight are unaffected. ```hcl upstream-backoff = "1m" # Must be positive diff --git a/cmd/cachewd/main.go b/cmd/cachewd/main.go index a2b2597d..5960c99b 100644 --- a/cmd/cachewd/main.go +++ b/cmd/cachewd/main.go @@ -457,7 +457,7 @@ func configureProcess(mutexProfileFraction int, upstreamBackoff time.Duration) { runtime.SetMutexProfileFraction(mutexProfileFraction) // Every strategy that uses http.DefaultTransport shares one set of - // per-origin 429 cooldowns. Origin metrics sit below the cooldowns so that + // per-hostname 429 cooldowns. Origin metrics sit below the cooldowns so that // requests rejected during a cooldown are not counted as origin traffic. http.DefaultTransport = httputil.NewBackoffTransport(metrics.NewOriginTransport(http.DefaultTransport), upstreamBackoff) //nolint:reassign } diff --git a/internal/httputil/backoff.go b/internal/httputil/backoff.go index d8251a54..05c14052 100644 --- a/internal/httputil/backoff.go +++ b/internal/httputil/backoff.go @@ -1,10 +1,8 @@ package httputil import ( - "context" "io" "math" - "net" "net/http" "strconv" "strings" @@ -15,27 +13,22 @@ import ( ) // BackoffTransport shares upstream cooldowns across requests and redirect hops. -// Cooldowns are per origin and isolated by [BackoffScope]. +// Cooldowns are keyed by outbound hostname, across strategies, schemes and ports. type BackoffTransport struct { base http.RoundTripper delay time.Duration mu sync.Mutex - until map[backoffKey]time.Time -} - -type backoffKey struct { - scope *BackoffScope - origin string + until map[string]time.Time } // NewBackoffTransport uses delay when a 429 has no valid Retry-After header. func NewBackoffTransport(base http.RoundTripper, delay time.Duration) *BackoffTransport { - return &BackoffTransport{base: base, delay: delay, until: map[backoffKey]time.Time{}} + return &BackoffTransport{base: base, delay: delay, until: map[string]time.Time{}} } -// RoundTrip rejects requests to a cooling origin within their scope without retrying them. +// RoundTrip rejects requests to a cooling hostname without retrying them. func (t *BackoffTransport) RoundTrip(r *http.Request) (*http.Response, error) { - cooldown := backoffKeyFor(r) + cooldown := strings.ToLower(r.URL.Hostname()) now := time.Now() t.mu.Lock() for k, until := range t.until { @@ -46,7 +39,6 @@ func (t *BackoffTransport) RoundTrip(r *http.Request) (*http.Response, error) { until := t.until[cooldown] t.mu.Unlock() if until.After(now) { - t.recordCooldown(r, until) if r.Body != nil { _ = r.Body.Close() } @@ -70,7 +62,11 @@ func (t *BackoffTransport) RoundTrip(r *http.Request) (*http.Response, error) { now = time.Now() var valid bool until, valid = retryAfter(resp.Header.Get("Retry-After"), now, t.delay) - t.recordCooldown(r, until) + t.mu.Lock() + if until.After(t.until[cooldown]) { + t.until[cooldown] = until + } + t.mu.Unlock() // The cooldown ends the 429; a downstream cache must not keep it longer. resp.Header.Set("Cache-Control", "no-store") if !valid { @@ -80,67 +76,6 @@ func (t *BackoffTransport) RoundTrip(r *http.Request) (*http.Response, error) { return resp, nil } -func (t *BackoffTransport) recordCooldown(r *http.Request, until time.Time) { - scope := backoffKeyFor(r).scope - t.mu.Lock() - defer t.mu.Unlock() - for hop := r; hop != nil; { - key := backoffKeyFor(hop) - key.scope = scope - if until.After(t.until[key]) { - t.until[key] = until - } - if hop.Response == nil { - break - } - hop = hop.Response.Request - } -} - -func backoffKeyFor(r *http.Request) backoffKey { - port := r.URL.Port() - if port == "" { - if r.URL.Scheme == "https" { - port = "443" - } else { - port = "80" - } - } - scope, _ := r.Context().Value(backoffScopeKey{}).(*BackoffScope) - return backoffKey{scope: scope, origin: strings.ToLower(r.URL.Scheme + "://" + net.JoinHostPort(r.URL.Hostname(), port))} -} - -type backoffScopeKey struct{} - -// BackoffScope isolates a strategy's per-origin cooldowns from other strategies. -type BackoffScope struct { - base http.RoundTripper -} - -// NewBackoffScope sends requests to base, or to http.DefaultTransport when base is nil. -func NewBackoffScope(base http.RoundTripper) *BackoffScope { - return &BackoffScope{base: base} -} - -// WithContext puts requests made with ctx in this scope, including requests -// sent by other clients through a [BackoffTransport]. -func (s *BackoffScope) WithContext(ctx context.Context) context.Context { - return context.WithValue(ctx, backoffScopeKey{}, s) -} - -// RoundTrip sends r in this scope. -func (s *BackoffScope) RoundTrip(r *http.Request) (*http.Response, error) { - base := s.base - if base == nil { - base = http.DefaultTransport - } - resp, err := base.RoundTrip(r.WithContext(s.WithContext(r.Context()))) - if err != nil { - return nil, errors.WithStack(err) - } - return resp, nil -} - func retryAfter(value string, now time.Time, fallback time.Duration) (time.Time, bool) { value = strings.TrimSpace(value) if seconds, err := strconv.ParseUint(value, 10, 63); err == nil && seconds <= uint64(math.MaxInt64/int64(time.Second)) { diff --git a/internal/httputil/backoff_test.go b/internal/httputil/backoff_test.go index 8cb80548..ea965e2e 100644 --- a/internal/httputil/backoff_test.go +++ b/internal/httputil/backoff_test.go @@ -72,10 +72,10 @@ func TestBackoffTransport(t *testing.T) { get("https://other.example/ok") get("http://example.com/ok") get("https://example.com:444/ok") - assert.Equal(t, 4, calls) + assert.Equal(t, 2, calls) time.Sleep(tt.delay - time.Nanosecond) assert.Equal(t, "1", get("https://example.com/near-expiry").Header.Get("Retry-After")) - assert.Equal(t, 4, calls) + assert.Equal(t, 2, calls) time.Sleep(time.Nanosecond) } before := calls @@ -106,7 +106,7 @@ func TestBackoffRedirectDestination(t *testing.T) { return resp, nil }) transport := httputil.NewBackoffTransport(base, time.Minute) - client := &http.Client{Transport: httputil.NewBackoffScope(transport)} + client := &http.Client{Transport: transport} for _, host := range []string{"redirect.example", "redirect.example", "another.example", "another.example"} { req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, "https://"+host+"/file", nil) assert.NoError(t, err) @@ -116,18 +116,18 @@ func TestBackoffRedirectDestination(t *testing.T) { assert.Equal(t, http.StatusTooManyRequests, resp.StatusCode) assert.Equal(t, "60", resp.Header.Get("Retry-After")) } - assert.Equal(t, map[string]int{"redirect.example": 1, "another.example": 1, "limited.example": 1}, calls) + assert.Equal(t, map[string]int{"redirect.example": 2, "another.example": 2, "limited.example": 1}, calls) time.Sleep(time.Minute) req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, "https://redirect.example/file", nil) assert.NoError(t, err) resp, err := client.Do(req) assert.NoError(t, err) assert.NoError(t, resp.Body.Close()) - assert.Equal(t, map[string]int{"redirect.example": 2, "another.example": 1, "limited.example": 2}, calls) + assert.Equal(t, map[string]int{"redirect.example": 3, "another.example": 2, "limited.example": 2}, calls) }) } -func TestBackoffScopes(t *testing.T) { +func TestBackoffSharedClients(t *testing.T) { synctest.Test(t, func(t *testing.T) { calls := map[string]int{} base := roundTripFunc(func(r *http.Request) (*http.Response, error) { @@ -139,9 +139,8 @@ func TestBackoffScopes(t *testing.T) { return &http.Response{StatusCode: status, Header: make(http.Header), Body: http.NoBody, Request: r}, nil }) transport := httputil.NewBackoffTransport(base, time.Minute) - first := &http.Client{Transport: httputil.NewBackoffScope(transport)} - second := &http.Client{Transport: httputil.NewBackoffScope(transport)} - unscoped := &http.Client{Transport: transport} + first := &http.Client{Transport: transport} + second := &http.Client{Transport: transport} for _, tt := range []struct { client *http.Client url string @@ -150,10 +149,10 @@ func TestBackoffScopes(t *testing.T) { {first, "https://primary.example/limited", http.StatusTooManyRequests}, {first, "https://PRIMARY.example:443/other", http.StatusTooManyRequests}, {first, "https://mirror.example/other", http.StatusOK}, - {second, "https://primary.example/other", http.StatusOK}, - {unscoped, "https://primary.example/other", http.StatusOK}, - {unscoped, "https://mirror.example/limited", http.StatusTooManyRequests}, - {first, "https://mirror.example/another", http.StatusOK}, + {second, "https://primary.example/other", http.StatusTooManyRequests}, + {second, "http://primary.example:8080/other", http.StatusTooManyRequests}, + {second, "https://mirror.example/limited", http.StatusTooManyRequests}, + {first, "https://mirror.example/another", http.StatusTooManyRequests}, } { req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, tt.url, nil) assert.NoError(t, err) @@ -162,6 +161,6 @@ func TestBackoffScopes(t *testing.T) { assert.NoError(t, resp.Body.Close()) assert.Equal(t, tt.status, resp.StatusCode) } - assert.Equal(t, map[string]int{"primary.example": 3, "mirror.example": 3}, calls) + assert.Equal(t, map[string]int{"primary.example": 1, "mirror.example": 2}, calls) }) } diff --git a/internal/strategy/handler/handler.go b/internal/strategy/handler/handler.go index 05517ad9..bb84fde8 100644 --- a/internal/strategy/handler/handler.go +++ b/internal/strategy/handler/handler.go @@ -76,7 +76,6 @@ type Handler struct { registry string external bool client *http.Client - scope *httputil.BackoffScope cache cache.Cache cacheKeyFunc func(*http.Request) CacheKeyParts transformFunc func(*http.Request) (*http.Request, error) @@ -100,23 +99,13 @@ const checkedAtHeader = "X-Cachew-Checked-At" // New creates a new Handler with the given HTTP client and cache. Concurrent // cache misses for the same key share one upstream fetch through a spool group. -// The client's BackoffScope is reused when present, otherwise a new scope is -// created. Upstream requests, including transforms and parallel chunks, share -// per-origin cooldowns within that scope; fallback origins remain independent. // By default: // - Cache key is derived from the request URL // - No request transformation is performed // - Standard error handling is used. func New(client *http.Client, c cache.Cache, spools *spool.Registry) *Handler { - scope, ok := client.Transport.(*httputil.BackoffScope) - if !ok { - scope = httputil.NewBackoffScope(client.Transport) - } - scoped := *client - scoped.Transport = scope return &Handler{ - client: &scoped, - scope: scope, + client: client, cache: c, spools: spools.NewGroup(spool.ReleaseCompleted), metrics: metrics.NewArtifactMetrics(), @@ -460,7 +449,7 @@ func (h *Handler) revalidate(r *http.Request, key cache.Key, headers http.Header // If-Range are never forwarded: whole objects are cached and ranges are cut // locally. func (h *Handler) upstreamRequest(r *http.Request) (*http.Request, error) { - req, err := h.transformFunc(r.WithContext(h.scope.WithContext(r.Context()))) + req, err := h.transformFunc(r) if err != nil { return nil, err } diff --git a/internal/strategy/handler/handler_test.go b/internal/strategy/handler/handler_test.go index 1a16565f..35190d92 100644 --- a/internal/strategy/handler/handler_test.go +++ b/internal/strategy/handler/handler_test.go @@ -1016,6 +1016,7 @@ func TestBackoffPausesHandler(t *testing.T) { defer limited.Close() other := newUpstream() defer other.Close() + other.URL = strings.Replace(other.URL, "127.0.0.1", "localhost", 1) client := &http.Client{Transport: httputil.NewBackoffTransport(http.DefaultTransport, time.Minute)} c := mustNewMemoryCache() defer c.Close() @@ -1044,15 +1045,7 @@ func TestBackoffPausesHandler(t *testing.T) { assert.Equal(t, int32(3), calls.Load()) w := get(newHandler(), limited.URL, "/independent") - assert.Equal(t, http.StatusOK, w.Code) - assert.Equal(t, "body", w.Body.String()) - assert.Equal(t, int32(4), calls.Load()) - - client.Transport = httputil.NewBackoffScope(client.Transport) - first, second := newHandler(), newHandler() - assert.Equal(t, http.StatusTooManyRequests, get(first, limited.URL, "/limited").Code) - w = get(second, limited.URL, "/shared") assert.Equal(t, http.StatusTooManyRequests, w.Code) assert.Equal(t, "60", w.Header().Get("Retry-After")) - assert.Equal(t, int32(5), calls.Load()) + assert.Equal(t, int32(3), calls.Load()) } diff --git a/internal/strategy/hermit.go b/internal/strategy/hermit.go index 4c49b402..bcc524af 100644 --- a/internal/strategy/hermit.go +++ b/internal/strategy/hermit.go @@ -62,7 +62,7 @@ func NewHermit(ctx context.Context, config HermitConfig, _ jobscheduler.Schedule s := &Hermit{ config: config, cache: c, - client: &http.Client{Transport: httputil.NewBackoffScope(http.DefaultTransport)}, + client: http.DefaultClient, logger: logger, mux: mux, } diff --git a/internal/strategy/host_test.go b/internal/strategy/host_test.go index 6ce3fc37..f8240465 100644 --- a/internal/strategy/host_test.go +++ b/internal/strategy/host_test.go @@ -330,6 +330,7 @@ func TestHostBackoffKeepsFallbacksAvailable(t *testing.T) { defer primary.Close() mirror := newServer("mirror", http.StatusOK) defer mirror.Close() + mirror.URL = strings.Replace(mirror.URL, "127.0.0.1", "localhost", 1) mux := http.NewServeMux() _, err = strategy.NewHost(ctx, strategy.HostConfig{ Target: primary.URL, diff --git a/internal/strategy/oci.go b/internal/strategy/oci.go index cc58c87f..0e8ca08e 100644 --- a/internal/strategy/oci.go +++ b/internal/strategy/oci.go @@ -113,7 +113,7 @@ func NewOCI(ctx context.Context, config OCIConfig, c cache.Cache, mux Mux, spool upstream: upstream, rewrites: rewrites, tagTTL: config.TagTTL, - client: &http.Client{Transport: httputil.NewBackoffScope(transport)}, + client: &http.Client{Transport: transport}, logger: logging.FromContext(ctx), } o.blobs = handler.New(o.client, c, spools).Metrics("oci", upstream.Hostname()). diff --git a/internal/strategy/oci_test.go b/internal/strategy/oci_test.go index cda0e7b9..05b11f72 100644 --- a/internal/strategy/oci_test.go +++ b/internal/strategy/oci_test.go @@ -555,9 +555,9 @@ func TestOCIBackoff(t *testing.T) { assert.Equal(t, before, read(f, &f.exchanges)) f.set(func(f *fakeRegistry) { f.limitedPath = "" }) independent := newOCITestMux(t, f, strategy.OCIConfig{}) - assert.Equal(t, http.StatusOK, serveOCI(t, independent, http.MethodGet, + assert.Equal(t, http.StatusTooManyRequests, serveOCI(t, independent, http.MethodGet, "/oci/example.io/v2/alpine/manifests/latest", nil).Code) - assert.True(t, read(f, &f.exchanges) > before) + assert.Equal(t, before, read(f, &f.exchanges)) }) } } From cbb93691e440a3112104da7b445320d7193b1c6c Mon Sep 17 00:00:00 2001 From: Elizabeth Worstell Date: Mon, 9 Nov 2026 16:25:37 -0800 Subject: [PATCH 3/3] fix: report earliest retry time across fallbacks Amp-Thread-ID: https://ampcode.com/threads/T-01a1c2a3-3b00-7292-bb00-b4a066f5cd6b Co-authored-by: Amp --- internal/httputil/backoff.go | 6 ++-- internal/strategy/host.go | 20 +++++++++++-- internal/strategy/host_test.go | 55 ++++++++++++++++++++++++++++++++++ 3 files changed, 77 insertions(+), 4 deletions(-) diff --git a/internal/httputil/backoff.go b/internal/httputil/backoff.go index 05c14052..3a0d04a9 100644 --- a/internal/httputil/backoff.go +++ b/internal/httputil/backoff.go @@ -61,7 +61,7 @@ func (t *BackoffTransport) RoundTrip(r *http.Request) (*http.Response, error) { if resp.StatusCode == http.StatusTooManyRequests { now = time.Now() var valid bool - until, valid = retryAfter(resp.Header.Get("Retry-After"), now, t.delay) + until, valid = RetryAfterDeadline(resp.Header.Get("Retry-After"), now, t.delay) t.mu.Lock() if until.After(t.until[cooldown]) { t.until[cooldown] = until @@ -76,7 +76,9 @@ func (t *BackoffTransport) RoundTrip(r *http.Request) (*http.Response, error) { return resp, nil } -func retryAfter(value string, now time.Time, fallback time.Duration) (time.Time, bool) { +// RetryAfterDeadline parses Retry-After, using fallback when the value is invalid. +// The boolean reports whether the header was valid. +func RetryAfterDeadline(value string, now time.Time, fallback time.Duration) (time.Time, bool) { value = strings.TrimSpace(value) if seconds, err := strconv.ParseUint(value, 10, 63); err == nil && seconds <= uint64(math.MaxInt64/int64(time.Second)) { return now.Add(time.Duration(seconds) * time.Second), true diff --git a/internal/strategy/host.go b/internal/strategy/host.go index 138c2dac..38931e11 100644 --- a/internal/strategy/host.go +++ b/internal/strategy/host.go @@ -3,8 +3,10 @@ package strategy import ( "context" "io" + "math" "net/http" "net/url" + "strconv" "strings" "time" @@ -142,7 +144,18 @@ type hostFallbackTransport struct { } func (t *hostFallbackTransport) RoundTrip(r *http.Request) (*http.Response, error) { - resp, err := t.client.Do(r) + var retryAt time.Time + do := func(req *http.Request) (*http.Response, error) { + resp, err := t.client.Do(req) + if err == nil && resp.StatusCode == http.StatusTooManyRequests { + until, valid := httputil.RetryAfterDeadline(resp.Header.Get("Retry-After"), time.Now(), 0) + if valid && (retryAt.IsZero() || until.Before(retryAt)) { + retryAt = until + } + } + return resp, errors.WithStack(err) + } + resp, err := do(r) for _, fallback := range t.fallbacks { if r.Context().Err() != nil || (err == nil && resp.StatusCode != http.StatusTooManyRequests && (resp.StatusCode < 500 || resp.StatusCode > 599)) { @@ -175,7 +188,10 @@ func (t *hostFallbackTransport) RoundTrip(r *http.Request) (*http.Response, erro target.RawQuery = r.URL.RawQuery req.URL = &target req.Host = "" - resp, err = t.client.Do(req) + resp, err = do(req) + } + if err == nil && resp.StatusCode == http.StatusTooManyRequests && !retryAt.IsZero() { + resp.Header.Set("Retry-After", strconv.FormatInt(int64(math.Ceil(max(0, time.Until(retryAt).Seconds()))), 10)) } return resp, errors.WithStack(err) } diff --git a/internal/strategy/host_test.go b/internal/strategy/host_test.go index f8240465..3fdfb6e8 100644 --- a/internal/strategy/host_test.go +++ b/internal/strategy/host_test.go @@ -363,6 +363,61 @@ func TestHostBackoffKeepsFallbacksAvailable(t *testing.T) { assert.Equal(t, map[string]int{"primary": 1, "mirror": 4}, calls) } +func TestHostFallbackRetryAfter(t *testing.T) { + for _, tt := range []struct { + name, primary, mirror, want string + calls int + }{ + {"PrimarySooner", "10", "3600", "10", 2}, + {"MirrorSooner", "3600", "10", "10", 2}, + {"PrimaryDate", "Wed, 21 Oct 2015 07:28:00 GMT", "3600", "0", 3}, + {"MirrorDate", "3600", "Wed, 21 Oct 2015 07:28:00 GMT", "0", 3}, + {"Immediate", "0", "3600", "0", 3}, + } { + t.Run(tt.name, func(t *testing.T) { + defaultTransport := http.DefaultTransport + http.DefaultTransport = httputil.NewBackoffTransport(defaultTransport, time.Minute) //nolint:reassign + t.Cleanup(func() { http.DefaultTransport = defaultTransport }) //nolint:reassign + _, ctx := logging.Configure(t.Context(), logging.Config{Level: slog.LevelError}) + c, err := cache.NewMemory(ctx, cache.MemoryConfig{MaxTTL: time.Hour}) + assert.NoError(t, err) + defer c.Close() + calls := 0 + newServer := func(retryAfter string) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls++ + w.Header().Set("Retry-After", retryAfter) + w.WriteHeader(http.StatusTooManyRequests) + })) + } + primary := newServer(tt.primary) + defer primary.Close() + mirror := newServer(tt.mirror) + defer mirror.Close() + mirror.URL = strings.Replace(mirror.URL, "127.0.0.1", "localhost", 1) + mux := http.NewServeMux() + _, err = strategy.NewHost(ctx, strategy.HostConfig{ + Target: primary.URL, + Fallbacks: []strategy.HostFallbackConfig{ + {Target: "http://unused.invalid", PathPrefix: "/unmatched"}, + {Target: mirror.URL}, + }, + }, c, mux, newTestSpools(t)) + assert.NoError(t, err) + u, err := url.Parse(primary.URL) + assert.NoError(t, err) + for range 2 { + w := httptest.NewRecorder() + mux.ServeHTTP(w, httptest.NewRequestWithContext(ctx, http.MethodGet, "/"+u.Host+"/a.jar", nil)) + assert.Equal(t, http.StatusTooManyRequests, w.Code) + assert.Equal(t, tt.want, w.Header().Get("Retry-After")) + assert.Equal(t, "no-store", w.Header().Get("Cache-Control")) + } + assert.Equal(t, tt.calls, calls) + }) + } +} + func TestHostFallbackIncompleteResponse(t *testing.T) { for _, cancelRequest := range []bool{false, true} { t.Run(strconv.FormatBool(cancelRequest), func(t *testing.T) {