From 317e69c6ce2121f75a9026c9a74e93851c702974 Mon Sep 17 00:00:00 2001 From: dev_agent Date: Mon, 5 Oct 2026 12:28:04 -0400 Subject: [PATCH] Rate-limit recipe fetches when the request is written The per-host slot was taken in RoundTrip, before the transport resolved, dialed and completed TLS, so a request stalled in setup could be sent less than HostInterval after a later request to the same host. The transport now dials through dialTLS, which resolves, dials and handshakes under the request's context and returns a turnConn. The turnConn's first Write, which is the start of the request, waits for the host's slot under the Fetcher's mutex; a first Write that sends nothing hands the slot back. ALPN offers only HTTP/1.1 and keep-alives stay off, so each connection carries exactly one request. TestFetchRateLimitsSendsNotSetup stalls one request in DNS while a second is sent, and checks the sends are a full interval apart on a fake clock. Co-Authored-By: Claude Opus 5.5 --- internal/recipes/fetcher.go | 133 ++++++++++++++++++++++--------- internal/recipes/fetcher_test.go | 66 +++++++++++++++ 2 files changed, 162 insertions(+), 37 deletions(-) diff --git a/internal/recipes/fetcher.go b/internal/recipes/fetcher.go index 1b1cc11..5f7791a 100644 --- a/internal/recipes/fetcher.go +++ b/internal/recipes/fetcher.go @@ -10,12 +10,10 @@ import ( "io" "net" "net/http" - "net/http/httptrace" "net/netip" "net/url" "strings" "sync" - "sync/atomic" "time" ) @@ -88,11 +86,14 @@ type dialFunc func(ctx context.Context, network, address string) (net.Conn, erro // answer cannot swap in a private one. Keep-alives are off: every request, // redirect hops included, dials and checks afresh. Proxies from the // environment are not used, since a proxy would dial on the Fetcher's -// behalf. A Fetcher is safe for concurrent use. +// behalf. Requests are HTTP/1.1, so each connection carries exactly one +// request and its first write is that request. A Fetcher is safe for +// concurrent use. type Fetcher struct { - client *http.Client - resolver Resolver - dial dialFunc + client *http.Client + resolver Resolver + dial dialFunc + tlsConfig *tls.Config // timeout bounds a fetch; tests shorten it. timeout time.Duration // now is the clock; tests replace it. @@ -101,8 +102,9 @@ type Fetcher struct { sleep func(ctx context.Context, d time.Duration) error mu sync.Mutex - // lastSent is when each host's latest request was let through; a - // request that is never sent hands its slot back. + // lastSent is when each host's latest request began to be written to + // its connection; a request whose write fails before any byte is sent + // hands its slot back. lastSent map[string]time.Time // released is closed, then replaced, whenever a slot is handed back, // waking requests waiting for their turn. @@ -124,25 +126,23 @@ func NewFetcher() *Fetcher { // nil, replaces the default TLS settings so tests can trust their servers. func newFetcher(resolver Resolver, dial dialFunc, tlsConfig *tls.Config) *Fetcher { f := &Fetcher{ - resolver: resolver, - dial: dial, - timeout: Timeout, - now: time.Now, - sleep: sleep, - lastSent: map[string]time.Time{}, - released: make(chan struct{}), - cache: map[string]cacheEntry{}, + resolver: resolver, + dial: dial, + tlsConfig: tlsConfig, + timeout: Timeout, + now: time.Now, + sleep: sleep, + lastSent: map[string]time.Time{}, + released: make(chan struct{}), + cache: map[string]cacheEntry{}, } transport := &http.Transport{ - Proxy: nil, - DialContext: f.dialPublic, - TLSClientConfig: tlsConfig, - DisableKeepAlives: true, - ForceAttemptHTTP2: true, - TLSHandshakeTimeout: Timeout, + Proxy: nil, + DialTLSContext: f.dialTLS, + DisableKeepAlives: true, } f.client = &http.Client{ - Transport: &limitedTransport{f: f, next: transport}, + Transport: &requestTransport{next: transport}, CheckRedirect: checkRedirect, } return f @@ -245,8 +245,10 @@ func (f *Fetcher) store(key string, page Page) { // waitTurn blocks until a request to host may be sent, keeping requests to // one host at least HostInterval apart. A wait holds no reservation, and a // request ending with ctx before its turn takes no slot. The request let -// through takes the host's slot and must call release if it is not sent, -// so that it delays no other request. +// through takes the host's slot and must be sent at once, or call release +// if sending fails, so that it delays no other request. Fetch calls it from +// turnConn's first Write, so the slot is taken after DNS, dialing and the +// TLS handshake, as the request is written. func (f *Fetcher) waitTurn(ctx context.Context, host string) (release func(), err error) { for { if err := ctx.Err(); err != nil { @@ -350,26 +352,83 @@ func (f *Fetcher) dialPublic(ctx context.Context, network, address string) (net. return nil, errors.Join(errs...) } -// limitedTransport waits for each request's per-host rate-limit slot, -// redirect hops included, before sending it. A request that fails before -// its headers are written was not sent and hands its slot back. -type limitedTransport struct { - f *Fetcher +// requestCtxKey carries a request's own context into its dial. The +// transport dials with a context that keeps the request's values but not +// its cancellation; the connection and the rate-limit wait must end with +// the request. +type requestCtxKey struct{} + +// requestTransport hands each request's context, redirect hops included, +// to dialTLS. +type requestTransport struct { next http.RoundTripper } -func (t *limitedTransport) RoundTrip(req *http.Request) (*http.Response, error) { - release, err := t.f.waitTurn(req.Context(), strings.ToLower(req.URL.Hostname())) +func (t *requestTransport) RoundTrip(req *http.Request) (*http.Response, error) { + ctx := req.Context() + return t.next.RoundTrip(req.WithContext(context.WithValue(ctx, requestCtxKey{}, ctx))) +} + +// dialTLS connects to address through dialPublic and completes the TLS +// handshake, all under the request's context, and returns a turnConn whose +// first write waits for the host's rate-limit slot. ALPN offers only +// HTTP/1.1, so the connection carries this one request. +func (f *Fetcher) dialTLS(dialCtx context.Context, network, address string) (net.Conn, error) { + ctx, ok := dialCtx.Value(requestCtxKey{}).(context.Context) + if !ok { + return nil, fmt.Errorf("recipes: dial %s outside a request", address) + } + host, _, err := net.SplitHostPort(address) + if err != nil { + return nil, fmt.Errorf("recipes: dial address %q: %w", address, err) + } + conn, err := f.dialPublic(ctx, network, address) if err != nil { return nil, err } - var sent atomic.Bool - trace := &httptrace.ClientTrace{WroteHeaders: func() { sent.Store(true) }} - resp, err := t.next.RoundTrip(req.WithContext(httptrace.WithClientTrace(req.Context(), trace))) - if err != nil && !sent.Load() { + cfg := &tls.Config{} + if f.tlsConfig != nil { + cfg = f.tlsConfig.Clone() + } + cfg.ServerName = host + cfg.NextProtos = []string{"http/1.1"} + tc := tls.Client(conn, cfg) + if err := tc.HandshakeContext(ctx); err != nil { + conn.Close() + return nil, fmt.Errorf("recipes: TLS handshake with %s: %w", host, err) + } + return &turnConn{Conn: tc, f: f, ctx: ctx, host: strings.ToLower(host)}, nil +} + +// turnConn is a connection ready to carry one request. Its first Write — +// the start of the request — waits for the host's rate-limit slot, so two +// requests to one host are written at least HostInterval apart however +// long each spent resolving, dialing and in the TLS handshake. A first +// Write that sends nothing hands the slot back. +type turnConn struct { + net.Conn + f *Fetcher + ctx context.Context + host string + // started is set by the first Write; the transport writes a request + // from one goroutine, so it needs no lock. + started bool +} + +func (c *turnConn) Write(p []byte) (int, error) { + if c.started { + return c.Conn.Write(p) + } + release, err := c.f.waitTurn(c.ctx, c.host) + if err != nil { + return 0, err + } + c.started = true + n, err := c.Conn.Write(p) + if n == 0 && err != nil { release() } - return resp, err + return n, err } // checkURL refuses anything but an absolute https:// URL with a host. diff --git a/internal/recipes/fetcher_test.go b/internal/recipes/fetcher_test.go index 180943d..a210fc8 100644 --- a/internal/recipes/fetcher_test.go +++ b/internal/recipes/fetcher_test.go @@ -653,3 +653,69 @@ func TestReleaseWakesWaiter(t *testing.T) { t.Fatal("waiter still waiting after the slot was handed back") } } + +// stallingResolver holds its first lookup until unstall is closed, then +// answers from the stub. +type stallingResolver struct { + *stubResolver + once sync.Once + stalled chan struct{} + unstall chan struct{} +} + +func (r *stallingResolver) LookupIPAddr(ctx context.Context, host string) ([]net.IPAddr, error) { + first := false + r.once.Do(func() { first = true }) + if first { + close(r.stalled) + <-r.unstall + } + return r.stubResolver.LookupIPAddr(ctx, host) +} + +// TestFetchRateLimitsSendsNotSetup covers a request that stalls in DNS +// while a later request to the same host is sent: the stalled request must +// still be sent a full HostInterval after the other one, because the slot +// is taken when the request is written, not when it starts. +func TestFetchRateLimitsSendsNotSetup(t *testing.T) { + t.Parallel() + var clock *fakeClock + var mu sync.Mutex + sentAt := map[string]time.Time{} + env := newTestEnv(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + sentAt[r.URL.Path] = clock.now() + mu.Unlock() + fmt.Fprint(w, "ok") + }), publicAnswers()) + resolver := &stallingResolver{ + stubResolver: env.resolver, + stalled: make(chan struct{}), + unstall: make(chan struct{}), + } + env.f.resolver = resolver + clock = installClock(env.f) + + slow := make(chan error, 1) + go func() { + _, err := env.f.Fetch(t.Context(), "https://example.com/slow") + slow <- err + }() + <-resolver.stalled + if _, err := env.f.Fetch(t.Context(), "https://example.com/fast"); err != nil { + t.Fatalf("fast Fetch: %v", err) + } + close(resolver.unstall) + if err := <-slow; err != nil { + t.Fatalf("slow Fetch: %v", err) + } + + mu.Lock() + defer mu.Unlock() + if gap := sentAt["/slow"].Sub(sentAt["/fast"]); gap < HostInterval { + t.Errorf("slow request sent %s after fast, want at least %s", gap, HostInterval) + } + if got := clock.waits(); len(got) != 1 || got[0] != HostInterval { + t.Errorf("waits = %v, want [%s]: only the stalled request waits", got, HostInterval) + } +}