From ea75ad13465109ef711cf82c36314272661c4549 Mon Sep 17 00:00:00 2001
From: Matthew Keeler <mkeeler@launchdarkly.com>
Date: Wed, 5 Aug 2026 16:08:19 -0400
Subject: [PATCH 1/5] perf: Deduplicate concurrent polling requests with a
 flight group

The streaming endpoints have always collapsed concurrent replay builds
through a singleflight group, but every polling request snapshotted and
serialized the store on its own. The two server-side full-payload
polling endpoints (GET /sdk/poll and GET /sdk/flags) now share one
payload build per environment through a flight group exposed by the
EnvContext, keyed by the FDv2 basis where the payload depends on it.
The per-context evaluation endpoints and the per-key PHP endpoints are
unchanged: their results rarely collide, so there is little duplicate
work to share.

The store and serialize spans belong to the one request that executes
the build; every request records relay.singleflight.shared on its
request span so a trace without those child spans is explainable.
---
 internal/relayenv/env_context.go             |   7 +
 internal/relayenv/env_context_impl.go        |  10 +
 internal/sharedtest/testclient/fake_store.go |  14 +
 internal/tracing/attributes.go               |   6 +
 relay/relay_endpoints.go                     | 359 +++++++++++--------
 relay/relay_endpoints_singleflight_test.go   | 198 ++++++++++
 6 files changed, 450 insertions(+), 144 deletions(-)
 create mode 100644 relay/relay_endpoints_singleflight_test.go

diff --git a/internal/relayenv/env_context.go b/internal/relayenv/env_context.go
index 0678efdb..d4e2c308 100644
--- a/internal/relayenv/env_context.go
+++ b/internal/relayenv/env_context.go
@@ -16,6 +16,7 @@ import (
 	"github.com/launchdarkly/ld-relay/v9/internal/events"
 	"github.com/launchdarkly/ld-relay/v9/internal/sdks"
 	"github.com/launchdarkly/ld-relay/v9/internal/streams"
+	"golang.org/x/sync/singleflight"
 
 	ldeval "github.com/launchdarkly/go-server-sdk-evaluation/v3"
 )
