perf: Deduplicate concurrent polling requests with a flight group

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
 

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)
 }

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)

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")
+		})
+	}
+}

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)

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,24 +134,42 @@ 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)
 		return out
 	}
 	go func() {
 		defer close(out)
-		event, err := r.getReplayEvent()
+		event, err := r.getReplayEvent(ctx)
 		if err == nil && event != nil {
 			out <- event
 		}
 	}()
 	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
 		}

internal/tracing/attributes.go

@@ -34,4 +34,16 @@ const (
 	StoreKeyKey      = attribute.Key("relay.store.key")
 	PayloadEventsKey = attribute.Key("relay.payload.events")
 	PayloadBytesKey  = attribute.Key("relay.payload.bytes")
+
+	// 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 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")
 )

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
+}