Cancel a shared flight when its last participant leaves

Reference-count flightCall so the fan-out context ends with the last
caller waiting on it, keeping cfg.Timeout as the upper bound. A leader
leaving with a follower still parked no longer disturbs the flight, and
a solo requester disconnecting releases the upstream sockets at once
instead of holding them for the whole timeout.

Return errFlightAbandoned from Do rather than inferring the abandon path
from the request context's sentinel, and read Age off the server's clock
so it matches the timestamp the cache stored.

README.md

@@ -177,7 +177,9 @@ backend later without further handler changes.
   upstream fan-out; the rest wait for it and share the result. That fan-out runs
   on its own context, bounded by `timeout`, so a client that disconnects can
   neither cancel nor fail the requests sharing its flight; a waiter whose own
-  client goes away leaves the flight running for the others.
+  client goes away leaves the flight running for the others. The flight is
+  cancelled once its last participant leaves, so a lone client disconnecting
+  releases the upstream connections straight away.
 - **Response headers** — every response on a cached path carries `X-Cache`
   (`hit` served from a fresh entry, `miss` built by this request, `stale` the
   expired-entry fallback) and `Age` in whole seconds since the served copy was

cache.go

@@ -3,6 +3,7 @@ package main
 import (
 	"container/list"
 	"context"
+	"errors"
 	"fmt"
 	"net/url"
 	"runtime/debug"
@@ -191,8 +192,18 @@ type flightCall struct {
 	done chan struct{}
 	resp cachedResponse
 	err  error
+
+	cancel context.CancelFunc
+	// participants is the number of callers still waiting on this flight,
+	// guarded by flightGroup.mu.
+	participants int
 }
 
+// errFlightAbandoned reports that this caller stopped waiting because its own
+// context ended. It says nothing about the flight, which may still be running
+// for other participants.
+var errFlightAbandoned = errors.New("abandoned the shared flight")
+
 // flightPanic is a panic from a flight's fn, reported to the leader and to every
 // waiter as an error so callers keep their error handling (stale fallback, 502)
 // instead of seeing a zero-value success.
@@ -207,43 +218,80 @@ func (p *flightPanic) Error() string {
 
 // Do returns fn's result and whether this caller shared another's in-flight run.
 //
-// ctx belongs to the caller alone. A waiter that gives up returns ctx.Err() and
-// leaves the flight running for everyone else; the leader ignores ctx entirely
-// and always runs fn to completion, so one participant walking away can neither
-// cancel nor fail the others. fn is therefore responsible for its own deadline.
-func (g *flightGroup) Do(ctx context.Context, key string, fn func() (cachedResponse, error)) (resp cachedResponse, err error, shared bool) {
+// fn runs on a context of the flight's own: detached from every caller's,
+// bounded by timeout, and cancelled once the last participant leaves. A caller
+// that gives up returns errFlightAbandoned and leaves the flight running for
+// whoever is still waiting, so one participant walking away can neither cancel
+// nor fail the others, while a flight nobody waits on any more is dropped at
+// once rather than holding upstream sockets until the timeout.
+func (g *flightGroup) Do(ctx context.Context, key string, timeout time.Duration, fn func(context.Context) (cachedResponse, error)) (resp cachedResponse, err error, shared bool) {
 	g.mu.Lock()
 	if g.calls == nil {
 		g.calls = make(map[string]*flightCall)
 	}
-	if c, ok := g.calls[key]; ok {
+	c, shared := g.calls[key]
+	if shared {
+		c.participants++
 		g.mu.Unlock()
-		select {
-		case <-c.done:
-			return c.resp, c.err, true
-		case <-ctx.Done():
-			return cachedResponse{}, ctx.Err(), true
-		}
+	} else {
+		flightCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), timeout)
+		c = &flightCall{done: make(chan struct{}), cancel: cancel, participants: 1}
+		g.calls[key] = c
+		g.mu.Unlock()
+		go g.run(flightCtx, key, c, fn)
 	}
-	c := &flightCall{done: make(chan struct{})}
-	g.calls[key] = c
-	g.mu.Unlock()
+	defer g.leave(key, c)
 
+	select {
+	case <-c.done:
+		return c.resp, c.err, shared
+	case <-ctx.Done():
+		return cachedResponse{}, fmt.Errorf("%w: %w", errFlightAbandoned, ctx.Err()), shared
+	}
+}
+
+func (g *flightGroup) run(ctx context.Context, key string, c *flightCall, fn func(context.Context) (cachedResponse, error)) {
 	defer func() {
 		if r := recover(); r != nil {
 			c.resp, c.err = cachedResponse{}, &flightPanic{value: r, stack: debug.Stack()}
-			resp, err = c.resp, c.err
 		}
 		// The key is released before the results are published, so a caller that
 		// arrives late leads a new flight instead of joining a finished one.
-		g.mu.Lock()
-		delete(g.calls, key)
-		g.mu.Unlock()
+		g.forget(key, c)
 		close(c.done)
 	}()
 
-	c.resp, c.err = fn()
-	return c.resp, c.err, false
+	c.resp, c.err = fn(ctx)
+}
+
+// leave drops one participant and, when it was the last, unregisters the key and
+// cancels the flight so a lone requester disconnecting aborts the fan-out
+// instead of pinning a socket per backend for the whole timeout. Unregistering
+// under the same lock that admits joiners keeps anyone from joining a flight
+// that is about to be cancelled.
+func (g *flightGroup) leave(key string, c *flightCall) {
+	g.mu.Lock()
+	c.participants--
+	last := c.participants == 0
+	if last {
+		g.unregister(key, c)
+	}
+	g.mu.Unlock()
+	if last {
+		c.cancel()
+	}
+}
+
+func (g *flightGroup) forget(key string, c *flightCall) {
+	g.mu.Lock()
+	g.unregister(key, c)
+	g.mu.Unlock()
+}
+
+func (g *flightGroup) unregister(key string, c *flightCall) {
+	if cur, ok := g.calls[key]; ok && cur == c {
+		delete(g.calls, key)
+	}
 }
 
 // staleTracker records stale fallbacks for /healthz. serving flips back to false

cache_test.go

@@ -5,11 +5,13 @@ import (
 	"encoding/json"
 	"errors"
 	"io"
+	"math/rand/v2"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"os"
 	"path/filepath"
+	"runtime"
 	"slices"
 	"strconv"
 	"strings"
@@ -19,6 +21,16 @@ import (
 	"time"
 )
 
+const (
+	// testFlightTimeout is long enough that a flight ending early can only be the
+	// refcount, never the deadline.
+	testFlightTimeout = 30 * time.Second
+
+	// disconnectAfter is long enough for the fan-out to reach the backends and
+	// short enough to leave the whole abort well inside cfg.Timeout.
+	disconnectAfter = 25 * time.Millisecond
+)
+
 // fakeClock drives the cache's TTL without sleeping.
 type fakeClock struct {
 	mu sync.Mutex
@@ -111,6 +123,7 @@ func newCachedServer(t *testing.T, cfg Config) (*Server, *fakeClock) {
 	}
 	clk := newFakeClock()
 	mc.now = clk.now
+	srv.now = clk.now
 	return srv, clk
 }
 
@@ -389,7 +402,7 @@ func TestFlightGroup_LeaderRunsOnce(t *testing.T) {
 	entered := make(chan struct{})
 	release := make(chan struct{})
 
-	run := func() (cachedResponse, error) {
+	run := func(context.Context) (cachedResponse, error) {
 		if calls.Add(1) == 1 {
 			close(entered)
 		}
@@ -402,7 +415,7 @@ func TestFlightGroup_LeaderRunsOnce(t *testing.T) {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			if _, _, s := g.Do(context.Background(), "k", run); s {
+			if _, _, s := g.Do(context.Background(), "k", testFlightTimeout, run); s {
 				shared.Add(1)
 			}
 		}()
@@ -419,7 +432,9 @@ func TestFlightGroup_LeaderRunsOnce(t *testing.T) {
 		t.Errorf("%d callers shared the flight, want 15", shared.Load())
 	}
 	// The key is released once the flight finishes.
-	if _, _, s := g.Do(context.Background(), "k", func() (cachedResponse, error) { return cachedResponse{Records: -1}, nil }); s {
+	if _, _, s := g.Do(context.Background(), "k", testFlightTimeout, func(context.Context) (cachedResponse, error) {
+		return cachedResponse{Records: -1}, nil
+	}); s {
 		t.Error("a later call must start its own flight")
 	}
 }
@@ -441,7 +456,7 @@ func TestFlightGroup_LeaderPanicFailsLeaderAndWaiters(t *testing.T) {
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		leader.resp, leader.err, leader.shared = g.Do(context.Background(), "k", func() (cachedResponse, error) {
+		leader.resp, leader.err, leader.shared = g.Do(context.Background(), "k", testFlightTimeout, func(context.Context) (cachedResponse, error) {
 			close(entered)
 			<-release
 			panic("build exploded")
@@ -453,7 +468,7 @@ func TestFlightGroup_LeaderPanicFailsLeaderAndWaiters(t *testing.T) {
 		wg.Add(1)
 		go func(i int) {
 			defer wg.Done()
-			waiters[i].resp, waiters[i].err, waiters[i].shared = g.Do(context.Background(), "k", func() (cachedResponse, error) {
+			waiters[i].resp, waiters[i].err, waiters[i].shared = g.Do(context.Background(), "k", testFlightTimeout, func(context.Context) (cachedResponse, error) {
 				t.Error("a waiter must not run its own fn")
 				return cachedResponse{}, nil
 			})
@@ -480,7 +495,7 @@ func TestFlightGroup_LeaderPanicFailsLeaderAndWaiters(t *testing.T) {
 	}
 
 	// The key is released on the panic path, so a later call leads its own flight.
-	resp, err, shared := g.Do(context.Background(), "k", func() (cachedResponse, error) {
+	resp, err, shared := g.Do(context.Background(), "k", testFlightTimeout, func(context.Context) (cachedResponse, error) {
 		return cachedResponse{Body: json.RawMessage(`[]`), Records: 0}, nil
 	})
 	if shared {
@@ -639,6 +654,117 @@ func TestHandler_LeaderDisconnectDoesNotFailFollowers(t *testing.T) {
 	if got := a.hitCount(nodesPath); got != 1 {
 		t.Errorf("backend a saw %d requests, want 1", got)
 	}
+
+	// The flight the leader started still populated the cache.
+	warmed := doGet(t, h, nodesPath, "")
+	if got := warmed.Header().Get(cacheStatusHeader); got != "hit" {
+		t.Errorf("%s = %q, want hit: the abandoned leader's flight must still warm the cache", cacheStatusHeader, got)
+	}
+	if got := a.hitCount(nodesPath); got != 1 {
+		t.Errorf("backend a saw %d requests after the cached read, want 1", got)
+	}
+	assertNoFlights(t, &srv.flights)
+}
+
+// blockingBackend parks every request until its own context ends, reporting when
+// that happened, so an abandoned fan-out is observable from upstream.
+func blockingBackend(t *testing.T, aborted chan<- time.Time) string {
+	t.Helper()
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		select {
+		case <-r.Context().Done():
+			aborted <- time.Now()
+		case <-time.After(10 * time.Second):
+		}
+	}))
+	t.Cleanup(srv.Close)
+	return srv.URL
+}
+
+// A requester with nobody else on its flight must take the fan-out down with it
+// rather than leave a socket per backend held open until cfg.Timeout.
+func TestHandler_SoloDisconnectAbortsFanOutPromptly(t *testing.T) {
+	aborted := make(chan time.Time, 4)
+	cfg := cacheTestConfig(blockingBackend(t, aborted), blockingBackend(t, aborted))
+	cfg.Timeout = 400 * time.Millisecond
+	srv, _ := newCachedServer(t, cfg)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go func() {
+		time.Sleep(disconnectAfter)
+		cancel()
+	}()
+
+	start := time.Now()
+	req := httptest.NewRequest(http.MethodGet, factsPath, nil).WithContext(ctx)
+	srv.Handler().ServeHTTP(httptest.NewRecorder(), req)
+	if elapsed := time.Since(start); elapsed >= cfg.Timeout/2 {
+		t.Errorf("handler returned after %s, want well under cfg.Timeout %s", elapsed, cfg.Timeout)
+	}
+
+	waitForReleases(t, aborted, len(cfg.Backends), start.Add(cfg.Timeout/2))
+	assertNoFlights(t, &srv.flights)
+}
+
+// Distinct cache keys do not collapse into one flight, so disconnecting clients
+// must not each hold len(backends) sockets for the whole timeout.
+func TestHandler_DisconnectedRequestsDoNotPinBackends(t *testing.T) {
+	aborted := make(chan time.Time, 128)
+	cfg := cacheTestConfig(blockingBackend(t, aborted), blockingBackend(t, aborted))
+	cfg.Timeout = 400 * time.Millisecond
+	srv, _ := newCachedServer(t, cfg)
+	h := srv.Handler()
+
+	// Keep-alive plumbing outlives the requests, so the transport is ours to shut
+	// down before counting goroutines.
+	transport := &http.Transport{}
+	srv.client.Transport = transport
+
+	baseline := runtime.NumGoroutine()
+	const callers = 25
+	var wg sync.WaitGroup
+	start := time.Now()
+	for i := range callers {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			go func() {
+				time.Sleep(disconnectAfter)
+				cancel()
+			}()
+			target := factsPath + "?query=" + url.QueryEscape(strconv.Itoa(i))
+			req := httptest.NewRequest(http.MethodGet, target, nil).WithContext(ctx)
+			h.ServeHTTP(httptest.NewRecorder(), req)
+		}(i)
+	}
+	wg.Wait()
+	if elapsed := time.Since(start); elapsed >= cfg.Timeout/2 {
+		t.Errorf("%d disconnecting callers took %s, want well under cfg.Timeout %s", callers, elapsed, cfg.Timeout)
+	}
+
+	// A caller cancelled before its fan-out was dispatched leaves the backend
+	// nothing to release, so one release per caller is the floor.
+	waitForReleases(t, aborted, callers, start.Add(cfg.Timeout/2))
+	assertNoFlights(t, &srv.flights)
+	transport.CloseIdleConnections()
+	assertGoroutinesSettle(t, baseline, 2)
+}
+
+// waitForReleases fails unless at least want backend requests were released by
+// cutoff, which is set well inside cfg.Timeout so only the flight going away can
+// have freed them.
+func waitForReleases(t *testing.T, aborted <-chan time.Time, want int, cutoff time.Time) {
+	t.Helper()
+	for got := 0; got < want; got++ {
+		select {
+		case <-aborted:
+		case <-time.After(time.Until(cutoff)):
+			t.Fatalf("%d of %d backend requests released before the cutoff, want all of them", got, want)
+		}
+	}
 }
 
 // A follower whose own client goes away must unpark rather than wait out the
@@ -654,7 +780,7 @@ func TestFlightGroup_WaiterAbandonsOnContextCancel(t *testing.T) {
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		leaderResp, leaderErr, _ = g.Do(context.Background(), "k", func() (cachedResponse, error) {
+		leaderResp, leaderErr, _ = g.Do(context.Background(), "k", testFlightTimeout, func(context.Context) (cachedResponse, error) {
 			close(entered)
 			<-release
 			return cachedResponse{Body: json.RawMessage(`[1]`), Records: 1}, nil
@@ -668,10 +794,7 @@ func TestFlightGroup_WaiterAbandonsOnContextCancel(t *testing.T) {
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		patientResp, patientErr, _ = g.Do(context.Background(), "k", func() (cachedResponse, error) {
-			t.Error("a waiter must not run its own fn")
-			return cachedResponse{}, nil
-		})
+		patientResp, patientErr, _ = g.Do(context.Background(), "k", testFlightTimeout, waiterMustNotBuild(t))
 	}()
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -680,10 +803,7 @@ func TestFlightGroup_WaiterAbandonsOnContextCancel(t *testing.T) {
 	var abandonedShared bool
 	go func() {
 		defer close(done)
-		_, abandoned, abandonedShared = g.Do(ctx, "k", func() (cachedResponse, error) {
-			t.Error("a waiter must not run its own fn")
-			return cachedResponse{}, nil
-		})
+		_, abandoned, abandonedShared = g.Do(ctx, "k", testFlightTimeout, waiterMustNotBuild(t))
 	}()
 
 	time.Sleep(100 * time.Millisecond)
@@ -693,8 +813,11 @@ func TestFlightGroup_WaiterAbandonsOnContextCancel(t *testing.T) {
 	case <-time.After(2 * time.Second):
 		t.Fatal("a waiter whose context was cancelled stayed parked on the leader")
 	}
+	if !errors.Is(abandoned, errFlightAbandoned) {
+		t.Errorf("abandoning waiter err = %v, want errFlightAbandoned", abandoned)
+	}
 	if !errors.Is(abandoned, context.Canceled) {
-		t.Errorf("abandoning waiter err = %v, want context.Canceled", abandoned)
+		t.Errorf("abandoning waiter err = %v, want it to carry context.Canceled", abandoned)
 	}
 	if !abandonedShared {
 		t.Error("the abandoning waiter did share the flight")
@@ -718,21 +841,22 @@ func TestFlightGroup_WaiterAbandonsOnContextCancel(t *testing.T) {
 	}
 }
 
-// Every participant walking away must still leave the flight bounded and the
-// group empty: nothing parked, nothing registered.
-func TestFlightGroup_AllCallersAbandon(t *testing.T) {
+// The last participant leaving must cancel the flight's context rather than let
+// it burn the whole timeout, and must leave the group empty.
+func TestFlightGroup_LastParticipantLeavingCancelsFlight(t *testing.T) {
 	var g flightGroup
 	entered := make(chan struct{})
-	release := make(chan struct{})
+	flightCancelled := make(chan struct{})
 
 	leaderCtx, cancelLeader := context.WithCancel(context.Background())
 	leaderDone := make(chan struct{})
 	go func() {
 		defer close(leaderDone)
-		_, _, _ = g.Do(leaderCtx, "k", func() (cachedResponse, error) {
+		_, _, _ = g.Do(leaderCtx, "k", testFlightTimeout, func(ctx context.Context) (cachedResponse, error) {
 			close(entered)
-			<-release
-			return cachedResponse{Records: -1}, nil
+			<-ctx.Done()
+			close(flightCancelled)
+			return cachedResponse{}, ctx.Err()
 		})
 	}()
 	<-entered
@@ -744,10 +868,7 @@ func TestFlightGroup_AllCallersAbandon(t *testing.T) {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			_, _, _ = g.Do(ctx, "k", func() (cachedResponse, error) {
-				t.Error("a waiter must not run its own fn")
-				return cachedResponse{}, nil
-			})
+			_, _, _ = g.Do(ctx, "k", testFlightTimeout, waiterMustNotBuild(t))
 		}()
 	}
 
@@ -762,13 +883,123 @@ func TestFlightGroup_AllCallersAbandon(t *testing.T) {
 		t.Fatal("waiters stayed parked after their contexts were cancelled")
 	}
 
-	close(release)
+	// testFlightTimeout is far longer, so only the refcount can have cancelled it.
+	select {
+	case <-flightCancelled:
+	case <-time.After(2 * time.Second):
+		t.Fatal("the flight ran on after its last participant left")
+	}
 	select {
 	case <-leaderDone:
 	case <-time.After(2 * time.Second):
 		t.Fatal("the leader goroutine leaked")
 	}
 
+	assertNoFlights(t, &g)
+}
+
+// A leader walking away while a waiter is still parked must leave the flight's
+// context untouched, so the waiter gets a real result.
+func TestFlightGroup_LeaderLeavingKeepsFlightAliveForWaiter(t *testing.T) {
+	var g flightGroup
+	entered := make(chan struct{})
+	release := make(chan struct{})
+
+	leaderCtx, cancelLeader := context.WithCancel(context.Background())
+	leaderDone := make(chan struct{})
+	var leaderErr error
+	go func() {
+		defer close(leaderDone)
+		_, leaderErr, _ = g.Do(leaderCtx, "k", testFlightTimeout, func(ctx context.Context) (cachedResponse, error) {
+			close(entered)
+			<-release
+			if err := ctx.Err(); err != nil {
+				return cachedResponse{}, err
+			}
+			return cachedResponse{Body: json.RawMessage(`[1]`), Records: 1}, nil
+		})
+	}()
+	<-entered
+
+	waiterDone := make(chan struct{})
+	var waiterResp cachedResponse
+	var waiterErr error
+	go func() {
+		defer close(waiterDone)
+		waiterResp, waiterErr, _ = g.Do(context.Background(), "k", testFlightTimeout, waiterMustNotBuild(t))
+	}()
+
+	time.Sleep(100 * time.Millisecond)
+	cancelLeader()
+	select {
+	case <-leaderDone:
+	case <-time.After(2 * time.Second):
+		t.Fatal("the leader stayed parked after its context was cancelled")
+	}
+	if !errors.Is(leaderErr, errFlightAbandoned) {
+		t.Errorf("leader err = %v, want errFlightAbandoned", leaderErr)
+	}
+
+	close(release)
+	select {
+	case <-waiterDone:
+	case <-time.After(2 * time.Second):
+		t.Fatal("the waiter never got a result")
+	}
+	if waiterErr != nil || waiterResp.Records != 1 {
+		t.Errorf("waiter = (%+v, %v), want the flight's result", waiterResp, waiterErr)
+	}
+
+	assertNoFlights(t, &g)
+}
+
+// Random cancellations across many keys must leave nothing registered and no
+// goroutines behind.
+func TestFlightGroup_HammerRandomCancellations(t *testing.T) {
+	var g flightGroup
+	baseline := runtime.NumGoroutine()
+
+	const callers = 400
+	var wg sync.WaitGroup
+	for i := range callers {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			if i%3 == 0 {
+				go func() {
+					time.Sleep(time.Duration(rand.IntN(3000)) * time.Microsecond)
+					cancel()
+				}()
+			}
+			key := strconv.Itoa(i % 17)
+			_, _, _ = g.Do(ctx, key, testFlightTimeout, func(ctx context.Context) (cachedResponse, error) {
+				select {
+				case <-ctx.Done():
+					return cachedResponse{}, ctx.Err()
+				case <-time.After(time.Duration(rand.IntN(3000)) * time.Microsecond):
+					return cachedResponse{Records: -1}, nil
+				}
+			})
+		}(i)
+	}
+	wg.Wait()
+
+	assertNoFlights(t, &g)
+	assertGoroutinesSettle(t, baseline, 2)
+}
+
+func waiterMustNotBuild(t *testing.T) func(context.Context) (cachedResponse, error) {
+	t.Helper()
+	return func(context.Context) (cachedResponse, error) {
+		t.Error("a waiter must not run its own fn")
+		return cachedResponse{}, nil
+	}
+}
+
+func assertNoFlights(t *testing.T, g *flightGroup) {
+	t.Helper()
 	g.mu.Lock()
 	remaining := len(g.calls)
 	g.mu.Unlock()
@@ -777,6 +1008,24 @@ func TestFlightGroup_AllCallersAbandon(t *testing.T) {
 	}
 }
 
+// Goroutines unwind after their caller returns, so settling is polled rather
+// than sampled once.
+func assertGoroutinesSettle(t *testing.T, baseline, slack int) {
+	t.Helper()
+	deadline := time.Now().Add(5 * time.Second)
+	for {
+		got := runtime.NumGoroutine()
+		if got <= baseline+slack {
+			return
+		}
+		if time.Now().After(deadline) {
+			t.Errorf("goroutines = %d, want back near the baseline of %d", got, baseline)
+			return
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+}
+
 func TestServeCached_CacheStatusHeaders(t *testing.T) {
 	stored := `[` + fact("h1", "role", "web", "") + `]`
 	a := newCountingBackend(t, map[string]string{factsPath: stored})
@@ -792,16 +1041,17 @@ func TestServeCached_CacheStatusHeaders(t *testing.T) {
 		t.Errorf("first request %s = %q, want 0", ageHeader, got)
 	}
 
+	clk.advance(7 * time.Second)
 	rec = doGet(t, h, factsPath, "")
 	if got := rec.Header().Get(cacheStatusHeader); got != "hit" {
 		t.Errorf("cached request %s = %q, want hit", cacheStatusHeader, got)
 	}
-	if _, err := strconv.Atoi(rec.Header().Get(ageHeader)); err != nil {
-		t.Errorf("cached request %s = %q, want whole seconds", ageHeader, rec.Header().Get(ageHeader))
+	if got := rec.Header().Get(ageHeader); got != "7" {
+		t.Errorf("cached request %s = %q, want 7", ageHeader, got)
 	}
 
 	// Past the TTL with every backend down, the stale fallback must say so.
-	clk.advance(31 * time.Second)
+	clk.advance(24 * time.Second)
 	a.setFail(true)
 	b.setFail(true)
 	rec = doGet(t, h, factsPath, "")
@@ -811,8 +1061,8 @@ func TestServeCached_CacheStatusHeaders(t *testing.T) {
 	if got := rec.Header().Get(cacheStatusHeader); got != "stale" {
 		t.Errorf("stale fallback %s = %q, want stale", cacheStatusHeader, got)
 	}
-	if got := rec.Header().Get(ageHeader); got == "" {
-		t.Errorf("stale fallback must carry an %s header", ageHeader)
+	if got := rec.Header().Get(ageHeader); got != "31" {
+		t.Errorf("stale fallback %s = %q, want 31 seconds since the entry was stored", ageHeader, got)
 	}
 	if got := strings.TrimSpace(rec.Body.String()); got != stored {
 		t.Errorf("stale body = %s, want %s", got, stored)

server.go

@@ -53,6 +53,9 @@ type Server struct {
 	flights    flightGroup
 	stale      staleTracker
 
+	// now is shared with the cache's clock so Age matches the stored timestamp.
+	now func() time.Time
+
 	// freshness cache (freshness merge only).
 	mu        sync.Mutex
 	freshData freshness
@@ -65,6 +68,7 @@ func NewServer(cfg Config, logger *log.Logger) *Server {
 		cfg:    cfg,
 		client: &http.Client{Timeout: cfg.Timeout},
 		log:    logger,
+		now:    time.Now,
 	}
 	if cfg.cacheEnabled() {
 		s.factsCache = newMemoryCache(cfg.FactsTTL, cfg.CacheBytes)
@@ -316,13 +320,11 @@ func (s *Server) serveCached(w http.ResponseWriter, r *http.Request, path string
 		stale = &ent
 	}
 
-	// The flight is shared, so it runs on a context detached from whichever
-	// request happened to lead it: one client disconnecting must not cancel the
-	// fan-out its followers are waiting on. cfg.Timeout keeps it bounded.
-	resp, err, _ := s.flights.Do(r.Context(), key, func() (cachedResponse, error) {
-		ctx, cancel := context.WithTimeout(context.WithoutCancel(r.Context()), s.flightTimeout())
-		defer cancel()
-
+	// The flight is shared, so it runs on its own context rather than the leading
+	// request's: one client disconnecting must not cancel the fan-out its
+	// followers are waiting on, and the flight ends as soon as the last of them
+	// goes. cfg.Timeout keeps it bounded.
+	resp, err, _ := s.flights.Do(r.Context(), key, s.flightTimeout(), func(ctx context.Context) (cachedResponse, error) {
 		built, buildErr := build(ctx)
 		if buildErr != nil {
 			return cachedResponse{}, buildErr
@@ -338,14 +340,13 @@ func (s *Server) serveCached(w http.ResponseWriter, r *http.Request, path string
 		return built, nil
 	})
 	if err != nil {
-		// This caller abandoned the flight because its own client went away; the
-		// flight itself is still running for everyone else and there is nobody
-		// left to write to.
-		if rerr := r.Context().Err(); rerr != nil && errors.Is(err, rerr) {
+		// This caller left the flight because its own client went away, so there
+		// is nobody to write to.
+		if errors.Is(err, errFlightAbandoned) {
 			return
 		}
 		if stale != nil {
-			s.stale.markStale(time.Now())
+			s.stale.markStale(s.now())
 			s.log.Printf("warning: serving stale %s from cache (stored %s): %v",
 				path, stale.StoredAt.UTC().Format(time.RFC3339), err)
 			s.writeStored(w, *stale, CacheStale)
@@ -355,7 +356,7 @@ func (s *Server) serveCached(w http.ResponseWriter, r *http.Request, path string
 		return
 	}
 	s.stale.markFresh()
-	setCacheHeaders(w, CacheMiss, time.Time{})
+	s.setCacheHeaders(w, CacheMiss, time.Time{})
 	writeCached(w, resp)
 }
 
@@ -375,14 +376,15 @@ func (s *Server) writeStored(w http.ResponseWriter, ent CacheEntry, status Cache
 		http.Error(w, "unreadable cache entry", http.StatusBadGateway)
 		return
 	}
-	setCacheHeaders(w, status, ent.StoredAt)
+	s.setCacheHeaders(w, status, ent.StoredAt)
 	writeCached(w, resp)
 }
 
 // setCacheHeaders labels a response from a cache-backed path: X-Cache is
 // hit/stale/miss and Age is whole seconds since the served copy was stored (0
-// for a response built by this request).
-func setCacheHeaders(w http.ResponseWriter, status CacheStatus, storedAt time.Time) {
+// for a response built by this request). It reads the same clock the cache
+// stamps entries with, so the two never disagree.
+func (s *Server) setCacheHeaders(w http.ResponseWriter, status CacheStatus, storedAt time.Time) {
 	label := "miss"
 	switch status {
 	case CacheFresh:
@@ -392,7 +394,7 @@ func setCacheHeaders(w http.ResponseWriter, status CacheStatus, storedAt time.Ti
 	}
 	age := 0
 	if !storedAt.IsZero() {
-		if secs := int(time.Since(storedAt).Seconds()); secs > 0 {
+		if secs := int(s.now().Sub(storedAt).Seconds()); secs > 0 {
 			age = secs
 		}
 	}