fix(router): add configurable SSE server write timeout

router-tests/events/kafka_sse_write_timeout_test.go

@@ -0,0 +1,248 @@
+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`)
+
+		xEnv.WaitForSubscriptionCount(1, EventWaitTimeout)
+		xEnv.KafkaPublishUntilReceived(topic,
+			`{"__typename":"Employee","id":2,"update":{"name":"recovery"}}`, 1, EventWaitTimeout)
+
+		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")
+		}
+	})
+}
+
+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")
+		}
+	}
+}

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 {

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,21 +288,24 @@ 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,
+		})
+		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,
 			})