diff --git a/router-tests/events/kafka_hydration_hang_test.go b/router-tests/events/kafka_hydration_hang_test.go new file mode 100644 index 0000000000..b3a5b8c79e --- /dev/null +++ b/router-tests/events/kafka_hydration_hang_test.go @@ -0,0 +1,182 @@ +package events_test + +import ( + "encoding/json" + "math" + "net/http" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/wundergraph/cosmo/router-tests/events" + "github.com/wundergraph/cosmo/router-tests/testenv" + "github.com/wundergraph/cosmo/router-tests/testutils" + "github.com/wundergraph/cosmo/router/core" + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + "github.com/wundergraph/cosmo/router/pkg/config" +) + +var ( + _ core.Module = (*blockingHydrationModule)(nil) + _ core.EnginePreOriginHandler = (*blockingHydrationModule)(nil) +) + +// blockingHydrationModule blocks exactly one employees hydration request until +// its request context is canceled. +type blockingHydrationModule struct { + armed *atomic.Bool + started chan struct{} + release chan struct{} + startedOnce *sync.Once +} + +func (m *blockingHydrationModule) Module() core.ModuleInfo { + return core.ModuleInfo{ + ID: "blockingHydrationModule", + Priority: math.MaxInt32, + New: func() core.Module { + return &blockingHydrationModule{ + armed: m.armed, + started: m.started, + release: m.release, + startedOnce: m.startedOnce, + } + }, + } +} + +func (m *blockingHydrationModule) OnOriginRequest(req *http.Request, ctx core.RequestContext) (*http.Request, *http.Response) { + subgraph := ctx.ActiveSubgraph(req) + if subgraph != nil && subgraph.Name == "employees" && m.armed.CompareAndSwap(true, false) { + m.startedOnce.Do(func() { close(m.started) }) + select { + case <-m.release: + case <-req.Context().Done(): + } + } + return req, nil +} + +// TestKafkaSubscriptionContinuesAfterHydrationHonorsCancellation proves that a +// request timeout lets the same WebSocket subscription receive an inline error +// and then a later event when the hydration operation honors cancellation. +func TestKafkaSubscriptionContinuesAfterHydrationHonorsCancellation(t *testing.T) { + if testing.Short() { + t.Skip("skipping Kafka integration test in short mode") + } + + recovered, receivedError := runKafkaHydrationTimeoutScenario(t, "employeeUpdated-hydration-canceled") + require.True(t, receivedError, "expected the timed-out event to emit an error over the subscription") + require.True(t, recovered, "expected the existing subscription to receive a later Kafka event") +} + +func runKafkaHydrationTimeoutScenario(t *testing.T, topic string) (recovered bool, receivedError bool) { + t.Helper() + + armed := &atomic.Bool{} + started := make(chan struct{}) + release := make(chan struct{}) + var releaseOnce sync.Once + t.Cleanup(func() { releaseOnce.Do(func() { close(release) }) }) + + module := &blockingHydrationModule{ + armed: armed, + started: started, + release: release, + startedOnce: &sync.Once{}, + } + + testenv.Run(t, &testenv.Config{ + RouterConfigJSONTemplate: testenv.ConfigWithEdfsKafkaJSONTemplate, + EnableKafka: true, + ModifyRouterConfig: func(routerConfig *nodev1.RouterConfig) { + overrideKafkaTopicsForField(t, routerConfig, "employeeUpdatedMyKafka", + []string{"employeeUpdated", "employeeUpdatedTwo"}, topic) + }, + RouterOptions: []core.Option{ + core.WithCustomModules(module), + core.WithSubgraphTransportOptions(core.NewSubgraphTransportOptions(config.TrafficShapingRules{ + All: config.GlobalSubgraphRequestRule{ + RequestTimeout: testutils.ToPtr(100 * time.Millisecond), + }, + })), + core.WithSubgraphRetryOptions(false, "", 0, 0, 0, "", nil), + }, + ModifyEngineExecutionConfiguration: func(cfg *config.EngineExecutionConfiguration) { + cfg.SubscriptionFetchTimeout = 100 * time.Millisecond + }, + }, func(t *testing.T, xEnv *testenv.Environment) { + events.KafkaEnsureTopicExists(t, xEnv, EventWaitTimeout, topic) + + conn := xEnv.InitGraphQLWebSocketConnection(nil, nil, nil) + defer conn.Close() + + require.NoError(t, testenv.WSWriteJSON(t, conn, &testenv.WebSocketMessage{ + ID: "1", + Type: "subscribe", + Payload: []byte(`{"query":"subscription { employeeUpdatedMyKafka(employeeID: 3) { id details { forename surname } } }"}`), + })) + + xEnv.WaitForSubscriptionCount(1, EventWaitTimeout) + xEnv.WaitForTriggerCount(1, EventWaitTimeout) + + // Warm up the Kafka pipeline before arming the blocking hydration request. + xEnv.KafkaPublishUntilReceived(topic, + `{"__typename":"Employee","id":1,"update":{"name":"warmup"}}`, 1, EventWaitTimeout) + + var message testenv.WebSocketMessage + require.NoError(t, testenv.WSReadJSON(t, conn, &message)) + require.Equal(t, "next", message.Type) + + armed.Store(true) + events.ProduceKafkaMessage(t, xEnv, EventWaitTimeout, topic, + `{"__typename":"Employee","id":2,"update":{"name":"blocked"}}`) + + select { + case <-started: + case <-time.After(EventWaitTimeout): + t.Fatal("timed out waiting for hydration request to block") + } + + // This record is queued behind the stuck hydration on current main. + events.ProduceKafkaMessage(t, xEnv, EventWaitTimeout, topic, + `{"__typename":"Employee","id":3,"update":{"name":"recovery"}}`) + + deadline := time.Now().Add(500 * time.Millisecond) + require.NoError(t, conn.SetReadDeadline(deadline)) + + for time.Now().Before(deadline) { + message = testenv.WebSocketMessage{} + if err := conn.ReadJSON(&message); err != nil { + break + } + + var payload struct { + Data struct { + Employee struct { + ID int `json:"id"` + } `json:"employeeUpdatedMyKafka"` + } `json:"data"` + Errors []json.RawMessage `json:"errors"` + } + if json.Unmarshal(message.Payload, &payload) != nil { + continue + } + if len(payload.Errors) != 0 { + receivedError = true + } + if payload.Data.Employee.ID == 3 { + recovered = true + break + } + } + + // Always release the intentionally stuck goroutine before an assertion can stop the test. + releaseOnce.Do(func() { close(release) }) + }) + + return recovered, receivedError +} diff --git a/router-tests/events/kafka_sse_write_timeout_test.go b/router-tests/events/kafka_sse_write_timeout_test.go new file mode 100644 index 0000000000..e8148a7557 --- /dev/null +++ b/router-tests/events/kafka_sse_write_timeout_test.go @@ -0,0 +1,249 @@ +package events_test + +import ( + "bufio" + "context" + "errors" + "net/http" + "os" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/wundergraph/cosmo/router-tests/events" + "github.com/wundergraph/cosmo/router-tests/testenv" + "github.com/wundergraph/cosmo/router/core" + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + "github.com/wundergraph/cosmo/router/pkg/config" +) + +const blockSSEWriteHeader = "X-Test-Block-SSE-Write" + +var ( + _ core.Module = (*blockingSSEWriterModule)(nil) + _ core.RouterOnRequestHandler = (*blockingSSEWriterModule)(nil) +) + +// blockingSSEWriterModule simulates a client that stops draining its SSE +// connection without closing it. The wrapped writer only returns when the +// router sets a write deadline or the test releases it during cleanup. +type blockingSSEWriterModule struct { + armed *atomic.Bool + writeStarted chan struct{} + startedOnce *sync.Once + release chan struct{} +} + +func (m *blockingSSEWriterModule) Module() core.ModuleInfo { + return core.ModuleInfo{ + ID: "blockingSSEWriterModule", + Priority: 1, + New: func() core.Module { + return &blockingSSEWriterModule{ + armed: m.armed, + writeStarted: m.writeStarted, + startedOnce: m.startedOnce, + release: m.release, + } + }, + } +} + +func (m *blockingSSEWriterModule) RouterOnRequest(ctx core.RequestContext, next http.Handler) { + if ctx.Request().Header.Get(blockSSEWriteHeader) != "true" { + next.ServeHTTP(ctx.ResponseWriter(), ctx.Request()) + return + } + + next.ServeHTTP(&deadlineBlockingResponseWriter{ + ResponseWriter: ctx.ResponseWriter(), + armed: m.armed, + writeStarted: m.writeStarted, + startedOnce: m.startedOnce, + release: m.release, + }, ctx.Request()) +} + +type deadlineBlockingResponseWriter struct { + http.ResponseWriter + armed *atomic.Bool + writeStarted chan struct{} + startedOnce *sync.Once + release chan struct{} + deadlineNanos atomic.Int64 +} + +func (w *deadlineBlockingResponseWriter) Write(data []byte) (int, error) { + if !w.armed.CompareAndSwap(true, false) { + return w.ResponseWriter.Write(data) + } + + w.startedOnce.Do(func() { close(w.writeStarted) }) + deadlineNanos := w.deadlineNanos.Load() + if deadlineNanos == 0 { + <-w.release + return 0, os.ErrDeadlineExceeded + } + + wait := time.Until(time.Unix(0, deadlineNanos)) + if wait <= 0 { + return 0, os.ErrDeadlineExceeded + } + timer := time.NewTimer(wait) + defer timer.Stop() + select { + case <-w.release: + return 0, os.ErrDeadlineExceeded + case <-timer.C: + return 0, os.ErrDeadlineExceeded + } +} + +func (w *deadlineBlockingResponseWriter) Flush() { + if flusher, ok := w.ResponseWriter.(http.Flusher); ok { + flusher.Flush() + } +} + +func (w *deadlineBlockingResponseWriter) FlushError() error { + if flusher, ok := w.ResponseWriter.(interface{ FlushError() error }); ok { + return flusher.FlushError() + } + w.Flush() + return nil +} + +func (w *deadlineBlockingResponseWriter) SetWriteDeadline(deadline time.Time) error { + w.deadlineNanos.Store(deadline.UnixNano()) + return nil +} + +func (w *deadlineBlockingResponseWriter) Unwrap() http.ResponseWriter { + return w.ResponseWriter +} + +func TestKafkaSubscriptionRecoversAfterSSEWriteTimeout(t *testing.T) { + if testing.Short() { + t.Skip("skipping Kafka integration test in short mode") + } + + const topic = "employeeUpdated-sse-write-timeout" + armed := &atomic.Bool{} + writeStarted := make(chan struct{}) + release := make(chan struct{}) + var releaseOnce sync.Once + t.Cleanup(func() { releaseOnce.Do(func() { close(release) }) }) + + module := &blockingSSEWriterModule{ + armed: armed, + writeStarted: writeStarted, + startedOnce: &sync.Once{}, + release: release, + } + + testenv.Run(t, &testenv.Config{ + RouterConfigJSONTemplate: testenv.ConfigWithEdfsKafkaJSONTemplate, + EnableKafka: true, + RouterOptions: []core.Option{core.WithCustomModules(module)}, + ModifyRouterConfig: func(routerConfig *nodev1.RouterConfig) { + overrideKafkaTopicsForField(t, routerConfig, "employeeUpdatedMyKafka", + []string{"employeeUpdated", "employeeUpdatedTwo"}, topic) + }, + ModifyEngineExecutionConfiguration: func(cfg *config.EngineExecutionConfiguration) { + cfg.SSEServerWriteTimeout = 100 * time.Millisecond + }, + }, func(t *testing.T, xEnv *testenv.Environment) { + events.KafkaEnsureTopicExists(t, xEnv, EventWaitTimeout, topic) + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + client := &http.Client{} + blockedResp := openSSESubscription(t, ctx, client, xEnv.GraphQLRequestURL(), true) + defer blockedResp.Body.Close() + healthyResp := openSSESubscription(t, ctx, client, xEnv.GraphQLRequestURL(), false) + defer healthyResp.Body.Close() + healthyReader := bufio.NewReader(healthyResp.Body) + + xEnv.WaitForSubscriptionCount(2, EventWaitTimeout) + xEnv.WaitForTriggerCount(1, EventWaitTimeout) + + armed.Store(true) + xEnv.KafkaPublishUntilReceived(topic, + `{"__typename":"Employee","id":1,"update":{"name":"blocked"}}`, 1, EventWaitTimeout) + + select { + case <-writeStarted: + case <-time.After(EventWaitTimeout): + t.Fatal("timed out waiting for the SSE write to block") + } + + require.Contains(t, readSSEData(t, healthyReader), `"id":1`) + + events.ProduceKafkaMessage(t, xEnv, EventWaitTimeout, topic, + `{"__typename":"Employee","id":2,"update":{"name":"recovery"}}`) + + recovery := make(chan string, 1) + go func() { + data, err := readSSEDataLine(healthyReader) + if err != nil { + recovery <- "error: " + err.Error() + return + } + recovery <- data + }() + + select { + case data := <-recovery: + require.Contains(t, data, `"id":2`) + case <-time.After(EventWaitTimeout): + t.Fatal("healthy subscription did not receive the queued event after the SSE write deadline") + } + + xEnv.WaitForSubscriptionCount(1, EventWaitTimeout) + }) +} + +func openSSESubscription(t *testing.T, ctx context.Context, client *http.Client, url string, blocked bool) *http.Response { + t.Helper() + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, + strings.NewReader(`{"query":"subscription { employeeUpdatedMyKafka(employeeID: 3) { id } }"}`)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + if blocked { + req.Header.Set(blockSSEWriteHeader, "true") + } + + resp, err := client.Do(req) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Equal(t, "text/event-stream", resp.Header.Get("Content-Type")) + return resp +} + +func readSSEData(t *testing.T, reader *bufio.Reader) string { + t.Helper() + data, err := readSSEDataLine(reader) + require.NoError(t, err) + return data +} + +func readSSEDataLine(reader *bufio.Reader) (string, error) { + for { + line, err := reader.ReadString('\n') + if err != nil { + return "", err + } + line = strings.TrimSpace(line) + if strings.HasPrefix(line, "data: ") { + return strings.TrimPrefix(line, "data: "), nil + } + if strings.HasPrefix(line, "event: complete") { + return "", errors.New("subscription completed before receiving data") + } + } +} diff --git a/router/core/graph_server.go b/router/core/graph_server.go index 3cc45d5f9e..e4876a514d 100644 --- a/router/core/graph_server.go +++ b/router/core/graph_server.go @@ -1806,6 +1806,7 @@ func (s *graphServer) buildGraphMux( SubgraphErrorPropagation: s.subgraphErrorPropagation, EngineLoaderHooks: loaderHooks, HeaderPropagation: s.headerPropagation, + SSEServerWriteTimeout: s.engineExecutionConfiguration.SSEServerWriteTimeout, } if s.redisClient != nil { diff --git a/router/core/graphql_handler.go b/router/core/graphql_handler.go index 4ef92da46b..aa05d63247 100644 --- a/router/core/graphql_handler.go +++ b/router/core/graphql_handler.go @@ -10,6 +10,7 @@ import ( "net/http" "strconv" "strings" + "time" otelmetric "go.opentelemetry.io/otel/metric" "go.opentelemetry.io/otel/trace" @@ -87,6 +88,7 @@ type HandlerOptions struct { EnableCostResponseHeaders bool ApolloSubscriptionMultipartPrintBoundary bool + SSEServerWriteTimeout time.Duration HeaderPropagation *HeaderPropagation } @@ -109,6 +111,7 @@ func NewGraphQLHandler(opts HandlerOptions) *GraphQLHandler { subgraphErrorPropagation: opts.SubgraphErrorPropagation, engineLoaderHooks: opts.EngineLoaderHooks, apolloSubscriptionMultipartPrintBoundary: opts.ApolloSubscriptionMultipartPrintBoundary, + sseServerWriteTimeout: opts.SSEServerWriteTimeout, headerPropagation: opts.HeaderPropagation, } return graphQLHandler @@ -143,6 +146,7 @@ type GraphQLHandler struct { enableCostResponseHeaders bool apolloSubscriptionMultipartPrintBoundary bool + sseServerWriteTimeout time.Duration } func (h *GraphQLHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { @@ -284,16 +288,20 @@ func (h *GraphQLHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { } case *plan.SubscriptionResponsePlan: var ( - writer resolve.SubscriptionResponseWriter - ok bool + writer resolve.SubscriptionResponseWriter + writerErr error ) h.setDebugCacheHeaders(w, reqCtx.operation) defer propagateSubgraphErrors(resolveCtx) - resolveCtx, writer, ok = GetSubscriptionResponseWriter(resolveCtx, r, w, h.apolloSubscriptionMultipartPrintBoundary) - if !ok { - reqCtx.logger.Error("unable to get subscription response writer", zap.Error(errCouldNotFlushResponse)) - trackFinalResponseError(r.Context(), errCouldNotFlushResponse) + resolveCtx, writer, writerErr = GetSubscriptionResponseWriter(resolveCtx, r, w, SubscriptionResponseWriterOptions{ + ApolloSubscriptionMultipartPrintBoundary: h.apolloSubscriptionMultipartPrintBoundary, + SSEWriteTimeout: h.sseServerWriteTimeout, + MetricStore: h.metricStore, + }) + if writerErr != nil { + reqCtx.logger.Error("unable to get subscription response writer", zap.Error(writerErr)) + trackFinalResponseError(r.Context(), writerErr) writeRequestErrors(writeRequestErrorsParams{ request: r, writer: w, diff --git a/router/core/http_transport_cancellation_test.go b/router/core/http_transport_cancellation_test.go new file mode 100644 index 0000000000..6e746e7aed --- /dev/null +++ b/router/core/http_transport_cancellation_test.go @@ -0,0 +1,53 @@ +package core + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// TestHTTPTransportHonorsRequestCancellation documents the cancellation +// contract relied on by subscription hydration. The router's concrete HTTP +// transport must return when its request context expires; otherwise a shared +// subscription trigger can remain blocked after the hydration timeout. +func TestHTTPTransportHonorsRequestCancellation(t *testing.T) { + requestStarted := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + close(requestStarted) + <-r.Context().Done() + })) + t.Cleanup(server.Close) + + ctx, cancel := context.WithCancel(context.Background()) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, server.URL, nil) + require.NoError(t, err) + + transport := newHTTPTransport(DefaultTransportRequestOptions(), nil, nil, "products", nil) + t.Cleanup(transport.CloseIdleConnections) + + done := make(chan error, 1) + go func() { + _, roundTripErr := transport.RoundTrip(req) + done <- roundTripErr + }() + + select { + case <-requestStarted: + case <-time.After(time.Second): + t.Fatal("request did not reach the test server") + } + cancel() + + select { + case roundTripErr := <-done: + require.Error(t, roundTripErr) + require.True(t, errors.Is(roundTripErr, context.Canceled), roundTripErr) + case <-time.After(time.Second): + t.Fatal("router HTTP transport did not return after request cancellation") + } +} diff --git a/router/core/subscription_response_writer.go b/router/core/subscription_response_writer.go index abe951a380..44026e9493 100644 --- a/router/core/subscription_response_writer.go +++ b/router/core/subscription_response_writer.go @@ -3,13 +3,20 @@ package core import ( "bytes" "context" + "errors" + "fmt" "io" "mime" "net/http" "strconv" "strings" + "time" + + "go.opentelemetry.io/otel/attribute" + otelmetric "go.opentelemetry.io/otel/metric" "github.com/wundergraph/astjson" + rmetric "github.com/wundergraph/cosmo/router/pkg/metric" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) @@ -31,16 +38,34 @@ type withFlushWriter interface { SubscriptionResponseWriter() resolve.SubscriptionResponseWriter } +type SubscriptionResponseWriterOptions struct { + ApolloSubscriptionMultipartPrintBoundary bool + SSEWriteTimeout time.Duration + MetricStore rmetric.Store +} + +type sseFrameType string + +const ( + sseFrameTypeHeaders sseFrameType = "headers" + sseFrameTypeHeartbeat sseFrameType = "heartbeat" + sseFrameTypeNext sseFrameType = "next" + sseFrameTypeComplete sseFrameType = "complete" +) + type HttpFlushWriter struct { - ctx context.Context - cancel context.CancelFunc - writer io.Writer - flusher http.Flusher - subscribeOnce bool - sse bool - multipart bool - buf *bytes.Buffer - firstMessage bool + ctx context.Context + cancel context.CancelFunc + writer io.Writer + flusher http.Flusher + responseControl *http.ResponseController + subscribeOnce bool + sse bool + multipart bool + buf *bytes.Buffer + firstMessage bool + sseWriteTimeout time.Duration + metricStore rmetric.Store // apolloSubscriptionMultipartPrintBoundary if set to true will send the multipart boundary at the end of the message to allow // misbehaving client (like apollo client) to read the message just sent before the next one or the heartbeat apolloSubscriptionMultipartPrintBoundary bool @@ -53,7 +78,10 @@ func (f *HttpFlushWriter) Complete() { return } if f.sse { - _, _ = f.writer.Write([]byte("event: complete\ndata: \n\n")) + _ = f.writeAndFlushSSE(sseFrameTypeComplete, func() error { + _, err := f.writer.Write([]byte("event: complete\ndata: \n\n")) + return err + }) } else if f.multipart { // Write the final boundary in the multipart response if f.apolloSubscriptionMultipartPrintBoundary { @@ -63,8 +91,10 @@ func (f *HttpFlushWriter) Complete() { } } - // Flush before closing the writer to ensure all data is sent - f.flusher.Flush() + if !f.sse { + // Flush before closing the writer to ensure all data is sent. + f.flusher.Flush() + } f.cancel() } @@ -85,12 +115,10 @@ func (f *HttpFlushWriter) Heartbeat() error { var heartbeat []byte if f.sse { heartbeat = []byte(":heartbeat\n\n") - - if _, err := f.writer.Write(heartbeat); err != nil { + return f.writeAndFlushSSE(sseFrameTypeHeartbeat, func() error { + _, err := f.writer.Write(heartbeat) return err - } - - f.flusher.Flush() + }) } else if f.multipart { if _, err := f.Write([]byte("{}")); err != nil { return err @@ -151,14 +179,22 @@ func (f *HttpFlushWriter) Flush() (err error) { } full := flushBreak + string(resp) + separation - _, err = f.writer.Write([]byte(full)) + if f.sse { + err = f.writeAndFlushSSE(sseFrameTypeNext, func() error { + _, writeErr := f.writer.Write([]byte(full)) + return writeErr + }) + } else { + _, err = f.writer.Write([]byte(full)) + if err == nil { + // Flush before closing the writer to ensure all data is sent. + f.flusher.Flush() + } + } if err != nil { return err } - // Flush before closing the writer to ensure all data is sent - f.flusher.Flush() - if f.subscribeOnce { defer f.cancel() } @@ -166,15 +202,74 @@ func (f *HttpFlushWriter) Flush() (err error) { return nil } -func GetSubscriptionResponseWriter(ctx *resolve.Context, r *http.Request, w http.ResponseWriter, apolloSubscriptionMultipartPrintBoundary bool) (*resolve.Context, resolve.SubscriptionResponseWriter, bool) { +func (f *HttpFlushWriter) writeAndFlushSSE(frameType sseFrameType, write func() error) (err error) { + if f.metricStore != nil && frameType != sseFrameTypeHeartbeat { + started := time.Now() + defer func() { + attrs := []attribute.KeyValue{attribute.String("wg.sse.frame_type", string(frameType))} + f.metricStore.MeasureSSEWriteDuration(context.WithoutCancel(f.ctx), time.Since(started), attrs, otelmetric.WithAttributes()) + }() + } + + if f.sseWriteTimeout > 0 { + if err = f.responseControl.SetWriteDeadline(time.Now().Add(f.sseWriteTimeout)); err != nil { + // Failing closed prevents a response writer without deadline support from + // reintroducing an unbounded shared-trigger stall. + err = fmt.Errorf("set SSE write deadline: %w", err) + f.measureSSEWriteFailure(frameType, err) + return err + } + } + + if err = write(); err != nil { + f.measureSSEWriteFailure(frameType, err) + return err + } + + err = f.responseControl.Flush() + if err != nil { + f.measureSSEWriteFailure(frameType, err) + } + return err +} + +func (f *HttpFlushWriter) measureSSEWriteFailure(frameType sseFrameType, err error) { + if f.metricStore == nil { + return + } + attrs := []attribute.KeyValue{ + attribute.String("wg.sse.frame_type", string(frameType)), + attribute.String("wg.sse.failure_reason", sseWriteFailureReason(err)), + } + f.metricStore.MeasureSSEWriteFailure(context.WithoutCancel(f.ctx), attrs, otelmetric.WithAttributes()) +} + +func sseWriteFailureReason(err error) string { + if errors.Is(err, context.DeadlineExceeded) { + return "timeout" + } + var timeout interface{ Timeout() bool } + if errors.As(err, &timeout) && timeout.Timeout() { + return "timeout" + } + if errors.Is(err, http.ErrNotSupported) { + return "deadline_unsupported" + } + if errors.Is(err, context.Canceled) { + return "client_disconnected" + } + return "other" +} + +func GetSubscriptionResponseWriter(ctx *resolve.Context, r *http.Request, w http.ResponseWriter, opts SubscriptionResponseWriterOptions) (*resolve.Context, resolve.SubscriptionResponseWriter, error) { if wfw, ok := w.(withFlushWriter); ok { - return ctx, wfw.SubscriptionResponseWriter(), true + return ctx, wfw.SubscriptionResponseWriter(), nil } wgParams := NegotiateSubscriptionParams(r, false) flusher, ok := w.(http.Flusher) if !ok { - return ctx, nil, false + return ctx, nil, errors.New("subscription response writer does not support flushing") } setSubscriptionHeaders(wgParams, r, w) @@ -182,12 +277,15 @@ func GetSubscriptionResponseWriter(ctx *resolve.Context, r *http.Request, w http flushWriter := &HttpFlushWriter{ writer: w, flusher: flusher, + responseControl: http.NewResponseController(w), sse: wgParams.UseSse, multipart: wgParams.UseMultipart, subscribeOnce: wgParams.SubscribeOnce, buf: &bytes.Buffer{}, firstMessage: true, - apolloSubscriptionMultipartPrintBoundary: apolloSubscriptionMultipartPrintBoundary, + sseWriteTimeout: opts.SSEWriteTimeout, + metricStore: opts.MetricStore, + apolloSubscriptionMultipartPrintBoundary: opts.ApolloSubscriptionMultipartPrintBoundary, } flushWriter.ctx, flushWriter.cancel = context.WithCancel(ctx.Context()) @@ -197,10 +295,17 @@ func GetSubscriptionResponseWriter(ctx *resolve.Context, r *http.Request, w http ctx.ExecutionOptions.SendHeartbeat = true // Flush the response head immediately so the client establishes the connection // before the first message, instead of blocking until one is streamed. - flusher.Flush() + if wgParams.UseSse { + if err := flushWriter.writeAndFlushSSE(sseFrameTypeHeaders, func() error { return nil }); err != nil { + flushWriter.cancel() + return ctx, nil, fmt.Errorf("flush initial SSE response headers: %w", err) + } + } else { + flusher.Flush() + } } - return ctx, flushWriter, true + return ctx, flushWriter, nil } func wrapMultipartMessage(resp []byte, wrapPayload bool) ([]byte, error) { diff --git a/router/core/subscription_response_writer_test.go b/router/core/subscription_response_writer_test.go index 02db6b7400..2f007a744b 100644 --- a/router/core/subscription_response_writer_test.go +++ b/router/core/subscription_response_writer_test.go @@ -2,16 +2,55 @@ package core import ( "context" + "errors" "net/http" "net/http/httptest" "net/url" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" + "go.opentelemetry.io/otel/attribute" + otelmetric "go.opentelemetry.io/otel/metric" + + routermetric "github.com/wundergraph/cosmo/router/pkg/metric" ) +type deadlineRecorder struct { + *httptest.ResponseRecorder + deadline time.Time + flushErr error +} + +type sseMetricSpy struct { + routermetric.Store + durationCalls int + failureCalls int +} + +func (s *sseMetricSpy) MeasureSSEWriteDuration(context.Context, time.Duration, []attribute.KeyValue, otelmetric.RecordOption) { + s.durationCalls++ +} + +func (s *sseMetricSpy) MeasureSSEWriteFailure(context.Context, []attribute.KeyValue, otelmetric.AddOption) { + s.failureCalls++ +} + +func (r *deadlineRecorder) SetWriteDeadline(deadline time.Time) error { + r.deadline = deadline + return nil +} + +func (r *deadlineRecorder) FlushError() error { + if r.flushErr != nil { + return r.flushErr + } + r.Flush() + return nil +} + func TestNegotiateSubscriptionParams(t *testing.T) { type args struct { r *http.Request @@ -137,10 +176,96 @@ func TestGetSubscriptionResponseWriter(t *testing.T) { req := httptest.NewRequest(http.MethodPost, "/graphql", nil) req.Header.Set("Accept", sseMimeType) - _, _, ok := GetSubscriptionResponseWriter(resolve.NewContext(context.Background()), req, recorder, false) - require.True(t, ok) + _, _, err := GetSubscriptionResponseWriter(resolve.NewContext(context.Background()), req, recorder, SubscriptionResponseWriterOptions{}) + require.NoError(t, err) assert.Equal(t, sseMimeType, recorder.Header().Get("Content-Type")) assert.True(t, recorder.Flushed, "expected the SSE response head to be flushed before any message is written") }) + + t.Run("sets and refreshes a deadline for SSE headers and data", func(t *testing.T) { + recorder := &deadlineRecorder{ResponseRecorder: httptest.NewRecorder()} + req := httptest.NewRequest(http.MethodPost, "/graphql", nil) + req.Header.Set("Accept", sseMimeType) + + _, writer, err := GetSubscriptionResponseWriter(resolve.NewContext(context.Background()), req, recorder, SubscriptionResponseWriterOptions{SSEWriteTimeout: time.Second}) + require.NoError(t, err) + headerDeadline := recorder.deadline + require.False(t, headerDeadline.IsZero()) + + time.Sleep(time.Millisecond) + _, err = writer.Write([]byte(`{"data":{"id":1}}`)) + require.NoError(t, err) + require.NoError(t, writer.Flush()) + assert.True(t, recorder.deadline.After(headerDeadline), "expected each SSE frame to refresh the write deadline") + }) + + t.Run("does not measure a successful heartbeat", func(t *testing.T) { + recorder := &deadlineRecorder{ResponseRecorder: httptest.NewRecorder()} + req := httptest.NewRequest(http.MethodPost, "/graphql", nil) + req.Header.Set("Accept", sseMimeType) + metrics := &sseMetricSpy{} + + _, writer, err := GetSubscriptionResponseWriter(resolve.NewContext(context.Background()), req, recorder, SubscriptionResponseWriterOptions{ + SSEWriteTimeout: time.Second, + MetricStore: metrics, + }) + require.NoError(t, err) + metrics.durationCalls = 0 + metrics.failureCalls = 0 + + require.NoError(t, writer.Heartbeat()) + assert.Zero(t, metrics.durationCalls) + assert.Zero(t, metrics.failureCalls) + }) + + t.Run("measures a successful data frame", func(t *testing.T) { + recorder := &deadlineRecorder{ResponseRecorder: httptest.NewRecorder()} + req := httptest.NewRequest(http.MethodPost, "/graphql", nil) + req.Header.Set("Accept", sseMimeType) + metrics := &sseMetricSpy{} + + _, writer, err := GetSubscriptionResponseWriter(resolve.NewContext(context.Background()), req, recorder, SubscriptionResponseWriterOptions{ + SSEWriteTimeout: time.Second, + MetricStore: metrics, + }) + require.NoError(t, err) + metrics.durationCalls = 0 + + _, err = writer.Write([]byte(`{"data":{"id":1}}`)) + require.NoError(t, err) + require.NoError(t, writer.Flush()) + assert.Equal(t, 1, metrics.durationCalls) + }) + + t.Run("reports a failed heartbeat", func(t *testing.T) { + recorder := &deadlineRecorder{ResponseRecorder: httptest.NewRecorder()} + req := httptest.NewRequest(http.MethodPost, "/graphql", nil) + req.Header.Set("Accept", sseMimeType) + metrics := &sseMetricSpy{} + + _, writer, err := GetSubscriptionResponseWriter(resolve.NewContext(context.Background()), req, recorder, SubscriptionResponseWriterOptions{ + SSEWriteTimeout: time.Second, + MetricStore: metrics, + }) + require.NoError(t, err) + metrics.failureCalls = 0 + flushErr := errors.New("flush failed") + recorder.flushErr = flushErr + + require.ErrorIs(t, writer.Heartbeat(), flushErr) + assert.Equal(t, 1, metrics.failureCalls) + }) + + t.Run("fails closed when an SSE deadline is configured but unsupported", func(t *testing.T) { + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/graphql", nil) + req.Header.Set("Accept", sseMimeType) + + _, writer, err := GetSubscriptionResponseWriter(resolve.NewContext(context.Background()), req, recorder, SubscriptionResponseWriterOptions{SSEWriteTimeout: time.Second}) + require.Error(t, err) + assert.ErrorIs(t, err, http.ErrNotSupported) + assert.ErrorContains(t, err, "set SSE write deadline") + assert.Nil(t, writer) + }) } diff --git a/router/internal/retrytransport/retry_transport.go b/router/internal/retrytransport/retry_transport.go index ca5f9dc9a4..efcbd9e506 100644 --- a/router/internal/retrytransport/retry_transport.go +++ b/router/internal/retrytransport/retry_transport.go @@ -142,8 +142,23 @@ func (rt *RetryHTTPTransport) RoundTrip(req *http.Request) (*http.Response, erro rt.RetryOptions.OnRetry(retries, req, resp, sleepDuration, err) } - // Wait for the specified duration - time.Sleep(sleepDuration) + // Stop waiting as soon as the caller's deadline is reached. A plain + // time.Sleep here can make retries outlive a canceled subscription fetch. + timer := time.NewTimer(sleepDuration) + select { + case <-timer.C: + case <-req.Context().Done(): + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + if resp != nil && resp.Body != nil { + _ = resp.Body.Close() + } + return nil, req.Context().Err() + } // drain the previous response before retrying rt.drainBody(resp, requestLogger) diff --git a/router/internal/retrytransport/retry_transport_test.go b/router/internal/retrytransport/retry_transport_test.go index 8b7c83e318..faf95d00d4 100644 --- a/router/internal/retrytransport/retry_transport_test.go +++ b/router/internal/retrytransport/retry_transport_test.go @@ -2,12 +2,14 @@ package retrytransport import ( "bytes" + "context" "errors" "fmt" "io" "net/http" "net/http/httptest" "strings" + "sync" "testing" "testing/synctest" "time" @@ -18,6 +20,121 @@ import ( "go.uber.org/zap" ) +func TestRetryWaitHonorsRequestCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + request, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.com", nil) + assert.NoError(t, err) + + attempts := 0 + retrying := make(chan struct{}) + transport := RetryHTTPTransport{ + RoundTripper: &MockTransport{handler: func(*http.Request) (*http.Response, error) { + attempts++ + return nil, errors.New("temporary failure") + }}, + getRequestLogger: func(*http.Request) *zap.Logger { return zap.NewNop() }, + RetryOptions: RetryOptions{ + MaxRetryCount: 5, + Interval: time.Minute, + MaxDuration: time.Minute, + ShouldRetry: simpleShouldRetry, + OnRetry: func(int, *http.Request, *http.Response, time.Duration, error) { + close(retrying) + }, + }, + } + + done := make(chan error, 1) + go func() { + _, roundTripErr := transport.RoundTrip(request) + done <- roundTripErr + }() + <-retrying + cancel() + + select { + case err := <-done: + assert.ErrorIs(t, err, context.Canceled) + case <-time.After(250 * time.Millisecond): + t.Fatal("retry backoff did not stop after request cancellation") + } + assert.Equal(t, 1, attempts) +} + +type blockingReadBody struct { + release chan struct{} + closed chan struct{} + releaseOnce sync.Once + closedOnce sync.Once +} + +func newBlockingReadBody() *blockingReadBody { + return &blockingReadBody{ + release: make(chan struct{}), + closed: make(chan struct{}), + } +} + +func (b *blockingReadBody) Read([]byte) (int, error) { + <-b.release + return 0, io.EOF +} + +func (b *blockingReadBody) Close() error { + b.closedOnce.Do(func() { close(b.closed) }) + b.releaseRead() + return nil +} + +func (b *blockingReadBody) releaseRead() { + b.releaseOnce.Do(func() { close(b.release) }) +} + +func TestRetryCancellationClosesBodyWithoutDraining(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + request, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.com", nil) + assert.NoError(t, err) + + body := newBlockingReadBody() + t.Cleanup(body.releaseRead) + retrying := make(chan struct{}) + transport := RetryHTTPTransport{ + RoundTripper: &MockTransport{handler: func(*http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusInternalServerError, Body: body}, nil + }}, + getRequestLogger: func(*http.Request) *zap.Logger { return zap.NewNop() }, + RetryOptions: RetryOptions{ + MaxRetryCount: 5, + Interval: time.Minute, + MaxDuration: time.Minute, + ShouldRetry: simpleShouldRetry, + OnRetry: func(int, *http.Request, *http.Response, time.Duration, error) { + close(retrying) + }, + }, + } + + done := make(chan error, 1) + go func() { + _, roundTripErr := transport.RoundTrip(request) + done <- roundTripErr + }() + <-retrying + cancel() + + select { + case err := <-done: + assert.ErrorIs(t, err, context.Canceled) + case <-time.After(250 * time.Millisecond): + t.Fatal("retry cancellation blocked while draining the response body") + } + select { + case <-body.closed: + case <-time.After(250 * time.Millisecond): + t.Fatal("retry cancellation did not close the response body") + } +} + const defaultMaxDuration = 100 * time.Second // simpleShouldRetry provides simple retry logic for testing the transport implementation diff --git a/router/pkg/config/config.go b/router/pkg/config/config.go index 0b1489a051..23fb2536d5 100644 --- a/router/pkg/config/config.go +++ b/router/pkg/config/config.go @@ -500,6 +500,7 @@ type EngineExecutionConfiguration struct { DisableVariablesRemapping bool `envDefault:"false" env:"ENGINE_DISABLE_VARIABLES_REMAPPING" yaml:"disable_variables_remapping"` EnableRequireFetchReasons bool `envDefault:"false" env:"ENGINE_ENABLE_REQUIRE_FETCH_REASONS" yaml:"enable_require_fetch_reasons"` SubscriptionFetchTimeout time.Duration `envDefault:"30s" env:"ENGINE_SUBSCRIPTION_FETCH_TIMEOUT" yaml:"subscription_fetch_timeout,omitempty"` + SSEServerWriteTimeout time.Duration `envDefault:"0s" env:"ENGINE_SSE_SERVER_WRITE_TIMEOUT" yaml:"sse_server_write_timeout,omitempty"` EnableDefer bool `envDefault:"false" env:"ENGINE_ENABLE_DEFER" yaml:"enable_defer"` // EnableMultiFetch merges entity fetches to the same subgraph that execute diff --git a/router/pkg/config/config.schema.json b/router/pkg/config/config.schema.json index ac6bf7d5ac..68fa6721ba 100644 --- a/router/pkg/config/config.schema.json +++ b/router/pkg/config/config.schema.json @@ -4139,6 +4139,12 @@ "default": "30s", "description": "The maximum time a subscription fetch can take before it is considered timed out. The period is specified as a string with a number and a unit, e.g. 10ms, 1s, 1m, 1h. The supported units are 'ms', 's', 'm', 'h'." }, + "sse_server_write_timeout": { + "type": "string", + "format": "go-duration", + "default": "0s", + "description": "The maximum time allowed for each downstream SSE write and flush. When exceeded, the affected SSE subscription is terminated so it cannot indefinitely block other subscriptions sharing a trigger. A value of 0s disables the deadline." + }, "enable_defer": { "type": "boolean", "default": false, diff --git a/router/pkg/config/fixtures/full.yaml b/router/pkg/config/fixtures/full.yaml index 6e33f4f779..9ac2eddc07 100644 --- a/router/pkg/config/fixtures/full.yaml +++ b/router/pkg/config/fixtures/full.yaml @@ -448,6 +448,7 @@ engine: websocket_client_write_timeout: 10s websocket_server_read_timeout: 5s websocket_server_write_timeout: 10s + sse_server_write_timeout: 10s websocket_server_poll_timeout: 1s websocket_server_conn_buffer_size: 128 websocket_client_read_limit: 1MB diff --git a/router/pkg/config/testdata/config_defaults.json b/router/pkg/config/testdata/config_defaults.json index 576760616c..80feb894ce 100644 --- a/router/pkg/config/testdata/config_defaults.json +++ b/router/pkg/config/testdata/config_defaults.json @@ -516,6 +516,7 @@ "DisableVariablesRemapping": false, "EnableRequireFetchReasons": false, "SubscriptionFetchTimeout": 30000000000, + "SSEServerWriteTimeout": 0, "EnableDefer": false, "EnableMultiFetch": false, "EnableScheduleFetches": false, diff --git a/router/pkg/config/testdata/config_full.json b/router/pkg/config/testdata/config_full.json index 1539deedab..b546969372 100644 --- a/router/pkg/config/testdata/config_full.json +++ b/router/pkg/config/testdata/config_full.json @@ -984,6 +984,7 @@ "DisableVariablesRemapping": false, "EnableRequireFetchReasons": false, "SubscriptionFetchTimeout": 30000000000, + "SSEServerWriteTimeout": 10000000000, "EnableDefer": false, "EnableMultiFetch": false, "EnableScheduleFetches": false, diff --git a/router/pkg/metric/measurements.go b/router/pkg/metric/measurements.go index 8b6105bad3..5d33a4d45b 100644 --- a/router/pkg/metric/measurements.go +++ b/router/pkg/metric/measurements.go @@ -88,6 +88,18 @@ func createMeasures(meter otelmetric.Meter, opts MetricOpts) (*Measurements, err h.upDownCounters[InFlightRequestsUpDownCounter] = inFlightRequestsGauge + sseWriteDuration, err := meter.Float64Histogram(SSEWriteDurationHistogram, SSEWriteDurationHistogramOptions...) + if err != nil { + return nil, fmt.Errorf("failed to create SSE write duration histogram: %w", err) + } + h.histograms[SSEWriteDurationHistogram] = sseWriteDuration + + sseWriteFailures, err := meter.Int64Counter(SSEWriteFailuresCounter, SSEWriteFailuresCounterOptions...) + if err != nil { + return nil, fmt.Errorf("failed to create SSE write failures counter: %w", err) + } + h.counters[SSEWriteFailuresCounter] = sseWriteFailures + operationPlanningTime, err := meter.Float64Histogram( OperationPlanningTime, OperationPlanningTimeHistogramOptions..., diff --git a/router/pkg/metric/metric_store.go b/router/pkg/metric/metric_store.go index aa74c911b1..0ba4939743 100644 --- a/router/pkg/metric/metric_store.go +++ b/router/pkg/metric/metric_store.go @@ -27,6 +27,9 @@ const ( CircuitBreakerStateGauge = "router.circuit_breaker.state" CircuitBreakerShortCircuitsCounter = "router.circuit_breaker.short_circuits" + SSEWriteDurationHistogram = "router.http.server.sse.write.duration" + SSEWriteFailuresCounter = "router.http.server.sse.write.failures" + SchemaFieldUsageCounter = "router.graphql.schema_field_usage" // Total field usage OperationPlanningTime = "router.graphql.operation.planning_time" // Time taken to plan the operation @@ -70,6 +73,13 @@ var ( InFlightRequestsUpDownCounterOptions = []otelmetric.Int64UpDownCounterOption{ otelmetric.WithDescription(InFlightRequestsUpDownCounterDescription), } + SSEWriteDurationHistogramOptions = []otelmetric.Float64HistogramOption{ + otelmetric.WithUnit("ms"), + otelmetric.WithDescription("SSE write and flush duration in milliseconds"), + } + SSEWriteFailuresCounterOptions = []otelmetric.Int64CounterOption{ + otelmetric.WithDescription("SSE writes that failed or exceeded their deadline"), + } // GraphQL operation metrics @@ -173,6 +183,8 @@ type ( MeasureCircuitBreakerShortCircuit(ctx context.Context, opts ...otelmetric.AddOption) MeasureOperationCostEstimated(ctx context.Context, cost int64, opts ...otelmetric.RecordOption) MeasureOperationCostActual(ctx context.Context, cost int64, opts ...otelmetric.RecordOption) + MeasureSSEWriteDuration(ctx context.Context, duration float64, opts ...otelmetric.RecordOption) + MeasureSSEWriteFailure(ctx context.Context, opts ...otelmetric.AddOption) Flush(ctx context.Context) error Shutdown() error } @@ -193,6 +205,8 @@ type ( SetCircuitBreakerState(ctx context.Context, state bool, sliceAttr []attribute.KeyValue, opt otelmetric.RecordOption) MeasureOperationCostEstimated(ctx context.Context, cost int64, sliceAttr []attribute.KeyValue, opt otelmetric.RecordOption) MeasureOperationCostActual(ctx context.Context, cost int64, sliceAttr []attribute.KeyValue, opt otelmetric.RecordOption) + MeasureSSEWriteDuration(ctx context.Context, duration time.Duration, sliceAttr []attribute.KeyValue, opt otelmetric.RecordOption) + MeasureSSEWriteFailure(ctx context.Context, sliceAttr []attribute.KeyValue, opt otelmetric.AddOption) Flush(ctx context.Context) error Shutdown(ctx context.Context) error } @@ -260,6 +274,39 @@ func (h *Metrics) MeasureInFlight(ctx context.Context, sliceAttr []attribute.Key } } +func (h *Metrics) MeasureSSEWriteFailure(ctx context.Context, sliceAttr []attribute.KeyValue, opt otelmetric.AddOption) { + h.measureAdd(ctx, sliceAttr, opt, func(provider Provider, ctx context.Context, opts ...otelmetric.AddOption) { + provider.MeasureSSEWriteFailure(ctx, opts...) + }) +} + +func (h *Metrics) MeasureSSEWriteDuration(ctx context.Context, duration time.Duration, sliceAttr []attribute.KeyValue, opt otelmetric.RecordOption) { + opts := []otelmetric.RecordOption{h.baseAttributesOpt, opt} + durationMs := float64(duration) / float64(time.Millisecond) + if len(sliceAttr) == 0 { + h.promRequestMetrics.MeasureSSEWriteDuration(ctx, durationMs, opts...) + } else { + explodeRecordInstrument(ctx, sliceAttr, func(ctx context.Context, newOpts ...otelmetric.RecordOption) { + h.promRequestMetrics.MeasureSSEWriteDuration(ctx, durationMs, append(newOpts, opts...)...) + }) + } + opts = append(opts, otelmetric.WithAttributes(sliceAttr...)) + h.otlpRequestMetrics.MeasureSSEWriteDuration(ctx, durationMs, opts...) +} + +func (h *Metrics) measureAdd(ctx context.Context, sliceAttr []attribute.KeyValue, opt otelmetric.AddOption, measure func(Provider, context.Context, ...otelmetric.AddOption)) { + opts := []otelmetric.AddOption{h.baseAttributesOpt, opt} + if len(sliceAttr) == 0 { + measure(h.promRequestMetrics, ctx, opts...) + } else { + explodeAddInstrument(ctx, sliceAttr, func(ctx context.Context, newOpts ...otelmetric.AddOption) { + measure(h.promRequestMetrics, ctx, append(newOpts, opts...)...) + }) + } + opts = append(opts, otelmetric.WithAttributes(sliceAttr...)) + measure(h.otlpRequestMetrics, ctx, opts...) +} + func (h *Metrics) MeasureRequestCount(ctx context.Context, sliceAttr []attribute.KeyValue, opt otelmetric.AddOption) { opts := []otelmetric.AddOption{h.baseAttributesOpt, opt} diff --git a/router/pkg/metric/metric_store_test.go b/router/pkg/metric/metric_store_test.go index c688b2cab5..e87fa5b528 100644 --- a/router/pkg/metric/metric_store_test.go +++ b/router/pkg/metric/metric_store_test.go @@ -416,6 +416,41 @@ func findMetricDataPoints(t *testing.T, rm metricdata.ResourceMetrics, name stri return nil } +func TestSSEWriteMetrics(t *testing.T) { + metricReader := metric.NewManualReader() + store := createTestStore(t, 0, metricReader) + ctx := context.Background() + attrs := []attribute.KeyValue{attribute.String("wg.sse.frame_type", "next")} + opt := otelmetric.WithAttributes(attrs...) + + store.MeasureSSEWriteDuration(ctx, 25*time.Millisecond, attrs, opt) + store.MeasureSSEWriteFailure(ctx, attrs, opt) + + var rm metricdata.ResourceMetrics + require.NoError(t, metricReader.Collect(ctx, &rm)) + require.EqualValues(t, 1, findMetricDataPoints(t, rm, SSEWriteFailuresCounter)[0].Value) + durationPoints := findFloatMetricDataPoints(t, rm, SSEWriteDurationHistogram) + require.Len(t, durationPoints, 1) + require.EqualValues(t, 1, durationPoints[0].Count) + require.Equal(t, 25.0, durationPoints[0].Sum) +} + +func findFloatMetricDataPoints(t *testing.T, rm metricdata.ResourceMetrics, name string) []metricdata.HistogramDataPoint[float64] { + t.Helper() + for _, sm := range rm.ScopeMetrics { + for _, m := range sm.Metrics { + if m.Name == name { + if histogram, ok := m.Data.(metricdata.Histogram[float64]); ok { + return histogram.DataPoints + } + t.Fatalf("metric %q has unexpected data type %T", name, m.Data) + } + } + } + t.Fatalf("metric %q not found", name) + return nil +} + // TestOperationCostMetrics tests that operation cost metrics are recorded correctly func TestOperationCostMetrics(t *testing.T) { t.Parallel() diff --git a/router/pkg/metric/noop_metrics.go b/router/pkg/metric/noop_metrics.go index 7783f8e499..3cde6ccdb4 100644 --- a/router/pkg/metric/noop_metrics.go +++ b/router/pkg/metric/noop_metrics.go @@ -70,6 +70,12 @@ func (n NoopMetrics) MeasureOperationCostEstimated(ctx context.Context, cost int func (n NoopMetrics) MeasureOperationCostActual(ctx context.Context, cost int64, sliceAttr []attribute.KeyValue, opt otelmetric.RecordOption) { } +func (n NoopMetrics) MeasureSSEWriteDuration(context.Context, time.Duration, []attribute.KeyValue, otelmetric.RecordOption) { +} + +func (n NoopMetrics) MeasureSSEWriteFailure(context.Context, []attribute.KeyValue, otelmetric.AddOption) { +} + func NewNoopMetrics() Store { return &NoopMetrics{} } diff --git a/router/pkg/metric/noop_stream_metrics.go b/router/pkg/metric/noop_stream_metrics.go index c312cc2472..a6dfc24259 100644 --- a/router/pkg/metric/noop_stream_metrics.go +++ b/router/pkg/metric/noop_stream_metrics.go @@ -2,12 +2,16 @@ package metric import ( "context" + "time" ) type NoopStreamMetricStore struct{} -func (n *NoopStreamMetricStore) Produce(ctx context.Context, event StreamsEvent) {} -func (n *NoopStreamMetricStore) Consume(ctx context.Context, event StreamsEvent) {} +func (n *NoopStreamMetricStore) Produce(ctx context.Context, event StreamsEvent) {} +func (n *NoopStreamMetricStore) Consume(ctx context.Context, event StreamsEvent) {} +func (n *NoopStreamMetricStore) DispatchStart(ctx context.Context, event StreamsEvent) {} +func (n *NoopStreamMetricStore) DispatchFinish(ctx context.Context, event StreamsEvent, duration time.Duration) { +} func (n *NoopStreamMetricStore) Flush(ctx context.Context) error { return nil } func (n *NoopStreamMetricStore) Shutdown(ctx context.Context) error { return nil } diff --git a/router/pkg/metric/oltp_stream_metric_store.go b/router/pkg/metric/oltp_stream_metric_store.go index 8d30c15364..e94645190d 100644 --- a/router/pkg/metric/oltp_stream_metric_store.go +++ b/router/pkg/metric/oltp_stream_metric_store.go @@ -46,3 +46,13 @@ func (o *otlpStreamEventMetrics) Produce(ctx context.Context, opts ...otelmetric func (o *otlpStreamEventMetrics) Consume(ctx context.Context, opts ...otelmetric.AddOption) { o.instruments.consumedMessages.Add(ctx, 1, opts...) } + +func (o *otlpStreamEventMetrics) Process(ctx context.Context, opts ...otelmetric.AddOption) { + o.instruments.processedMessages.Add(ctx, 1, opts...) +} +func (o *otlpStreamEventMetrics) DispatchInFlight(ctx context.Context, delta int64, opts ...otelmetric.AddOption) { + o.instruments.dispatchInFlight.Add(ctx, delta, opts...) +} +func (o *otlpStreamEventMetrics) DispatchDuration(ctx context.Context, duration float64, opts ...otelmetric.RecordOption) { + o.instruments.dispatchDuration.Record(ctx, duration, opts...) +} diff --git a/router/pkg/metric/otlp_metric_store.go b/router/pkg/metric/otlp_metric_store.go index c80b5c0b8f..218f89b0c6 100644 --- a/router/pkg/metric/otlp_metric_store.go +++ b/router/pkg/metric/otlp_metric_store.go @@ -85,6 +85,14 @@ func (h *OtlpMetricStore) MeasureRequestCount(ctx context.Context, opts ...otelm } } +func (h *OtlpMetricStore) MeasureSSEWriteDuration(ctx context.Context, duration float64, opts ...otelmetric.RecordOption) { + h.measurements.histograms[SSEWriteDurationHistogram].Record(ctx, duration, opts...) +} + +func (h *OtlpMetricStore) MeasureSSEWriteFailure(ctx context.Context, opts ...otelmetric.AddOption) { + h.measurements.counters[SSEWriteFailuresCounter].Add(ctx, 1, opts...) +} + func (h *OtlpMetricStore) MeasureCircuitBreakerShortCircuit(ctx context.Context, opts ...otelmetric.AddOption) { if !h.circuitBreakerEnabled { return diff --git a/router/pkg/metric/prom_metric_store.go b/router/pkg/metric/prom_metric_store.go index 2988f0ef70..a830d96bd5 100644 --- a/router/pkg/metric/prom_metric_store.go +++ b/router/pkg/metric/prom_metric_store.go @@ -90,6 +90,14 @@ func (h *PromMetricStore) MeasureRequestCount(ctx context.Context, opts ...otelm } } +func (h *PromMetricStore) MeasureSSEWriteDuration(ctx context.Context, duration float64, opts ...otelmetric.RecordOption) { + h.measurements.histograms[SSEWriteDurationHistogram].Record(ctx, duration, opts...) +} + +func (h *PromMetricStore) MeasureSSEWriteFailure(ctx context.Context, opts ...otelmetric.AddOption) { + h.measurements.counters[SSEWriteFailuresCounter].Add(ctx, 1, opts...) +} + func (h *PromMetricStore) MeasureCircuitBreakerShortCircuit(ctx context.Context, opts ...otelmetric.AddOption) { if !h.circuitBreakerEnabled { return diff --git a/router/pkg/metric/prom_stream_metric_store.go b/router/pkg/metric/prom_stream_metric_store.go index 30309f2444..371bdcf37b 100644 --- a/router/pkg/metric/prom_stream_metric_store.go +++ b/router/pkg/metric/prom_stream_metric_store.go @@ -47,6 +47,16 @@ func (p *promStreamEventMetrics) Consume(ctx context.Context, opts ...otelmetric p.instruments.consumedMessages.Add(ctx, 1, opts...) } +func (p *promStreamEventMetrics) Process(ctx context.Context, opts ...otelmetric.AddOption) { + p.instruments.processedMessages.Add(ctx, 1, opts...) +} +func (p *promStreamEventMetrics) DispatchInFlight(ctx context.Context, delta int64, opts ...otelmetric.AddOption) { + p.instruments.dispatchInFlight.Add(ctx, delta, opts...) +} +func (p *promStreamEventMetrics) DispatchDuration(ctx context.Context, duration float64, opts ...otelmetric.RecordOption) { + p.instruments.dispatchDuration.Record(ctx, duration, opts...) +} + func (p *promStreamEventMetrics) Flush(ctx context.Context) error { return p.meterProvider.ForceFlush(ctx) } diff --git a/router/pkg/metric/stream_measurements.go b/router/pkg/metric/stream_measurements.go index a5e1dadfb1..1c774be28b 100644 --- a/router/pkg/metric/stream_measurements.go +++ b/router/pkg/metric/stream_measurements.go @@ -7,8 +7,11 @@ import ( ) const ( - messagingSentMessages = "router.streams.sent.messages" - messagingConsumedMessages = "router.streams.received.messages" + messagingSentMessages = "router.streams.sent.messages" + messagingConsumedMessages = "router.streams.received.messages" + messagingProcessedMessages = "router.streams.processed.messages" + messagingDispatchInFlight = "router.streams.dispatch.in_flight" + messagingDispatchDuration = "router.streams.dispatch.duration" ) var ( @@ -18,11 +21,17 @@ var ( messagingConsumedMessagesOptions = []otelmetric.Int64CounterOption{ otelmetric.WithDescription("Number of stream consumed messages"), } + messagingProcessedMessagesOptions = []otelmetric.Int64CounterOption{otelmetric.WithDescription("Number of stream messages whose subscription dispatch completed")} + messagingDispatchInFlightOptions = []otelmetric.Int64UpDownCounterOption{otelmetric.WithDescription("Number of stream messages currently dispatching to subscriptions")} + messagingDispatchDurationOptions = []otelmetric.Float64HistogramOption{otelmetric.WithUnit("ms"), otelmetric.WithDescription("Duration of stream message dispatch to subscriptions")} ) type eventInstruments struct { - producedMessages otelmetric.Int64Counter - consumedMessages otelmetric.Int64Counter + producedMessages otelmetric.Int64Counter + consumedMessages otelmetric.Int64Counter + processedMessages otelmetric.Int64Counter + dispatchInFlight otelmetric.Int64UpDownCounter + dispatchDuration otelmetric.Float64Histogram } func newStreamEventInstruments(meter otelmetric.Meter) (*eventInstruments, error) { @@ -41,9 +50,24 @@ func newStreamEventInstruments(meter otelmetric.Meter) (*eventInstruments, error if err != nil { return nil, fmt.Errorf("failed to create received messages counter: %w", err) } + processedCounter, err := meter.Int64Counter(messagingProcessedMessages, messagingProcessedMessagesOptions...) + if err != nil { + return nil, fmt.Errorf("failed to create processed messages counter: %w", err) + } + dispatchInFlight, err := meter.Int64UpDownCounter(messagingDispatchInFlight, messagingDispatchInFlightOptions...) + if err != nil { + return nil, fmt.Errorf("failed to create dispatch in flight counter: %w", err) + } + dispatchDuration, err := meter.Float64Histogram(messagingDispatchDuration, messagingDispatchDurationOptions...) + if err != nil { + return nil, fmt.Errorf("failed to create dispatch duration histogram: %w", err) + } return &eventInstruments{ - producedMessages: producedCounter, - consumedMessages: consumedCounter, + producedMessages: producedCounter, + consumedMessages: consumedCounter, + processedMessages: processedCounter, + dispatchInFlight: dispatchInFlight, + dispatchDuration: dispatchDuration, }, nil } diff --git a/router/pkg/metric/stream_metric_store.go b/router/pkg/metric/stream_metric_store.go index 361f49388d..04dbba0976 100644 --- a/router/pkg/metric/stream_metric_store.go +++ b/router/pkg/metric/stream_metric_store.go @@ -3,6 +3,7 @@ package metric import ( "context" "fmt" + "time" "go.opentelemetry.io/otel/attribute" otelmetric "go.opentelemetry.io/otel/metric" @@ -27,17 +28,23 @@ type StreamsEvent struct { ProviderType ProviderType // The messaging system type that are supported ErrorType string // Optional error type, e.g., "publish_error" or "receive_error". If empty, the attribute is not set DestinationName string // The name of the destination queue / topic / channel + RootFieldName string // The GraphQL subscription root field } // StreamMetricProvider is the interface that wraps the basic Event metric methods. type StreamMetricProvider interface { Produce(ctx context.Context, opts ...otelmetric.AddOption) Consume(ctx context.Context, opts ...otelmetric.AddOption) + Process(ctx context.Context, opts ...otelmetric.AddOption) + DispatchInFlight(ctx context.Context, delta int64, opts ...otelmetric.AddOption) + DispatchDuration(ctx context.Context, duration float64, opts ...otelmetric.RecordOption) } type StreamMetricStore interface { Produce(ctx context.Context, event StreamsEvent) Consume(ctx context.Context, event StreamsEvent) + DispatchStart(ctx context.Context, event StreamsEvent) + DispatchFinish(ctx context.Context, event StreamsEvent, duration time.Duration) } // StreamMetrics is the store for Event (Kafka/Redis/NATS) metrics. @@ -80,27 +87,42 @@ func (e *StreamMetrics) withAttrs(attrs ...attribute.KeyValue) otelmetric.AddOpt } func (e *StreamMetrics) Produce(ctx context.Context, event StreamsEvent) { - attrs := []attribute.KeyValue{ - otel.WgStreamOperationName.String(event.StreamOperationName), - otel.WgProviderType.String(string(event.ProviderType)), - } - if event.ErrorType != "" { - attrs = append(attrs, otel.WgErrorType.String(event.ErrorType)) - } - if event.ProviderId != "" { - attrs = append(attrs, otel.WgProviderId.String(event.ProviderId)) - } - if event.DestinationName != "" { - attrs = append(attrs, otel.WgDestinationName.String(event.DestinationName)) + e.recordAdd(ctx, event, func(provider StreamMetricProvider, ctx context.Context, opt otelmetric.AddOption) { + provider.Produce(ctx, opt) + }) +} + +func (e *StreamMetrics) Consume(ctx context.Context, event StreamsEvent) { + e.recordAdd(ctx, event, func(provider StreamMetricProvider, ctx context.Context, opt otelmetric.AddOption) { + provider.Consume(ctx, opt) + }) +} + +func (e *StreamMetrics) DispatchStart(ctx context.Context, event StreamsEvent) { + e.recordAdd(ctx, event, func(provider StreamMetricProvider, ctx context.Context, opt otelmetric.AddOption) { + provider.DispatchInFlight(ctx, 1, opt) + }) +} + +func (e *StreamMetrics) DispatchFinish(ctx context.Context, event StreamsEvent, duration time.Duration) { + attrs := e.eventAttrs(event) + addOpt := e.withAttrs(attrs...) + recordOpt := otelmetric.WithAttributes(append(append([]attribute.KeyValue{}, e.baseAttributes...), attrs...)...) + for _, provider := range e.providers { + provider.DispatchInFlight(ctx, -1, addOpt) + provider.Process(ctx, addOpt) + provider.DispatchDuration(ctx, float64(duration)/float64(time.Millisecond), recordOpt) } - opt := e.withAttrs(attrs...) +} +func (e *StreamMetrics) recordAdd(ctx context.Context, event StreamsEvent, record func(StreamMetricProvider, context.Context, otelmetric.AddOption)) { + opt := e.withAttrs(e.eventAttrs(event)...) for _, provider := range e.providers { - provider.Produce(ctx, opt) + record(provider, ctx, opt) } } -func (e *StreamMetrics) Consume(ctx context.Context, event StreamsEvent) { +func (e *StreamMetrics) eventAttrs(event StreamsEvent) []attribute.KeyValue { attrs := []attribute.KeyValue{ otel.WgStreamOperationName.String(event.StreamOperationName), otel.WgProviderType.String(string(event.ProviderType)), @@ -114,10 +136,8 @@ func (e *StreamMetrics) Consume(ctx context.Context, event StreamsEvent) { if event.DestinationName != "" { attrs = append(attrs, otel.WgDestinationName.String(event.DestinationName)) } - - opt := e.withAttrs(attrs...) - - for _, provider := range e.providers { - provider.Consume(ctx, opt) + if event.RootFieldName != "" { + attrs = append(attrs, otel.WgGraphQLFieldName.String(event.RootFieldName)) } + return attrs } diff --git a/router/pkg/pubsub/kafka/adapter.go b/router/pkg/pubsub/kafka/adapter.go index 3ee51437f4..5a96405f0a 100644 --- a/router/pkg/pubsub/kafka/adapter.go +++ b/router/pkg/pubsub/kafka/adapter.go @@ -48,7 +48,8 @@ type ProviderAdapter struct { } type PollerOpts struct { - providerId string + providerId string + rootFieldName string } // topicPoller polls the Kafka topic for new records and calls the updateTriggers function. @@ -102,12 +103,16 @@ func (p *ProviderAdapter) topicPoller(ctx context.Context, client *kgo.Client, u headers[header.Key] = header.Value } - p.streamMetricStore.Consume(ctx, metric.StreamsEvent{ + streamEvent := metric.StreamsEvent{ ProviderId: pollerOpts.providerId, StreamOperationName: kafkaReceive, ProviderType: metric.ProviderTypeKafka, DestinationName: r.Topic, - }) + RootFieldName: pollerOpts.rootFieldName, + } + p.streamMetricStore.Consume(ctx, streamEvent) + p.streamMetricStore.DispatchStart(ctx, streamEvent) + dispatchStarted := time.Now() updater.Update([]datasource.StreamEvent{ &Event{ @@ -118,6 +123,7 @@ func (p *ProviderAdapter) topicPoller(ctx context.Context, client *kgo.Client, u }, }, }) + p.streamMetricStore.DispatchFinish(context.WithoutCancel(ctx), streamEvent, time.Since(dispatchStarted)) } } } @@ -173,7 +179,7 @@ func (p *ProviderAdapter) Subscribe(ctx context.Context, conf datasource.Subscri stop := context.AfterFunc(p.ctx, cancel) defer stop() - err := p.topicPoller(pollerCtx, client, updater, PollerOpts{providerId: conf.ProviderID()}) + err := p.topicPoller(pollerCtx, client, updater, PollerOpts{providerId: conf.ProviderID(), rootFieldName: conf.RootFieldName()}) if err != nil { if errors.Is(err, errClientClosed) || errors.Is(err, context.Canceled) { log.Debug("poller canceled", zap.Error(err)) diff --git a/router/pkg/pubsub/nats/adapter.go b/router/pkg/pubsub/nats/adapter.go index a8ba5c3c7e..cb432e512c 100644 --- a/router/pkg/pubsub/nats/adapter.go +++ b/router/pkg/pubsub/nats/adapter.go @@ -151,12 +151,16 @@ func (p *ProviderAdapter) Subscribe(ctx context.Context, cfg datasource.Subscrip log.Debug("subscription update", zap.String("message_subject", msg.Subject()), zap.ByteString("data", msg.Data())) - p.streamMetricStore.Consume(p.ctx, metric.StreamsEvent{ + streamEvent := metric.StreamsEvent{ ProviderId: subConf.ProviderID(), StreamOperationName: natsReceive, ProviderType: metric.ProviderTypeNats, DestinationName: msg.Subject(), - }) + RootFieldName: subConf.RootFieldName(), + } + p.streamMetricStore.Consume(p.ctx, streamEvent) + p.streamMetricStore.DispatchStart(p.ctx, streamEvent) + dispatchStarted := time.Now() updater.Update([]datasource.StreamEvent{ &Event{evt: &MutableEvent{ @@ -164,6 +168,7 @@ func (p *ProviderAdapter) Subscribe(ctx context.Context, cfg datasource.Subscrip Headers: map[string][]string(msg.Headers()), }}, }) + p.streamMetricStore.DispatchFinish(context.WithoutCancel(p.ctx), streamEvent, time.Since(dispatchStarted)) // Acknowledge the message after it has been processed ackErr := msg.Ack() @@ -202,18 +207,23 @@ func (p *ProviderAdapter) Subscribe(ctx context.Context, cfg datasource.Subscrip select { case msg := <-msgChan: log.Debug("subscription update", zap.String("message_subject", msg.Subject), zap.ByteString("data", msg.Data)) - p.streamMetricStore.Consume(p.ctx, metric.StreamsEvent{ + streamEvent := metric.StreamsEvent{ ProviderId: subConf.ProviderID(), StreamOperationName: natsReceive, ProviderType: metric.ProviderTypeNats, DestinationName: msg.Subject, - }) + RootFieldName: subConf.RootFieldName(), + } + p.streamMetricStore.Consume(p.ctx, streamEvent) + p.streamMetricStore.DispatchStart(p.ctx, streamEvent) + dispatchStarted := time.Now() updater.Update([]datasource.StreamEvent{ &Event{evt: &MutableEvent{ Data: msg.Data, Headers: map[string][]string(msg.Header), }}, }) + p.streamMetricStore.DispatchFinish(context.WithoutCancel(p.ctx), streamEvent, time.Since(dispatchStarted)) case <-p.ctx.Done(): // When the application context is done, we stop the subscriptions for _, subscription := range subscriptions { diff --git a/router/pkg/pubsub/redis/adapter.go b/router/pkg/pubsub/redis/adapter.go index 606a473e96..25f043ca9b 100644 --- a/router/pkg/pubsub/redis/adapter.go +++ b/router/pkg/pubsub/redis/adapter.go @@ -137,10 +137,7 @@ func (p *ProviderAdapter) Subscribe(ctx context.Context, conf datasource.Subscri } } - p.closeWg.Add(1) - - go func() { - defer p.closeWg.Done() + p.closeWg.Go(func() { defer cleanup() for { @@ -155,17 +152,22 @@ func (p *ProviderAdapter) Subscribe(ctx context.Context, conf datasource.Subscri return } log.Debug("subscription update", zap.String("message_channel", msg.Channel), zap.String("data", msg.Payload)) - p.streamMetricStore.Consume(ctx, metric.StreamsEvent{ + streamEvent := metric.StreamsEvent{ ProviderId: conf.ProviderID(), StreamOperationName: redisReceive, ProviderType: metric.ProviderTypeRedis, DestinationName: msg.Channel, - }) + RootFieldName: conf.RootFieldName(), + } + p.streamMetricStore.Consume(ctx, streamEvent) + p.streamMetricStore.DispatchStart(ctx, streamEvent) + dispatchStarted := time.Now() updater.Update([]datasource.StreamEvent{ &Event{evt: &MutableEvent{ Data: []byte(msg.Payload), }}, }) + p.streamMetricStore.DispatchFinish(context.WithoutCancel(ctx), streamEvent, time.Since(dispatchStarted)) case <-p.ctx.Done(): // When the application context is done, we stop the subscription if it is not already done log.Debug("application context done, stopping subscription") @@ -176,7 +178,7 @@ func (p *ProviderAdapter) Subscribe(ctx context.Context, conf datasource.Subscri return } } - }() + }) return nil }