feat(router): observe subscription delivery outcomes

router/core/graphql_handler.go

@@ -477,14 +477,25 @@ func (h *GraphQLHandler) writeError(ctx *resolve.Context, err error, res *resolv
 	}
 
 	requestLogger := reqContext.logger
+	errorKind := getErrorType(err)
+	observer, observesSubscriptions := h.engineStats.(statistics.SubscriptionObserver)
+	if _, isSubscription := w.(resolve.SubscriptionResponseWriter); isSubscription && observesSubscriptions {
+		reason := subscriptionResolutionErrorReason(errorKind)
+		observer.SubscriptionResolutionError(reason)
+		requestLogger.Warn("Subscription resolution failed",
+			zap.String("reason", string(reason)),
+			zap.Bool("terminal", terminal || isTerminalSubscriptionError(err)),
+			zap.Error(err),
+		)
+	}
 
 	httpWriter, isHttpResponseWriter := w.(http.ResponseWriter)
 	response := GraphQLErrorResponse{
 		Errors: make([]graphqlError, 1),
 		Data:   nil,
 	}
 
-	switch getErrorType(err) {
+	switch errorKind {
 	case errorTypeMergeResult:
 		var errMerge resolve.ErrMergeResult
 		if !errors.As(err, &errMerge) {
@@ -638,6 +649,29 @@ func (h *GraphQLHandler) writeError(ctx *resolve.Context, err error, res *resolv
 	}
 }
 
+func subscriptionResolutionErrorReason(kind errorType) statistics.SubscriptionResolutionErrorReason {
+	switch kind {
+	case errorTypeRateLimit:
+		return statistics.SubscriptionResolutionErrorRateLimit
+	case errorTypeUnauthorized:
+		return statistics.SubscriptionResolutionErrorAuthorization
+	case errorTypeContextCanceled:
+		return statistics.SubscriptionResolutionErrorContextCanceled
+	case errorTypeContextTimeout:
+		return statistics.SubscriptionResolutionErrorTimeout
+	case errorTypeUpgradeFailed, errorTypeEDFS:
+		return statistics.SubscriptionResolutionErrorFetch
+	case errorTypeStreamsHandlerError:
+		return statistics.SubscriptionResolutionErrorHandler
+	case errorTypeEDFSInvalidMessage, errorTypeInvalidWsSubprotocol:
+		return statistics.SubscriptionResolutionErrorInvalidMessage
+	case errorTypeMergeResult:
+		return statistics.SubscriptionResolutionErrorResolve
+	default:
+		return statistics.SubscriptionResolutionErrorUnknown
+	}
+}
+
 func (h *GraphQLHandler) setDebugCacheHeaders(w http.ResponseWriter, opCtx *operationContext) {
 	if h.enableCacheResponseHeaders {
 		if opCtx.normalizationCacheHit {

router/core/graphql_handler_observability_test.go

@@ -0,0 +1,26 @@
+package core
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/wundergraph/cosmo/router/pkg/statistics"
+)
+
+func TestSubscriptionResolutionErrorReason(t *testing.T) {
+	t.Parallel()
+	tests := map[errorType]statistics.SubscriptionResolutionErrorReason{
+		errorTypeUnauthorized:        statistics.SubscriptionResolutionErrorAuthorization,
+		errorTypeContextCanceled:     statistics.SubscriptionResolutionErrorContextCanceled,
+		errorTypeContextTimeout:      statistics.SubscriptionResolutionErrorTimeout,
+		errorTypeUpgradeFailed:       statistics.SubscriptionResolutionErrorFetch,
+		errorTypeEDFS:                statistics.SubscriptionResolutionErrorFetch,
+		errorTypeStreamsHandlerError: statistics.SubscriptionResolutionErrorHandler,
+		errorTypeEDFSInvalidMessage:  statistics.SubscriptionResolutionErrorInvalidMessage,
+		errorTypeMergeResult:         statistics.SubscriptionResolutionErrorResolve,
+		errorTypeUnknown:             statistics.SubscriptionResolutionErrorUnknown,
+	}
+	for kind, want := range tests {
+		require.Equal(t, want, subscriptionResolutionErrorReason(kind))
+	}
+}

router/core/websocket.go

@@ -669,8 +669,12 @@ func (rw *websocketResponseWriter) Complete() {
 	}
 	err := rw.protocol.Complete(rw.id)
 	if err != nil {
-		rw.logger.Debug("Sending complete message", zap.Error(err))
+		reason := webSocketWriteFailureReason(err)
+		rw.observeFrame("complete", "none", "failure", reason)
+		rw.logger.Warn("Sending WebSocket complete message failed", zap.String("reason", reason), zap.Error(err))
+		return
 	}
+	rw.observeFrame("complete", "none", "success", "none")
 }
 
 // Heartbeat is a no-op function for WebSocket subscriptions.
@@ -700,15 +704,22 @@ func (rw *websocketResponseWriter) Error(data []byte) {
 		errors = json.RawMessage(`[{"message":"Unable to subscribe"}]`)
 	}
 	if err := rw.protocol.WriteGraphQLErrors(rw.id, errors, nil); err != nil {
-		rw.logger.Debug("Sending error message", zap.Error(err))
+		reason := webSocketWriteFailureReason(err)
+		rw.observeFrame("terminal_error", "errors_only", "failure", reason)
+		rw.logger.Warn("Sending terminal GraphQL error frame failed", zap.String("reason", reason), zap.Error(err))
 		return
 	}
+	rw.observeFrame("terminal_error", "errors_only", "success", "none")
 	// subscriptions-transport-ws clients rely on an explicit "complete" to end
 	// the stream after a data+errors frame. graphql-transport-ws treats the
 	// "error" frame as terminal per spec, so no follow-up is needed there.
 	if rw.protocol.Subprotocol() == wsproto.SubscriptionsTransportWSSubprotocol {
 		if err := rw.protocol.Complete(rw.id); err != nil {
-			rw.logger.Debug("Sending complete after error", zap.Error(err))
+			reason := webSocketWriteFailureReason(err)
+			rw.observeFrame("complete", "none", "failure", reason)
+			rw.logger.Warn("Sending WebSocket complete after error failed", zap.String("reason", reason), zap.Error(err))
+		} else {
+			rw.observeFrame("complete", "none", "success", "none")
 		}
 	}
 }
@@ -728,6 +739,7 @@ func (rw *websocketResponseWriter) Flush() error {
 				"response_headers": rw.header,
 			})
 			if err != nil {
+				rw.observeFrame("data", webSocketPayloadType(payload), "failure", "serialization_error")
 				rw.logger.Warn("Serializing response headers", zap.Error(err))
 				return err
 			}
@@ -746,12 +758,56 @@ func (rw *websocketResponseWriter) Flush() error {
 		err = rw.protocol.WriteGraphQLData(rw.id, payload, extensions)
 		rw.buf.Reset()
 		if err != nil {
+			reason := webSocketWriteFailureReason(err)
+			rw.observeFrame("data", webSocketPayloadType(payload), "failure", reason)
+			rw.logger.Warn("Sending GraphQL data frame failed", zap.String("reason", reason), zap.Error(err))
 			return err
 		}
+		rw.observeFrame("data", webSocketPayloadType(payload), "success", "none")
 	}
 	return nil
 }
 
