diff --git a/README.md b/README.md index e242a3d7..7eeb875d 100644 --- a/README.md +++ b/README.md @@ -219,14 +219,15 @@ spool { ## Upstream Rate Limits -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. +After an upstream returns 429, cachewd pauses HTTP requests to that outbound hostname across strategies, schemes, and ports until a cooldown ends. Requests made with a GitHub App installation token, an App JWT, or the static GitHub token cool down separately, because GitHub limits each of those credentials on its own: one installation's 429 does not pause another installation, token minting, or anonymous requests. All other requests share their hostname's cooldown. A redirect to another host drops the credential, so that hop also shares the hostname's cooldown. Each redirect destination is checked independently. Cache hits and other hostnames remain available. Requests that one strategy forwards to another cachew route, such as Hermit's GitHub releases, are exempt; the target route applies the cooldown to the real upstream. 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 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 and apply to clients that use the shared HTTP transport, not subprocess traffic. Other replicas and requests already in flight are unaffected. +The cooldown uses the upstream's `Retry-After` (seconds or HTTP date) when it is valid, and `upstream-backoff` otherwise, and never lasts longer than `upstream-backoff-max`. Each cooldown that starts or is extended is logged and counted in `cachew.upstream.cooldowns`; local 429s are counted in `cachew.upstream.cooldown_rejections`. 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 +upstream-backoff = "1m" # Must be positive +upstream-backoff-max = "15m" # Raised to upstream-backoff when lower ``` ## Authorization (OPA) diff --git a/cmd/cachewd/main.go b/cmd/cachewd/main.go index 5960c99b..561c186b 100644 --- a/cmd/cachewd/main.go +++ b/cmd/cachewd/main.go @@ -48,6 +48,7 @@ type GlobalConfig struct { Bind string `hcl:"bind" default:"127.0.0.1:8080" help:"Bind address for the server."` URL string `hcl:"url" default:"http://127.0.0.1:8080/" help:"Base URL for cachewd."` UpstreamBackoff time.Duration `hcl:"upstream-backoff,optional" default:"1m" help:"Cooldown after an upstream 429 with missing or invalid Retry-After; must be positive."` + UpstreamBackoffMax time.Duration `hcl:"upstream-backoff-max,optional" default:"15m" help:"Longest cooldown after an upstream 429, whatever its Retry-After; raised to upstream-backoff when lower."` ShutdownReadinessDelay time.Duration `hcl:"shutdown-readiness-delay,optional" default:"5s" help:"Delay between flipping readiness to 503 on SIGTERM and starting graceful shutdown."` // ShutdownTimeout must be less than the pod's terminationGracePeriodSeconds // (minus ShutdownReadinessDelay) or the kubelet will SIGKILL before Shutdown returns. @@ -105,7 +106,7 @@ func main() { // Flipped on SIGTERM so /_readiness fails before the listener closes. var shuttingDown atomic.Bool - configureProcess(cli.MutexProfileFraction, globalConfig.UpstreamBackoff) + configureProcess(logger, cli.MutexProfileFraction, globalConfig.UpstreamBackoff, globalConfig.UpstreamBackoffMax) // Register the OpenTelemetry tracer provider. tracing.New is a // no-op when OTEL_EXPORTER_OTLP_ENDPOINT is unset, so deployments @@ -449,7 +450,7 @@ func newServer( } // configureProcess applies process-wide runtime and HTTP transport settings. -func configureProcess(mutexProfileFraction int, upstreamBackoff time.Duration) { +func configureProcess(logger *slog.Logger, mutexProfileFraction int, upstreamBackoff, upstreamBackoffMax time.Duration) { // Enable mutex contention sampling so /admin/pprof/mutex returns // useful data. The pprof endpoints themselves are mounted on the // main HTTP listener under /admin/pprof/* via the net/http/pprof @@ -457,9 +458,9 @@ func configureProcess(mutexProfileFraction int, upstreamBackoff time.Duration) { runtime.SetMutexProfileFraction(mutexProfileFraction) // Every strategy that uses http.DefaultTransport shares one set of - // per-hostname 429 cooldowns. Origin metrics sit below the cooldowns so that + // per-hostname and per-credential 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 + http.DefaultTransport = httputil.NewBackoffTransport(metrics.NewOriginTransport(http.DefaultTransport), upstreamBackoff, upstreamBackoffMax, logger) //nolint:reassign } // loadGlobalConfig unmarshals the global config from HCL, using a two-pass @@ -505,6 +506,8 @@ func loadGlobalConfig(ast *hcl.AST) (GlobalConfig, map[string]string, error) { if cfg.UpstreamBackoff <= 0 { return cfg, nil, errors.New("upstream-backoff must be positive") } + // Keeps configs that set only a long upstream-backoff valid. + cfg.UpstreamBackoffMax = max(cfg.UpstreamBackoffMax, cfg.UpstreamBackoff) return cfg, envars, nil } diff --git a/internal/githubapp/ratelimit.go b/internal/githubapp/ratelimit.go index 9bbc38c2..7282a479 100644 --- a/internal/githubapp/ratelimit.go +++ b/internal/githubapp/ratelimit.go @@ -98,7 +98,7 @@ func (p *RateLimitPoller) poll(ctx context.Context) error { func (p *RateLimitPoller) pollInstallation(ctx context.Context, app *appState, installationID string) error { ctx, cancel := context.WithTimeout(ctx, rateLimitPollTimeout) defer cancel() - token, owner, err := app.getInstallationToken(ctx, installationID) + token, installation, err := app.getInstallationToken(ctx, installationID) if err != nil { return errors.Wrap(err, "get installation token") } @@ -107,9 +107,10 @@ func (p *RateLimitPoller) pollInstallation(ctx context.Context, app *appState, i return errors.Wrap(err, "create rate-limit request") } req.Header.Set("Authorization", "Bearer "+token) + req = req.WithContext(WithInstallationCooldown(req.Context(), installationID)) req.Header.Set("Accept", "application/vnd.github+json") req.Header.Set("X-Github-Api-Version", "2022-11-28") - resp, err := app.httpClient.Do(WithAPIMetrics(req, owner, "get_rate_limit")) + resp, err := app.httpClient.Do(WithAPIMetrics(req, installation.Owner, "get_rate_limit")) if err != nil { return errors.Wrap(err, "fetch rate limit") } @@ -126,7 +127,7 @@ func (p *RateLimitPoller) pollInstallation(ctx context.Context, app *appState, i for resource, values := range result.Resources { attrs := metric.WithAttributes( attribute.String("app", app.appID), attribute.String("installation_id", installationID), - attribute.String("owner", owner), attribute.String("resource", resource), + attribute.String("owner", installation.Owner), attribute.String("resource", resource), ) p.remaining.Record(ctx, values.Remaining, attrs) p.limit.Record(ctx, values.Limit, attrs) diff --git a/internal/githubapp/ratelimit_test.go b/internal/githubapp/ratelimit_test.go index b1dd8cae..f66ab522 100644 --- a/internal/githubapp/ratelimit_test.go +++ b/internal/githubapp/ratelimit_test.go @@ -161,10 +161,10 @@ func TestRateLimitPollerRetainsOwnerAcrossAliases(t *testing.T) { poller := NewRateLimitPoller(&TokenManager{apps: []*appState{app}}, 0) for _, alias := range []string{"Example-Org", "old-name", "renamed-org"} { app.rememberInstallation(alias, "123", alias) - token, owner, err := app.getToken(pollerContext(t), alias) + token, installation, err := app.getToken(pollerContext(t), alias) assert.NoError(t, err) assert.Equal(t, "token-111", token) - assert.Equal(t, "example-org", owner) + assert.Equal(t, Installation{ID: "123", Owner: "example-org"}, installation) assert.Equal(t, map[string]string{"123": "example-org"}, app.knownInstallations()) assert.NoError(t, poller.poll(pollerContext(t))) points := collectMetrics(t, reader)["github.ratelimit_remaining"] diff --git a/internal/githubapp/tokens.go b/internal/githubapp/tokens.go index 9879f02f..e284c0cb 100644 --- a/internal/githubapp/tokens.go +++ b/internal/githubapp/tokens.go @@ -13,6 +13,7 @@ import ( "github.com/alecthomas/errors" + "github.com/block/cachew/internal/httputil" "github.com/block/cachew/internal/logging" ) @@ -110,14 +111,21 @@ func newTokenManager(configs []Config, logger *slog.Logger, apiBase string, http // caches the result, and falls back to the fallback org's token for orgs where // the app is not installed. func (tm *TokenManager) GetTokenForOrg(ctx context.Context, org string) (string, error) { - token, _, err := tm.GetTokenForOrgWithOwner(ctx, org) + token, _, err := tm.GetTokenForOrgWithInstallation(ctx, org) return token, err } -// GetTokenForOrgWithOwner also returns the owner whose installation minted the token. -func (tm *TokenManager) GetTokenForOrgWithOwner(ctx context.Context, org string) (string, string, error) { +// Installation is the GitHub App installation that minted a token. +type Installation struct { + // ID is unique across GitHub Apps. + ID string + Owner string +} + +// GetTokenForOrgWithInstallation also returns the installation that minted the token. +func (tm *TokenManager) GetTokenForOrgWithInstallation(ctx context.Context, org string) (string, Installation, error) { if tm == nil { - return "", "", errors.New("token manager not initialized") + return "", Installation{}, errors.New("token manager not initialized") } logger := logging.FromContext(ctx) @@ -151,12 +159,12 @@ func (tm *TokenManager) GetTokenForOrgWithOwner(ctx context.Context, org string) if tm.fallbackApp != nil { logger.InfoContext(ctx, "Using fallback org token", "requested_org", org, "fallback_org", tm.fallbackOrg) if _, err := tm.fallbackApp.lookupInstallationID(ctx, tm.fallbackOrg); err != nil { - return "", "", errors.Wrap(err, "resolve fallback installation") + return "", Installation{}, errors.Wrap(err, "resolve fallback installation") } return tm.fallbackApp.getToken(ctx, tm.fallbackOrg) } - return "", "", errors.Errorf("no GitHub App installation found for org: %s", org) + return "", Installation{}, errors.Errorf("no GitHub App installation found for org: %s", org) } // GetTokenForURL extracts the org from a GitHub URL and returns an installation token. @@ -172,23 +180,24 @@ func (tm *TokenManager) GetTokenForURL(ctx context.Context, url string) (string, return tm.GetTokenForOrg(ctx, org) } -func (a *appState) getToken(ctx context.Context, org string) (string, string, error) { +func (a *appState) getToken(ctx context.Context, org string) (string, Installation, error) { a.installationMu.RLock() installationID := a.installationCache[org] a.installationMu.RUnlock() if installationID == "" { - return "", "", errors.Errorf("no installation ID for org: %s", org) + return "", Installation{}, errors.Errorf("no installation ID for org: %s", org) } return a.getInstallationToken(ctx, installationID) } -func (a *appState) getInstallationToken(ctx context.Context, installationID string) (string, string, error) { +func (a *appState) getInstallationToken(ctx context.Context, installationID string) (string, Installation, error) { a.installationMu.RLock() owner := a.installationOwners[installationID] a.installationMu.RUnlock() if owner == "" { - return "", "", errors.Errorf("no owner for installation: %s", installationID) + return "", Installation{}, errors.Errorf("no owner for installation: %s", installationID) } + installation := Installation{ID: installationID, Owner: owner} logger := logging.FromContext(ctx).With("org", owner, "app_id", a.appID) a.mu.RLock() @@ -197,14 +206,14 @@ func (a *appState) getInstallationToken(ctx context.Context, installationID stri if exists && time.Now().Add(a.cacheConfig.RefreshBuffer).Before(cached.expiresAt) { logger.DebugContext(ctx, "Using cached GitHub App token") - return cached.token, owner, nil + return cached.token, installation, nil } logger.DebugContext(ctx, "Fetching new GitHub App installation token", "installation_id", installationID) token, expiresAt, err := a.fetchInstallationToken(ctx, owner, installationID) if err != nil { - return "", "", errors.Wrap(err, "fetch installation token") + return "", Installation{}, errors.Wrap(err, "fetch installation token") } a.mu.Lock() @@ -216,7 +225,7 @@ func (a *appState) getInstallationToken(ctx context.Context, installationID stri logger.InfoContext(ctx, "GitHub App token refreshed", "expires_at", expiresAt) - return token, owner, nil + return token, installation, nil } func (a *appState) fetchInstallationToken(ctx context.Context, org, installationID string) (string, time.Time, error) { @@ -233,6 +242,7 @@ func (a *appState) fetchInstallationToken(ctx context.Context, org, installation req.Header.Set("Accept", "application/vnd.github+json") req.Header.Set("Authorization", "Bearer "+jwt) + req = req.WithContext(httputil.WithCooldownScope(req.Context(), "github-app:"+a.appID)) req.Header.Set("X-Github-Api-Version", "2022-11-28") resp, err := a.httpClient.Do(WithAPIMetrics(req, org, "create_installation_token")) @@ -285,6 +295,7 @@ func (a *appState) lookupInstallationID(ctx context.Context, org string) (string req.Header.Set("Accept", "application/vnd.github+json") req.Header.Set("Authorization", "Bearer "+jwt) + req = req.WithContext(httputil.WithCooldownScope(req.Context(), "github-app:"+a.appID)) req.Header.Set("X-Github-Api-Version", "2022-11-28") resp, err := a.httpClient.Do(WithAPIMetrics(req, "unknown", "lookup_installation")) diff --git a/internal/githubapp/tokens_test.go b/internal/githubapp/tokens_test.go index f7a637c8..02f6dad9 100644 --- a/internal/githubapp/tokens_test.go +++ b/internal/githubapp/tokens_test.go @@ -149,17 +149,17 @@ func TestGetTokenForOrgDynamicDiscovery(t *testing.T) { ctx := logging.ContextWithLogger(t.Context(), slog.Default()) t.Run("FallbackBeforeDiscovery", func(t *testing.T) { - token, owner, err := tm.GetTokenForOrgWithOwner(ctx, "unknown-org") + token, installation, err := tm.GetTokenForOrgWithInstallation(ctx, "unknown-org") assert.NoError(t, err) assert.Equal(t, "ghs_squareup_token", token) - assert.Equal(t, "squareup", owner) + assert.Equal(t, githubapp.Installation{ID: "12345", Owner: "squareup"}, installation) }) t.Run("DiscoverAndCacheInstallation", func(t *testing.T) { - token, owner, err := tm.GetTokenForOrgWithOwner(ctx, "squareup") + token, installation, err := tm.GetTokenForOrgWithInstallation(ctx, "squareup") assert.NoError(t, err) assert.Equal(t, "ghs_squareup_token", token) - assert.Equal(t, "squareup", owner) + assert.Equal(t, githubapp.Installation{ID: "12345", Owner: "squareup"}, installation) // Second call should use cache token, err = tm.GetTokenForOrg(ctx, "squareup") @@ -168,10 +168,10 @@ func TestGetTokenForOrgDynamicDiscovery(t *testing.T) { }) t.Run("DiscoverNewOrg", func(t *testing.T) { - token, owner, err := tm.GetTokenForOrgWithOwner(ctx, "AfterpayTouch") + token, installation, err := tm.GetTokenForOrgWithInstallation(ctx, "AfterpayTouch") assert.NoError(t, err) assert.Equal(t, "ghs_afterpay_token", token) - assert.Equal(t, "AfterpayTouch", owner) + assert.Equal(t, githubapp.Installation{ID: "67890", Owner: "AfterpayTouch"}, installation) }) t.Run("FallbackForUnknownOrg", func(t *testing.T) { @@ -230,10 +230,10 @@ func TestGetTokenForOrgRetainsInstallationOwner(t *testing.T) { slog.Default(), server.URL, server.Client()) ctx := logging.ContextWithLogger(t.Context(), slog.Default()) for _, alias := range []string{"Example-Org", "example-org", "renamed-org"} { - token, owner, err := manager.GetTokenForOrgWithOwner(ctx, alias) + token, installation, err := manager.GetTokenForOrgWithInstallation(ctx, alias) assert.NoError(t, err) assert.Equal(t, "test-token", token) - assert.Equal(t, "example-org", owner) + assert.Equal(t, githubapp.Installation{ID: "123", Owner: "example-org"}, installation) } assert.Equal(t, int64(1), mints.Load()) } diff --git a/internal/githubapp/transport.go b/internal/githubapp/transport.go index 0275d965..acb08437 100644 --- a/internal/githubapp/transport.go +++ b/internal/githubapp/transport.go @@ -10,6 +10,7 @@ import ( "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/metric" + "github.com/block/cachew/internal/httputil" "github.com/block/cachew/internal/logging" "github.com/block/cachew/internal/metrics" ) @@ -42,6 +43,12 @@ func NewAPIClient(base *http.Client) *http.Client { return &client } +// WithInstallationCooldown scopes upstream cooldowns to a GitHub App +// installation, since GitHub limits each installation separately. +func WithInstallationCooldown(ctx context.Context, installationID string) context.Context { + return httputil.WithCooldownScope(ctx, "github-installation:"+installationID) +} + // WithAPIMetrics attributes an API request to its credential owner and operation. func WithAPIMetrics(req *http.Request, owner, operation string) *http.Request { call := apiCall{host: req.URL.Host, owner: owner, operation: operation} diff --git a/internal/httputil/backoff.go b/internal/httputil/backoff.go index 3a0d04a9..d692d55a 100644 --- a/internal/httputil/backoff.go +++ b/internal/httputil/backoff.go @@ -1,7 +1,9 @@ package httputil import ( + "context" "io" + "log/slog" "math" "net/http" "strconv" @@ -10,25 +12,85 @@ import ( "time" "github.com/alecthomas/errors" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" ) // BackoffTransport shares upstream cooldowns across requests and redirect hops. -// Cooldowns are keyed by outbound hostname, across strategies, schemes and ports. +// Cooldowns are keyed by outbound hostname, across strategies, schemes and +// ports, and by the scope set with [WithCooldownScope]. type BackoffTransport struct { - base http.RoundTripper - delay time.Duration - mu sync.Mutex - until map[string]time.Time + base http.RoundTripper + delay time.Duration + maxDelay time.Duration + logger *slog.Logger + cooldowns metric.Int64Counter + rejections metric.Int64Counter + mu sync.Mutex + 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[string]time.Time{}} +// NewBackoffTransport uses delay when a 429 has no valid Retry-After header, +// never cools down for longer than maxDelay, and reports cooldowns to logger, +// because outbound requests do not always carry a logger in their context. +func NewBackoffTransport(base http.RoundTripper, delay, maxDelay time.Duration, logger *slog.Logger) *BackoffTransport { + meter := otel.Meter("cachew.upstream") + cooldowns, err := meter.Int64Counter("cachew.upstream.cooldowns", metric.WithUnit("{cooldown}"), + metric.WithDescription("Upstream 429 responses that started or extended a cooldown")) + if err != nil { + panic(err) + } + rejections, err := meter.Int64Counter("cachew.upstream.cooldown_rejections", metric.WithUnit("{request}"), + metric.WithDescription("Outbound requests answered with a local 429 during a cooldown")) + if err != nil { + panic(err) + } + return &BackoffTransport{base: base, delay: delay, maxDelay: maxDelay, logger: logger, + cooldowns: cooldowns, rejections: rejections, until: map[string]time.Time{}} +} + +type noCooldownKey struct{} +type cooldownScopeKey struct{} + +// WithCooldownScope gives requests their own cooldown on each host. Set it +// where a request gets a credential whose rate limit is separate from the +// host's other callers, such as a GitHub App installation token. Requests +// without a scope, or without an Authorization header, share their host's +// cooldown. +func WithCooldownScope(ctx context.Context, scope string) context.Context { + return context.WithValue(ctx, cooldownScopeKey{}, scope) +} + +// CooldownScope returns the scope set by [WithCooldownScope], if any. +func CooldownScope(ctx context.Context) (string, bool) { + scope, ok := ctx.Value(cooldownScopeKey{}).(string) + return scope, ok +} + +// WithoutCooldown exempts requests that cachew forwards to one of its own +// routes. That route applies cooldowns to the real upstream, and a 429 from it +// must not cool down every other route behind the same address. +func WithoutCooldown(ctx context.Context) context.Context { + return context.WithValue(ctx, noCooldownKey{}, true) } // RoundTrip rejects requests to a cooling hostname without retrying them. func (t *BackoffTransport) RoundTrip(r *http.Request) (*http.Response, error) { + if exempt, _ := r.Context().Value(noCooldownKey{}).(bool); exempt { + resp, err := t.base.RoundTrip(r) + return resp, errors.WithStack(err) + } cooldown := strings.ToLower(r.URL.Hostname()) + scope, scoped := CooldownScope(r.Context()) + // A redirect to another host drops the credential that the scope stands for. + scoped = scoped && r.Header.Get("Authorization") != "" + if scoped { + cooldown += "#" + scope + } else { + scope = "" + } + attrs := metric.WithAttributes(attribute.Bool("scoped", scoped)) now := time.Now() t.mu.Lock() for k, until := range t.until { @@ -39,6 +101,7 @@ func (t *BackoffTransport) RoundTrip(r *http.Request) (*http.Response, error) { until := t.until[cooldown] t.mu.Unlock() if until.After(now) { + t.rejections.Add(r.Context(), 1, attrs) if r.Body != nil { _ = r.Body.Close() } @@ -62,14 +125,28 @@ func (t *BackoffTransport) RoundTrip(r *http.Request) (*http.Response, error) { now = time.Now() var valid bool until, valid = RetryAfterDeadline(resp.Header.Get("Retry-After"), now, t.delay) + capped := until.After(now.Add(t.maxDelay)) + if capped { + until = now.Add(t.maxDelay) + } t.mu.Lock() - if until.After(t.until[cooldown]) { + previous := t.until[cooldown] + if until.After(previous) { t.until[cooldown] = until } t.mu.Unlock() + if until.After(previous) { + event := "extended" + if !previous.After(now) { + event = "started" + } + t.cooldowns.Add(r.Context(), 1, metric.WithAttributes(attribute.Bool("scoped", scoped), attribute.String("event", event))) + t.logger.WarnContext(r.Context(), "Upstream rate limited, cooling down", "host", r.URL.Hostname(), "scope", scope, + "event", event, "until", until, "capped", capped) + } // The cooldown ends the 429; a downstream cache must not keep it longer. resp.Header.Set("Cache-Control", "no-store") - if !valid { + if !valid || capped { resp.Header.Set("Retry-After", strconv.FormatInt(int64(math.Ceil(until.Sub(now).Seconds())), 10)) } } diff --git a/internal/httputil/backoff_test.go b/internal/httputil/backoff_test.go index ea965e2e..ae0388b1 100644 --- a/internal/httputil/backoff_test.go +++ b/internal/httputil/backoff_test.go @@ -2,6 +2,7 @@ package httputil_test import ( "io" + "log/slog" "net/http" "strconv" "strings" @@ -45,7 +46,7 @@ func TestBackoffTransport(t *testing.T) { return &http.Response{StatusCode: http.StatusTooManyRequests, Header: http.Header{"Retry-After": {value}, "Cache-Control": {"public, max-age=3600"}}, Body: io.NopCloser(strings.NewReader("origin limit")), Request: r}, nil }) - transport := httputil.NewBackoffTransport(base, 7*time.Second) + transport := httputil.NewBackoffTransport(base, 7*time.Second, time.Hour, slog.New(slog.DiscardHandler)) client := &http.Client{Transport: transport} get := func(url string) *http.Response { req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, url, nil) @@ -105,7 +106,7 @@ func TestBackoffRedirectDestination(t *testing.T) { } return resp, nil }) - transport := httputil.NewBackoffTransport(base, time.Minute) + transport := httputil.NewBackoffTransport(base, time.Minute, time.Hour, slog.New(slog.DiscardHandler)) 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) @@ -138,7 +139,7 @@ func TestBackoffSharedClients(t *testing.T) { } return &http.Response{StatusCode: status, Header: make(http.Header), Body: http.NoBody, Request: r}, nil }) - transport := httputil.NewBackoffTransport(base, time.Minute) + transport := httputil.NewBackoffTransport(base, time.Minute, time.Hour, slog.New(slog.DiscardHandler)) first := &http.Client{Transport: transport} second := &http.Client{Transport: transport} for _, tt := range []struct { @@ -164,3 +165,132 @@ func TestBackoffSharedClients(t *testing.T) { assert.Equal(t, map[string]int{"primary.example": 1, "mirror.example": 2}, calls) }) } + +func TestBackoffCooldownScope(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + calls := map[string]int{} + base := roundTripFunc(func(r *http.Request) (*http.Response, error) { + scope, _ := httputil.CooldownScope(r.Context()) + calls[scope]++ + 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 + }) + client := &http.Client{Transport: httputil.NewBackoffTransport(base, time.Minute, time.Hour, slog.New(slog.DiscardHandler))} + for _, tt := range []struct { + path, scope string + status int + }{ + {"/limited", "installation-a", http.StatusTooManyRequests}, + {"/other", "installation-a", http.StatusTooManyRequests}, + {"/other", "installation-b", http.StatusOK}, + {"/other", "", http.StatusOK}, + {"/limited", "", http.StatusTooManyRequests}, + {"/other", "", http.StatusTooManyRequests}, + {"/other", "installation-b", http.StatusOK}, + } { + ctx := t.Context() + if tt.scope != "" { + ctx = httputil.WithCooldownScope(ctx, tt.scope) + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://api.example"+tt.path, nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer rotating-"+strconv.Itoa(len(calls))) + resp, err := client.Do(req) + assert.NoError(t, err) + assert.NoError(t, resp.Body.Close()) + assert.Equal(t, tt.status, resp.StatusCode, "%s %s", tt.scope, tt.path) + } + assert.Equal(t, map[string]int{"installation-a": 1, "installation-b": 2, "": 2}, calls) + }) +} + +func TestBackoffScopeDroppedOnCrossHostRedirect(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + cdnCalls := 0 + base := roundTripFunc(func(r *http.Request) (*http.Response, error) { + if r.URL.Host == "api.example" { + return &http.Response{StatusCode: http.StatusFound, Header: http.Header{"Location": {"https://cdn.example/asset"}}, + Body: http.NoBody, Request: r}, nil + } + cdnCalls++ + return &http.Response{StatusCode: http.StatusTooManyRequests, Header: make(http.Header), Body: http.NoBody, Request: r}, nil + }) + client := &http.Client{Transport: httputil.NewBackoffTransport(base, time.Minute, time.Hour, slog.New(slog.DiscardHandler))} + for _, scope := range []string{"installation-a", "installation-b"} { + req, err := http.NewRequestWithContext(httputil.WithCooldownScope(t.Context(), scope), http.MethodGet, "https://api.example/asset", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+scope) + resp, err := client.Do(req) + assert.NoError(t, err) + assert.NoError(t, resp.Body.Close()) + assert.Equal(t, http.StatusTooManyRequests, resp.StatusCode, scope) + } + assert.Equal(t, 1, cdnCalls, "the anonymous redirect hop shares the CDN host's cooldown") + }) +} + +func TestBackoffMaxDelay(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + calls := 0 + base := roundTripFunc(func(r *http.Request) (*http.Response, error) { + calls++ + return &http.Response{StatusCode: http.StatusTooManyRequests, Header: http.Header{"Retry-After": {"86400"}}, Body: http.NoBody, Request: r}, nil + }) + client := &http.Client{Transport: httputil.NewBackoffTransport(base, time.Minute, 10*time.Minute, slog.New(slog.DiscardHandler))} + get := func() *http.Response { + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, "https://example.com/file", nil) + assert.NoError(t, err) + resp, err := client.Do(req) + assert.NoError(t, err) + assert.NoError(t, resp.Body.Close()) + return resp + } + assert.Equal(t, "600", get().Header.Get("Retry-After")) + assert.Equal(t, "600", get().Header.Get("Retry-After")) + assert.Equal(t, 1, calls) + time.Sleep(10 * time.Minute) + get() + assert.Equal(t, 2, calls) + }) +} + +func TestBackoffWithoutCooldown(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + calls := 0 + base := roundTripFunc(func(r *http.Request) (*http.Response, error) { + calls++ + 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 + }) + client := &http.Client{Transport: httputil.NewBackoffTransport(base, time.Minute, time.Hour, slog.New(slog.DiscardHandler))} + for _, tt := range []struct { + path string + exempt bool + status int + }{ + {"/limited", true, http.StatusTooManyRequests}, + {"/other", false, http.StatusOK}, + {"/limited", false, http.StatusTooManyRequests}, + {"/other", true, http.StatusOK}, + {"/other", false, http.StatusTooManyRequests}, + } { + ctx := t.Context() + if tt.exempt { + ctx = httputil.WithoutCooldown(ctx) + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://127.0.0.1:8080"+tt.path, nil) + assert.NoError(t, err) + resp, err := client.Do(req) + assert.NoError(t, err) + assert.NoError(t, resp.Body.Close()) + assert.Equal(t, tt.status, resp.StatusCode, "%s exempt=%v", tt.path, tt.exempt) + } + assert.Equal(t, 4, calls) + }) +} diff --git a/internal/strategy/git/git.go b/internal/strategy/git/git.go index a513d7ca..63350800 100644 --- a/internal/strategy/git/git.go +++ b/internal/strategy/git/git.go @@ -220,10 +220,11 @@ func New( parts := strings.Split(strings.TrimPrefix(req.URL.Path, "/"), "/") if len(parts) >= 1 && parts[0] != "" { org := parts[0] - token, err := s.tokenManager.GetTokenForOrg(req.Context(), org) + token, installation, err := s.tokenManager.GetTokenForOrgWithInstallation(req.Context(), org) if err == nil && token != "" { // Inject token as Basic auth with "x-access-token" username req.SetBasicAuth("x-access-token", token) + *req = *req.WithContext(githubapp.WithInstallationCooldown(req.Context(), installation.ID)) logger.DebugContext(req.Context(), "Injecting GitHub App auth into upstream request", "org", org) } } diff --git a/internal/strategy/github_releases.go b/internal/strategy/github_releases.go index ef64c131..79871b89 100644 --- a/internal/strategy/github_releases.go +++ b/internal/strategy/github_releases.go @@ -117,19 +117,21 @@ func (g *GitHubReleases) newGitHubRequest(ctx context.Context, reqURL, accept, o // Try GitHub App authentication first, fall back to static token owner := org if g.tokenManager != nil && org != "" { - token, resolvedOwner, err := g.tokenManager.GetTokenForOrgWithOwner(ctx, org) + token, installation, err := g.tokenManager.GetTokenForOrgWithInstallation(ctx, org) if err != nil { logging.FromContext(ctx).WarnContext(ctx, "Failed to get GitHub App token, falling back to static token", "org", org, "error", err) } else if token != "" { req.Header.Set("Authorization", "Bearer "+token) - owner = resolvedOwner + req = req.WithContext(githubapp.WithInstallationCooldown(req.Context(), installation.ID)) + owner = installation.Owner } } // Fall back to static token if GitHub App not used or failed if req.Header.Get("Authorization") == "" && g.config.Token != "" { req.Header.Set("Authorization", "Bearer "+g.config.Token) + req = req.WithContext(httputil.WithCooldownScope(req.Context(), "github-token")) } return githubapp.WithAPIMetrics(req, owner, operation), nil diff --git a/internal/strategy/github_releases_test.go b/internal/strategy/github_releases_test.go index 2fd5090b..356bacc9 100644 --- a/internal/strategy/github_releases_test.go +++ b/internal/strategy/github_releases_test.go @@ -130,7 +130,7 @@ func setupTest(t *testing.T, config strategy.GitHubReleasesConfig) (*mockGitHubS http.DefaultTransport = httputil.NewBackoffTransport(metrics.NewOriginTransport(&testTransport{ //nolint:reassign // Required for testing backend: mock.server, originalTransport: originalTransport, - }), time.Minute) + }), time.Minute, time.Hour, slog.New(slog.DiscardHandler)) _, ctx := logging.Configure(context.Background(), logging.Config{Level: slog.LevelError}) memCache, err := cache.NewMemory(ctx, cache.MemoryConfig{MaxTTL: time.Hour}) diff --git a/internal/strategy/handler/handler.go b/internal/strategy/handler/handler.go index bb84fde8..d30c0fe0 100644 --- a/internal/strategy/handler/handler.go +++ b/internal/strategy/handler/handler.go @@ -129,7 +129,9 @@ func (h *Handler) Metrics(strategy, registry string) *Handler { return h } -// InternalForward excludes delegation to another cache route from external-origin and fill metrics. +// InternalForward marks delegation to another cache route. Its requests are +// excluded from external-origin and fill metrics and from upstream cooldowns, +// which the target route applies to the real upstream. func (h *Handler) InternalForward() *Handler { h.external = false return h @@ -453,7 +455,14 @@ func (h *Handler) upstreamRequest(r *http.Request) (*http.Request, error) { if err != nil { return nil, err } - req = req.Clone(r.Context()) + ctx := r.Context() + if !h.external { + ctx = httputil.WithoutCooldown(ctx) + } + if scope, ok := httputil.CooldownScope(req.Context()); ok { + ctx = httputil.WithCooldownScope(ctx, scope) + } + req = req.Clone(ctx) // Forward safe headers from the original request, without overwriting headers set by transform. forwardable := httputil.FilterHeaders(r.Header, slices.Concat(httputil.HopByHopHeaders, []string{"Range", "If-Range"})...) for key, values := range forwardable { diff --git a/internal/strategy/handler/handler_test.go b/internal/strategy/handler/handler_test.go index 35190d92..ec00e41e 100644 --- a/internal/strategy/handler/handler_test.go +++ b/internal/strategy/handler/handler_test.go @@ -1017,12 +1017,20 @@ func TestBackoffPausesHandler(t *testing.T) { 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)} + client := &http.Client{Transport: httputil.NewBackoffTransport(http.DefaultTransport, time.Minute, time.Hour, slog.New(slog.DiscardHandler))} c := mustNewMemoryCache() defer c.Close() newHandler := func() *handler.Handler { return handler.New(client, c, newTestSpools(t)).Transform(func(r *http.Request) (*http.Request, error) { - return http.NewRequestWithContext(r.Context(), http.MethodGet, r.URL.Query().Get("upstream")+r.URL.Path, nil) + ctx := r.Context() + if scope := r.URL.Query().Get("scope"); scope != "" { + ctx = httputil.WithCooldownScope(ctx, scope) + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, r.URL.Query().Get("upstream")+r.URL.Path, nil) + if err == nil && r.URL.Query().Has("scope") { + req.Header.Set("Authorization", "Bearer token") + } + return req, err }) } h := newHandler() @@ -1048,4 +1056,8 @@ func TestBackoffPausesHandler(t *testing.T) { assert.Equal(t, http.StatusTooManyRequests, w.Code) assert.Equal(t, "60", w.Header().Get("Retry-After")) assert.Equal(t, int32(3), calls.Load()) + + w = serve(t, h, http.MethodGet, "http://example.com/scoped?scope=installation&upstream="+limited.URL, nil) + assert.Equal(t, http.StatusOK, w.Code, "a scoped request has its own cooldown") + assert.Equal(t, int32(4), calls.Load()) } diff --git a/internal/strategy/handler/instrumentation_test.go b/internal/strategy/handler/instrumentation_test.go index f7a590ea..a762eefa 100644 --- a/internal/strategy/handler/instrumentation_test.go +++ b/internal/strategy/handler/instrumentation_test.go @@ -396,7 +396,7 @@ func TestArtifactCooldownIsNotOriginTraffic(t *testing.T) { w.WriteHeader(http.StatusTooManyRequests) })) defer origin.Close() - client := &http.Client{Transport: httputil.NewBackoffTransport(metrics.NewOriginTransport(origin.Client().Transport), time.Minute)} + client := &http.Client{Transport: httputil.NewBackoffTransport(metrics.NewOriginTransport(origin.Client().Transport), time.Minute, time.Hour, slog.New(slog.DiscardHandler))} c := mustNewMemoryCache() defer c.Close() h := handler.New(client, c, newTestSpools(t)).Metrics("host", "registry.example").Transform(func(r *http.Request) (*http.Request, error) { diff --git a/internal/strategy/hermit_test.go b/internal/strategy/hermit_test.go index 1bbd1705..873965d0 100644 --- a/internal/strategy/hermit_test.go +++ b/internal/strategy/hermit_test.go @@ -50,7 +50,7 @@ func TestHermitBackoff(t *testing.T) { defer backend.Close() original := http.DefaultTransport http.DefaultTransport = httputil.NewBackoffTransport( //nolint:reassign - &mockTransport{backend: backend, originalTransport: original}, time.Minute) + &mockTransport{backend: backend, originalTransport: original}, time.Minute, time.Hour, slog.New(slog.DiscardHandler)) 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}) @@ -80,6 +80,47 @@ func TestHermitBackoff(t *testing.T) { assert.Equal(t, 3, calls) } +func TestHermitInternalForwardSkipsCooldown(t *testing.T) { + httpTransportMutexHermit.Lock() + defer httpTransportMutexHermit.Unlock() + + var paths []string + loopback := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + paths = append(paths, r.URL.Path) + if r.URL.Path == "/github.com/org/repo/releases/download/v1/limited" { + w.Header().Set("Retry-After", "60") + w.WriteHeader(http.StatusTooManyRequests) + return + } + _, _ = w.Write([]byte("release")) + })) + defer loopback.Close() + original := http.DefaultTransport + http.DefaultTransport = httputil.NewBackoffTransport( //nolint:reassign + &mockTransport{backend: loopback, originalTransport: original}, time.Minute, time.Hour, slog.New(slog.DiscardHandler)) + 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: "http://127.0.0.1:8080/github.com"}, + nil, c, mux, newTestSpools(t)) + assert.NoError(t, err) + for _, tt := range []struct { + path string + status int + }{ + {"/hermit/github.com/org/repo/releases/download/v1/limited", http.StatusTooManyRequests}, + {"/hermit/github.com/other/tool/releases/download/v2/file", 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) + } + assert.Equal(t, []string{"/github.com/org/repo/releases/download/v1/limited", "/github.com/other/tool/releases/download/v2/file"}, paths) +} + 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 3fdfb6e8..4003af68 100644 --- a/internal/strategy/host_test.go +++ b/internal/strategy/host_test.go @@ -306,8 +306,8 @@ func TestHostFallbackOrder(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 + http.DefaultTransport = httputil.NewBackoffTransport(defaultTransport, time.Minute, time.Hour, slog.New(slog.DiscardHandler)) //nolint:reassign + t.Cleanup(func() { http.DefaultTransport = defaultTransport }) //nolint:reassign _, ctx := logging.Configure(context.Background(), logging.Config{Level: slog.LevelError}) c, err := cache.NewMemory(ctx, cache.MemoryConfig{MaxTTL: time.Hour}) assert.NoError(t, err) @@ -376,8 +376,8 @@ func TestHostFallbackRetryAfter(t *testing.T) { } { 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 + http.DefaultTransport = httputil.NewBackoffTransport(defaultTransport, time.Minute, time.Hour, slog.New(slog.DiscardHandler)) //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) diff --git a/internal/strategy/oci_test.go b/internal/strategy/oci_test.go index 05b11f72..fc068c82 100644 --- a/internal/strategy/oci_test.go +++ b/internal/strategy/oci_test.go @@ -526,8 +526,8 @@ func TestOCIBackoff(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 + http.DefaultTransport = httputil.NewBackoffTransport(defaultTransport, time.Minute, time.Hour, slog.New(slog.DiscardHandler)) //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)