diff --git a/README.md b/README.md index f9772345..e8ecec05 100644 --- a/README.md +++ b/README.md @@ -1630,7 +1630,7 @@ Failing checks include an `"error"` field. Storage failures also include a `"ste `circuit_breakers` reports the state of each upstream's artifact-fetch circuit breaker (`"open"` or `"closed"`), keyed by upstream host — or by the `hostless-url-` placeholder described under [Monitoring](#monitoring) where the fetch URL has no host to read. The key is omitted until the proxy has fetched an artifact from at least one upstream, and a host appears only once a breaker has been created for it. Breakers trip after repeated upstream failures and retry the upstream after an exponential backoff. While one is open, artifact downloads for that host return HTTP 502 on a cache miss without contacting the upstream; already-cached artifacts are still served from storage, since the cache is checked before the fetcher. A breaker is reported as `"open"` throughout its backoff, including the half-open window in which it admits one probe request to test recovery. Breaker state is per process and in memory, so a restart clears it, but a restart is not needed for recovery: the backoff keeps retrying for as long as the breaker is open, so it closes on its own once the upstream serves again. -`upstream_holds` lists upstream hosts that answered the proxy with HTTP 429, or 503 with a `Retry-After`, and the time until which the proxy sends them nothing. Every request to a held host, metadata and artifacts alike, fails locally for the duration the upstream asked for (capped at one hour), or for one minute doubling up to fifteen when the upstream gave no `Retry-After`. Cached metadata is served stale during a hold, and anything uncached returns HTTP 503 with the remaining `Retry-After` so a client can pace itself too. Metadata requests do not go through the circuit breaker, so without the hold a rate-limited upstream would keep receiving one request per cache miss, which is exactly the traffic it is asking the proxy to stop sending. The first non-limited response after a hold ends clears it, and `proxy_upstream_rate_limit_holds_total` counts the holds placed. The key is omitted while no host is held. +`upstream_holds` lists upstream hosts that answered the proxy with HTTP 429, or 503 with a `Retry-After`, and the time until which the proxy sends them nothing. Every request to a held host, metadata and artifacts alike, fails locally for the duration the upstream asked for (capped at one hour), or for one minute doubling up to fifteen when the upstream gave no `Retry-After`. Cached metadata is served stale during a hold, and anything uncached returns HTTP 503 with the remaining `Retry-After` so a client can pace itself too. Metadata requests do not go through the circuit breaker, so without the hold a rate-limited upstream would keep receiving one request per cache miss, which is exactly the traffic it is asking the proxy to stop sending. The first non-limited response after a hold ends clears it, and `proxy_upstream_rate_limit_holds_total` counts the holds placed. Holds are recorded in the database, so instances sharing one honour each other's holds and list them here too. The key is omitted while no host is held. An open breaker does **not** set `status` to `"error"` or change the HTTP status code: it reports a specific upstream refusing to serve, not this proxy being unfit to receive traffic, and failing the readiness probe over one unhealthy upstream would pull the pod out of rotation for every other ecosystem too. Use `proxy_circuit_breaker_state` for alerting on it. diff --git a/docs/configuration.md b/docs/configuration.md index 5a877e38..959eb45b 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -207,7 +207,9 @@ Instances coordinate through PostgreSQL advisory locks: - Eviction, reclaim, retention and the Gradle build cache sweep run on one instance at a time. An instance that finds a job running elsewhere skips that pass. - A [scheduled mirror](#scheduled-refresh) runs on whichever instance reaches it first in each interval. -Cache misses are fetched from upstream once across instances. The instance that fetches holds a lease on the miss in the `fetch_leases` table and renews it while it downloads. Another instance missing the same file waits until it is stored and serves the stored copy; `proxy_fetch_lease_waits_total` counts those waits. If the fetch fails, the instances waiting on it stop waiting together: each serves its stale metadata copy if it has one, and otherwise fetches itself, so a failing upstream costs every instance one attempt at the same time rather than one after another. A lease of an instance that stopped expires after 30 seconds. Metadata is shared this way only with [`cache_metadata`](#metadata-caching) and a `metadata_ttl` above zero, since without a stored copy there is nothing to wait for. A cache hit costs no extra query; a miss costs one. +Cache misses are fetched from upstream once across instances. The instance that fetches holds a lease on the miss in the `fetch_leases` table and renews it while it downloads. Another instance missing the same file waits until it is stored and serves the stored copy; `proxy_fetch_lease_waits_total` counts those waits. If the fetch fails, the instances waiting on it stop waiting together: each serves its stale metadata copy if it has one, and otherwise fetches itself, so a failing upstream costs every instance one attempt at the same time rather than one after another. When the fetch failed because the upstream rate limited it, the waiters without a stale copy answer 503 with the upstream's `Retry-After` instead of fetching. A lease of an instance that stopped expires after 30 seconds. Metadata is shared this way only with [`cache_metadata`](#metadata-caching) and a `metadata_ttl` above zero, since without a stored copy there is nothing to wait for. A cache hit costs no extra query; a miss costs one. + +Rate-limit holds (see `upstream_holds` under [Health Check](../README.md#health-check)) are shared through the `upstream_holds` table. When one instance gets a 429, or a 503 with `Retry-After`, it records the host and the end of the hold, and the other instances stop sending to that host too instead of each waiting for its own 429. Each instance rereads the table at most every two seconds, so another instance can still send for up to two seconds after a hold starts. A hold is also recorded with SQLite, where it then outlasts a restart. Each instance needs its own [`git.cache_dir`](#git-repositories). diff --git a/internal/database/database_test.go b/internal/database/database_test.go index 8b29cede..f8956a63 100644 --- a/internal/database/database_test.go +++ b/internal/database/database_test.go @@ -613,7 +613,7 @@ func createTestPostgresDB(t *testing.T) *DB { // Drop and recreate every table CreateSchema creates for clean test state; // leftover migration records make the next CreateSchema fail on the // migrations primary key. - tables := []string{"artifacts", "pending_deletes", "versions", "packages", "vulnerabilities", "metadata_cache", "metadata_dates", "job_runs", "fetch_leases", "migrations", "schema_info"} + tables := []string{"artifacts", "pending_deletes", "versions", "packages", "vulnerabilities", "metadata_cache", "metadata_dates", "job_runs", "fetch_leases", "upstream_holds", "migrations", "schema_info"} for _, table := range tables { _, _ = db.Exec("DROP TABLE IF EXISTS " + table + " CASCADE") } diff --git a/internal/database/schema.go b/internal/database/schema.go index 74cc00fb..b2fef631 100644 --- a/internal/database/schema.go +++ b/internal/database/schema.go @@ -147,6 +147,11 @@ CREATE TABLE IF NOT EXISTS metadata_dates ( PRIMARY KEY (ecosystem, proxy_url, name) ); +CREATE TABLE IF NOT EXISTS upstream_holds ( + host TEXT NOT NULL PRIMARY KEY, + until_ms INTEGER NOT NULL +); + CREATE TABLE IF NOT EXISTS migrations ( name TEXT NOT NULL PRIMARY KEY, applied_at DATETIME NOT NULL @@ -283,6 +288,11 @@ CREATE TABLE IF NOT EXISTS metadata_dates ( PRIMARY KEY (ecosystem, proxy_url, name) ); +CREATE TABLE IF NOT EXISTS upstream_holds ( + host TEXT NOT NULL PRIMARY KEY, + until_ms BIGINT NOT NULL +); + CREATE TABLE IF NOT EXISTS migrations ( name TEXT NOT NULL PRIMARY KEY, applied_at TIMESTAMP NOT NULL @@ -447,6 +457,7 @@ var migrations = []migration{ {"014_add_metadata_last_accessed", migrateAddMetadataLastAccessed}, {"015_add_job_runs", migrateAddJobRuns}, {"016_add_fetch_leases", migrateAddFetchLeases}, + {"018_add_upstream_holds", migrateAddUpstreamHolds}, } // isTableNotFound returns true if the error indicates a missing table. @@ -782,6 +793,23 @@ func migrateAddFetchLeases(db *DB) error { return nil } +// migrateAddUpstreamHolds creates the table through which instances sharing +// the database honour each other's upstream rate-limit holds. +func migrateAddUpstreamHolds(db *DB) error { + millis := "INTEGER" + if db.dialect == DialectPostgres { + millis = "BIGINT" + } + query := fmt.Sprintf(`CREATE TABLE IF NOT EXISTS upstream_holds ( + host TEXT NOT NULL PRIMARY KEY, + until_ms %s NOT NULL + )`, millis) + if _, err := db.Exec(query); err != nil { + return fmt.Errorf("creating upstream_holds table: %w", err) + } + return nil +} + // migrateAddPendingDeletesQueuedAtIndex indexes the column reclaim orders // the queue by, which grows large after a retention sweep. func migrateAddPendingDeletesQueuedAtIndex(db *DB) error { diff --git a/internal/database/upstream_holds.go b/internal/database/upstream_holds.go new file mode 100644 index 00000000..6f842ca7 --- /dev/null +++ b/internal/database/upstream_holds.go @@ -0,0 +1,58 @@ +package database + +import ( + "fmt" + "time" +) + +// Upstream holds let instances sharing a database honour a rate-limit hold +// any one of them placed, so each does not need its own 429 before it stops +// sending. Times are the instances' own clocks in Unix milliseconds; a few +// seconds of skew does not matter next to holds of a minute or more. + +type upstreamHoldRow struct { + Host string `db:"host"` + UntilMS int64 `db:"until_ms"` +} + +// PutUpstreamHold records that requests to host are held until until. A hold +// already recorded for longer is kept. +func (db *DB) PutUpstreamHold(host string, until time.Time) error { + _, err := db.Exec(db.Rebind(` + INSERT INTO upstream_holds (host, until_ms) VALUES (?, ?) + ON CONFLICT (host) DO UPDATE SET until_ms = CASE + WHEN excluded.until_ms > upstream_holds.until_ms THEN excluded.until_ms + ELSE upstream_holds.until_ms END + `), host, until.UnixMilli()) + if err != nil { + return fmt.Errorf("recording upstream hold: %w", err) + } + return nil +} + +// UpstreamHolds returns the holds that have not ended by now, and deletes +// those that have. +func (db *DB) UpstreamHolds(now time.Time) (map[string]time.Time, error) { + var rows []upstreamHoldRow + if err := db.Select(&rows, `SELECT host, until_ms FROM upstream_holds`); err != nil { + return nil, fmt.Errorf("reading upstream holds: %w", err) + } + holds := make(map[string]time.Time, len(rows)) + expired := false + for _, row := range rows { + until := time.UnixMilli(row.UntilMS) + if now.Before(until) { + holds[row.Host] = until + } else { + expired = true + } + } + // Only write when there is something to delete, so the frequent reads + // stay reads. + if expired { + if _, err := db.Exec(db.Rebind(`DELETE FROM upstream_holds WHERE until_ms <= ?`), now.UnixMilli()); err != nil { + return holds, fmt.Errorf("deleting expired upstream holds: %w", err) + } + } + return holds, nil +} diff --git a/internal/database/upstream_holds_test.go b/internal/database/upstream_holds_test.go new file mode 100644 index 00000000..5d6617cc --- /dev/null +++ b/internal/database/upstream_holds_test.go @@ -0,0 +1,48 @@ +package database + +import ( + "testing" + "time" +) + +func TestUpstreamHolds(t *testing.T) { + runWithBothDatabases(t, func(t *testing.T, db *DB) { + now := time.UnixMilli(time.Now().UnixMilli()) + if err := db.PutUpstreamHold("registry.npmjs.org", now.Add(time.Minute)); err != nil { + t.Fatal(err) + } + // A shorter hold reported later keeps the longer one. + if err := db.PutUpstreamHold("registry.npmjs.org", now.Add(30*time.Second)); err != nil { + t.Fatal(err) + } + if err := db.PutUpstreamHold("pypi.org", now.Add(10*time.Second)); err != nil { + t.Fatal(err) + } + holds, err := db.UpstreamHolds(now) + if err != nil { + t.Fatal(err) + } + if len(holds) != 2 || !holds["registry.npmjs.org"].Equal(now.Add(time.Minute)) { + t.Fatalf("holds = %v", holds) + } + + // A longer one replaces it. + if err := db.PutUpstreamHold("registry.npmjs.org", now.Add(time.Hour)); err != nil { + t.Fatal(err) + } + holds, err = db.UpstreamHolds(now.Add(20 * time.Second)) + if err != nil { + t.Fatal(err) + } + if len(holds) != 1 || !holds["registry.npmjs.org"].Equal(now.Add(time.Hour)) { + t.Fatalf("holds after pypi.org ended = %v", holds) + } + var rows int + if err := db.Get(&rows, "SELECT COUNT(*) FROM upstream_holds"); err != nil { + t.Fatal(err) + } + if rows != 1 { + t.Errorf("rows = %d, want the ended hold deleted", rows) + } + }) +} diff --git a/internal/handler/handler.go b/internal/handler/handler.go index 2e0d688c..513a5319 100644 --- a/internal/handler/handler.go +++ b/internal/handler/handler.go @@ -26,7 +26,6 @@ import ( "github.com/git-pkgs/proxy/internal/scanner" "github.com/git-pkgs/proxy/internal/storage" "github.com/git-pkgs/purl" - registryclient "github.com/git-pkgs/registries/client" "github.com/git-pkgs/registries/fetch" "github.com/opencontainers/go-digest" ) @@ -195,7 +194,7 @@ type Proxy struct { // RateLimits is the gate under the upstream clients. Artifact fetches // check it before calling Fetcher, so held hosts never reach its breaker. - RateLimits *registryclient.RateLimitGate + RateLimits UpstreamHolds // HostedNPMScopes and HostedComposerVendors name namespaces served from // pushed packages alone, without asking the upstream. @@ -584,7 +583,7 @@ func (p *Proxy) fetchAndCache(ctx context.Context, ecosystem, name, version, fil } } - if err := p.RateLimits.Held(info.URL); err != nil { + if err := p.upstreamHeld(info.URL); err != nil { return artifacts.Artifact{}, "", err } p.Logger.Info("fetching from upstream", @@ -874,7 +873,12 @@ func (p *Proxy) runSharedFetch(ctx context.Context, key string, f *inflightFetch } // With leaseFailed, release is a no-op and this instance fetches on its // own, so a failing upstream costs each instance one attempt in - // parallel. + // parallel. A rate-limited holder has published its hold by now, and + // rereading it makes commit refuse the fetch with the holder's + // Retry-After instead of sending one more request to the limited host. + if outcome == leaseFailed { + p.reloadUpstreamHolds() + } failed := true defer func() { release(failed) }() @@ -1441,7 +1445,10 @@ func (p *Proxy) coalescedMetadataMiss(ctx context.Context, ecosystem, cacheKey, return *hit case leaseFailed: // Serve the stale copy at once, as the holder did, rather - // than trying the failing upstream again. Without one, fetch. + // than trying the failing upstream again. Without one, fetch, + // unless the holder was rate limited: its hold, read here, + // refuses the fetch with its Retry-After. + p.reloadUpstreamHolds() if entry != nil { if res := p.staleMetadata(fetchCtx, ecosystem, cacheKey, validate, entry, errFetchFailedElsewhere); res.err == nil { return res @@ -2202,7 +2209,7 @@ func (p *Proxy) discardStaleArtifact(pkgPURL, versionPURL, filename, upstreamHas } func (p *Proxy) fetchAndCacheFromURL(ctx context.Context, ecosystem, name, version, filename, pkgPURL, versionPURL, downloadURL string, headers http.Header, upstreamHash, declared string) (artifacts.Artifact, string, error) { - if err := p.RateLimits.Held(downloadURL); err != nil { + if err := p.upstreamHeld(downloadURL); err != nil { return artifacts.Artifact{}, "", err } p.Logger.Info("fetching from upstream", @@ -2269,7 +2276,7 @@ func (p *Proxy) fetchUpstream(ctx context.Context, downloadURL string, headers h // after the bytes have been sent, and an unterminated chunked response is the // only way left to make the client reject them (see serveArtifact). func (p *Proxy) streamFromUpstream(ctx context.Context, ecosystem, name, version, filename, versionPURL, upstreamURL, upstreamHash, declared string, fetchArtifact func(context.Context) (*fetch.Artifact, error)) (*CacheResult, error) { - if err := p.RateLimits.Held(upstreamURL); err != nil { + if err := p.upstreamHeld(upstreamURL); err != nil { return nil, err } p.Logger.Info("streaming from upstream", diff --git a/internal/handler/shared_holds_test.go b/internal/handler/shared_holds_test.go new file mode 100644 index 00000000..5536c00c --- /dev/null +++ b/internal/handler/shared_holds_test.go @@ -0,0 +1,170 @@ +package handler + +import ( + "context" + "errors" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "path/filepath" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/git-pkgs/proxy/internal/database" + upstreamhttp "github.com/git-pkgs/proxy/internal/httpclient" + "github.com/git-pkgs/proxy/internal/metrics" + registryclient "github.com/git-pkgs/registries/client" + "github.com/git-pkgs/registries/fetch" + "github.com/prometheus/client_golang/prometheus/testutil" +) + +// limitedUpstream answers every request with 429 and a Retry-After of ten +// minutes. Requests wait for release, which is open unless a test closes it +// later. +type limitedUpstream struct { + server *httptest.Server + calls atomic.Int32 + entered chan struct{} + release chan struct{} + released sync.Once +} + +func newLimitedUpstream(t *testing.T, blocked bool) *limitedUpstream { + u := &limitedUpstream{entered: make(chan struct{}, 8), release: make(chan struct{})} + if !blocked { + u.open() + } + t.Cleanup(u.open) + u.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + u.calls.Add(1) + u.entered <- struct{}{} + <-u.release + w.Header().Set("Retry-After", "600") + w.WriteHeader(http.StatusTooManyRequests) + })) + t.Cleanup(u.server.Close) + return u +} + +func (u *limitedUpstream) open() { u.released.Do(func() { close(u.release) }) } + +func (u *limitedUpstream) artifactURL() string { return u.server.URL + "/left-pad-1.3.0.tgz" } + +// useSharedGate routes p's upstream requests through its own rate limit gate, +// sharing holds through db as a replica does. +func useSharedGate(p *Proxy, db *database.DB, u *limitedUpstream) *upstreamhttp.SharedRateLimitGate { + gate := upstreamhttp.NewSharedRateLimitGate(registryclient.NewRateLimitGate(u.server.Client().Transport), db, nil) + client := &http.Client{Transport: gate} + p.HTTPClient = client + p.Fetcher = fetch.NewFetcher(fetch.WithHTTPClient(client), fetch.WithBaseDelay(time.Millisecond)) + p.RateLimits = gate + return gate +} + +func fetchLimited(ctx context.Context, p *Proxy, u *limitedUpstream) error { + res, err := p.GetOrFetchArtifactFromURL(ctx, "npm", "left-pad", "1.3.0", "left-pad-1.3.0.tgz", u.artifactURL()) + if err == nil { + _ = res.Reader.Close() + } + return err +} + +// requireHeld checks that err is answered like any held upstream request, +// with 503 and a Retry-After. +func requireHeld(t *testing.T, name string, err error) { + t.Helper() + w := httptest.NewRecorder() + if !writeUpstreamHeld(w, err) { + t.Fatalf("%s: err = %v, want a held upstream", name, err) + } + if w.Code != http.StatusServiceUnavailable || w.Header().Get("Retry-After") == "" { + t.Fatalf("%s: status %d, Retry-After %q; want 503 with Retry-After", name, w.Code, w.Header().Get("Retry-After")) + } +} + +// TestSharedHoldStopsOtherInstance checks that once one instance is rate +// limited, another sharing its database refuses requests to that host +// without sending them, until the hold ends. +func TestSharedHoldStopsOtherInstance(t *testing.T) { + path := filepath.Join(t.TempDir(), "shared.db") + dbA, err := database.Create(path) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = dbA.Close() }) + dbB, err := database.Open(path) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = dbB.Close() }) + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + store := newMockStorage() + a := NewProxy(dbA, store, &mockFetcher{}, fetch.NewResolver(), logger) + b := NewProxy(dbB, store, &mockFetcher{}, fetch.NewResolver(), logger) + + u := newLimitedUpstream(t, false) + useSharedGate(a, dbA, u) + gateB := useSharedGate(b, dbB, u) + gateB.Refresh = 0 + + if err := fetchLimited(context.Background(), a, u); err == nil { + t.Fatal("instance a: want the rate limit error") + } + if got := u.calls.Load(); got != 1 { + t.Fatalf("upstream requests after a = %d, want 1", got) + } + + requireHeld(t, "instance b artifact", fetchLimited(context.Background(), b, u)) + _, _, err = b.FetchOrCacheMetadata(context.Background(), "npm", "left-pad", u.server.URL+"/left-pad", "application/json") + if !errors.Is(err, registryclient.ErrUpstreamHeld) { + t.Errorf("instance b metadata: err = %v, want a held upstream", err) + } + if got := u.calls.Load(); got != 1 { + t.Errorf("upstream requests after b = %d, want 1", got) + } + + // The hold ends when the upstream said, and is then deleted. + holds, err := dbB.UpstreamHolds(time.Now().Add(11 * time.Minute)) + if err != nil || len(holds) != 0 { + t.Fatalf("holds after expiry = %v, %v; want none", holds, err) + } + if err := b.upstreamHeld(u.artifactURL()); err != nil { + t.Errorf("instance b after expiry: %v, want no hold", err) + } +} + +// TestFetchLeaseRateLimitedHolderStopsWaiters checks that instances waiting +// on a fetch that was rate limited answer with the holder's Retry-After +// instead of each sending one more request to the limited host. +func TestFetchLeaseRateLimitedHolderStopsWaiters(t *testing.T) { + a, b, db := sharedInstances(t) + u := newLimitedUpstream(t, true) + useSharedGate(a, db, u) + gateB := useSharedGate(b, db, u) + // b has read the shared holds just before, and would not read them + // again for an hour on its own. + gateB.Refresh = time.Hour + if err := b.upstreamHeld(u.artifactURL()); err != nil { + t.Fatal(err) + } + + waits := testutil.ToFloat64(metrics.FetchLeaseWaits.WithLabelValues(metrics.CacheKindArtifact)) + resA := make(chan error, 1) + go func() { resA <- fetchLimited(context.Background(), a, u) }() + <-u.entered + resB := make(chan error, 1) + go func() { resB <- fetchLimited(context.Background(), b, u) }() + waitForLeaseWaits(t, metrics.CacheKindArtifact, waits) + u.open() + + if err := <-resA; err == nil { + t.Fatal("instance a: want the rate limit error") + } + requireHeld(t, "instance b", <-resB) + if got := u.calls.Load(); got != 1 { + t.Errorf("upstream requests = %d, want 1", got) + } +} diff --git a/internal/handler/upstream_held.go b/internal/handler/upstream_held.go index 1097e7ef..57ac84d1 100644 --- a/internal/handler/upstream_held.go +++ b/internal/handler/upstream_held.go @@ -25,3 +25,29 @@ func writeUpstreamHeld(w http.ResponseWriter, err error) bool { http.Error(w, held.Error(), http.StatusServiceUnavailable) return true } + +// UpstreamHolds reports whether requests to a URL's host are held by rate +// limiting. *registryclient.RateLimitGate and the shared gate in +// internal/httpclient implement it. +type UpstreamHolds interface { + // Held returns a *registryclient.HeldError while rawURL's host is held. + Held(rawURL string) error +} + +// upstreamHeld returns the hold on rawURL's host, if any. +func (p *Proxy) upstreamHeld(rawURL string) error { + if p.RateLimits == nil { + return nil + } + return p.RateLimits.Held(rawURL) +} + +// reloadUpstreamHolds rereads holds shared with other instances, when the +// gate shares them. It is called after another instance's fetch failed, +// since a 429 there will have placed a hold this instance may not have read +// yet. +func (p *Proxy) reloadUpstreamHolds() { + if r, ok := p.RateLimits.(interface{ Reload() }); ok { + r.Reload() + } +} diff --git a/internal/httpclient/shared_holds.go b/internal/httpclient/shared_holds.go new file mode 100644 index 00000000..175931b0 --- /dev/null +++ b/internal/httpclient/shared_holds.go @@ -0,0 +1,213 @@ +package httpclient + +import ( + "log/slog" + "net/http" + "net/url" + "strings" + "sync" + "time" + + registryclient "github.com/git-pkgs/registries/client" +) + +// DefaultHoldRefresh is how long a SharedRateLimitGate trusts its copy of +// the shared holds before reading them again. It bounds both the database +// reads, at one per interval per instance, and how long another instance +// may keep sending after a hold is placed. +const DefaultHoldRefresh = 2 * time.Second + +// HoldStore keeps rate-limit holds where every instance sharing it sees them. +type HoldStore interface { + PutUpstreamHold(host string, until time.Time) error + UpstreamHolds(now time.Time) (map[string]time.Time, error) +} + +// SharedRateLimitGate adds holds shared with other instances to a +// RateLimitGate. Holds the gate places are published to the store, and +// requests to a host another instance holds are refused like the gate's own, +// with a HeldError. Without a store it behaves as the gate alone. +type SharedRateLimitGate struct { + gate *registryclient.RateLimitGate + store HoldStore + logger *slog.Logger + now func() time.Time + + // Refresh is how long the shared holds read from the store are used + // before reading them again. Set it before use. + Refresh time.Duration + + mu sync.Mutex + shared map[string]time.Time + // published holds the holds this instance wrote, so a read that started + // before a write does not drop it. + published map[string]time.Time + loaded time.Time + loading bool +} + +// NewSharedRateLimitGate wraps gate, publishing its holds to store and +// honouring those in store. It takes over gate.OnHold, calling the callback +// set before. A nil store shares nothing. +func NewSharedRateLimitGate(gate *registryclient.RateLimitGate, store HoldStore, logger *slog.Logger) *SharedRateLimitGate { + s := &SharedRateLimitGate{ + gate: gate, + store: store, + logger: logger, + now: time.Now, + Refresh: DefaultHoldRefresh, + shared: make(map[string]time.Time), + published: make(map[string]time.Time), + } + if store != nil { + onHold := gate.OnHold + gate.OnHold = func(event registryclient.HoldEvent) { + s.publish(event.Host, event.Until) + if onHold != nil { + onHold(event) + } + } + } + return s +} + +// RoundTrip implements http.RoundTripper. +func (s *SharedRateLimitGate) RoundTrip(req *http.Request) (*http.Response, error) { + if err := s.sharedHeld(hostKey(req.URL)); err != nil { + return nil, err + } + return s.gate.RoundTrip(req) +} + +// Held returns a *registryclient.HeldError when requests to rawURL's host are +// held by this instance or another, and nil otherwise. It is safe to call on +// a nil gate. +func (s *SharedRateLimitGate) Held(rawURL string) error { + if s == nil { + return nil + } + if err := s.gate.Held(rawURL); err != nil { + return err + } + u, err := url.Parse(rawURL) + if err != nil { + return nil + } + return s.sharedHeld(hostKey(u)) +} + +// Holds returns the hosts held by this instance or another, and when each +// hold ends. It is safe to call on a nil gate. +func (s *SharedRateLimitGate) Holds() map[string]time.Time { + if s == nil { + return nil + } + holds := s.gate.Holds() + if s.store == nil { + return holds + } + s.load(false) + now := s.now() + s.mu.Lock() + defer s.mu.Unlock() + for host, until := range s.shared { + if now.Before(until) && until.After(holds[host]) { + if holds == nil { + holds = make(map[string]time.Time) + } + holds[host] = until + } + } + return holds +} + +// Reload reads the shared holds now rather than when the copy runs out. A +// caller that knows another instance has just failed, say on a 429, uses it +// to see that instance's hold at once. +func (s *SharedRateLimitGate) Reload() { + if s == nil || s.store == nil { + return + } + s.load(true) +} + +func (s *SharedRateLimitGate) sharedHeld(host string) error { + if s.store == nil { + return nil + } + s.load(false) + now := s.now() + s.mu.Lock() + until, ok := s.shared[host] + s.mu.Unlock() + if !ok || !now.Before(until) { + return nil + } + return ®istryclient.HeldError{Host: host, Until: until, RetryAfter: until.Sub(now)} +} + +// load refreshes the copy of the shared holds once it is older than Refresh, +// or at once when force is set. One caller reads the store while concurrent +// callers go on with the copy they have; a forced read does not wait for +// one already running, which may have started before the hold it is after. +func (s *SharedRateLimitGate) load(force bool) { + now := s.now() + s.mu.Lock() + if !force && (s.loading || now.Sub(s.loaded) < s.Refresh) { + s.mu.Unlock() + return + } + s.loading = true + s.mu.Unlock() + + holds, err := s.store.UpstreamHolds(now) + + s.mu.Lock() + defer s.mu.Unlock() + s.loading = false + // A failed read keeps the last copy and waits a full interval before + // trying again, so a database outage costs no more queries than usual. + s.loaded = now + if err != nil { + if s.logger != nil { + s.logger.Warn("reading shared upstream holds failed", "error", err) + } + if holds == nil { + return + } + } + for host, until := range s.published { + if !now.Before(until) { + delete(s.published, host) + } else if until.After(holds[host]) { + holds[host] = until + } + } + s.shared = holds +} + +// publish records a hold this instance placed, locally and in the store. The +// responses already in flight when a hold starts each report it again; only a +// later end is written. +func (s *SharedRateLimitGate) publish(host string, until time.Time) { + s.mu.Lock() + if !until.After(s.published[host]) { + s.mu.Unlock() + return + } + s.published[host] = until + if until.After(s.shared[host]) { + s.shared[host] = until + } + s.mu.Unlock() + if err := s.store.PutUpstreamHold(host, until); err != nil && s.logger != nil { + s.logger.Warn("publishing upstream hold failed", "host", host, "error", err) + } +} + +func hostKey(u *url.URL) string { + if u == nil { + return "" + } + return strings.ToLower(u.Host) +} diff --git a/internal/httpclient/shared_holds_test.go b/internal/httpclient/shared_holds_test.go new file mode 100644 index 00000000..31b9d794 --- /dev/null +++ b/internal/httpclient/shared_holds_test.go @@ -0,0 +1,125 @@ +package httpclient + +import ( + "errors" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + "time" + + registryclient "github.com/git-pkgs/registries/client" +) + +// memoryHoldStore is a HoldStore shared by gates in one test, counting reads. +type memoryHoldStore struct { + mu sync.Mutex + holds map[string]time.Time + reads atomic.Int32 +} + +func (m *memoryHoldStore) PutUpstreamHold(host string, until time.Time) error { + m.mu.Lock() + defer m.mu.Unlock() + if until.After(m.holds[host]) { + m.holds[host] = until + } + return nil +} + +func (m *memoryHoldStore) UpstreamHolds(now time.Time) (map[string]time.Time, error) { + m.reads.Add(1) + m.mu.Lock() + defer m.mu.Unlock() + out := make(map[string]time.Time) + for host, until := range m.holds { + if now.Before(until) { + out[host] = until + } + } + return out, nil +} + +func TestSharedRateLimitGateSharesHolds(t *testing.T) { + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls.Add(1) + w.Header().Set("Retry-After", "120") + w.WriteHeader(http.StatusTooManyRequests) + })) + defer upstream.Close() + + store := &memoryHoldStore{holds: make(map[string]time.Time)} + a := NewSharedRateLimitGate(registryclient.NewRateLimitGate(upstream.Client().Transport), store, nil) + b := NewSharedRateLimitGate(registryclient.NewRateLimitGate(upstream.Client().Transport), store, nil) + clock := time.Now() + b.now = func() time.Time { return clock } + + // b reads the store before a is rate limited, and keeps that copy for + // Refresh. + if err := b.Held(upstream.URL); err != nil { + t.Fatalf("b held before any 429: %v", err) + } + resp, err := a.RoundTrip(httptest.NewRequest(http.MethodGet, upstream.URL, nil)) + if err != nil || resp.StatusCode != http.StatusTooManyRequests { + t.Fatalf("a: %v, %v; want the 429", resp, err) + } + _ = resp.Body.Close() + if err := b.Held(upstream.URL); err != nil { + t.Errorf("b refreshed before Refresh ran out: %v", err) + } + + clock = clock.Add(DefaultHoldRefresh) + _, err = b.RoundTrip(httptest.NewRequest(http.MethodGet, upstream.URL, nil)) + var held *registryclient.HeldError + if !errors.As(err, &held) || held.RetryAfter <= 100*time.Second { + t.Fatalf("b after refresh: %v, want a hold of about two minutes", err) + } + if calls.Load() != 1 { + t.Errorf("upstream requests = %d, want 1", calls.Load()) + } + if holds := b.Holds(); len(holds) != 1 { + t.Errorf("b.Holds() = %v, want a's hold", holds) + } + // One read by a before its request, two by b: the first and the one + // after Refresh. + if n := store.reads.Load(); n != 3 { + t.Errorf("store reads = %d, want 3", n) + } + + // Once the hold ends, b sends again. + clock = clock.Add(2 * time.Minute) + if err := b.Held(upstream.URL); err != nil { + t.Errorf("b after the hold ended: %v", err) + } +} + +func TestSharedRateLimitGateReload(t *testing.T) { + store := &memoryHoldStore{holds: make(map[string]time.Time)} + g := NewSharedRateLimitGate(registryclient.NewRateLimitGate(nil), store, nil) + g.Refresh = time.Hour + if err := g.Held("https://registry.example/x"); err != nil { + t.Fatal(err) + } + _ = store.PutUpstreamHold("registry.example", time.Now().Add(time.Minute)) + if err := g.Held("https://registry.example/x"); err != nil { + t.Fatalf("held before reload: %v", err) + } + g.Reload() + if !errors.Is(g.Held("https://registry.example/x"), registryclient.ErrUpstreamHeld) { + t.Error("not held after reload") + } +} + +func TestSharedRateLimitGateWithoutStore(t *testing.T) { + g := NewSharedRateLimitGate(registryclient.NewRateLimitGate(nil), nil, nil) + if err := g.Held("https://registry.example/x"); err != nil { + t.Fatal(err) + } + g.Reload() + var nilGate *SharedRateLimitGate + if nilGate.Held("https://registry.example/x") != nil || nilGate.Holds() != nil { + t.Error("nil gate holds") + } +} diff --git a/internal/server/breakers_test.go b/internal/server/breakers_test.go index 3a9e7640..6cfb5114 100644 --- a/internal/server/breakers_test.go +++ b/internal/server/breakers_test.go @@ -12,6 +12,7 @@ import ( "testing" "time" + upstreamhttp "github.com/git-pkgs/proxy/internal/httpclient" "github.com/git-pkgs/proxy/internal/metrics" registryclient "github.com/git-pkgs/registries/client" "github.com/git-pkgs/registries/fetch" @@ -364,7 +365,7 @@ func TestHealthEndpoint_ReportsUpstreamHolds(t *testing.T) { if resp, err := gate.RoundTrip(req); err == nil { _ = resp.Body.Close() } - ts.server.rateLimits = gate + ts.server.rateLimits = upstreamhttp.NewSharedRateLimitGate(gate, nil, nil) w := httptest.NewRecorder() ts.handler.ServeHTTP(w, httptest.NewRequest("GET", "/health", nil)) diff --git a/internal/server/registry_policies_test.go b/internal/server/registry_policies_test.go index 881356a3..4fb1942f 100644 --- a/internal/server/registry_policies_test.go +++ b/internal/server/registry_policies_test.go @@ -13,6 +13,7 @@ import ( "time" "github.com/git-pkgs/proxy/internal/config" + upstreamhttp "github.com/git-pkgs/proxy/internal/httpclient" registryclient "github.com/git-pkgs/registries/client" "github.com/git-pkgs/registries/fetch" ) @@ -59,7 +60,7 @@ func TestMirrorProxySharesHoldAfterBearerRetry(t *testing.T) { t.Fatalf("metadata fetch = %v, upstream calls %d", err, calls.Load()) } var held *registryclient.HeldError - if !errors.As(proxy.RateLimits.Held(upstream.URL), &held) || held.RetryAfter < 59*time.Minute || len(proxy.RateLimits.Holds()) != 1 { + if !errors.As(proxy.RateLimits.Held(upstream.URL), &held) || held.RetryAfter < 59*time.Minute || len(proxy.RateLimits.(*upstreamhttp.SharedRateLimitGate).Holds()) != 1 { t.Fatalf("shared hold = %+v", held) } } diff --git a/internal/server/server.go b/internal/server/server.go index 3dc2f9be..4a7738fc 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -120,7 +120,7 @@ type Server struct { accessLog *accesslog.Logger ecr *ecrTokens breakers *breakerMonitor - rateLimits *registryclient.RateLimitGate + rateLimits *upstreamhttp.SharedRateLimitGate ecoStats ecosystemStatsCache metaStats metadataStatsCache sources sourceTracker @@ -571,8 +571,14 @@ func (s *Server) newUpstreamClients(artifactTimeout time.Duration) (metadataClie s.logger.Info("upstream rate limit cleared", "host", host) } } - s.rateLimits = gate - authTransport := upstreamhttp.NewTransport(gate, upstreamhttp.AuthFunc(s.authForURL)) + // Holds go through the database so instances sharing it stop sending to + // a host as soon as any one of them is rate limited. + var holdStore upstreamhttp.HoldStore + if s.db != nil { + holdStore = s.db + } + s.rateLimits = upstreamhttp.NewSharedRateLimitGate(gate, holdStore, s.logger) + authTransport := upstreamhttp.NewTransport(s.rateLimits, upstreamhttp.AuthFunc(s.authForURL)) metadata := *safeClient metadata.Timeout = s.cfg.ParseHTTPTimeout() metadata.Transport = authTransport