+func (rw *websocketResponseWriter) observeFrame(frameType, payloadType, result, reason string) {
+	observer, ok := rw.stats.(statistics.SubscriptionObserver)
+	if !ok {
+		return
+	}
+	observer.WebSocketFrame(statistics.WebSocketFrameObservation{
+		FrameType:   frameType,
+		PayloadType: payloadType,
+		Result:      result,
+		Reason:      reason,
+	})
+}
+
+func webSocketPayloadType(payload []byte) string {
+	data := gjson.GetBytes(payload, "data")
+	hasData := data.Exists() && data.Type != gjson.Null
+	hasErrors := gjson.GetBytes(payload, "errors").Type == gjson.JSON
+	switch {
+	case hasData && hasErrors:
+		return "data_with_errors"
+	case hasErrors:
+		return "errors_only"
+	case hasData:
+		return "data"
+	default:
+		return "none"
+	}
+}
+
+func webSocketWriteFailureReason(err error) string {
+	switch {
+	case errors.Is(err, context.Canceled):
+		return "context_canceled"
+	case errors.Is(err, net.ErrClosed), errors.Is(err, syscall.EPIPE), errors.Is(err, syscall.ECONNRESET):
+		return "client_disconnected"
+	default:
+		return "protocol_error"
+	}
+}
+
 func (rw *websocketResponseWriter) SubscriptionResponseWriter() resolve.SubscriptionResponseWriter {
 	return rw
 }

