diff --git a/docs-website/router/configuration.mdx b/docs-website/router/configuration.mdx index ae004c1ff..d625e0b40 100644 --- a/docs-website/router/configuration.mdx +++ b/docs-website/router/configuration.mdx @@ -2149,6 +2149,7 @@ Configure the GraphQL Execution Engine of the Router. | ENGINE_ENABLE_NET_POLL | enable_net_poll | | Enables the more efficient poll implementation for the server-side WebSocket handler of the router. This is only available on Linux and MacOS. On Windows or when the host system is limited, the default synchronous implementation is used. Has no effect on the router's upstream connections to subgraphs. | true | | ENGINE_WEBSOCKET_SERVER_READ_TIMEOUT | websocket_server_read_timeout | | Read timeout on the server-side WebSocket handler (router accepting clients). Specified as a Go duration string, e.g. `10ms`, `1s`, `1m`. | 5s | | ENGINE_WEBSOCKET_SERVER_WRITE_TIMEOUT | websocket_server_write_timeout | | Write timeout on the server-side WebSocket handler (router accepting clients). | 10s | +| ENGINE_SSE_SERVER_WRITE_TIMEOUT | sse_server_write_timeout | | Write timeout for server-side SSE responses (router writing to clients). When exceeded, the router closes the affected subscription. Set to `0s` to disable. | 10s | | ENGINE_WEBSOCKET_SERVER_POLL_TIMEOUT | websocket_server_poll_timeout | | The timeout for the poll loop of the server-side WebSocket handler. The period is specified as a string with a number and a unit. | 1s | | ENGINE_WEBSOCKET_SERVER_CONN_BUFFER_SIZE | websocket_server_conn_buffer_size | | The buffer size for the poll buffer of the server-side WebSocket handler. The buffer size determines how many connections can be handled in one loop. | 128 | | ENGINE_WEBSOCKET_CLIENT_WRITE_TIMEOUT | websocket_client_write_timeout | | The timeout for the websocket write of the WebSocket client implementation. | 10s | diff --git a/docs-website/router/custom-modules.mdx b/docs-website/router/custom-modules.mdx index 404f6b83b..2439ada9d 100644 --- a/docs-website/router/custom-modules.mdx +++ b/docs-website/router/custom-modules.mdx @@ -338,6 +338,29 @@ func (m *CustomModule) RouterOnRequest(ctx core.RequestContext, next http.Handle } ``` +#### Streaming responses + +Custom response writers used for subscriptions must implement `http.Flusher`. + +To retain SSE write timeouts, they must also either implement `SetWriteDeadline(time.Time) error` **or** expose the wrapped response writer: + +```go +// required for streaming to function +func (w *headerCapturingWriter) Flush() { + w.ResponseWriter.(http.Flusher).Flush() +} + +// Forward write deadlines directly +func (w *headerCapturingWriter) SetWriteDeadline(deadline time.Time) error { + return http.NewResponseController(w.ResponseWriter).SetWriteDeadline(deadline) +} + +// or expose the wrapped writer. +func (w *headerCapturingWriter) Unwrap() http.ResponseWriter { + return w.ResponseWriter +} +``` + ### Request Handler lifecycle The current module handler allow to intercept and modify request / response subgraphs. diff --git a/router-tests/subscriptions/http_subscriptions_test.go b/router-tests/subscriptions/http_subscriptions_test.go index 5ab3cd960..e02e7c8cb 100644 --- a/router-tests/subscriptions/http_subscriptions_test.go +++ b/router-tests/subscriptions/http_subscriptions_test.go @@ -3,11 +3,14 @@ package integration import ( "bufio" "bytes" + "context" "errors" "fmt" "io" "net/http" + "os" "strings" + "sync/atomic" "testing" "time" @@ -59,10 +62,10 @@ func readMultipartPrefix(reader *bufio.Reader) error { return nil } -func TestHeartbeats(t *testing.T) { +func TestHTTPMultipartSubscriptions(t *testing.T) { subscriptionHeartbeatInterval := time.Millisecond * 300 - t.Run("should work correctly for multipart", func(t *testing.T) { + t.Run("send heartbeats while waiting for data", func(t *testing.T) { testenv.Run(t, &testenv.Config{ RouterOptions: []core.Option{ core.WithSubscriptionHeartbeatInterval(subscriptionHeartbeatInterval), @@ -159,8 +162,12 @@ func TestHeartbeats(t *testing.T) { assert.Equal(t, 6, dataIdx, "expected 6 data messages") }) }) +} + +func TestSSESubscriptions(t *testing.T) { + subscriptionHeartbeatInterval := time.Millisecond * 300 - t.Run("should work correctly for sse", func(t *testing.T) { + t.Run("send heartbeats while waiting for data", func(t *testing.T) { testenv.Run(t, &testenv.Config{ RouterOptions: []core.Option{ core.WithSubscriptionHeartbeatInterval(subscriptionHeartbeatInterval), @@ -240,7 +247,7 @@ func TestHeartbeats(t *testing.T) { }) }) - t.Run("should write an error on sse", func(t *testing.T) { + t.Run("write upstream subscription errors", func(t *testing.T) { testenv.Run(t, &testenv.Config{ RouterOptions: []core.Option{ core.WithSubscriptionHeartbeatInterval(subscriptionHeartbeatInterval), @@ -303,14 +310,284 @@ func TestHeartbeats(t *testing.T) { }) }) }) + + testSSEWriteTimeout(t) + testSSENonFlusherWriter(t) +} + +const blockSSEWriteHeader = "X-Test-Block-SSE-Write" + +var ( + _ core.Module = (*blockingSSEWriterModule)(nil) + _ core.RouterOnRequestHandler = (*blockingSSEWriterModule)(nil) +) + +type blockingSSEWriteState struct { + armed atomic.Bool + writeStarted chan struct{} + writeDone chan struct{} + release chan struct{} +} + +type blockingSSEWriterModule struct { + state *blockingSSEWriteState +} + +func (m *blockingSSEWriterModule) Module() core.ModuleInfo { + return core.ModuleInfo{ + ID: "blockingSSEWriterModule", + Priority: 1, + New: func() core.Module { + return &blockingSSEWriterModule{state: m.state} + }, + } +} + +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(), + state: m.state, + }, ctx.Request()) } -func TestNonFlusherWriterSubscriptionError(t *testing.T) { - t.Parallel() +type deadlineBlockingResponseWriter struct { + http.ResponseWriter + state *blockingSSEWriteState + deadlineNanos atomic.Int64 +} + +func (w *deadlineBlockingResponseWriter) Write(data []byte) (int, error) { + if !w.state.armed.CompareAndSwap(true, false) { + return w.ResponseWriter.Write(data) + } - t.Run("subscription error when writer cannot flush", func(t *testing.T) { - t.Parallel() + close(w.state.writeStarted) + defer close(w.state.writeDone) + deadlineNanos := w.deadlineNanos.Load() + if deadlineNanos == 0 { + <-w.state.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.state.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 { + if deadline.IsZero() { + w.deadlineNanos.Store(0) + return nil + } + w.deadlineNanos.Store(deadline.UnixNano()) + return nil +} + +func (w *deadlineBlockingResponseWriter) Unwrap() http.ResponseWriter { + return w.ResponseWriter +} + +func testSSEWriteTimeout(t *testing.T) { + t.Run("remain writable after being idle longer than the write timeout", func(t *testing.T) { + const ( + sseWriteTimeout = 100 * time.Millisecond + eventIntervalMilliseconds = 500 + eventWaitTimeout = 5 * time.Second + ) + + testenv.Run(t, &testenv.Config{ + RouterOptions: []core.Option{ + core.WithSubscriptionHeartbeatInterval(time.Minute), + }, + // TLS enables HTTP/2, where an expired SSE write deadline fails the stream. + TLSConfig: config.TLSConfiguration{ + Server: config.TLSServerConfiguration{ + Enabled: true, + CertFile: "../testdata/tls/cert.pem", + KeyFile: "../testdata/tls/key.pem", + }, + }, + ModifyEngineExecutionConfiguration: func(cfg *config.EngineExecutionConfiguration) { + cfg.SSEServerWriteTimeout = sseWriteTimeout + }, + }, func(t *testing.T, xEnv *testenv.Environment) { + ctx, cancel := context.WithTimeout(t.Context(), eventWaitTimeout) + defer cancel() + + response := openCountEmpSSESubscription( + t, + ctx, + xEnv.RouterClient, + xEnv.GraphQLRequestURL(), + false, + eventIntervalMilliseconds, + ) + defer func() { + _ = response.Body.Close() + }() + require.Equal(t, 2, response.ProtoMajor) + reader := bufio.NewReader(response.Body) + + require.JSONEq(t, `{"data":{"countEmp":0}}`, readSSEData(t, reader)) + require.JSONEq(t, `{"data":{"countEmp":1}}`, readSSEData(t, reader)) + }) + }) + + t.Run("remove a blocked subscriber after write timeout while a healthy subscriber continues", func(t *testing.T) { + const ( + sseWriteTimeout = time.Second + eventWaitTimeout = 5 * time.Second + ) + + state := &blockingSSEWriteState{ + writeStarted: make(chan struct{}), + writeDone: make(chan struct{}), + release: make(chan struct{}), + } + + testenv.Run(t, &testenv.Config{ + RouterOptions: []core.Option{ + core.WithCustomModules(&blockingSSEWriterModule{state: state}), + core.WithSubscriptionHeartbeatInterval(time.Minute), + }, + ModifyEngineExecutionConfiguration: func(cfg *config.EngineExecutionConfiguration) { + cfg.SSEServerWriteTimeout = sseWriteTimeout + }, + }, func(t *testing.T, xEnv *testenv.Environment) { + defer close(state.release) + + ctx, cancel := context.WithTimeout(t.Context(), eventWaitTimeout) + defer cancel() + + client := &http.Client{} + blockedResponse := openCountEmpSSESubscription(t, ctx, client, xEnv.GraphQLRequestURL(), true, 250) + defer func() { + _ = blockedResponse.Body.Close() + }() + healthyResponse := openCountEmpSSESubscription(t, ctx, client, xEnv.GraphQLRequestURL(), false, 250) + defer func() { + _ = healthyResponse.Body.Close() + }() + healthyReader := bufio.NewReader(healthyResponse.Body) + + xEnv.WaitForSubscriptionCount(2, eventWaitTimeout) + xEnv.WaitForTriggerCount(1, eventWaitTimeout) + xEnv.RequireTriggerCount(1) + + readSSEData(t, healthyReader) + state.armed.Store(true) + + select { + case <-state.writeStarted: + case <-time.After(eventWaitTimeout): + t.Fatal("timed out waiting for the SSE write to block") + } + + beforeTimeout := readSSEData(t, healthyReader) + + select { + case <-state.writeDone: + case <-time.After(sseWriteTimeout + time.Second): + t.Fatal("blocked SSE write did not return after its deadline") + } + + xEnv.WaitForSubscriptionCount(1, eventWaitTimeout) + afterTimeout := readSSEData(t, healthyReader) + require.NotEqual(t, beforeTimeout, afterTimeout) + }) + }) +} + +func openCountEmpSSESubscription( + t *testing.T, + ctx context.Context, + client *http.Client, + url string, + blocked bool, + intervalMilliseconds int, +) *http.Response { + t.Helper() + + request, err := http.NewRequestWithContext( + ctx, + http.MethodPost, + url, + strings.NewReader(fmt.Sprintf( + `{"query":"subscription { countEmp(max: 20, intervalMilliseconds: %d) }"}`, + intervalMilliseconds, + )), + ) + require.NoError(t, err) + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Accept", "text/event-stream") + if blocked { + request.Header.Set(blockSSEWriteHeader, "true") + } + + response, err := client.Do(request) + require.NoError(t, err) + require.Equal(t, http.StatusOK, response.StatusCode) + require.Equal(t, "text/event-stream", response.Header.Get("Content-Type")) + return response +} + +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") + } + } +} + +func testSSENonFlusherWriter(t *testing.T) { + t.Run("return an error when the response writer cannot flush", func(t *testing.T) { cfg := config.Config{ Graph: config.Graph{}, Modules: map[string]interface{}{ @@ -340,8 +617,14 @@ func TestNonFlusherWriterSubscriptionError(t *testing.T) { body, err := io.ReadAll(resp.Body) require.NoError(t, err) - require.Contains(t, string(body), "errors") - require.Contains(t, string(body), "could not flush response") + require.Equal( + t, + `event: next +data: {"errors":[{"message":"subscription response writer does not support flushing"}]} + +`, + string(body), + ) }) }) } diff --git a/router/core/errors.go b/router/core/errors.go index 8b11963c0..c6b2a64bb 100644 --- a/router/core/errors.go +++ b/router/core/errors.go @@ -299,6 +299,17 @@ func writeRequestErrors(params writeRequestErrorsParams) { } params.logger.Error("Error writing response", zap.Error(err)) } + return + } + + if wgRequestParams.UseSse { + if _, err := params.writer.Write([]byte("\n\n")); err != nil && params.logger != nil { + if rErrors.IsBrokenPipe(err) { + params.logger.Warn("Broken pipe, error writing response", zap.Error(err)) + return + } + params.logger.Error("Error writing response", zap.Error(err)) + } } } diff --git a/router/core/graph_server.go b/router/core/graph_server.go index 1bb23b37a..e949c1d79 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.responseCache != nil { diff --git a/router/core/graphql_handler.go b/router/core/graphql_handler.go index 37db7c3dc..893fed9d9 100644 --- a/router/core/graphql_handler.go +++ b/router/core/graphql_handler.go @@ -90,6 +90,7 @@ type HandlerOptions struct { EnableCostResponseHeaders bool ApolloSubscriptionMultipartPrintBoundary bool + SSEServerWriteTimeout time.Duration HeaderPropagation *HeaderPropagation ResponseCache caching.Cache @@ -116,6 +117,7 @@ func NewGraphQLHandler(opts HandlerOptions) *GraphQLHandler { subgraphErrorPropagation: opts.SubgraphErrorPropagation, engineLoaderHooks: opts.EngineLoaderHooks, apolloSubscriptionMultipartPrintBoundary: opts.ApolloSubscriptionMultipartPrintBoundary, + sseServerWriteTimeout: opts.SSEServerWriteTimeout, headerPropagation: opts.HeaderPropagation, responseCacheStore: opts.ResponseCache, responseCacheFallbackTTL: opts.ResponseCacheFallbackTTL, @@ -177,6 +179,7 @@ type GraphQLHandler struct { enableCostResponseHeaders bool apolloSubscriptionMultipartPrintBoundary bool + sseServerWriteTimeout time.Duration } func (h *GraphQLHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { @@ -330,21 +333,25 @@ 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, + Logger: reqCtx.logger, + }) + 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, statusCode: http.StatusInternalServerError, - requestErrors: graphqlerrors.RequestErrorsFromError(errCouldNotFlushResponse), + requestErrors: graphqlerrors.RequestErrorsFromError(writerErr), logger: reqCtx.logger, headerPropagation: h.headerPropagation, }) diff --git a/router/core/subscription_response_writer.go b/router/core/subscription_response_writer.go index abe951a38..f38badb43 100644 --- a/router/core/subscription_response_writer.go +++ b/router/core/subscription_response_writer.go @@ -3,14 +3,18 @@ package core import ( "bytes" "context" + "errors" + "fmt" "io" "mime" "net/http" "strconv" "strings" + "time" "github.com/wundergraph/astjson" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" + "go.uber.org/zap" ) const ( @@ -31,16 +35,25 @@ type withFlushWriter interface { SubscriptionResponseWriter() resolve.SubscriptionResponseWriter } +type SubscriptionResponseWriterOptions struct { + ApolloSubscriptionMultipartPrintBoundary bool + SSEWriteTimeout time.Duration + Logger *zap.Logger +} + 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 + logger *zap.Logger // 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 +66,10 @@ func (f *HttpFlushWriter) Complete() { return } if f.sse { - _, _ = f.writer.Write([]byte("event: complete\ndata: \n\n")) + _ = f.writeAndFlushSSE(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 +79,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 +103,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(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 +167,22 @@ func (f *HttpFlushWriter) Flush() (err error) { } full := flushBreak + string(resp) + separation - _, err = f.writer.Write([]byte(full)) + if f.sse { + err = f.writeAndFlushSSE(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 +190,45 @@ 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(write func() error) (err error) { + if f.sseWriteTimeout > 0 { + if deadlineErr := f.responseControl.SetWriteDeadline(time.Now().Add(f.sseWriteTimeout)); deadlineErr != nil { + if !errors.Is(deadlineErr, http.ErrNotSupported) { + return fmt.Errorf("set SSE write deadline: %w", deadlineErr) + } + + f.sseWriteTimeout = 0 + if f.logger != nil { + f.logger.Warn( + "SSE write timeout disabled because response writer does not support write deadlines", + zap.Error(deadlineErr), + ) + } + } else { + defer func() { + if clearErr := f.responseControl.SetWriteDeadline(time.Time{}); clearErr != nil { + err = errors.Join(err, fmt.Errorf("clear SSE write deadline: %w", clearErr)) + } + }() + } + } + + if err := write(); err != nil { + return err + } + + return f.responseControl.Flush() +} + +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 +236,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, + logger: opts.Logger, + apolloSubscriptionMultipartPrintBoundary: opts.ApolloSubscriptionMultipartPrintBoundary, } flushWriter.ctx, flushWriter.cancel = context.WithCancel(ctx.Context()) @@ -197,10 +254,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(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 02db6b740..4c041d09f 100644 --- a/router/core/subscription_response_writer_test.go +++ b/router/core/subscription_response_writer_test.go @@ -2,16 +2,47 @@ 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.uber.org/zap" + "go.uber.org/zap/zaptest/observer" ) +type deadlineRecorder struct { + *httptest.ResponseRecorder + deadlines []time.Time + deadlineErr error + clearDeadlineErr error + flushErr error +} + +func (r *deadlineRecorder) SetWriteDeadline(deadline time.Time) error { + if deadline.IsZero() && r.clearDeadlineErr != nil { + return r.clearDeadlineErr + } + if r.deadlineErr != nil { + return r.deadlineErr + } + r.deadlines = append(r.deadlines, 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 +168,108 @@ 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 clears a deadline for every SSE write and flush", 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) + require.Len(t, recorder.deadlines, 2, "expected the initial header flush deadline to be set and cleared") + assert.False(t, recorder.deadlines[0].IsZero()) + assert.True(t, recorder.deadlines[1].IsZero()) + + _, err = writer.Write([]byte(`{"data":{"id":1}}`)) + require.NoError(t, err) + require.NoError(t, writer.Flush()) + require.Len(t, recorder.deadlines, 4, "expected the data frame deadline to be set and cleared") + assert.False(t, recorder.deadlines[2].Before(recorder.deadlines[0])) + assert.True(t, recorder.deadlines[3].IsZero()) + + require.NoError(t, writer.Heartbeat()) + require.Len(t, recorder.deadlines, 6, "expected the heartbeat deadline to be set and cleared") + assert.True(t, recorder.deadlines[5].IsZero()) + + writer.Complete() + require.Len(t, recorder.deadlines, 8, "expected the completion frame deadline to be set and cleared") + assert.True(t, recorder.deadlines[7].IsZero()) + }) + + t.Run("propagates an SSE flush error", 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{}) + require.NoError(t, err) + + flushErr := errors.New("flush failed") + recorder.flushErr = flushErr + require.ErrorIs(t, writer.Heartbeat(), flushErr) + }) + + t.Run("propagates an SSE deadline error", 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) + + deadlineErr := errors.New("deadline failed") + recorder.deadlineErr = deadlineErr + err = writer.Heartbeat() + assert.ErrorIs(t, err, deadlineErr) + assert.ErrorContains(t, err, "set SSE write deadline") + }) + + t.Run("returns an error when clearing an SSE deadline fails", func(t *testing.T) { + clearDeadlineErr := errors.New("clear deadline failed") + recorder := &deadlineRecorder{ + ResponseRecorder: httptest.NewRecorder(), + clearDeadlineErr: clearDeadlineErr, + } + 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, clearDeadlineErr) + assert.ErrorContains(t, err, "clear SSE write deadline") + assert.Nil(t, writer) + }) + + t.Run("disables the SSE timeout when write deadlines are unsupported", func(t *testing.T) { + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/graphql", nil) + req.Header.Set("Accept", sseMimeType) + logCore, logs := observer.New(zap.WarnLevel) + + _, writer, err := GetSubscriptionResponseWriter(resolve.NewContext(context.Background()), req, recorder, SubscriptionResponseWriterOptions{ + SSEWriteTimeout: time.Second, + Logger: zap.New(logCore), + }) + require.NoError(t, err) + require.NotNil(t, writer) + require.NoError(t, writer.Heartbeat()) + assert.True(t, recorder.Flushed) + assert.Equal(t, 1, logs.FilterMessage("SSE write timeout disabled because response writer does not support write deadlines").Len()) + }) + + t.Run("does not require deadline support when the timeout is disabled", 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{}) + require.NoError(t, err) + assert.NotNil(t, writer) + }) } diff --git a/router/pkg/config/config.go b/router/pkg/config/config.go index 211f2f2df..b3c293a96 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:"10s" 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 @@ -1780,6 +1781,10 @@ func LoadConfig(configFilePaths []string) (*LoadResult, error) { } } + if cfg.Config.EngineExecutionConfiguration.SSEServerWriteTimeout < 0 { + return nil, errors.New("engine.sse_server_write_timeout must be greater or equal to 0s") + } + // Post-process the config if cfg.Config.DevelopmentMode { cfg.Config.JSONLog = false diff --git a/router/pkg/config/config.schema.json b/router/pkg/config/config.schema.json index 88b844f27..592770624 100644 --- a/router/pkg/config/config.schema.json +++ b/router/pkg/config/config.schema.json @@ -4253,6 +4253,15 @@ "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", + "duration": { + "minimum": "0s" + }, + "default": "10s", + "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/config_test.go b/router/pkg/config/config_test.go index 26e6aa03c..9e4f8ba8c 100644 --- a/router/pkg/config/config_test.go +++ b/router/pkg/config/config_test.go @@ -478,6 +478,36 @@ telemetry: require.Equal(t, "at '/telemetry/tracing/exporters/0/export_timeout': duration must be less or equal than 2m0s", js.Causes[0].Error()) } +func TestSSEServerWriteTimeoutRejectsNegativeValues(t *testing.T) { + t.Run("config file", func(t *testing.T) { + f := createTempFileFromFixture(t, ` +version: "1" + +graph: + token: "token" + +engine: + sse_server_write_timeout: -1s +`) + + _, err := LoadConfig([]string{f}) + require.ErrorContains(t, err, "duration must be greater or equal than 0s") + }) + + t.Run("environment variable", func(t *testing.T) { + t.Setenv("ENGINE_SSE_SERVER_WRITE_TIMEOUT", "-1s") + f := createTempFileFromFixture(t, ` +version: "1" + +graph: + token: "token" +`) + + _, err := LoadConfig([]string{f}) + require.EqualError(t, err, "engine.sse_server_write_timeout must be greater or equal to 0s") + }) +} + func TestLoadFullConfig(t *testing.T) { t.Parallel() diff --git a/router/pkg/config/fixtures/full.yaml b/router/pkg/config/fixtures/full.yaml index 5636fe5c8..9957ac453 100644 --- a/router/pkg/config/fixtures/full.yaml +++ b/router/pkg/config/fixtures/full.yaml @@ -450,6 +450,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/json_schema.go b/router/pkg/config/json_schema.go index 46bc4432a..dc0675f6f 100644 --- a/router/pkg/config/json_schema.go +++ b/router/pkg/config/json_schema.go @@ -27,8 +27,10 @@ const ( ) type duration struct { - min time.Duration - max time.Duration + min time.Duration + max time.Duration + hasMin bool + hasMax bool } func (d duration) Validate(ctx *jsonschema.ValidatorContext, v any) { @@ -51,7 +53,7 @@ func (d duration) Validate(ctx *jsonschema.ValidatorContext, v any) { return } - if d.min > 0 { + if d.hasMin { if duration < d.min { ctx.AddError(&validationErrorKind{ fmt.Sprintf("duration must be greater or equal than %s", d.min), @@ -61,7 +63,7 @@ func (d duration) Validate(ctx *jsonschema.ValidatorContext, v any) { } } - if d.max > 0 { + if d.hasMax { if duration > d.max { ctx.AddError(&validationErrorKind{ fmt.Sprintf("duration must be less or equal than %s", d.max), @@ -118,23 +120,25 @@ func compileDuration(ctx *jsonschema.CompilerContext, m map[string]any) (jsonsch var minDuration, maxDuration time.Duration var err error - minDurationString, ok := mapVal["minimum"].(string) - if ok { + minDurationString, hasMin := mapVal["minimum"].(string) + if hasMin { minDuration, err = time.ParseDuration(minDurationString) if err != nil { return nil, err } } - maxDurationString, ok := mapVal["maximum"].(string) - if ok { + maxDurationString, hasMax := mapVal["maximum"].(string) + if hasMax { maxDuration, err = time.ParseDuration(maxDurationString) if err != nil { return nil, err } } return duration{ - min: minDuration, - max: maxDuration, + min: minDuration, + max: maxDuration, + hasMin: hasMin, + hasMax: hasMax, }, nil } diff --git a/router/pkg/config/testdata/config_defaults.json b/router/pkg/config/testdata/config_defaults.json index bffa8b9f1..2c1caf810 100644 --- a/router/pkg/config/testdata/config_defaults.json +++ b/router/pkg/config/testdata/config_defaults.json @@ -535,6 +535,7 @@ "DisableVariablesRemapping": false, "EnableRequireFetchReasons": false, "SubscriptionFetchTimeout": 30000000000, + "SSEServerWriteTimeout": 10000000000, "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 accbb14d3..63e8940f6 100644 --- a/router/pkg/config/testdata/config_full.json +++ b/router/pkg/config/testdata/config_full.json @@ -1003,6 +1003,7 @@ "DisableVariablesRemapping": false, "EnableRequireFetchReasons": false, "SubscriptionFetchTimeout": 30000000000, + "SSEServerWriteTimeout": 10000000000, "EnableDefer": false, "EnableMultiFetch": false, "EnableScheduleFetches": false,