diff --git a/internal/recipes/fetcher.go b/internal/recipes/fetcher.go index 3c680a5..1b1cc11 100644 --- a/internal/recipes/fetcher.go +++ b/internal/recipes/fetcher.go @@ -10,10 +10,12 @@ import ( "io" "net" "net/http" + "net/http/httptrace" "net/netip" "net/url" "strings" "sync" + "sync/atomic" "time" ) @@ -99,9 +101,12 @@ 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; only - // requests actually sent advance it. + // lastSent is when each host's latest request was let through; a + // request that is never 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. + released chan struct{} cache map[string]cacheEntry } @@ -125,6 +130,7 @@ func newFetcher(resolver Resolver, dial dialFunc, tlsConfig *tls.Config) *Fetche now: time.Now, sleep: sleep, lastSent: map[string]time.Time{}, + released: make(chan struct{}), cache: map[string]cacheEntry{}, } transport := &http.Transport{ @@ -237,11 +243,15 @@ 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: the -// host's slot advances only when a request is let through, so a wait that -// ends with ctx does not delay later requests. -func (f *Fetcher) waitTurn(ctx context.Context, host string) error { +// 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. +func (f *Fetcher) waitTurn(ctx context.Context, host string) (release func(), err error) { for { + if err := ctx.Err(); err != nil { + return nil, err + } f.mu.Lock() now := f.now() last, sent := f.lastSent[host] @@ -249,15 +259,55 @@ func (f *Fetcher) waitTurn(ctx context.Context, host string) error { if !sent || wait <= 0 { f.lastSent[host] = now f.mu.Unlock() - return nil + return func() { f.release(host, now, last, sent) }, nil } + released := f.released f.mu.Unlock() - if err := f.sleep(ctx, wait); err != nil { - return err + if err := f.waitOrRelease(ctx, wait, released); err != nil { + return nil, err } } } +// release hands back the slot taken at taken for a request that was not +// sent, restoring the host's previous send time, and wakes waiting requests. +func (f *Fetcher) release(host string, taken, prev time.Time, hadPrev bool) { + f.mu.Lock() + defer f.mu.Unlock() + if !f.lastSent[host].Equal(taken) { + // A later request already holds the slot. + return + } + if hadPrev { + f.lastSent[host] = prev + } else { + delete(f.lastSent, host) + } + close(f.released) + f.released = make(chan struct{}) +} + +// waitOrRelease sleeps for d, returning early without error when a slot is +// handed back so the caller can look again. +func (f *Fetcher) waitOrRelease(ctx context.Context, d time.Duration, released <-chan struct{}) error { + sleepCtx, cancel := context.WithCancel(ctx) + defer cancel() + go func() { + select { + case <-released: + cancel() + case <-sleepCtx.Done(): + } + }() + err := f.sleep(sleepCtx, d) + select { + case <-released: + return nil + default: + return err + } +} + // dialPublic resolves address's host once, refuses unless every answer is // public, then dials the checked addresses in turn. func (f *Fetcher) dialPublic(ctx context.Context, network, address string) (net.Conn, error) { @@ -301,17 +351,25 @@ func (f *Fetcher) dialPublic(ctx context.Context, network, address string) (net. } // limitedTransport waits for each request's per-host rate-limit slot, -// redirect hops included, before sending it. +// 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 next http.RoundTripper } func (t *limitedTransport) RoundTrip(req *http.Request) (*http.Response, error) { - if err := t.f.waitTurn(req.Context(), strings.ToLower(req.URL.Hostname())); err != nil { + release, err := t.f.waitTurn(req.Context(), strings.ToLower(req.URL.Hostname())) + if err != nil { return nil, err } - return t.next.RoundTrip(req) + 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() { + release() + } + return resp, 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 de72215..180943d 100644 --- a/internal/recipes/fetcher_test.go +++ b/internal/recipes/fetcher_test.go @@ -592,3 +592,64 @@ func TestFetchCanceledWaitHoldsNoSlot(t *testing.T) { t.Errorf("waits = %v, want none: canceled waits held slots", got) } } + +// TestFetchUnsentRequestHoldsNoSlot covers requests to a host with no +// previous request that are never sent — canceled before Fetch, or failing +// to connect — and checks that the next request is not delayed. +func TestFetchUnsentRequestHoldsNoSlot(t *testing.T) { + t.Parallel() + const unroutableIP = "93.184.215.99" + answers := map[string][][]string{"example.com": {{unroutableIP}, {publicIP}}} + env := newTestEnv(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, "ok") + }), answers) + clock := installClock(env.f) + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + if _, err := env.f.Fetch(ctx, "https://example.com/canceled"); !errors.Is(err, context.Canceled) { + t.Fatalf("canceled Fetch error = %v, want context.Canceled", err) + } + if _, err := env.f.Fetch(t.Context(), "https://example.com/unroutable"); err == nil || !strings.Contains(err.Error(), "no route") { + t.Fatalf("unroutable Fetch error = %v, want a dial failure", err) + } + if _, err := env.f.Fetch(t.Context(), "https://example.com/next"); err != nil { + t.Fatalf("next Fetch: %v", err) + } + if got := clock.waits(); len(got) != 0 { + t.Errorf("waits = %v, want none: unsent requests held slots", got) + } + if got := env.dials(); len(got) != 2 { + t.Errorf("dials = %v, want the unroutable and the next request only", got) + } +} + +// TestReleaseWakesWaiter checks that a request waiting behind a slot that +// is handed back goes ahead at once instead of sleeping out the interval. +func TestReleaseWakesWaiter(t *testing.T) { + t.Parallel() + f := newFetcher(&stubResolver{}, nil, nil) + release, err := f.waitTurn(t.Context(), "example.com") + if err != nil { + t.Fatalf("first waitTurn: %v", err) + } + done := make(chan error, 1) + start := time.Now() + go func() { + _, err := f.waitTurn(t.Context(), "example.com") + done <- err + }() + time.Sleep(50 * time.Millisecond) + release() + select { + case err := <-done: + if err != nil { + t.Fatalf("waiting waitTurn: %v", err) + } + if waited := time.Since(start); waited >= HostInterval { + t.Errorf("waiter took %s, want less than %s", waited, HostInterval) + } + case <-time.After(HostInterval / 2): + t.Fatal("waiter still waiting after the slot was handed back") + } +}