Detach the shared fan-out from its leader's request context

A single flight is built by whichever request arrived first, but every request on that key waits for it. Running the fan-out on the leader's cancelable request context hands the leader's disconnect to followers whose own connections are healthy: they get 502 all backends failed. Waiters also parked on a WaitGroup, so a follower whose own client went away stayed blocked until the leader finished. Run the flight on a context detached from the leader's request and bounded by the configured timeout, and pass that context into build so the fan-out uses it. Give Do a context so a waiter can abandon a flight it no longer needs; the leader ignores it and always runs fn to completion, keeping the cache warm for the others. A caller that abandons on its own cancellation writes no response.
2026-09-05 21:29:20 +10:00
parent c0aab83c20
commit c77298584d
3 changed files with 250 additions and 31 deletions
@@ -188,7 +188,7 @@ type flightGroup struct {
 }

 type flightCall struct {
-	wg   sync.WaitGroup
+	done chan struct{}
 	resp cachedResponse
 	err  error
 }
@@ -206,18 +206,26 @@ func (p *flightPanic) Error() string {
 }

 // Do returns fn's result and whether this caller shared another's in-flight run.
-func (g *flightGroup) Do(key string, fn func() (cachedResponse, error)) (resp cachedResponse, err error, shared bool) {
+//
+// 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) {
 	g.mu.Lock()
 	if g.calls == nil {
 		g.calls = make(map[string]*flightCall)
 	}
 	if c, ok := g.calls[key]; ok {
 		g.mu.Unlock()
-		c.wg.Wait()
-		return c.resp, c.err, true
+		select {
+		case <-c.done:
+			return c.resp, c.err, true
+		case <-ctx.Done():
+			return cachedResponse{}, ctx.Err(), true
+		}
 	}
-	c := &flightCall{}
-	c.wg.Add(1)
+	c := &flightCall{done: make(chan struct{})}
 	g.calls[key] = c
 	g.mu.Unlock()

@@ -226,11 +234,12 @@ func (g *flightGroup) Do(key string, fn func() (cachedResponse, error)) (resp ca
 			c.resp, c.err = cachedResponse{}, &flightPanic{value: r, stack: debug.Stack()}
 			resp, err = c.resp, c.err
 		}
-		// Done only after the results are stored, so waiters read them.
-		c.wg.Done()
+		// 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()
+		close(c.done)
 	}()

 	c.resp, c.err = fn()