router/core/websocket_observability_test.go

@@ -0,0 +1,86 @@
+package core
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/wundergraph/cosmo/router/internal/wsproto"
+	"github.com/wundergraph/cosmo/router/pkg/statistics"
+	"go.uber.org/zap"
+)
+
+type observabilityTestProtocol struct {
+	dataErr     error
+	errorErr    error
+	completeErr error
+}
+
+func (p *observabilityTestProtocol) Subprotocol() string                    { return wsproto.GraphQLWSSubprotocol }
+func (p *observabilityTestProtocol) Initialize() (json.RawMessage, error)   { return nil, nil }
+func (p *observabilityTestProtocol) ReadMessage() (*wsproto.Message, error) { return nil, nil }
+func (p *observabilityTestProtocol) Pong(*wsproto.Message) error            { return nil }
+func (p *observabilityTestProtocol) WriteGraphQLData(string, json.RawMessage, json.RawMessage) error {
+	return p.dataErr
+}
+func (p *observabilityTestProtocol) WriteGraphQLErrors(string, json.RawMessage, json.RawMessage) error {
+	return p.errorErr
+}
+func (p *observabilityTestProtocol) Complete(string) error { return p.completeErr }
+
+func TestWebsocketResponseWriterObservesFrameOutcomes(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name        string
+		payload     string
+		protocolErr error
+		wantResult  string
+		wantPayload string
+	}{
+		{name: "data success", payload: `{"data":{"value":1}}`, wantResult: "success", wantPayload: "data"},
+		{name: "data with errors success", payload: `{"data":{"value":1},"errors":[{"message":"partial"}]}`, wantResult: "success", wantPayload: "data_with_errors"},
+		{name: "errors only success", payload: `{"errors":[{"message":"failed"}],"data":null}`, wantResult: "success", wantPayload: "errors_only"},
+		{name: "write failure", payload: `{"data":{"value":1}}`, protocolErr: errors.New("write failed"), wantResult: "failure", wantPayload: "data"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			stats := statistics.NewEngineStats(t.Context(), zap.NewNop(), false)
+			writer := newWebsocketResponseWriter("1", &observabilityTestProtocol{dataErr: tt.protocolErr}, true, zap.NewNop(), stats, nil)
+			_, err := writer.Write([]byte(tt.payload))
+			require.NoError(t, err)
+			flushErr := writer.Flush()
+			if tt.protocolErr == nil {
+				require.NoError(t, flushErr)
+			} else {
+				require.ErrorIs(t, flushErr, tt.protocolErr)
+			}
+
+			report := stats.GetReport()
+			require.Len(t, report.WebSocketFrames, 1)
+			require.Equal(t, statistics.WebSocketFrameObservation{
+				FrameType: "data", PayloadType: tt.wantPayload, Result: tt.wantResult,
+				Reason: map[bool]string{true: "protocol_error", false: "none"}[tt.protocolErr != nil],
+			}, report.WebSocketFrames[0].Observation)
+			require.Equal(t, uint64(1), report.WebSocketFrames[0].Count)
+		})
+	}
+}
+
+func TestWebsocketResponseWriterObservesTerminalErrorWriteFailure(t *testing.T) {
+	t.Parallel()
+	stats := statistics.NewEngineStats(context.Background(), zap.NewNop(), false)
+	writer := newWebsocketResponseWriter("1", &observabilityTestProtocol{errorErr: errors.New("client gone")}, true, zap.NewNop(), stats, nil)
+
+	writer.Error([]byte(`{"errors":[{"message":"timeout"}]}`))
+
+	report := stats.GetReport()
+	require.Len(t, report.WebSocketFrames, 1)
+	require.Equal(t, statistics.WebSocketFrameObservation{
+		FrameType: "terminal_error", PayloadType: "errors_only", Result: "failure", Reason: "protocol_error",
+	}, report.WebSocketFrames[0].Observation)
+}