@@ -118,6 +119,12 @@ type EnvContext interface {
 	// This supports the modern V2 delivery protocol.
 	GetStreamHandlerV2(streams.StreamProvider, credential.SDKCredential) http.Handler
 
+	// GetPollingFlightGroup returns the singleflight group that the polling endpoints use to
+	// deduplicate concurrent requests whose response payloads would be identical. The group is
+	// scoped to this environment, so callers only need to key on the endpoint and any request
+	// parameters that change the payload.
+	GetPollingFlightGroup() *singleflight.Group
+
 	// GetEventDispatcher returns the object that proxies events for this environment.
 	GetEventDispatcher() *events.EventDispatcher
 
diff --git a/internal/relayenv/env_context_impl.go b/internal/relayenv/env_context_impl.go
index 88fb6b53..aa25871b 100644
--- a/internal/relayenv/env_context_impl.go
+++ b/internal/relayenv/env_context_impl.go
@@ -30,6 +30,7 @@ import (
 	"github.com/launchdarkly/go-server-sdk/v7/subsystems"
 	"github.com/launchdarkly/go-server-sdk/v7/subsystems/ldstoreimpl"
 	"github.com/launchdarkly/go-server-sdk/v7/subsystems/ldstoretypes"
+	"golang.org/x/sync/singleflight"
 )
 
 // LogNameMode is used in NewEnvContext to determine whether the environment's log messages should be
@@ -123,6 +124,11 @@ type envContextImpl struct {
 	connectionMapper          ConnectionMapper
 	offline                   bool
 	closed                    bool
+
+	// pollFlightGroup deduplicates concurrent polling requests for this environment; it is
+	// internally synchronized and needs no zero-value setup. It is deliberately not guarded by
+	// mu: the field itself is never reassigned.
+	pollFlightGroup singleflight.Group
 }
 
 // Implementation of the DataStoreQueries interface that the streams package uses as an abstraction of
@@ -668,6 +674,10 @@ func (c *envContextImpl) GetStreamHandlerV2(streamProvider streams.StreamProvide
 	return h
 }
 
+func (c *envContextImpl) GetPollingFlightGroup() *singleflight.Group {
+	return &c.pollFlightGroup
+}
+
 func invalidStreamHandler(w http.ResponseWriter, req *http.Request) {
 	w.WriteHeader(http.StatusNotFound)
 }
diff --git a/internal/sharedtest/testclient/fake_store.go b/internal/sharedtest/testclient/fake_store.go
index 07f1bbf2..191e9082 100644
--- a/internal/sharedtest/testclient/fake_store.go
+++ b/internal/sharedtest/testclient/fake_store.go
@@ -13,6 +13,14 @@ type FakeStore struct {
 	collections []ldstoretypes.Collection
 	selector    subsystems.Selector
 	mu          sync.Mutex
+
+	// SnapshotHook, if non-nil, is called at the start of every Snapshot call, before the
+	// store's lock is taken. Tests set it (before sharing the store across goroutines) to
+	// observe or block concurrent snapshot reads.
+	SnapshotHook func()
+
+	// GetAllHook is SnapshotHook for GetAll calls.
+	GetAllHook func()
 }
 
 func NewFakeStore(collections []ldstoretypes.Collection) *FakeStore {
@@ -152,6 +160,9 @@ func (s *FakeStore) Get(kind ldstoretypes.DataKind, key string) (ldstoretypes.It
 }
 
 func (s *FakeStore) GetAll(kind ldstoretypes.DataKind) ([]ldstoretypes.KeyedItemDescriptor, error) {
+	if s.GetAllHook != nil {
+		s.GetAllHook()
+	}
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	result := []ldstoretypes.KeyedItemDescriptor{}
@@ -170,6 +181,9 @@ func (s *FakeStore) IsInitialized() bool {
 }
 
 func (s *FakeStore) Snapshot() (map[ldstoretypes.DataKind][]ldstoretypes.KeyedItemDescriptor, subsystems.Selector, error) {
+	if s.SnapshotHook != nil {
+		s.SnapshotHook()
+	}
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	result := make(map[ldstoretypes.DataKind][]ldstoretypes.KeyedItemDescriptor)
diff --git a/internal/tracing/attributes.go b/internal/tracing/attributes.go
index c7eafa74..0fd64489 100644
--- a/internal/tracing/attributes.go
+++ b/internal/tracing/attributes.go
@@ -34,4 +34,10 @@ const (
 	StoreKeyKey      = attribute.Key("relay.store.key")
 	PayloadEventsKey = attribute.Key("relay.payload.events")
 	PayloadBytesKey  = attribute.Key("relay.payload.bytes")
+
+	// SingleflightSharedKey reports, on a polling endpoint's request span, whether the
+	// response payload build was shared with concurrent requests through the environment's
+	// polling flight group. When it is true and the request's trace has no store or serialize
+	// child spans, another request's trace carries them.
+	SingleflightSharedKey = attribute.Key("relay.singleflight.shared")
 )
diff --git a/relay/relay_endpoints.go b/relay/relay_endpoints.go
index 4359f756..175363bc 100644
--- a/relay/relay_endpoints.go
+++ b/relay/relay_endpoints.go
@@ -1,6 +1,7 @@
 package relay
 
 import (
+	"context"
 	"crypto/sha1" //nolint:gosec // we're not using SHA1 for encryption, just for generating an insecure hash
 	"encoding/hex"
 	"encoding/json"
@@ -187,13 +188,70 @@ type payloadEvent struct {
 	EventData any    `json:"data"`
 }
 
+// pollResult is what one execution of a polling flight-group closure produces: the exact response
+// body and the value the Etag header is derived from. Concurrent requests that share a flight all
+// receive the same pollResult; nothing in it is specific to any one request.
+type pollResult struct {
+	data []byte
+	etag string
+}
+
+// Flight-group keys for the polling endpoints. The flight group is scoped to one environment, so
+// these only need to distinguish the endpoints from each other, plus any request parameter that
+// changes the payload (which the handlers append).
+const (
+	serverSidePollFlightKey     = "sdk-poll"
+	serverSideAllFlagsFlightKey = "sdk-flags"
+)
+
 // Server-side SDK polling endpoint: app.ld.com/sdk/poll/
 func pollHandlerV2(w http.ResponseWriter, req *http.Request) {
 	clientCtx := middleware.GetEnvContextInfo(req.Context())
 	tr := tracing.Tracer()
 
-	_, storeSpan := tr.Start(req.Context(), tracing.SpanStoreSnapshot)
-	collection, selector, err := clientCtx.Env.GetStore().Snapshot()
+	basis := req.URL.Query().Get("basis")
+
+	// Concurrent requests would each take a store snapshot and serialize an identical payload;
+	// the flight group runs that work once and hands every waiting request the same result.
+	// The payload depends on the caller's basis -- a basis matching the current selector state
+	// gets an "up-to-date" response, anything else gets a full transfer -- so only requests
+	// with the same basis may share a result, and the basis is part of the key.
+	//
+	// The snapshot and serialize spans belong to the one request that executes the closure;
+	// every request records on its own request span whether the payload build was shared.
+	data, err, shared := clientCtx.Env.GetPollingFlightGroup().Do(serverSidePollFlightKey+":"+basis, func() (any, error) {
+		result, err := buildServerSidePollPayload(req.Context(), tr, clientCtx.Env, basis)
+		if err != nil {
+			return nil, err
+		}
+		return result, nil
+	})
+	trace.SpanFromContext(req.Context()).SetAttributes(tracing.SingleflightSharedKey.Bool(shared))
+	if err != nil {
+		w.WriteHeader(http.StatusInternalServerError)
+		return
+	}
+
+	// panic if it's not a pollResult - as this should be impossible
+	result := data.(pollResult)
+
+	traceWriteResponse(tr, req, func() (int, error) {
+		return writeCacheableJSONResponse(w, req, clientCtx.Env, result.data, result.etag)
+	})
+}
+
+// buildServerSidePollPayload takes the store snapshot and serializes the FDv2 polling payload for
+// pollHandlerV2. It runs inside the environment's polling flight group, so it must not touch any
+// one request's ResponseWriter; failures are logged and traced here (once per flight, not once
+// per waiting request) and returned as a plain error that every sharing request maps to a 500.
+func buildServerSidePollPayload(
+	ctx context.Context,
+	tr trace.Tracer,
+	env relayenv.EnvContext,
+	basis string,
+) (pollResult, error) {
+	_, storeSpan := tr.Start(ctx, tracing.SpanStoreSnapshot)
+	collection, selector, err := env.GetStore().Snapshot()
 	if err != nil {
 		storeSpan.RecordError(err)
 		storeSpan.SetStatus(codes.Error, err.Error())
@@ -201,138 +259,123 @@ func pollHandlerV2(w http.ResponseWriter, req *http.Request) {
 	storeSpan.End()
 
 	if err != nil {
-		clientCtx.Env.GetLogger().Error("error reading feature store", "error", err)
-		w.WriteHeader(http.StatusInternalServerError)
-		return
+		env.GetLogger().Error("error reading feature store", "error", err)
+		return pollResult{}, err
 	} else if collection == nil || !selector.IsDefined() {
-		clientCtx.Env.GetLogger().Error("snapshot selector is not defined; no data to return")
-		w.WriteHeader(http.StatusInternalServerError)
-		return
+		err := errors.New("snapshot selector is not defined; no data to return")
+		env.GetLogger().Error(err.Error())
+		return pollResult{}, err
 	}
 
-	payloadJSON, ok := func() ([]byte, bool) {
-		_, span := tr.Start(req.Context(), tracing.SpanSerializePayload)
-		defer span.End()
+	_, span := tr.Start(ctx, tracing.SpanSerializePayload)
+	defer span.End()
 
-		numItems := 2
-		if len(collection) > 0 {
-			for _, keyedItems := range collection {
-				numItems += len(keyedItems)
-			}
+	numItems := 2
+	if len(collection) > 0 {
+		for _, keyedItems := range collection {
+			numItems += len(keyedItems)
 		}
+	}
 
-		pollingPayload := pollingPayload{
-			Events: make([]payloadEvent, 0, numItems),
-		}
+	pollingPayload := pollingPayload{
+		Events: make([]payloadEvent, 0, numItems),
+	}
 
-		basis := req.URL.Query().Get("basis")
-		if selector.IsDefined() && basis != "" && selector.State() == basis {
-			pollingPayload.Events = append(pollingPayload.Events, payloadEvent{
-				Event: "server-intent",
-				EventData: subsystems.ServerIntent{Payload: subsystems.Payload{
-					ID:     selector.State(),
-					Target: selector.Version(),
-					Code:   subsystems.IntentNone,
-					Reason: "up-to-date",
-				}},
-			})
-		} else {
-			pollingPayload.Events = append(pollingPayload.Events, payloadEvent{
-				Event: "server-intent",
-				EventData: subsystems.ServerIntent{Payload: subsystems.Payload{
-					ID:     selector.State(),
-					Target: selector.Version(),
-					Code:   subsystems.IntentTransferFull,
-					Reason: "cant-catchup",
-				}},
-			})
-			for kind, keyedItems := range collection {
-				for _, keyedItem := range keyedItems {
-					if keyedItem.Item.Item == nil {
-						continue // this should not happen, but just in case
+	if selector.IsDefined() && basis != "" && selector.State() == basis {
+		pollingPayload.Events = append(pollingPayload.Events, payloadEvent{
+			Event: "server-intent",
+			EventData: subsystems.ServerIntent{Payload: subsystems.Payload{
+				ID:     selector.State(),
+				Target: selector.Version(),
+				Code:   subsystems.IntentNone,
+				Reason: "up-to-date",
+			}},
+		})
+	} else {
+		pollingPayload.Events = append(pollingPayload.Events, payloadEvent{
+			Event: "server-intent",
+			EventData: subsystems.ServerIntent{Payload: subsystems.Payload{
+				ID:     selector.State(),
+				Target: selector.Version(),
+				Code:   subsystems.IntentTransferFull,
+				Reason: "cant-catchup",
+			}},
+		})
+		for kind, keyedItems := range collection {
+			for _, keyedItem := range keyedItems {
+				if keyedItem.Item.Item == nil {
+					continue // this should not happen, but just in case
+				}
+				switch kind {
+				case ldstoreimpl.Features():
+					if flag, ok := keyedItem.Item.Item.(*ldmodel.FeatureFlag); ok {
+						writer := jwriter.NewWriter()
+						ldmodel.MarshalFeatureFlagToJSONWriter(*flag, &writer)
+
+						pollingPayload.Events = append(pollingPayload.Events, payloadEvent{
+							Event: "put-object",
+							EventData: subsystems.PutObject{
+								Version: keyedItem.Item.Version,
+								Kind:    subsystems.FlagKind,
+								Key:     keyedItem.Key,
+								Object:  writer.Bytes(),
+							},
+						})
+					} else {
+						err := errors.New("error casting keyed item to feature flag")
+						env.GetLogger().Error(err.Error())
+						span.RecordError(err)
+						span.SetStatus(codes.Error, err.Error())
+						return pollResult{}, err
 					}
-					switch kind {
-					case ldstoreimpl.Features():
-						if flag, ok := keyedItem.Item.Item.(*ldmodel.FeatureFlag); ok {
-							writer := jwriter.NewWriter()
-							ldmodel.MarshalFeatureFlagToJSONWriter(*flag, &writer)
-
-							pollingPayload.Events = append(pollingPayload.Events, payloadEvent{
-								Event: "put-object",
-								EventData: subsystems.PutObject{
-									Version: keyedItem.Item.Version,
-									Kind:    subsystems.FlagKind,
-									Key:     keyedItem.Key,
-									Object:  writer.Bytes(),
-								},
-							})
-						} else {
-							err := errors.New("error casting keyed item to feature flag")
-							clientCtx.Env.GetLogger().Error(err.Error())
-							span.RecordError(err)
-							span.SetStatus(codes.Error, err.Error())
-							w.WriteHeader(http.StatusInternalServerError)
-							return nil, false
-						}
-					case ldstoreimpl.Segments():
-						if segment, ok := keyedItem.Item.Item.(*ldmodel.Segment); ok {
-							writer := jwriter.NewWriter()
-							ldmodel.MarshalSegmentToJSONWriter(*segment, &writer)
-
-							pollingPayload.Events = append(pollingPayload.Events, payloadEvent{
-								Event: "put-object",
-								EventData: subsystems.PutObject{
-									Version: keyedItem.Item.Version,
-									Kind:    subsystems.SegmentKind,
-									Key:     keyedItem.Key,
-									Object:  writer.Bytes(),
-								},
-							})
-						} else {
-							err := errors.New("error casting keyed item to feature segment")
-							clientCtx.Env.GetLogger().Error(err.Error())
-							span.RecordError(err)
-							span.SetStatus(codes.Error, err.Error())
-							w.WriteHeader(http.StatusInternalServerError)
-							return nil, false
-						}
-					default:
-						err := errors.New("unexpected data kind in store snapshot")
-						clientCtx.Env.GetLogger().Error(err.Error(), "kind", kind)
+				case ldstoreimpl.Segments():
+					if segment, ok := keyedItem.Item.Item.(*ldmodel.Segment); ok {
+						writer := jwriter.NewWriter()
+						ldmodel.MarshalSegmentToJSONWriter(*segment, &writer)
+
+						pollingPayload.Events = append(pollingPayload.Events, payloadEvent{
+							Event: "put-object",
+							EventData: subsystems.PutObject{
+								Version: keyedItem.Item.Version,
+								Kind:    subsystems.SegmentKind,
+								Key:     keyedItem.Key,
+								Object:  writer.Bytes(),
+							},
+						})
+					} else {
+						err := errors.New("error casting keyed item to feature segment")
+						env.GetLogger().Error(err.Error())
 						span.RecordError(err)
 						span.SetStatus(codes.Error, err.Error())
-						w.WriteHeader(http.StatusInternalServerError)
-						return nil, false
+						return pollResult{}, err
 					}
+				default:
+					err := errors.New("unexpected data kind in store snapshot")
+					env.GetLogger().Error(err.Error(), "kind", kind)
+					span.RecordError(err)
+					span.SetStatus(codes.Error, err.Error())
+					return pollResult{}, err
 				}
 			}
-			pollingPayload.Events = append(pollingPayload.Events, payloadEvent{
-				Event:     "payload-transferred",
-				EventData: selector,
-			})
-		}
-
-		data, err := json.Marshal(pollingPayload)
-		if err != nil {
-			span.RecordError(err)
-			span.SetStatus(codes.Error, err.Error())
-			clientCtx.Env.GetLogger().Error("error marshaling polling response", "error", err)
-			w.WriteHeader(http.StatusInternalServerError)
-			return nil, false
 		}
-		span.SetAttributes(
-			tracing.PayloadEventsKey.Int(len(pollingPayload.Events)),
-			tracing.PayloadBytesKey.Int(len(data)),
-		)
-		return data, true
-	}()
-	if !ok {
-		return
+		pollingPayload.Events = append(pollingPayload.Events, payloadEvent{
+			Event:     "payload-transferred",
+			EventData: selector,
+		})
 	}
 
-	traceWriteResponse(tr, req, func() (int, error) {
-		return writeCacheableJSONResponse(w, req, clientCtx.Env, payloadJSON, selector.State())
-	})
+	data, err := json.Marshal(pollingPayload)
+	if err != nil {
+		span.RecordError(err)
+		span.SetStatus(codes.Error, err.Error())
+		env.GetLogger().Error("error marshaling polling response", "error", err)
+		return pollResult{}, err
+	}
+	span.SetAttributes(
+		tracing.PayloadEventsKey.Int(len(pollingPayload.Events)),
+		tracing.PayloadBytesKey.Int(len(data)),
+	)
+	return pollResult{data: data, etag: selector.State()}, nil
 }
 
 // FDv2 client-side polling endpoint that evaluates flags against a context.
@@ -508,8 +551,39 @@ func pollAllFlagsHandler(w http.ResponseWriter, req *http.Request) {
 	clientCtx := middleware.GetEnvContextInfo(req.Context())
 	tr := tracing.Tracer()
 
-	_, storeSpan := tr.Start(req.Context(), tracing.SpanStoreGetAll)
-	data, err := clientCtx.Env.GetStore().GetAll(ldstoreimpl.Features())
+	// Concurrent requests would each read every flag and serialize an identical map; the
+	// flight group runs that work once and hands every waiting request the same result. The
+	// payload and Etag depend only on the store contents, so a single key covers all callers.
+	// Span placement follows pollHandlerV2: the store and serialize spans belong to the one
+	// request that executes the closure.
+	data, err, shared := clientCtx.Env.GetPollingFlightGroup().Do(serverSideAllFlagsFlightKey, func() (any, error) {
+		result, err := buildAllFlagsPayload(req.Context(), tr, clientCtx.Env)
+		if err != nil {
+			return nil, err
+		}
+		return result, nil
+	})
+	trace.SpanFromContext(req.Context()).SetAttributes(tracing.SingleflightSharedKey.Bool(shared))
+	if err != nil {
+		w.WriteHeader(http.StatusInternalServerError)
+		return
+	}
+
+	// panic if it's not a pollResult - as this should be impossible
+	result := data.(pollResult)
+
+	traceWriteResponse(tr, req, func() (int, error) {
+		return writeCacheableJSONResponse(w, req, clientCtx.Env, result.data, result.etag)
+	})
+}
+
+// buildAllFlagsPayload reads every flag and serializes the all-flags map for pollAllFlagsHandler.
+// It runs inside the environment's polling flight group, so it must not touch any one request's
+// ResponseWriter; failures are logged and traced here (once per flight, not once per waiting
+// request) and returned as a plain error that every sharing request maps to a 500.
+func buildAllFlagsPayload(ctx context.Context, tr trace.Tracer, env relayenv.EnvContext) (pollResult, error) {
+	_, storeSpan := tr.Start(ctx, tracing.SpanStoreGetAll)
+	data, err := env.GetStore().GetAll(ldstoreimpl.Features())
 	if err != nil {
 		storeSpan.RecordError(err)
 		storeSpan.SetStatus(codes.Error, err.Error())
@@ -517,32 +591,26 @@ func pollAllFlagsHandler(w http.ResponseWriter, req *http.Request) {
 	storeSpan.End()
 
 	if err != nil {
-		clientCtx.Env.GetLogger().Error("error reading feature store", "error", err)
-		w.WriteHeader(http.StatusInternalServerError)
-		return
+		env.GetLogger().Error("error reading feature store", "error", err)
+		return pollResult{}, err
 	}
-	respData, etag := func() ([]byte, string) {
-		_, span := tr.Start(req.Context(), tracing.SpanSerializePayload)
-		defer span.End()
 
-		respData, flagCount := serializeFlagsAsMap(data)
-		// Compute an overall Etag for the data set by hashing flag keys and versions
-		hash := sha1.New()                                                         //nolint:gosec // just used for insecure hashing
-		sort.Slice(data, func(i, j int) bool { return data[i].Key < data[j].Key }) // makes the hash deterministic
-		for _, item := range data {
-			_, _ = io.WriteString(hash, fmt.Sprintf("%s:%d", item.Key, item.Item.Version))
-		}
-		etag := hex.EncodeToString(hash.Sum(nil))[:15]
-		span.SetAttributes(
-			tracing.FlagCountKey.Int(flagCount),
-			tracing.PayloadBytesKey.Int(len(respData)),
-		)
-		return respData, etag
-	}()
+	_, span := tr.Start(ctx, tracing.SpanSerializePayload)
+	defer span.End()
 
-	traceWriteResponse(tr, req, func() (int, error) {
-		return writeCacheableJSONResponse(w, req, clientCtx.Env, respData, etag)
-	})
+	respData, flagCount := serializeFlagsAsMap(data)
+	// Compute an overall Etag for the data set by hashing flag keys and versions
+	hash := sha1.New()                                                         //nolint:gosec // just used for insecure hashing
+	sort.Slice(data, func(i, j int) bool { return data[i].Key < data[j].Key }) // makes the hash deterministic
+	for _, item := range data {
+		_, _ = io.WriteString(hash, fmt.Sprintf("%s:%d", item.Key, item.Item.Version))
+	}
+	etag := hex.EncodeToString(hash.Sum(nil))[:15]
+	span.SetAttributes(
+		tracing.FlagCountKey.Int(flagCount),
+		tracing.PayloadBytesKey.Int(len(respData)),
+	)
+	return pollResult{data: respData, etag: etag}, nil
 }
 
 // PHP SDK polling endpoint for a flag: app.ld.com/sdk/flags/{key}
@@ -868,7 +936,10 @@ func writeCacheableJSONResponse(w http.ResponseWriter, req *http.Request, client
 	w.Header().Set("Etag", etag)
 	w.WriteHeader(http.StatusOK)
 
-	_, err := w.Write(bytes)
+	// Not XSS: every caller passes the output of a JSON encoder, and the response is served as
+	// application/json. The taint analysis only sees that a request parameter (the polling basis)
+	// influenced which payload was built.
+	_, err := w.Write(bytes) //nolint:gosec
 	return http.StatusOK, err
 }
 
diff --git a/relay/relay_endpoints_singleflight_test.go b/relay/relay_endpoints_singleflight_test.go
new file mode 100644
index 00000000..f0a65fae
--- /dev/null
+++ b/relay/relay_endpoints_singleflight_test.go
@@ -0,0 +1,198 @@
+package relay
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/launchdarkly/ld-relay/v9/internal/relayenv"
+	"github.com/launchdarkly/ld-relay/v9/internal/sharedtest/testclient"
+	"github.com/launchdarkly/ld-relay/v9/internal/sharedtest/testenv"
+	"github.com/launchdarkly/ld-relay/v9/internal/tracing"
+
+	"github.com/launchdarkly/go-server-sdk/v7/subsystems/ldstoreimpl"
+	"github.com/launchdarkly/go-server-sdk/v7/subsystems/ldstoretypes"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// doTracedPollRequest runs one pre-routed request against handler on the calling goroutine,
+// wrapped in its own root span the way otelmux wraps a routed request, and returns the response.
+func doTracedPollRequest(env relayenv.EnvContext, handler http.HandlerFunc, rawQuery string) *httptest.ResponseRecorder {
+	req := buildPreRoutedRequest("GET", nil, make(http.Header), nil, env)
+	req.URL.RawQuery = rawQuery
+	ctx, span := tracing.Tracer().Start(req.Context(), "test.request")
+	req = req.WithContext(ctx)
+	w := httptest.NewRecorder()
+	handler(w, req)
+	span.End()
+	return w
+}
+
+func storeWithOneFlag() *testclient.FakeStore {
+	return testclient.NewFakeStore([]ldstoretypes.Collection{
+		{Kind: ldstoreimpl.Features(), Items: []ldstoretypes.KeyedItemDescriptor{liveFlag("flag-one")}},
+		{Kind: ldstoreimpl.Segments(), Items: []ldstoretypes.KeyedItemDescriptor{}},
+	})
+}
+
+// blockFirstRead installs hook-driven gating on a store read: the first read signals `entered` and
+// then blocks until `release` is closed; later reads pass straight through. The returned counter
+// reports how many reads happened. It is installed after the environment is built so that only the
+// handlers' reads are observed.
+func blockFirstRead(setHook func(func()), entered chan<- struct{}, release <-chan struct{}) *atomic.Int32 {
+	var reads atomic.Int32
+	setHook(func() {
+		if reads.Add(1) == 1 {
+			close(entered)
+			<-release
+		}
+	})
+	return &reads
+}
+
+// TestPollingHandlersShareOnePayloadBuildAcrossConcurrentRequests blocks one request inside the
+// store read and then issues more identical requests. Because the first request cannot leave the
+// flight group until the test releases it, the followers join its flight: the store is read and
+// the payload serialized exactly once, and every request receives the same response.
+func TestPollingHandlersShareOnePayloadBuildAcrossConcurrentRequests(t *testing.T) {
+	cases := []struct {
+		name          string
+		handler       http.HandlerFunc
+		setReadHook   func(*testclient.FakeStore, func())
+		storeSpanName string
+	}{
+		{
+			name:          "pollHandlerV2 GET /sdk/poll",
+			handler:       pollHandlerV2,
+			setReadHook:   func(s *testclient.FakeStore, hook func()) { s.SnapshotHook = hook },
+			storeSpanName: tracing.SpanStoreSnapshot,
+		},
+		{
+			name:          "pollAllFlagsHandler GET /sdk/flags",
+			handler:       pollAllFlagsHandler,
+			setReadHook:   func(s *testclient.FakeStore, hook func()) { s.GetAllHook = hook },
+			storeSpanName: tracing.SpanStoreGetAll,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			recorder := installSpanRecorder(t)
+
+			store := storeWithOneFlag()
+			env := envWithStore(store)
+
+			entered := make(chan struct{})
+			release := make(chan struct{})
+			reads := blockFirstRead(func(hook func()) { tc.setReadHook(store, hook) }, entered, release)
+
+			const followers = 3
+			responses := make(chan *httptest.ResponseRecorder, followers+1)
+
+			// The leader enters the store read and blocks there, holding its flight open.
+			go func() { responses <- doTracedPollRequest(env, tc.handler, "") }()
+			<-entered
+
+			// The followers reach the flight group while the leader's flight is still open, so
+			// they join it instead of reading the store themselves. The sleep is generous time
+			// for them to get from goroutine start to the flight-group call (a short pure-CPU
+			// path); a follower that somehow had not arrived before the release would read the
+			// store itself and loudly fail the read-count assertion below.
+			for range followers {
+				go func() { responses <- doTracedPollRequest(env, tc.handler, "") }()
+			}
+			time.Sleep(100 * time.Millisecond)
+			close(release)
+
+			var bodies, etags []string
+			for range followers + 1 {
+				w := <-responses
+				require.Equal(t, http.StatusOK, w.Code)
+				bodies = append(bodies, w.Body.String())
+				etags = append(etags, w.Header().Get("Etag"))
+			}
+			for i := 1; i <= followers; i++ {
+				assert.Equal(t, bodies[0], bodies[i], "every request should receive the same payload")
+				assert.Equal(t, etags[0], etags[i], "every request should receive the same Etag")
+			}
+			assert.Equal(t, int32(1), reads.Load(), "the store should be read once, not once per request")
+
+			spans := recorder.Ended()
+			assert.Len(t, spansNamed(spans, tc.storeSpanName), 1)
+			assert.Len(t, spansNamed(spans, tracing.SpanSerializePayload), 1,
+				"one flight builds one payload, so there should be exactly one serialize span")
+			assert.Len(t, spansNamed(spans, tracing.SpanWriteResponse), followers+1,
+				"every request should still write its own response")
+
+			// Every request's own span reports that the payload build was shared.
+			roots := spansNamed(spans, "test.request")
+			require.Len(t, roots, followers+1)
+			for _, root := range roots {
+				shared, ok := spanAttrs(root)[tracing.SingleflightSharedKey]
+				require.True(t, ok, "the request span should report whether the payload build was shared")
+				assert.True(t, shared.AsBool())
+			}
+		})
+	}
+}
+
+// TestPollHandlerV2DoesNotShareAcrossDifferentBasisValues proves the basis is part of the flight
+// key: while a no-basis request is still blocked inside its flight, a request whose basis matches
+// the store's selector state completes on its own with the small "up-to-date" payload -- which is
+// only possible if it did not join the blocked flight.
+func TestPollHandlerV2DoesNotShareAcrossDifferentBasisValues(t *testing.T) {
+	store := storeWithOneFlag()
+	env := envWithStore(store)
+
+	entered := make(chan struct{})
+	release := make(chan struct{})
+	reads := blockFirstRead(func(hook func()) { store.SnapshotHook = hook }, entered, release)
+
+	leaderResp := make(chan *httptest.ResponseRecorder, 1)
+	go func() { leaderResp <- doTracedPollRequest(env, pollHandlerV2, "") }()
+	<-entered
+
+	// NewFakeStore's selector state is "initial-state", so this request is up to date.
+	upToDate := doTracedPollRequest(env, pollHandlerV2, "basis=initial-state")
+	require.Equal(t, http.StatusOK, upToDate.Code)
+	assert.Equal(t, int32(2), reads.Load(), "a request with a different basis must run a flight of its own")
+	assert.Contains(t, upToDate.Body.String(), "up-to-date")
+	assert.NotContains(t, upToDate.Body.String(), "put-object")
+
+	close(release)
+	leader := <-leaderResp
+	require.Equal(t, http.StatusOK, leader.Code)
+	assert.Contains(t, leader.Body.String(), "cant-catchup")
+	assert.Contains(t, leader.Body.String(), "put-object")
+}
+
+// TestPollingSingleflightAttributeIsFalseForALoneRequest checks the other side of the shared
+// attribute: a request that shares its flight with nobody reports shared=false.
+func TestPollingSingleflightAttributeIsFalseForALoneRequest(t *testing.T) {
+	recorder := installSpanRecorder(t)
+	env := testenv.MakeTestContextWithData()
+
+	for _, tc := range []struct {
+		name    string
+		handler http.HandlerFunc
+	}{
+		{"pollHandlerV2", pollHandlerV2},
+		{"pollAllFlagsHandler", pollAllFlagsHandler},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			recorder.Reset()
+
+			w := doTracedPollRequest(env, tc.handler, "")
+			require.Equal(t, http.StatusOK, w.Code)
+
+			root := requireSpan(t, recorder.Ended(), "test.request")
+			shared, ok := spanAttrs(root)[tracing.SingleflightSharedKey]
+			require.True(t, ok, "the request span should always report whether the payload build was shared")
+			assert.False(t, shared.AsBool())
+		})
+	}
+}

From 1bc37bbfbfa1d05096391c426794f166330317d7 Mon Sep 17 00:00:00 2001
From: Matthew Keeler <mkeeler@launchdarkly.com>
Date: Thu, 6 Aug 2026 08:23:20 -0400
Subject: [PATCH 2/5] feat: Annotate waiting polling requests with their
 flight-group wait time

A request that received its payload from a flight another request was
executing now records how long it waited as relay.singleflight.wait_ms
on its request span. The executing request carries no wait attribute --
it did not wait, and its time is visible as the store and serialize
child spans -- so the attribute's presence alone identifies a request
that waited, and slow waits are queryable.

Both polling handlers now resolve their flights through a shared
runPollingFlight helper that owns the Do call and both span
annotations. That indirection also breaks the taint chain behind the
gosec G705 false positive on writeCacheableJSONResponse, so its nolint
directive is removed.
---
 internal/tracing/attributes.go             |  7 +++
 relay/relay_endpoints.go                   | 70 ++++++++++++++--------
 relay/relay_endpoints_singleflight_test.go | 37 +++++++++---
 3 files changed, 81 insertions(+), 33 deletions(-)

diff --git a/internal/tracing/attributes.go b/internal/tracing/attributes.go
index 0fd64489..6bb34b2b 100644
--- a/internal/tracing/attributes.go
+++ b/internal/tracing/attributes.go
@@ -40,4 +40,11 @@ const (
 	// polling flight group. When it is true and the request's trace has no store or serialize
 	// child spans, another request's trace carries them.
 	SingleflightSharedKey = attribute.Key("relay.singleflight.shared")
+
+	// SingleflightWaitMSKey reports, on the request span of a polling request that received
+	// its payload from a flight another request was already executing, how many milliseconds
+	// it spent waiting for that flight. It is absent from the request that executed the build:
+	// that request did not wait, and its time is carried by the store and serialize child
+	// spans.
+	SingleflightWaitMSKey = attribute.Key("relay.singleflight.wait_ms")
 )
diff --git a/relay/relay_endpoints.go b/relay/relay_endpoints.go
index 175363bc..2b1cf065 100644
--- a/relay/relay_endpoints.go
+++ b/relay/relay_endpoints.go
@@ -204,6 +204,44 @@ const (
 	serverSideAllFlagsFlightKey = "sdk-flags"
 )
 
+// runPollingFlight resolves one polling request through the environment's flight group and
+// annotates the request's span with how that went: relay.singleflight.shared reports whether the
+// result was handed to multiple requests, and relay.singleflight.wait_ms reports, only on a
+// request that waited for a flight another request was already executing, how long it waited.
+// The executing request carries no wait attribute; it did not wait, and its time is visible as
+// the child spans that build starts on its context.
+func runPollingFlight(
+	req *http.Request,
+	env relayenv.EnvContext,
+	key string,
+	build func() (pollResult, error),
+) (pollResult, error) {
+	executed := false
+	start := time.Now()
+	data, err, shared := env.GetPollingFlightGroup().Do(key, func() (any, error) {
+		executed = true
+		result, err := build()
+		if err != nil {
+			return nil, err
+		}
+		return result, nil
+	})
+
+	span := trace.SpanFromContext(req.Context())
+	span.SetAttributes(tracing.SingleflightSharedKey.Bool(shared))
+	if !executed {
+		span.SetAttributes(tracing.SingleflightWaitMSKey.Float64(
+			float64(time.Since(start)) / float64(time.Millisecond)))
+	}
+
+	if err != nil {
+		return pollResult{}, err
+	}
+
+	// panic if it's not a pollResult - as this should be impossible
+	return data.(pollResult), nil
+}
+
 // Server-side SDK polling endpoint: app.ld.com/sdk/poll/
 func pollHandlerV2(w http.ResponseWriter, req *http.Request) {
 	clientCtx := middleware.GetEnvContextInfo(req.Context())
@@ -218,23 +256,16 @@ func pollHandlerV2(w http.ResponseWriter, req *http.Request) {
 	// with the same basis may share a result, and the basis is part of the key.
 	//
 	// The snapshot and serialize spans belong to the one request that executes the closure;
-	// every request records on its own request span whether the payload build was shared.
-	data, err, shared := clientCtx.Env.GetPollingFlightGroup().Do(serverSidePollFlightKey+":"+basis, func() (any, error) {
-		result, err := buildServerSidePollPayload(req.Context(), tr, clientCtx.Env, basis)
-		if err != nil {
-			return nil, err
-		}
-		return result, nil
+	// every request records on its own request span whether the payload build was shared, and
+	// how long it waited if another request built it.
+	result, err := runPollingFlight(req, clientCtx.Env, serverSidePollFlightKey+":"+basis, func() (pollResult, error) {
+		return buildServerSidePollPayload(req.Context(), tr, clientCtx.Env, basis)
 	})
-	trace.SpanFromContext(req.Context()).SetAttributes(tracing.SingleflightSharedKey.Bool(shared))
 	if err != nil {
 		w.WriteHeader(http.StatusInternalServerError)
 		return
 	}
 
-	// panic if it's not a pollResult - as this should be impossible
-	result := data.(pollResult)
-
 	traceWriteResponse(tr, req, func() (int, error) {
 		return writeCacheableJSONResponse(w, req, clientCtx.Env, result.data, result.etag)
 	})
@@ -556,22 +587,14 @@ func pollAllFlagsHandler(w http.ResponseWriter, req *http.Request) {
 	// payload and Etag depend only on the store contents, so a single key covers all callers.
 	// Span placement follows pollHandlerV2: the store and serialize spans belong to the one
 	// request that executes the closure.
-	data, err, shared := clientCtx.Env.GetPollingFlightGroup().Do(serverSideAllFlagsFlightKey, func() (any, error) {
-		result, err := buildAllFlagsPayload(req.Context(), tr, clientCtx.Env)
-		if err != nil {
-			return nil, err
-		}
-		return result, nil
+	result, err := runPollingFlight(req, clientCtx.Env, serverSideAllFlagsFlightKey, func() (pollResult, error) {
+		return buildAllFlagsPayload(req.Context(), tr, clientCtx.Env)
 	})
-	trace.SpanFromContext(req.Context()).SetAttributes(tracing.SingleflightSharedKey.Bool(shared))
 	if err != nil {
 		w.WriteHeader(http.StatusInternalServerError)
 		return
 	}
 
-	// panic if it's not a pollResult - as this should be impossible
-	result := data.(pollResult)
-
 	traceWriteResponse(tr, req, func() (int, error) {
 		return writeCacheableJSONResponse(w, req, clientCtx.Env, result.data, result.etag)
 	})
@@ -936,10 +959,7 @@ func writeCacheableJSONResponse(w http.ResponseWriter, req *http.Request, client
 	w.Header().Set("Etag", etag)
 	w.WriteHeader(http.StatusOK)
 
-	// Not XSS: every caller passes the output of a JSON encoder, and the response is served as
-	// application/json. The taint analysis only sees that a request parameter (the polling basis)
-	// influenced which payload was built.
-	_, err := w.Write(bytes) //nolint:gosec
+	_, err := w.Write(bytes)
 	return http.StatusOK, err
 }
 
diff --git a/relay/relay_endpoints_singleflight_test.go b/relay/relay_endpoints_singleflight_test.go
index f0a65fae..94e6da85 100644
--- a/relay/relay_endpoints_singleflight_test.go
+++ b/relay/relay_endpoints_singleflight_test.go
@@ -123,19 +123,35 @@ func TestPollingHandlersShareOnePayloadBuildAcrossConcurrentRequests(t *testing.
 
 			spans := recorder.Ended()
 			assert.Len(t, spansNamed(spans, tc.storeSpanName), 1)
-			assert.Len(t, spansNamed(spans, tracing.SpanSerializePayload), 1,
-				"one flight builds one payload, so there should be exactly one serialize span")
+			serialize := requireSpan(t, spans, tracing.SpanSerializePayload)
 			assert.Len(t, spansNamed(spans, tracing.SpanWriteResponse), followers+1,
 				"every request should still write its own response")
 
-			// Every request's own span reports that the payload build was shared.
+			// Every request's own span reports that the payload build was shared. The one
+			// request that executed the build -- the one whose trace carries the serialize
+			// span -- did not wait and must record no wait time; every follower must record
+			// how long it waited.
 			roots := spansNamed(spans, "test.request")
 			require.Len(t, roots, followers+1)
+			executedRoots, waitedRoots := 0, 0
 			for _, root := range roots {
-				shared, ok := spanAttrs(root)[tracing.SingleflightSharedKey]
+				attrs := spanAttrs(root)
+				shared, ok := attrs[tracing.SingleflightSharedKey]
 				require.True(t, ok, "the request span should report whether the payload build was shared")
 				assert.True(t, shared.AsBool())
+
+				wait, waited := attrs[tracing.SingleflightWaitMSKey]
+				if root.SpanContext().TraceID() == serialize.SpanContext().TraceID() {
+					executedRoots++
+					assert.False(t, waited, "the request that built the payload did not wait")
+				} else {
+					waitedRoots++
+					require.True(t, waited, "a request that shared another's payload build should record its wait")
+					assert.Positive(t, wait.AsFloat64())
+				}
 			}
+			assert.Equal(t, 1, executedRoots)
+			assert.Equal(t, followers, waitedRoots)
 		})
 	}
 }
@@ -170,9 +186,10 @@ func TestPollHandlerV2DoesNotShareAcrossDifferentBasisValues(t *testing.T) {
 	assert.Contains(t, leader.Body.String(), "put-object")
 }
 
-// TestPollingSingleflightAttributeIsFalseForALoneRequest checks the other side of the shared
-// attribute: a request that shares its flight with nobody reports shared=false.
-func TestPollingSingleflightAttributeIsFalseForALoneRequest(t *testing.T) {
+// TestPollingSingleflightAttributesForALoneRequest checks the other side of the flight-group
+// annotations: a request that shares its flight with nobody reports shared=false, and records no
+// wait time because it executed the build itself.
+func TestPollingSingleflightAttributesForALoneRequest(t *testing.T) {
 	recorder := installSpanRecorder(t)
 	env := testenv.MakeTestContextWithData()
 
@@ -190,9 +207,13 @@ func TestPollingSingleflightAttributeIsFalseForALoneRequest(t *testing.T) {
 			require.Equal(t, http.StatusOK, w.Code)
 
 			root := requireSpan(t, recorder.Ended(), "test.request")
-			shared, ok := spanAttrs(root)[tracing.SingleflightSharedKey]
+			attrs := spanAttrs(root)
+			shared, ok := attrs[tracing.SingleflightSharedKey]
 			require.True(t, ok, "the request span should always report whether the payload build was shared")
 			assert.False(t, shared.AsBool())
+
+			_, waited := attrs[tracing.SingleflightWaitMSKey]
+			assert.False(t, waited, "a request that built its own payload should record no wait time")
 		})
 	}
 }

From 898cca960e09e33c7eb814dfed59777be6e607c8 Mon Sep 17 00:00:00 2001
From: Matthew Keeler <mkeeler@launchdarkly.com>
Date: Thu, 6 Aug 2026 09:42:31 -0400
Subject: [PATCH 3/5] feat: Annotate streaming replay flights with the
 singleflight telemetry

The streaming repositories' replay flight groups now record the same
information the polling endpoints do: relay.singleflight.shared on the
subscribing request's span, plus relay.singleflight.wait_ms when a
replay waited on a flight another subscriber was already executing.
The annotation logic moves into tracing.SingleflightDo so both sides
record identical information by construction.

The flags-only repository never implemented ReplayWithContext, so it
had no request context to annotate; it now advertises context support
the same way the main server-side repository does.
---
 .../streams/stream_flight_telemetry_test.go   |  73 +++++++++++
 .../streams/stream_provider_server_side.go    |  17 ++-
 .../stream_provider_server_side_flags.go      |  26 +++-
 internal/tracing/attributes.go                |  17 ++-
 internal/tracing/singleflight.go              |  38 ++++++
 internal/tracing/singleflight_test.go         | 119 ++++++++++++++++++
 relay/relay_endpoints.go                      |  26 ++--
 7 files changed, 279 insertions(+), 37 deletions(-)
 create mode 100644 internal/streams/stream_flight_telemetry_test.go
 create mode 100644 internal/tracing/singleflight.go
 create mode 100644 internal/tracing/singleflight_test.go

diff --git a/internal/streams/stream_flight_telemetry_test.go b/internal/streams/stream_flight_telemetry_test.go
new file mode 100644
index 00000000..e1da9cfa
--- /dev/null
+++ b/internal/streams/stream_flight_telemetry_test.go
@@ -0,0 +1,73 @@
+package streams
+
+import (
+	"context"
+	"log/slog"
+	"testing"
+	"time"
+
+	"github.com/launchdarkly/ld-relay/v9/internal/tracing"
+
+	"github.com/launchdarkly/eventsource"
+	"github.com/launchdarkly/go-server-sdk-evaluation/v3/ldmodel"
+
+	helpers "github.com/launchdarkly/go-test-helpers/v3"
+	"go.opentelemetry.io/otel/attribute"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestReplayAnnotatesSubscriberSpanWithFlightTelemetry checks that every repository whose replay
+// goes through a flight group reports the flight-group telemetry (relay.singleflight.shared, and
+// relay.singleflight.wait_ms when a replay waits on another's) on the subscribing request's span,
+// matching what the polling endpoints record. The shared/waiting semantics themselves are covered
+// by the tracing package's SingleflightDo tests; these subtests prove each replay call site hands
+// the subscriber's context through.
+func TestReplayAnnotatesSubscriberSpanWithFlightTelemetry(t *testing.T) {
+	store := makeMockStore([]ldmodel.FeatureFlag{testFlag1}, []ldmodel.Segment{testSegment1})
+
+	repos := []struct {
+		name string
+		repo eventsource.RepositoryWithContext
+	}{
+		{"server-side v1", &serverSideEnvStreamRepository{store: store, logger: slog.Default()}},
+		{"server-side v2", &serverSideEnvStreamRepository{store: store, logger: slog.Default(), isV2: true}},
+		{"server-side flags only", &serverSideFlagsOnlyEnvStreamRepository{store: store, logger: slog.Default()}},
+	}
+
+	for _, tc := range repos {
+		t.Run(tc.name, func(t *testing.T) {
+			recorder := tracetest.NewSpanRecorder()
+			provider := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder))
+			t.Cleanup(func() { _ = provider.Shutdown(context.Background()) })
+
+			ctx, span := provider.Tracer("test").Start(context.Background(), "test.request")
+			eventCh := tc.repo.ReplayWithContext(ctx, "", "")
+			for {
+				if _, ok, closed := helpers.TryReceive(eventCh, time.Second); closed {
+					break
+				} else if !ok {
+					require.Fail(t, "timed out waiting for replayed event (channel was not closed)")
+				}
+			}
+			span.End()
+
+			ended := recorder.Ended()
+			require.Len(t, ended, 1)
+			attrs := make(map[attribute.Key]attribute.Value)
+			for _, kv := range ended[0].Attributes() {
+				attrs[kv.Key] = kv.Value
+			}
+
+			shared, ok := attrs[tracing.SingleflightSharedKey]
+			require.True(t, ok, "the subscriber's span should report whether the replay build was shared")
+			assert.False(t, shared.AsBool(), "a lone replay shares with nobody")
+
+			_, waited := attrs[tracing.SingleflightWaitMSKey]
+			assert.False(t, waited, "a lone replay builds its own payload, so it should record no wait")
+		})
+	}
+}
diff --git a/internal/streams/stream_provider_server_side.go b/internal/streams/stream_provider_server_side.go
index c8639813..17179828 100644
--- a/internal/streams/stream_provider_server_side.go
+++ b/internal/streams/stream_provider_server_side.go
@@ -8,6 +8,7 @@ import (
 
 	"github.com/launchdarkly/go-jsonstream/v3/jwriter"
 	"github.com/launchdarkly/ld-relay/v9/internal/sdkauth"
+	"github.com/launchdarkly/ld-relay/v9/internal/tracing"
 
 	"github.com/launchdarkly/ld-relay/v9/config"
 	"golang.org/x/sync/singleflight"
@@ -208,9 +209,9 @@ func (r *serverSideEnvStreamRepository) replay(ctx context.Context, id string) c
 		if r.isV2 {
 			// See the note in HandlerV2 about how we use the Last-Event-ID header to
 			// pass the basis.
-			events, err = r.getReplayEventsV2(id)
+			events, err = r.getReplayEventsV2(ctx, id)
 		} else {
-			events, err = r.getReplayEventsV1()
+			events, err = r.getReplayEventsV1(ctx)
 		}
 
 		if err != nil {
@@ -231,8 +232,10 @@ func (r *serverSideEnvStreamRepository) replay(ctx context.Context, id string) c
 }
 
 // getReplayEvent will return a ServerSidePutEvent with all the data needed for a Replay.
-func (r *serverSideEnvStreamRepository) getReplayEventsV1() ([]eventsource.Event, error) {
-	data, err, _ := r.flightGroup.Do("getReplayEventV1", func() (interface{}, error) {
+// The context is only used for telemetry: the subscribing request's span is annotated with how
+// the flight resolved (refer to tracing.SingleflightDo).
+func (r *serverSideEnvStreamRepository) getReplayEventsV1(ctx context.Context) ([]eventsource.Event, error) {
+	data, err := tracing.SingleflightDo(ctx, &r.flightGroup, "getReplayEventV1", func() (interface{}, error) {
 		snapshot, _, err := r.store.Snapshot()
 		if err != nil {
 			r.logger.Error("error getting all flags", "error", err)
@@ -260,12 +263,14 @@ func (r *serverSideEnvStreamRepository) getReplayEventsV1() ([]eventsource.Event
 	return []eventsource.Event{event}, nil
 }
 
-func (r *serverSideEnvStreamRepository) getReplayEventsV2(basis string) ([]eventsource.Event, error) {
+// getReplayEventsV2 is getReplayEventsV1 for the FDv2 protocol; the context serves the same
+// telemetry-only purpose.
+func (r *serverSideEnvStreamRepository) getReplayEventsV2(ctx context.Context, basis string) ([]eventsource.Event, error) {
 	// The result depends on the caller's basis: a client whose basis matches the current
 	// selector state gets an "up-to-date" event, while any other client gets a full data
 	// transfer. Only requests with the same basis may share a result, so the basis must be
 	// part of the key.
-	data, err, _ := r.flightGroup.Do("getReplayEventV2:"+basis, func() (interface{}, error) {
+	data, err := tracing.SingleflightDo(ctx, &r.flightGroup, "getReplayEventV2:"+basis, func() (interface{}, error) {
 		snapshot, selector, err := r.store.Snapshot()
 		if err != nil {
 			r.logger.Error("error getting all flags", "error", err)
diff --git a/internal/streams/stream_provider_server_side_flags.go b/internal/streams/stream_provider_server_side_flags.go
index 2079473e..e7a8758e 100644
--- a/internal/streams/stream_provider_server_side_flags.go
+++ b/internal/streams/stream_provider_server_side_flags.go
@@ -1,11 +1,13 @@
 package streams
 
 import (
+	"context"
 	"log/slog"
 	"net/http"
 	"sync"
 
 	"github.com/launchdarkly/ld-relay/v9/internal/sdkauth"
+	"github.com/launchdarkly/ld-relay/v9/internal/tracing"
 
 	"github.com/launchdarkly/ld-relay/v9/config"
 	"golang.org/x/sync/singleflight"
@@ -132,7 +134,25 @@ func (e *serverSideFlagsOnlyEnvStreamProvider) Close() {
 	}
 }
 
+// Ensure the repository advertises context support so the eventsource server calls
+// ReplayWithContext (and thus hands over the subscribing request's context) rather than Replay.
+var _ eventsource.RepositoryWithContext = (*serverSideFlagsOnlyEnvStreamRepository)(nil)
+
+// Replay satisfies the eventsource.Repository interface. It delegates to replay with a background
+// context; in practice the eventsource server prefers ReplayWithContext (below) whenever the
+// repository implements it, so this context-less path is only a fallback.
 func (r *serverSideFlagsOnlyEnvStreamRepository) Replay(channel, id string) chan eventsource.Event {
+	return r.replay(context.Background())
+}
+
+// ReplayWithContext satisfies the eventsource.RepositoryWithContext interface. The context is
+// only used for telemetry: the subscribing request's span is annotated with how the replay's
+// flight resolved (refer to tracing.SingleflightDo).
+func (r *serverSideFlagsOnlyEnvStreamRepository) ReplayWithContext(ctx context.Context, channel, id string) <-chan eventsource.Event {
+	return r.replay(ctx)
+}
+
+func (r *serverSideFlagsOnlyEnvStreamRepository) replay(ctx context.Context) chan eventsource.Event {
 	out := make(chan eventsource.Event)
 	if !r.store.IsInitialized() { // See serverSideEnvStreamRepository.Replay
 		close(out)
@@ -140,7 +160,7 @@ func (r *serverSideFlagsOnlyEnvStreamRepository) Replay(channel, id string) chan
 	}
 	go func() {
 		defer close(out)
-		event, err := r.getReplayEvent()
+		event, err := r.getReplayEvent(ctx)
 		if err == nil && event != nil {
 			out <- event
 		}
@@ -148,8 +168,8 @@ func (r *serverSideFlagsOnlyEnvStreamRepository) Replay(channel, id string) chan
 	return out
 }
 
-func (r *serverSideFlagsOnlyEnvStreamRepository) getReplayEvent() (eventsource.Event, error) {
-	data, err, _ := r.flightGroup.Do("getReplayEvent", func() (interface{}, error) {
+func (r *serverSideFlagsOnlyEnvStreamRepository) getReplayEvent(ctx context.Context) (eventsource.Event, error) {
+	data, err := tracing.SingleflightDo(ctx, &r.flightGroup, "getReplayEvent", func() (interface{}, error) {
 		if !r.store.IsInitialized() {
 			return nil, nil
 		}
diff --git a/internal/tracing/attributes.go b/internal/tracing/attributes.go
index 6bb34b2b..0a7dca17 100644
--- a/internal/tracing/attributes.go
+++ b/internal/tracing/attributes.go
@@ -35,16 +35,15 @@ const (
 	PayloadEventsKey = attribute.Key("relay.payload.events")
 	PayloadBytesKey  = attribute.Key("relay.payload.bytes")
 
-	// SingleflightSharedKey reports, on a polling endpoint's request span, whether the
-	// response payload build was shared with concurrent requests through the environment's
-	// polling flight group. When it is true and the request's trace has no store or serialize
-	// child spans, another request's trace carries them.
+	// SingleflightSharedKey reports, on the request span of a polling request or an SSE
+	// replay, whether the payload build was shared with concurrent requests through a flight
+	// group. When it is true and the request's trace shows no sign of the build itself,
+	// another request's trace carries it.
 	SingleflightSharedKey = attribute.Key("relay.singleflight.shared")
 
-	// SingleflightWaitMSKey reports, on the request span of a polling request that received
-	// its payload from a flight another request was already executing, how many milliseconds
-	// it spent waiting for that flight. It is absent from the request that executed the build:
-	// that request did not wait, and its time is carried by the store and serialize child
-	// spans.
+	// SingleflightWaitMSKey reports, on the request span of a request that received its
+	// payload from a flight another request was already executing, how many milliseconds it
+	// spent waiting for that flight. It is absent from the request that executed the build:
+	// that request did not wait.
 	SingleflightWaitMSKey = attribute.Key("relay.singleflight.wait_ms")
 )
diff --git a/internal/tracing/singleflight.go b/internal/tracing/singleflight.go
new file mode 100644
index 00000000..289df207
--- /dev/null
+++ b/internal/tracing/singleflight.go
@@ -0,0 +1,38 @@
+package tracing
+
+import (
+	"context"
+	"time"
+
+	"go.opentelemetry.io/otel/trace"
+	"golang.org/x/sync/singleflight"
+)
+
+// SingleflightDo runs fn through group under key and annotates the span in ctx with how the
+// flight resolved: SingleflightSharedKey reports whether the result was handed to multiple
+// callers, and SingleflightWaitMSKey records, only on a caller that waited for a flight another
+// caller was already executing, how long it waited. The caller that executed fn records no wait
+// time: it did not wait, and its work is expected to show up as the child spans fn starts on
+// that caller's context.
+func SingleflightDo(
+	ctx context.Context,
+	group *singleflight.Group,
+	key string,
+	fn func() (any, error),
+) (any, error) {
+	executed := false
+	start := time.Now()
+	data, err, shared := group.Do(key, func() (any, error) {
+		executed = true
+		return fn()
+	})
+
+	span := trace.SpanFromContext(ctx)
+	span.SetAttributes(SingleflightSharedKey.Bool(shared))
+	if !executed {
+		span.SetAttributes(SingleflightWaitMSKey.Float64(
+			float64(time.Since(start)) / float64(time.Millisecond)))
+	}
+
+	return data, err
+}
diff --git a/internal/tracing/singleflight_test.go b/internal/tracing/singleflight_test.go
new file mode 100644
index 00000000..a8afc77a
--- /dev/null
+++ b/internal/tracing/singleflight_test.go
@@ -0,0 +1,119 @@
+package tracing
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"go.opentelemetry.io/otel/attribute"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
+	"golang.org/x/sync/singleflight"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func attrsOf(s sdktrace.ReadOnlySpan) map[attribute.Key]attribute.Value {
+	m := make(map[attribute.Key]attribute.Value)
+	for _, kv := range s.Attributes() {
+		m[kv.Key] = kv.Value
+	}
+	return m
+}
+
+func TestSingleflightDoAnnotatesALoneCaller(t *testing.T) {
+	recorder := tracetest.NewSpanRecorder()
+	provider := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder))
+	t.Cleanup(func() { _ = provider.Shutdown(context.Background()) })
+
+	var group singleflight.Group
+	ctx, span := provider.Tracer("test").Start(context.Background(), "caller")
+	data, err := SingleflightDo(ctx, &group, "key", func() (any, error) { return "result", nil })
+	span.End()
+
+	require.NoError(t, err)
+	assert.Equal(t, "result", data)
+
+	ended := recorder.Ended()
+	require.Len(t, ended, 1)
+	attrs := attrsOf(ended[0])
+
+	shared, ok := attrs[SingleflightSharedKey]
+	require.True(t, ok, "the span should always report whether the flight was shared")
+	assert.False(t, shared.AsBool())
+
+	_, waited := attrs[SingleflightWaitMSKey]
+	assert.False(t, waited, "the caller executed the function itself, so it should record no wait")
+}
+
+func TestSingleflightDoAnnotatesSharedCallers(t *testing.T) {
+	recorder := tracetest.NewSpanRecorder()
+	provider := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder))
+	t.Cleanup(func() { _ = provider.Shutdown(context.Background()) })
+	tracer := provider.Tracer("test")
+
+	var group singleflight.Group
+	entered := make(chan struct{})
+	release := make(chan struct{})
+
+	run := func(name string, fn func() (any, error)) chan error {
+		done := make(chan error, 1)
+		go func() {
+			ctx, span := tracer.Start(context.Background(), name)
+			data, err := SingleflightDo(ctx, &group, "key", fn)
+			span.End()
+			if err == nil && data != "result" {
+				err = assert.AnError
+			}
+			done <- err
+		}()
+		return done
+	}
+
+	// The leader enters the function and blocks there, holding its flight open.
+	leaderDone := run("leader", func() (any, error) {
+		close(entered)
+		<-release
+		return "result", nil
+	})
+	<-entered
+
+	// The follower joins the leader's flight as long as it calls Do before the leader's
+	// function returns, which cannot happen until release is closed. The sleep is generous
+	// time for the goroutine to get there; if it somehow arrived late, its own function would
+	// run and the error below would fail the test loudly.
+	followerDone := run("follower", func() (any, error) {
+		return nil, assert.AnError
+	})
+	time.Sleep(100 * time.Millisecond)
+	close(release)
+
+	require.NoError(t, <-leaderDone)
+	require.NoError(t, <-followerDone)
+
+	ended := recorder.Ended()
+	require.Len(t, ended, 2)
+	for _, s := range ended {
+		attrs := attrsOf(s)
+		shared, ok := attrs[SingleflightSharedKey]
+		require.True(t, ok, "the span should always report whether the flight was shared")
+		assert.True(t, shared.AsBool())
+
+		wait, waited := attrs[SingleflightWaitMSKey]
+		switch s.Name() {
+		case "leader":
+			assert.False(t, waited, "the leader executed the function, so it should record no wait")
+		case "follower":
+			require.True(t, waited, "the follower waited on the leader's flight, so it should record its wait")
+			assert.Positive(t, wait.AsFloat64())
+		}
+	}
+}
+
+func TestSingleflightDoToleratesASpanlessContext(t *testing.T) {
+	var group singleflight.Group
+	data, err := SingleflightDo(context.Background(), &group, "key", func() (any, error) { return "result", nil })
+	require.NoError(t, err)
+	assert.Equal(t, "result", data)
+}
diff --git a/relay/relay_endpoints.go b/relay/relay_endpoints.go
index 2b1cf065..66fd5adf 100644
--- a/relay/relay_endpoints.go
+++ b/relay/relay_endpoints.go
@@ -204,36 +204,21 @@ const (
 	serverSideAllFlagsFlightKey = "sdk-flags"
 )
 
-// runPollingFlight resolves one polling request through the environment's flight group and
-// annotates the request's span with how that went: relay.singleflight.shared reports whether the
-// result was handed to multiple requests, and relay.singleflight.wait_ms reports, only on a
-// request that waited for a flight another request was already executing, how long it waited.
-// The executing request carries no wait attribute; it did not wait, and its time is visible as
-// the child spans that build starts on its context.
+// runPollingFlight resolves one polling request through the environment's flight group,
+// annotating the request's span with how the flight resolved (refer to tracing.SingleflightDo).
 func runPollingFlight(
 	req *http.Request,
 	env relayenv.EnvContext,
 	key string,
 	build func() (pollResult, error),
 ) (pollResult, error) {
-	executed := false
-	start := time.Now()
-	data, err, shared := env.GetPollingFlightGroup().Do(key, func() (any, error) {
-		executed = true
+	data, err := tracing.SingleflightDo(req.Context(), env.GetPollingFlightGroup(), key, func() (any, error) {
 		result, err := build()
 		if err != nil {
 			return nil, err
 		}
 		return result, nil
 	})
-
-	span := trace.SpanFromContext(req.Context())
-	span.SetAttributes(tracing.SingleflightSharedKey.Bool(shared))
-	if !executed {
-		span.SetAttributes(tracing.SingleflightWaitMSKey.Float64(
-			float64(time.Since(start)) / float64(time.Millisecond)))
-	}
-
 	if err != nil {
 		return pollResult{}, err
 	}
@@ -959,7 +944,10 @@ func writeCacheableJSONResponse(w http.ResponseWriter, req *http.Request, client
 	w.Header().Set("Etag", etag)
 	w.WriteHeader(http.StatusOK)
 
-	_, err := w.Write(bytes)
+	// Not XSS: every caller passes the output of a JSON encoder, and the response is served as
+	// application/json. The taint analysis only sees that a request parameter (the polling basis)
+	// influenced which payload was built.
+	_, err := w.Write(bytes) //nolint:gosec
 	return http.StatusOK, err
 }
 

From dbdeacc64eea04c42d483e8115ecb242b66553cc Mon Sep 17 00:00:00 2001
From: Matthew Keeler <mkeeler@launchdarkly.com>
Date: Fri, 7 Aug 2026 10:20:50 -0400
Subject: [PATCH 4/5] fix: Honor subscriber cancellation in the flags-only
 replay

Review feedback: ReplayWithContext's comment claimed the context was
only used for telemetry, but cancellation-on-disconnect is the reason
the eventsource RepositoryWithContext interface exists. The flags-only
replay now behaves like the main server-side repository: it skips the
payload build if the subscriber is already gone and abandons the send
on disconnect instead of relying on the eventsource server draining
the channel on its behalf.

The flight functions' telemetry-only comments now explain why they do
not cancel: a flight in progress may be shared with other waiting
subscribers, so disconnect handling belongs to replay's send loop.
---
 .../streams/stream_provider_server_side.go    |  9 +++--
 .../stream_provider_server_side_flags.go      | 28 +++++++++++---
 .../stream_provider_server_side_flags_test.go | 37 +++++++++++++++++++
 3 files changed, 65 insertions(+), 9 deletions(-)

diff --git a/internal/streams/stream_provider_server_side.go b/internal/streams/stream_provider_server_side.go
index 17179828..35108a36 100644
--- a/internal/streams/stream_provider_server_side.go
+++ b/internal/streams/stream_provider_server_side.go
@@ -232,8 +232,11 @@ func (r *serverSideEnvStreamRepository) replay(ctx context.Context, id string) c
 }
 
 // getReplayEvent will return a ServerSidePutEvent with all the data needed for a Replay.
-// The context is only used for telemetry: the subscribing request's span is annotated with how
-// the flight resolved (refer to tracing.SingleflightDo).
+// Within this function the context is used only for telemetry -- the subscribing request's span
+// is annotated with how the flight resolved (refer to tracing.SingleflightDo). A flight in
+// progress is deliberately never abandoned on cancellation: its result may be shared with other
+// subscribers still waiting on it. Disconnect handling belongs to the caller, in replay's send
+// loop.
 func (r *serverSideEnvStreamRepository) getReplayEventsV1(ctx context.Context) ([]eventsource.Event, error) {
 	data, err := tracing.SingleflightDo(ctx, &r.flightGroup, "getReplayEventV1", func() (interface{}, error) {
 		snapshot, _, err := r.store.Snapshot()
@@ -264,7 +267,7 @@ func (r *serverSideEnvStreamRepository) getReplayEventsV1(ctx context.Context) (
 }
 
 // getReplayEventsV2 is getReplayEventsV1 for the FDv2 protocol; the context serves the same
-// telemetry-only purpose.
+// telemetry-only purpose there, with cancellation likewise left to the caller.
 func (r *serverSideEnvStreamRepository) getReplayEventsV2(ctx context.Context, basis string) ([]eventsource.Event, error) {
 	// The result depends on the caller's basis: a client whose basis matches the current
 	// selector state gets an "up-to-date" event, while any other client gets a full data
diff --git a/internal/streams/stream_provider_server_side_flags.go b/internal/streams/stream_provider_server_side_flags.go
index e7a8758e..92889a55 100644
--- a/internal/streams/stream_provider_server_side_flags.go
+++ b/internal/streams/stream_provider_server_side_flags.go
@@ -135,7 +135,7 @@ func (e *serverSideFlagsOnlyEnvStreamProvider) Close() {
 }
 
 // Ensure the repository advertises context support so the eventsource server calls
-// ReplayWithContext (and thus hands over the subscribing request's context) rather than Replay.
+// ReplayWithContext (and thus propagates the connection's lifetime) rather than Replay.
 var _ eventsource.RepositoryWithContext = (*serverSideFlagsOnlyEnvStreamRepository)(nil)
 
 // Replay satisfies the eventsource.Repository interface. It delegates to replay with a background
@@ -145,9 +145,11 @@ func (r *serverSideFlagsOnlyEnvStreamRepository) Replay(channel, id string) chan
 	return r.replay(context.Background())
 }
 
-// ReplayWithContext satisfies the eventsource.RepositoryWithContext interface. The context is
-// only used for telemetry: the subscribing request's span is annotated with how the replay's
-// flight resolved (refer to tracing.SingleflightDo).
+// ReplayWithContext satisfies the eventsource.RepositoryWithContext interface. The eventsource
+// server passes the subscribing request's context, which is cancelled when the SDK client
+// disconnects. This lets the producer goroutine below stop immediately on disconnect instead of
+// blocking on a send whose reader has gone away; the same context also carries the request span
+// that the replay's flight-group telemetry is recorded on (refer to tracing.SingleflightDo).
 func (r *serverSideFlagsOnlyEnvStreamRepository) ReplayWithContext(ctx context.Context, channel, id string) <-chan eventsource.Event {
 	return r.replay(ctx)
 }
@@ -160,9 +162,23 @@ func (r *serverSideFlagsOnlyEnvStreamRepository) replay(ctx context.Context) cha
 	}
 	go func() {
 		defer close(out)
+		select {
+		case <-ctx.Done():
+			// The subscriber already disconnected; don't bother building a payload nobody will read.
+			r.logger.Info("subscriber disconnected before replay started; skipping replay")
+			return
+		default:
+		}
 		event, err := r.getReplayEvent(ctx)
-		if err == nil && event != nil {
-			out <- event
+		if err != nil || event == nil {
+			return
+		}
+		select {
+		case out <- event:
+		case <-ctx.Done():
+			// The subscriber disconnected before consuming the replay; stop producing so this
+			// goroutine and its payload are released promptly instead of leaking.
+			r.logger.Info("subscriber disconnected mid-replay; stopping replay")
 		}
 	}()
 	return out
diff --git a/internal/streams/stream_provider_server_side_flags_test.go b/internal/streams/stream_provider_server_side_flags_test.go
index 22895656..6bebcdea 100644
--- a/internal/streams/stream_provider_server_side_flags_test.go
+++ b/internal/streams/stream_provider_server_side_flags_test.go
@@ -1,6 +1,7 @@
 package streams
 
 import (
+	"context"
 	"log/slog"
 	"testing"
 	"time"
@@ -15,6 +16,7 @@ import (
 	"github.com/launchdarkly/go-server-sdk/v7/subsystems/ldstoreimpl"
 	"github.com/launchdarkly/go-server-sdk/v7/subsystems/ldstoretypes"
 
+	helpers "github.com/launchdarkly/go-test-helpers/v3"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -237,3 +239,38 @@ func TestStreamProviderServerSideFlagsOnly(t *testing.T) {
 		})
 	})
 }
+
+func TestFlagsOnlyReplayWithContextStopsWhenSubscriberCancels(t *testing.T) {
+	// A subscriber that disconnects mid-replay cancels the request context. The producer must
+	// stop sending promptly instead of blocking forever on a send that nobody will receive.
+	// This mirrors the equivalent test for serverSideEnvStreamRepository.
+	snapshotReturned := make(chan struct{}, 1)
+	store := newMockStoreQueries()
+	store.setupSnapshotFn(func() (map[ldstoretypes.DataKind][]ldstoretypes.KeyedItemDescriptor, subsystems.Selector, error) {
+		defer func() { snapshotReturned <- struct{}{} }()
+		return map[ldstoretypes.DataKind][]ldstoretypes.KeyedItemDescriptor{
+			ldstoreimpl.Features(): {ldstoretypes.KeyedItemDescriptor{Key: testFlag1.Key, Item: sharedtest.FlagDesc(testFlag1)}},
+			ldstoreimpl.Segments(): {},
+		}, subsystems.NoSelector(), nil
+	})
+	repo := &serverSideFlagsOnlyEnvStreamRepository{store: store, logger: slog.Default()}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	eventCh := repo.ReplayWithContext(ctx, "", "")
+
+	// Wait until the producer has computed the event and is parked on its (unbuffered, unread)
+	// send, then cancel without ever consuming it.
+	<-snapshotReturned
+	time.Sleep(50 * time.Millisecond)
+	cancel()
+
+	// Let the producer observe cancellation while no receiver exists: its select then has only
+	// the ctx.Done case ready, so it must exit without delivering anything. Only then attach a
+	// receiver. Asserting closed-with-no-event on the first receive is what makes this test fail
+	// against a producer that ignores the context -- such a producer would be rescued by the
+	// receive, deliver its event, and only then close the channel.
+	time.Sleep(50 * time.Millisecond)
+	_, ok, closed := helpers.TryReceive(eventCh, time.Second)
+	require.False(t, ok, "producer delivered an event after cancellation")
+	require.True(t, closed, "producer did not stop after context cancellation (channel never closed)")
+}

From c6da085f7c45b533f8a96e634d676b98db201a7d Mon Sep 17 00:00:00 2001
From: Matthew Keeler <mkeeler@launchdarkly.com>
Date: Fri, 7 Aug 2026 10:55:30 -0400
Subject: [PATCH 5/5] feat: Show the flight-group wait as a span instead of a
 timeline gap

A request that waited on another request's payload build used to show
that wait as unexplained empty space between its child spans, with only
the relay.singleflight.wait_ms attribute to account for it. The wait is
now also emitted as a relay.singleflight.wait child span covering
exactly the waiting window, so trace timelines read directly.

The span is back-dated: whether a caller waited (rather than executed)
is only known once the flight resolves, so it cannot be opened
beforehand without giving the executing caller a bogus wait span. It
comes from the provider owning the surrounding request span, so it
lands wherever that span is recorded.
---
 internal/tracing/attributes.go             |  4 +++-
 internal/tracing/singleflight.go           | 16 ++++++++++---
 internal/tracing/singleflight_test.go      | 28 +++++++++++++++++++---
 relay/relay_endpoints_singleflight_test.go |  9 +++++++
 4 files changed, 50 insertions(+), 7 deletions(-)

diff --git a/internal/tracing/attributes.go b/internal/tracing/attributes.go
index 0a7dca17..2e1727d9 100644
--- a/internal/tracing/attributes.go
+++ b/internal/tracing/attributes.go
@@ -23,6 +23,7 @@ const (
 	SpanEventsDispatch   = "relay.events.dispatch"
 	SpanSerializePayload = "relay.payload.serialize"
 	SpanWriteResponse    = "relay.response.write"
+	SpanSingleflightWait = "relay.singleflight.wait"
 )
 
 // Relay-specific span attribute keys.
@@ -44,6 +45,7 @@ const (
 	// SingleflightWaitMSKey reports, on the request span of a request that received its
 	// payload from a flight another request was already executing, how many milliseconds it
 	// spent waiting for that flight. It is absent from the request that executed the build:
-	// that request did not wait.
+	// that request did not wait. The same window is also visible in the trace timeline as a
+	// SpanSingleflightWait child span.
 	SingleflightWaitMSKey = attribute.Key("relay.singleflight.wait_ms")
 )
diff --git a/internal/tracing/singleflight.go b/internal/tracing/singleflight.go
index 289df207..d5401f2b 100644
--- a/internal/tracing/singleflight.go
+++ b/internal/tracing/singleflight.go
@@ -11,8 +11,10 @@ import (
 // SingleflightDo runs fn through group under key and annotates the span in ctx with how the
 // flight resolved: SingleflightSharedKey reports whether the result was handed to multiple
 // callers, and SingleflightWaitMSKey records, only on a caller that waited for a flight another
-// caller was already executing, how long it waited. The caller that executed fn records no wait
-// time: it did not wait, and its work is expected to show up as the child spans fn starts on
+// caller was already executing, how long it waited. That wait is additionally emitted as a
+// SpanSingleflightWait child span covering exactly the waiting window, so a trace's timeline
+// shows a labeled bar instead of an unexplained gap. The caller that executed fn records
+// neither: it did not wait, and its work is expected to show up as the child spans fn starts on
 // that caller's context.
 func SingleflightDo(
 	ctx context.Context,
@@ -30,8 +32,16 @@ func SingleflightDo(
 	span := trace.SpanFromContext(ctx)
 	span.SetAttributes(SingleflightSharedKey.Bool(shared))
 	if !executed {
+		end := time.Now()
 		span.SetAttributes(SingleflightWaitMSKey.Float64(
-			float64(time.Since(start)) / float64(time.Millisecond)))
+			float64(end.Sub(start)) / float64(time.Millisecond)))
+		// The wait span is back-dated: whether this caller waited (rather than executed) is
+		// only known once Do returns, so it cannot be opened beforehand without also giving
+		// the executing caller a bogus wait span. It comes from the provider that owns the
+		// surrounding span, so it lands wherever that span is recorded.
+		tr := span.TracerProvider().Tracer(TracerName)
+		_, waitSpan := tr.Start(ctx, SpanSingleflightWait, trace.WithTimestamp(start))
+		waitSpan.End(trace.WithTimestamp(end))
 	}
 
 	return data, err
diff --git a/internal/tracing/singleflight_test.go b/internal/tracing/singleflight_test.go
index a8afc77a..bd70eb9f 100644
--- a/internal/tracing/singleflight_test.go
+++ b/internal/tracing/singleflight_test.go
@@ -36,7 +36,7 @@ func TestSingleflightDoAnnotatesALoneCaller(t *testing.T) {
 	assert.Equal(t, "result", data)
 
 	ended := recorder.Ended()
-	require.Len(t, ended, 1)
+	require.Len(t, ended, 1, "a caller that executed the function should produce no wait span")
 	attrs := attrsOf(ended[0])
 
 	shared, ok := attrs[SingleflightSharedKey]
@@ -93,9 +93,17 @@ func TestSingleflightDoAnnotatesSharedCallers(t *testing.T) {
 	require.NoError(t, <-followerDone)
 
 	ended := recorder.Ended()
-	require.Len(t, ended, 2)
+	require.Len(t, ended, 3, "expected the leader span, the follower span, and the follower's wait span")
+
+	var followerWaitMS float64
+	var waitSpan, followerSpan sdktrace.ReadOnlySpan
 	for _, s := range ended {
 		attrs := attrsOf(s)
+		if s.Name() == SpanSingleflightWait {
+			waitSpan = s
+			continue
+		}
+
 		shared, ok := attrs[SingleflightSharedKey]
 		require.True(t, ok, "the span should always report whether the flight was shared")
 		assert.True(t, shared.AsBool())
@@ -105,10 +113,24 @@ func TestSingleflightDoAnnotatesSharedCallers(t *testing.T) {
 		case "leader":
 			assert.False(t, waited, "the leader executed the function, so it should record no wait")
 		case "follower":
+			followerSpan = s
 			require.True(t, waited, "the follower waited on the leader's flight, so it should record its wait")
-			assert.Positive(t, wait.AsFloat64())
+			followerWaitMS = wait.AsFloat64()
+			assert.Positive(t, followerWaitMS)
 		}
 	}
+
+	// The follower's wait is also a span, back-dated to cover exactly the waiting window, so
+	// the trace timeline shows the wait instead of a gap.
+	require.NotNil(t, waitSpan, "the follower's wait should be visible as a span")
+	require.NotNil(t, followerSpan)
+	assert.Equal(t, followerSpan.SpanContext().TraceID(), waitSpan.SpanContext().TraceID(),
+		"the wait span should live in the follower's trace")
+	assert.Equal(t, followerSpan.SpanContext().SpanID(), waitSpan.Parent().SpanID(),
+		"the wait span should be a child of the follower's span")
+	spanMS := float64(waitSpan.EndTime().Sub(waitSpan.StartTime())) / float64(time.Millisecond)
+	assert.InDelta(t, followerWaitMS, spanMS, 0.001,
+		"the wait span should cover the same window the wait_ms attribute reports")
 }
 
 func TestSingleflightDoToleratesASpanlessContext(t *testing.T) {
diff --git a/relay/relay_endpoints_singleflight_test.go b/relay/relay_endpoints_singleflight_test.go
index 94e6da85..f054e493 100644
--- a/relay/relay_endpoints_singleflight_test.go
+++ b/relay/relay_endpoints_singleflight_test.go
@@ -127,6 +127,15 @@ func TestPollingHandlersShareOnePayloadBuildAcrossConcurrentRequests(t *testing.
 			assert.Len(t, spansNamed(spans, tracing.SpanWriteResponse), followers+1,
 				"every request should still write its own response")
 
+			// Each follower's wait shows up as a span in its own trace -- never in the
+			// winner's, whose time is the store and serialize spans instead.
+			waitSpans := spansNamed(spans, tracing.SpanSingleflightWait)
+			assert.Len(t, waitSpans, followers, "every waiting request should show its wait as a span")
+			for _, w := range waitSpans {
+				assert.NotEqual(t, serialize.SpanContext().TraceID(), w.SpanContext().TraceID(),
+					"the request that built the payload must not also show a wait span")
+			}
+
 			// Every request's own span reports that the payload build was shared. The one
 			// request that executed the build -- the one whose trace carries the serialize
 			// span -- did not wait and must record no wait time; every follower must record