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/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 c7eafa74..0a7dca17 100644 --- a/internal/tracing/attributes.go +++ b/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") ) 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 4359f756..66fd5adf 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,86 @@ 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" +) + +// 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) { + data, err := tracing.SingleflightDo(req.Context(), env.GetPollingFlightGroup(), key, func() (any, error) { + result, err := build() + if err != nil { + return nil, err + } + return result, nil + }) + 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()) 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, 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) + }) + if err != nil { + w.WriteHeader(http.StatusInternalServerError) + return + } + + 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 +275,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 +567,31 @@ 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. + result, err := runPollingFlight(req, clientCtx.Env, serverSideAllFlagsFlightKey, func() (pollResult, error) { + return buildAllFlagsPayload(req.Context(), tr, clientCtx.Env) + }) + if err != nil { + w.WriteHeader(http.StatusInternalServerError) + return + } + + 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 +599,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 +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 } diff --git a/relay/relay_endpoints_singleflight_test.go b/relay/relay_endpoints_singleflight_test.go new file mode 100644 index 00000000..94e6da85 --- /dev/null +++ b/relay/relay_endpoints_singleflight_test.go @@ -0,0 +1,219 @@ +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) + 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. 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 { + 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) + }) + } +} + +// 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") +} + +// 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() + + 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") + 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") + }) + } +}