Add artifact limits and shared rate-limit transport

README.md

@@ -296,7 +296,7 @@ The default client includes:
 
 - 30 second timeout
 - 5 retries with exponential backoff (50ms base, 10% jitter)
-- Automatic retry on 429 and 5xx responses
+- Automatic retry on 429 without a usable Retry-After and on 5xx responses
 
 Custom client via the top-level package:
 
@@ -329,6 +329,19 @@ body, err := c.GetBody(ctx, "https://crates.io/api/v1/crates/serde")
 statusCode, err := c.Head(ctx, "https://registry.npmjs.org/lodash")
 ```
 
+`client.NewRateLimitGate(transport)` optionally holds requests by host after a
+429, or a 503 with Retry-After. The first response passes through; later requests
+fail locally with `*client.HeldError`, matching `client.ErrUpstreamHeld`.
+`Held(url)` supports prechecks and `Holds()` reports active holds. Share one gate
+between metadata and artifact HTTP clients to apply the same holds to both.
+
+The gate accepts delay seconds and HTTP dates, caps header delays at one hour,
+and doubles missing or unusable delays from one minute up to fifteen minutes.
+`OnHold` and `OnRelease` callbacks support caller-owned logging and metrics;
+configure them before concurrent use. Metadata clients return held errors and
+429 responses with usable Retry-After immediately. `RateLimitError.RetryAfter`
+reports seconds, rounded up for HTTP dates and capped to fit a 32-bit integer.
+
 ## Artifact Downloads (`fetch/`)
 
 The `fetch` sub-package provides streaming artifact downloads with retry, circuit breaking, DNS caching, and URL resolution.
@@ -360,6 +373,14 @@ io.Copy(dst, artifact.Body)
 
 The fetcher uses DNS caching (5-minute refresh), connection pooling, and a 5-minute timeout suited for large artifacts. It retries on rate limits and server errors with exponential backoff and jitter.
 
+Use `fetch.WithMaxArtifactSize(bytes)` to reject oversized declared bodies before
+reading and enforce the same limit while streaming unknown-length bodies. Errors
+match `fetch.ErrArtifactTooLarge`. Bodies exactly at the limit are accepted;
+detecting overflow consumes at most one extra byte. Zero or negative limits leave
+downloads unlimited. The option also covers observed fetches: rejected bodies
+never produce a completed observation or content digests. HEAD requests still
+report the upstream size, and size rejections do not trip the circuit breaker.
+
 ### Observing artifact responses
 
 Use `FetchObserved` when the response metadata and downloaded content digests need to be retained:

client/client.go

@@ -4,11 +4,11 @@ package client
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"io"
 	"math"
 	"math/rand"
 	"net/http"
-	"strconv"
 	"time"
 
 	"github.com/git-pkgs/registries/safehttp"
@@ -90,6 +90,13 @@ func (c *Client) GetBody(ctx context.Context, url string) ([]byte, error) {
 		}
 
 		lastErr = err
+		if errors.Is(err, ErrUpstreamHeld) {
+			return nil, err
+		}
+		var rateErr *RateLimitError
+		if errors.As(err, &rateErr) {
+			return nil, err
+		}
 
 		var httpErr *HTTPError
 		if ok := isHTTPError(err, &httpErr); ok {
@@ -133,10 +140,12 @@ func (c *Client) doRequest(ctx context.Context, url string) ([]byte, error) {
 			Body:       string(body),
 		}
 		if resp.StatusCode == http.StatusTooManyRequests {
-			if retryAfter := resp.Header.Get("Retry-After"); retryAfter != "" {
-				if seconds, err := strconv.Atoi(retryAfter); err == nil {
-					return nil, &RateLimitError{RetryAfter: seconds}
+			if delay, ok := parseRetryAfter(resp.Header.Get("Retry-After"), time.Now()); ok {
+				seconds := delay / time.Second
+				if delay%time.Second != 0 {
+					seconds++
 				}
+				return nil, &RateLimitError{RetryAfter: int(seconds)}
 			}
 		}
 		return nil, httpErr

client/ratelimit.go

@@ -0,0 +1,202 @@
+package client
+
+import (
+	"errors"
+	"fmt"
+	"net/http"
+	"net/url"
+	"strings"
+	"sync"
+	"time"
+)
+
+const (
+	holdWithoutHeader    = time.Minute
+	maxHoldWithoutHeader = 15 * time.Minute //nolint:mnd // doubling stops here
+	// maxHold caps what an upstream's Retry-After can ask for.
+	maxHold = time.Hour
+)
+
+// ErrUpstreamHeld marks a request refused locally during a rate-limit hold.
+var ErrUpstreamHeld = errors.New("upstream requests are held by rate limiting")
+
+// HeldError is returned in place of a response while a host is held.
+type HeldError struct {
+	Host       string
+	Until      time.Time
+	RetryAfter time.Duration
+}
+
+func (e *HeldError) Error() string {
+	return fmt.Sprintf("%s is rate limiting requests, holding for %s", e.Host, e.RetryAfter.Round(time.Second))
+}
+
+func (e *HeldError) Is(target error) bool { return target == ErrUpstreamHeld }
+
+type hostHold struct {
+	until    time.Time
+	fallback time.Duration
+}
+
+// HoldEvent describes a rate-limit response, including any longer active hold.
+type HoldEvent struct {
+	Host       string
+	Until      time.Time
+	RetryAfter time.Duration
+	FromHeader bool
+}
+
+// RateLimitGate holds a host after a 429 or a 503 with Retry-After. Limited
+// responses pass through; requests during a hold return HeldError. The first
+// non-limited response after expiry releases the hold.
+type RateLimitGate struct {
+	base http.RoundTripper
+	// Set callbacks before use. They run outside the lock and may run concurrently.
+	OnHold    func(HoldEvent)
+	OnRelease func(host string)
+	now       func() time.Time
+
+	mu    sync.Mutex
+	holds map[string]*hostHold
+}
+
+// NewRateLimitGate wraps base. A nil base uses http.DefaultTransport.
+func NewRateLimitGate(base http.RoundTripper) *RateLimitGate {
+	if base == nil {
+		base = http.DefaultTransport
+	}
+	return &RateLimitGate{
+		base:  base,
+		now:   time.Now,
+		holds: make(map[string]*hostHold),
+	}
+}
+
+// RoundTrip implements http.RoundTripper.
+func (g *RateLimitGate) RoundTrip(req *http.Request) (*http.Response, error) {
+	host := hostKey(req.URL)
+	if until, held := g.heldUntil(host); held {
+		return nil, &HeldError{Host: host, Until: until, RetryAfter: until.Sub(g.now())}
+	}
+
+	resp, err := g.base.RoundTrip(req)
+	if err != nil {
+		return nil, err
+	}
+
+	retryAfter := resp.Header.Get("Retry-After")
+	switch {
+	case resp.StatusCode == http.StatusTooManyRequests,
+		resp.StatusCode == http.StatusServiceUnavailable && retryAfter != "":
+		g.hold(host, retryAfter)
+	default:
+		g.release(host)
+	}
+	return resp, nil
+}
+
+// Held returns a *HeldError when requests to rawURL's host are being held,
+// and nil otherwise. It is safe to call on a nil gate.
+func (g *RateLimitGate) Held(rawURL string) error {
+	if g == nil {
+		return nil
+	}
+	u, err := url.Parse(rawURL)
+	if err != nil {
+		return nil
+	}
+	host := hostKey(u)
+	until, held := g.heldUntil(host)
+	if !held {
+		return nil
+	}
+	return &HeldError{Host: host, Until: until, RetryAfter: until.Sub(g.now())}
+}
+
+// Holds returns the hosts currently held and when each hold ends. It is
+// safe to call on a nil gate.
+func (g *RateLimitGate) Holds() map[string]time.Time {
+	if g == nil {
+		return nil
+	}
+	now := g.now()
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	held := make(map[string]time.Time)
+	for host, h := range g.holds {
+		if now.Before(h.until) {
+			held[host] = h.until
+		}
+	}
+	return held
+}
+
+func (g *RateLimitGate) heldUntil(host string) (time.Time, bool) {
+	now := g.now()
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	h, ok := g.holds[host]
+	if !ok || !now.Before(h.until) {
+		return time.Time{}, false
+	}
+	return h.until, true
+}
+
+func (g *RateLimitGate) hold(host, header string) {
+	now := g.now()
+
+	g.mu.Lock()
+	h, ok := g.holds[host]
+	if !ok {
+		h = &hostHold{}
+		g.holds[host] = h
+	}
+	// Responses already in flight when a hold starts all report the same
+	// rate limiting; only a 429 after the hold ran out is a further strike.
+	if !now.Before(h.until) {
+		if h.fallback == 0 {
+			h.fallback = holdWithoutHeader
+		} else {
+			h.fallback = min(h.fallback, maxHoldWithoutHeader/backoffBase) * backoffBase
+		}
+	}
+	d, fromHeader := parseRetryAfter(header, now)
+	if !fromHeader {
+		d = h.fallback
+	}
+	d = min(d, maxHold)
+	until := now.Add(d)
+	if until.After(h.until) {
+		h.until = until
+	}
+	event := HoldEvent{Host: host, Until: h.until, RetryAfter: h.until.Sub(now), FromHeader: fromHeader}
+	g.mu.Unlock()
+
+	if g.OnHold != nil {
+		g.OnHold(event)
+	}
+}
+
+// release drops host's hold once it has ended. Nothing is sent during a
+// hold, so a success arriving while one is active was sent before it began.
+func (g *RateLimitGate) release(host string) {
+	now := g.now()
+	g.mu.Lock()
+	h, had := g.holds[host]
+	if had && now.Before(h.until) {
+		g.mu.Unlock()
+		return
+	}
+	delete(g.holds, host)
+	g.mu.Unlock()
+	if had && g.OnRelease != nil {
+		g.OnRelease(host)
+	}
+}
+
+func hostKey(u *url.URL) string {
+	if u == nil {
+		return ""
+	}
+	return strings.ToLower(u.Host)
+}