@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -400,7 +401,7 @@ func TestFlightGroup_LeaderRunsOnce(t *testing.T) {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			if _, _, s := g.Do("k", run); s {
+			if _, _, s := g.Do(context.Background(), "k", run); s {
 				shared.Add(1)
 			}
 		}()
@@ -417,7 +418,7 @@ 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("k", func() (cachedResponse, error) { return cachedResponse{Records: -1}, nil }); s {
+	if _, _, s := g.Do(context.Background(), "k", func() (cachedResponse, error) { return cachedResponse{Records: -1}, nil }); s {
 		t.Error("a later call must start its own flight")
 	}
 }
@@ -439,7 +440,7 @@ func TestFlightGroup_LeaderPanicFailsLeaderAndWaiters(t *testing.T) {
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		leader.resp, leader.err, leader.shared = g.Do("k", func() (cachedResponse, error) {
+		leader.resp, leader.err, leader.shared = g.Do(context.Background(), "k", func() (cachedResponse, error) {
 			close(entered)
 			<-release
 			panic("build exploded")
@@ -451,7 +452,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("k", func() (cachedResponse, error) {
+			waiters[i].resp, waiters[i].err, waiters[i].shared = g.Do(context.Background(), "k", func() (cachedResponse, error) {
 				t.Error("a waiter must not run its own fn")
 				return cachedResponse{}, nil
 			})
@@ -478,7 +479,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("k", func() (cachedResponse, error) {
+	resp, err, shared := g.Do(context.Background(), "k", func() (cachedResponse, error) {
 		return cachedResponse{Body: json.RawMessage(`[]`), Records: 0}, nil
 	})
 	if shared {
@@ -494,8 +495,8 @@ func TestServeCached_BuildPanicIs502ThenStale(t *testing.T) {
 	b := newCountingBackend(t, map[string]string{factsPath: `[]`})
 	srv, clk := newCachedServer(t, cacheTestConfig(a.srv.URL, b.srv.URL))

-	panicBuild := func() (cachedResponse, error) { panic("build exploded") }
-	serve := func(build func() (cachedResponse, error)) *httptest.ResponseRecorder {
+	panicBuild := func(context.Context) (cachedResponse, error) { panic("build exploded") }
+	serve := func(build func(context.Context) (cachedResponse, error)) *httptest.ResponseRecorder {
 		rec := httptest.NewRecorder()
 		req := httptest.NewRequest(http.MethodGet, factsPath, nil)
 		srv.serveCached(rec, req, factsPath, nil, build)
@@ -509,7 +510,7 @@ func TestServeCached_BuildPanicIs502ThenStale(t *testing.T) {
 	}

 	stored := `[` + fact("h1", "role", "web", "") + `]`
-	rec = serve(func() (cachedResponse, error) {
+	rec = serve(func(context.Context) (cachedResponse, error) {
 		return cachedResponse{Body: json.RawMessage(stored), Records: 1}, nil
 	})
 	if rec.Code != http.StatusOK {
@@ -539,7 +540,7 @@ func TestServeCached_BuildPanicFailsConcurrentWaiters(t *testing.T) {

 	entered := make(chan struct{})
 	release := make(chan struct{})
-	serve := func(build func() (cachedResponse, error)) *httptest.ResponseRecorder {
+	serve := func(build func(context.Context) (cachedResponse, error)) *httptest.ResponseRecorder {
 		rec := httptest.NewRecorder()
 		req := httptest.NewRequest(http.MethodGet, factsPath, nil)
 		srv.serveCached(rec, req, factsPath, nil, build)
@@ -551,7 +552,7 @@ func TestServeCached_BuildPanicFailsConcurrentWaiters(t *testing.T) {
 	var leaderRec *httptest.ResponseRecorder
 	go func() {
 		defer wg.Done()
-		leaderRec = serve(func() (cachedResponse, error) {
+		leaderRec = serve(func(context.Context) (cachedResponse, error) {
 			close(entered)
 			<-release
 			panic("build exploded")
@@ -565,7 +566,7 @@ func TestServeCached_BuildPanicFailsConcurrentWaiters(t *testing.T) {
 		wg.Add(1)
 		go func(i int) {
 			defer wg.Done()
-			recs[i] = serve(func() (cachedResponse, error) {
+			recs[i] = serve(func(context.Context) (cachedResponse, error) {
 				t.Error("a waiter must not run its own build")
 				return cachedResponse{}, nil
 			})
@@ -586,6 +587,195 @@ func TestServeCached_BuildPanicFailsConcurrentWaiters(t *testing.T) {
 	}
 }

+// The leader's client disconnecting mid-fetch must not fail the followers whose
+// own connections are healthy.
+func TestHandler_LeaderDisconnectDoesNotFailFollowers(t *testing.T) {
+	body := `[` + node("h1", "2026-01-01T00:00:00.000Z") + `]`
+	a := newCountingBackend(t, map[string]string{nodesPath: body})
+	b := newCountingBackend(t, map[string]string{nodesPath: `[]`})
+	srv, _ := newCachedServer(t, cacheTestConfig(a.srv.URL, b.srv.URL))
+	h := srv.Handler()
+
+	release := make(chan struct{})
+	a.setBlock(release)
+	b.setBlock(release)
+
+	leaderCtx, cancelLeader := context.WithCancel(context.Background())
+	defer cancelLeader()
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		req := httptest.NewRequest(http.MethodGet, nodesPath, nil).WithContext(leaderCtx)
+		h.ServeHTTP(httptest.NewRecorder(), req)
+	}()
+
+	// Only once the leader is inside the fan-out does a second caller join its
+	// flight rather than starting one of its own.
+	waitFor(t, func() bool { return a.hitCount(nodesPath) >= 1 })
+
+	follower := httptest.NewRecorder()
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		h.ServeHTTP(follower, httptest.NewRequest(http.MethodGet, nodesPath, nil))
+	}()
+
+	time.Sleep(250 * time.Millisecond)
+	cancelLeader()
+	time.Sleep(100 * time.Millisecond)
+	close(release)
+	wg.Wait()
+
+	if follower.Code != http.StatusOK {
+		t.Fatalf("follower status %d (%s), want 200: a healthy client must not inherit the leader's cancellation",
+			follower.Code, follower.Body.String())
+	}
+	if got := strings.TrimSpace(follower.Body.String()); got != body {
+		t.Errorf("follower body = %s, want %s", got, body)
+	}
+	if got := a.hitCount(nodesPath); got != 1 {
+		t.Errorf("backend a saw %d requests, want 1", got)
+	}
+}
+
+// A follower whose own client goes away must unpark rather than wait out the
+// leader, and must not disturb the flight the others are sharing.
+func TestFlightGroup_WaiterAbandonsOnContextCancel(t *testing.T) {
+	var g flightGroup
+	entered := make(chan struct{})
+	release := make(chan struct{})
+
+	var leaderResp cachedResponse
+	var leaderErr error
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		leaderResp, leaderErr, _ = g.Do(context.Background(), "k", func() (cachedResponse, error) {
+			close(entered)
+			<-release
+			return cachedResponse{Body: json.RawMessage(`[1]`), Records: 1}, nil
+		})
+	}()
+	<-entered
+
+	// A patient waiter proves the flight survives the abandoning one.
+	var patientResp cachedResponse
+	var patientErr error
+	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
+		})
+	}()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan struct{})
+	var abandoned error
+	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
+		})
+	}()
+
+	time.Sleep(100 * time.Millisecond)
+	cancel()
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("a waiter whose context was cancelled stayed parked on the leader")
+	}
+	if !errors.Is(abandoned, context.Canceled) {
+		t.Errorf("abandoning waiter err = %v, want context.Canceled", abandoned)
+	}
+	if !abandonedShared {
+		t.Error("the abandoning waiter did share the flight")
+	}
+
+	close(release)
+	wg.Wait()
+
+	if leaderErr != nil || leaderResp.Records != 1 {
+		t.Errorf("leader = (%+v, %v), want a clean success", leaderResp, leaderErr)
+	}
+	if patientErr != nil || patientResp.Records != 1 {
+		t.Errorf("patient waiter = (%+v, %v), want the leader's result", patientResp, patientErr)
+	}
+
+	g.mu.Lock()
+	remaining := len(g.calls)
+	g.mu.Unlock()
+	if remaining != 0 {
+		t.Errorf("%d flights left registered, want 0", remaining)
+	}
+}
+
+// Every participant walking away must still leave the flight bounded and the
+// group empty: nothing parked, nothing registered.
+func TestFlightGroup_AllCallersAbandon(t *testing.T) {
+	var g flightGroup
+	entered := make(chan struct{})
+	release := 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) {
+			close(entered)
+			<-release
+			return cachedResponse{Records: -1}, nil
+		})
+	}()
+	<-entered
+
+	const waiters = 8
+	ctx, cancel := context.WithCancel(context.Background())
+	var wg sync.WaitGroup
+	for range waiters {
+		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
+			})
+		}()
+	}
+
+	time.Sleep(100 * time.Millisecond)
+	cancel()
+	cancelLeader()
+	waitDone := make(chan struct{})
+	go func() { wg.Wait(); close(waitDone) }()
+	select {
+	case <-waitDone:
+	case <-time.After(2 * time.Second):
+		t.Fatal("waiters stayed parked after their contexts were cancelled")
+	}
+
+	close(release)
+	select {
+	case <-leaderDone:
+	case <-time.After(2 * time.Second):
+		t.Fatal("the leader goroutine leaked")
+	}
+
+	g.mu.Lock()
+	remaining := len(g.calls)
+	g.mu.Unlock()
+	if remaining != 0 {
+		t.Errorf("%d flights left registered, want 0", remaining)
+	}
+}
+
 func TestHandler_HealthzReportsCacheState(t *testing.T) {
 	a := newCountingBackend(t, map[string]string{factsPath: `[` + fact("h1", "role", "web", "") + `]`})
 	b := newCountingBackend(t, map[string]string{factsPath: `[]`})
@@ -142,8 +142,8 @@ func isReportSubResource(path string) bool {

 func (s *Server) serveMerged(w http.ResponseWriter, r *http.Request, path string, merge func([]backendResult) []json.RawMessage) {
 	params := queryParams(r.URL.Query().Get("query"))
-	s.serveCached(w, r, path, params, func() (cachedResponse, error) {
-		alive, err := s.aliveResults(r.Context(), path, params)
+	s.serveCached(w, r, path, params, func(ctx context.Context) (cachedResponse, error) {
+		alive, err := s.aliveResults(ctx, path, params)
 		if err != nil {
 			return cachedResponse{}, err
 		}
@@ -162,8 +162,8 @@ func (s *Server) serveUnion(w http.ResponseWriter, r *http.Request, path string,

 	// Keyed on the request's own params, not the upstream ones: upstreamParams
 	// folds offset into limit, so different windows would collide.
-	s.serveCached(w, r, path, in, func() (cachedResponse, error) {
-		alive, err := s.aliveResults(r.Context(), path, page.upstreamParams(in))
+	s.serveCached(w, r, path, in, func(ctx context.Context) (cachedResponse, error) {
+		alive, err := s.aliveResults(ctx, path, page.upstreamParams(in))
 		if err != nil {
 			return cachedResponse{}, err
 		}
@@ -215,8 +215,8 @@ func (s *Server) serveSummed(w http.ResponseWriter, r *http.Request, path string
 		return
 	}

-	s.serveCached(w, r, path, in, func() (cachedResponse, error) {
-		alive, err := s.aliveResults(r.Context(), path, page.upstreamParams(in))
+	s.serveCached(w, r, path, in, func(ctx context.Context) (cachedResponse, error) {
+		alive, err := s.aliveResults(ctx, path, page.upstreamParams(in))
 		if err != nil {
 			return cachedResponse{}, err
 		}
@@ -285,10 +285,10 @@ type cachedResponse struct {
 // fan-out — and stores the result. A build failure falls back to a stale entry
 // when one exists; that is the only path on which stale data is served. Paths
 // with no cache configured run build directly, unchanged.
-func (s *Server) serveCached(w http.ResponseWriter, r *http.Request, path string, params url.Values, build func() (cachedResponse, error)) {
+func (s *Server) serveCached(w http.ResponseWriter, r *http.Request, path string, params url.Values, build func(context.Context) (cachedResponse, error)) {
 	cache, enabled := s.cacheFor(path, params)
 	if !enabled {
-		resp, err := build()
+		resp, err := build(r.Context())
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusBadGateway)
 			return
@@ -311,8 +311,14 @@ func (s *Server) serveCached(w http.ResponseWriter, r *http.Request, path string
 		stale = &ent
 	}

-	resp, err, _ := s.flights.Do(key, func() (cachedResponse, error) {
-		built, buildErr := build()
+	// 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()
+
+		built, buildErr := build(ctx)
 		if buildErr != nil {
 			return cachedResponse{}, buildErr
 		}
@@ -321,13 +327,18 @@ func (s *Server) serveCached(w http.ResponseWriter, r *http.Request, path string
 			s.log.Printf("warning: encoding cache entry for %s failed: %v", key, marshalErr)
 			return built, nil
 		}
-		// The leader's request may be cancelled while followers still wait.
-		if putErr := cache.Put(context.WithoutCancel(r.Context()), key, body); putErr != nil {
+		if putErr := cache.Put(ctx, key, body); putErr != nil {
 			s.log.Printf("warning: cache store for %s failed: %v", key, putErr)
 		}
 		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) {
+			return
+		}
 		if stale != nil {
 			s.stale.markStale(time.Now())
 			s.log.Printf("warning: serving stale %s from cache (stored %s): %v",
@@ -342,6 +353,15 @@ func (s *Server) serveCached(w http.ResponseWriter, r *http.Request, path string
 	writeCached(w, resp)
 }

+// http.Client reads a zero Timeout as "no deadline", but it would expire a
+// context immediately, so an unset value falls back to the default.
+func (s *Server) flightTimeout() time.Duration {
+	if s.cfg.Timeout > 0 {
+		return s.cfg.Timeout
+	}
+	return defaultTimeout
+}
+
 func (s *Server) writeStored(w http.ResponseWriter, body []byte) {
 	var resp cachedResponse
 	if err := json.Unmarshal(body, &resp); err != nil {