[Diagnostics] Improve WebSocket handling

src/HotChocolate/AspNetCore/src/AspNetCore.Pipeline/Instrumentation/AggregateServerDiagnosticEventListener.cs

@@ -1,4 +1,6 @@
 using Microsoft.AspNetCore.Http;
+using HotChocolate.AspNetCore.Subscriptions;
+using HotChocolate.AspNetCore.Subscriptions.Protocols;
 using HotChocolate.Language;
 
 namespace HotChocolate.AspNetCore.Instrumentation;
@@ -111,6 +113,16 @@ public IDisposable WebSocketSession(HttpContext context)
         return new AggregateActivityScope(scopes);
     }
 
+    public void WebSocketConnectionInitialized(
+        ISocketSession session,
+        IOperationMessagePayload connectionInitMessage)
+    {
+        for (var i = 0; i < _listeners.Length; i++)
+        {
+            _listeners[i].WebSocketConnectionInitialized(session, connectionInitMessage);
+        }
+    }
+
     public void WebSocketSessionError(HttpContext context, Exception exception)
     {
         for (var i = 0; i < _listeners.Length; i++)

src/HotChocolate/AspNetCore/src/AspNetCore.Pipeline/Instrumentation/IServerDiagnosticEvents.cs

@@ -1,4 +1,6 @@
 using Microsoft.AspNetCore.Http;
+using HotChocolate.AspNetCore.Subscriptions;
+using HotChocolate.AspNetCore.Subscriptions.Protocols;
 using HotChocolate.Language;
 
 namespace HotChocolate.AspNetCore.Instrumentation;
@@ -138,6 +140,21 @@ void StartOperationBatchRequest(
     /// </returns>
     IDisposable WebSocketSession(HttpContext context);
 
+    /// <summary>
+    /// Called within the <see cref="WebSocketSession"/> scope once the client's connection
+    /// initialization message has been accepted.
+    /// </summary>
+    /// <param name="session">
+    /// The GraphQL socket session.
+    /// </param>
+    /// <param name="connectionInitMessage">
+    /// The connection initialization message sent by the client. Its payload is only valid
+    /// for the duration of this call.
+    /// </param>
+    void WebSocketConnectionInitialized(
+        ISocketSession session,
+        IOperationMessagePayload connectionInitMessage);
+
     /// <summary>
     /// Called within the <see cref="WebSocketSession"/> scope and signals
     /// that an error occurred that terminated the session.

src/HotChocolate/AspNetCore/src/AspNetCore.Pipeline/Instrumentation/NoopServerDiagnosticEventListener.cs

@@ -1,4 +1,6 @@
 using Microsoft.AspNetCore.Http;
+using HotChocolate.AspNetCore.Subscriptions;
+using HotChocolate.AspNetCore.Subscriptions.Protocols;
 using HotChocolate.Language;
 using static HotChocolate.AspNetCore.Instrumentation.ServerDiagnosticEventListener;
 
@@ -41,6 +43,12 @@ public void ParserErrors(HttpContext context, IReadOnlyList<IError> errors)
 
     public IDisposable WebSocketSession(HttpContext context) => EmptyScope;
 
+    public void WebSocketConnectionInitialized(
+        ISocketSession session,
+        IOperationMessagePayload connectionInitMessage)
+    {
+    }
+
     public void WebSocketSessionError(HttpContext context, Exception exception)
     {
     }

src/HotChocolate/AspNetCore/src/AspNetCore.Pipeline/Instrumentation/ServerDiagnosticEventListener.cs

@@ -1,4 +1,6 @@
 using Microsoft.AspNetCore.Http;
+using HotChocolate.AspNetCore.Subscriptions;
+using HotChocolate.AspNetCore.Subscriptions.Protocols;
 using HotChocolate.Language;
 
 namespace HotChocolate.AspNetCore.Instrumentation;
@@ -64,6 +66,13 @@ public virtual IDisposable FormatHttpResponse(HttpContext context, OperationResu
     public virtual IDisposable WebSocketSession(HttpContext context)
         => EmptyScope;
 
+    /// <inheritdoc />
+    public virtual void WebSocketConnectionInitialized(
+        ISocketSession session,
+        IOperationMessagePayload connectionInitMessage)
+    {
+    }
+
     /// <inheritdoc />
     public virtual void WebSocketSessionError(HttpContext context, Exception exception)
     {

src/HotChocolate/AspNetCore/src/AspNetCore.Pipeline/Subscriptions/Protocols/Apollo/ApolloSubscriptionProtocolHandler.cs

@@ -3,6 +3,7 @@
 using System.Runtime.InteropServices;
 using System.Text.Json;
 using HotChocolate.AspNetCore.Formatters;
+using HotChocolate.AspNetCore.Instrumentation;
 using HotChocolate.Buffers;
 using HotChocolate.Language;
 using HotChocolate.Text.Json;
@@ -18,13 +19,16 @@ internal sealed class ApolloSubscriptionProtocolHandler : IProtocolHandler
 {
     private readonly ISocketSessionInterceptor _interceptor;
     private readonly IWebSocketPayloadFormatter _formatter;
+    private readonly IServerDiagnosticEvents _diagnosticEvents;
 
     public ApolloSubscriptionProtocolHandler(
         ISocketSessionInterceptor interceptor,
-        IWebSocketPayloadFormatter formatter)
+        IWebSocketPayloadFormatter formatter,
+        IServerDiagnosticEvents diagnosticEvents)
     {
         _interceptor = interceptor;
         _formatter = formatter;
+        _diagnosticEvents = diagnosticEvents;
     }
 
     public string Name => GraphQL_WS;
@@ -113,6 +117,7 @@ await _interceptor.OnConnectAsync(
             if (connectionStatus.Accepted)
             {
                 ((WebSocketConnection)connection).IsConnected = true;
+                _diagnosticEvents.WebSocketConnectionInitialized(session, operationMessageObj);
                 await SendConnectionAcceptMessage(
                     session,
                     connectionStatus.Extensions,

src/HotChocolate/AspNetCore/src/AspNetCore.Pipeline/Subscriptions/Protocols/GraphQLOverWebSocket/GraphQLOverWebSocketProtocolHandler.cs

@@ -3,6 +3,7 @@
 using System.Runtime.InteropServices;
 using System.Text.Json;
 using HotChocolate.AspNetCore.Formatters;
+using HotChocolate.AspNetCore.Instrumentation;
 using HotChocolate.Buffers;
 using HotChocolate.Language;
 using HotChocolate.AspNetCore.Utilities;
@@ -19,7 +20,8 @@ internal sealed class GraphQLOverWebSocketProtocolHandler(
     IWebSocketPayloadFormatter formatter,
     IDocumentCache documentCache,
     IDocumentHashProvider documentHashProvider,
-    ParserOptions parserOptions)
+    ParserOptions parserOptions,
+    IServerDiagnosticEvents diagnosticEvents)
     : IGraphQLOverWebSocketProtocolHandler
 {
     public string Name => GraphQL_Transport_WS;
@@ -128,6 +130,7 @@ await interceptor.OnConnectAsync(
             if (connectionStatus.Accepted)
             {
                 ((WebSocketConnection)connection).IsConnected = true;
+                diagnosticEvents.WebSocketConnectionInitialized(session, operationMessageObj);
                 await SendConnectionAcceptMessage(
                     session,
                     connectionStatus.Extensions,

src/HotChocolate/AspNetCore/src/AspNetCore/Extensions/HotChocolateAspNetCoreServiceCollectionExtensions.Subscriptions.cs

@@ -2,6 +2,7 @@
 using Microsoft.Extensions.DependencyInjection.Extensions;
 using HotChocolate.AspNetCore;
 using HotChocolate.AspNetCore.Formatters;
+using HotChocolate.AspNetCore.Instrumentation;
 using HotChocolate.AspNetCore.Subscriptions.Protocols;
 using HotChocolate.AspNetCore.Subscriptions.Protocols.Apollo;
 using HotChocolate.AspNetCore.Subscriptions.Protocols.GraphQLOverWebSocket;
@@ -86,7 +87,8 @@ private static IRequestExecutorBuilder AddApolloProtocol(
             s => s.AddSingleton<IProtocolHandler>(
                 sp => new ApolloSubscriptionProtocolHandler(
                     sp.GetRequiredService<ISocketSessionInterceptor>(),
-                    sp.GetRequiredService<IWebSocketPayloadFormatter>())));
+                    sp.GetRequiredService<IWebSocketPayloadFormatter>(),
+                    sp.GetRequiredService<IServerDiagnosticEvents>())));
 
     private static IRequestExecutorBuilder AddGraphQLOverWebSocketProtocol(
         this IRequestExecutorBuilder builder)
@@ -97,7 +99,8 @@ private static IRequestExecutorBuilder AddGraphQLOverWebSocketProtocol(
                     sp.GetRequiredService<IWebSocketPayloadFormatter>(),
                     sp.GetRequiredService<IDocumentCache>(),
                     sp.GetRequiredService<IDocumentHashProvider>(),
-                    sp.GetRequiredService<ParserOptions>())));
+                    sp.GetRequiredService<ParserOptions>(),
+                    sp.GetRequiredService<IServerDiagnosticEvents>())));
 
     /// <summary>
     /// Adds a custom WebSocket payload formatter to the DI.

src/HotChocolate/Core/src/Types/Types/Descriptors/Conventions/DescriptorContext.cs

@@ -6,6 +6,7 @@
 using HotChocolate.Resolvers;
 using HotChocolate.Types.Relay;
 using HotChocolate.Utilities;
+using EmptyServiceProvider = HotChocolate.Utilities.EmptyServiceProvider;
 using ThrowHelper = HotChocolate.Utilities.ThrowHelper;
 
 namespace HotChocolate.Types.Descriptors;

src/HotChocolate/Core/test/Types.Tests/Resolvers/ResolverCompilerTests.cs

@@ -9,8 +9,8 @@
 using HotChocolate.Tests;
 using HotChocolate.Types;
 using HotChocolate.Types.Descriptors;
-using HotChocolate.Utilities;
 using Moq;
+using EmptyServiceProvider = HotChocolate.Utilities.EmptyServiceProvider;
 
 namespace HotChocolate.Resolvers;
 

src/HotChocolate/Diagnostics/src/Diagnostics.Core/ActivityEnricherBase.cs

@@ -1,6 +1,8 @@
 using System.Diagnostics;
 using Microsoft.AspNetCore.Http;
 using HotChocolate.AspNetCore.Instrumentation;
+using HotChocolate.AspNetCore.Subscriptions;
+using HotChocolate.AspNetCore.Subscriptions.Protocols;
 using HotChocolate.Execution;
 using HotChocolate.Language;
 
@@ -12,6 +14,11 @@ namespace HotChocolate.Diagnostics;
 /// </summary>
 public abstract class ActivityEnricherBase
 {
+    public virtual void OnWebSocketConnectionInitialized(
+        ISocketSession session,
+        IOperationMessagePayload connectionInitMessage)
+    { }
+
     public virtual void EnrichExecuteHttpRequest(
         HttpContext httpContext,
         HttpRequestKind kind,

src/HotChocolate/Diagnostics/src/Diagnostics/Listeners/ActivityExecutionDiagnosticListener.cs

@@ -25,15 +25,18 @@ public override IDisposable ExecuteRequest(RequestContext context)
 
         if (options.SkipExecuteRequest)
         {
-            if (!options.SkipExecuteHttpRequest
-                && context.Features.TryGet<HttpContext>(out var httpContext)
-                && httpContext.Features.Get<ExecuteHttpRequestSpan>() is { } httpRequestSpan)
+            if (options.SkipExecuteHttpRequest
+                || !context.Features.TryGet<HttpContext>(out var httpContext))
             {
-                httpContextActivity = httpRequestSpan.Activity;
+                return EmptyScope;
             }
-            else
+
+            // A GraphQL over HTTP request reuses the transport span as the root request
+            // span instead of emitting a second one. A transport that has no span of its
+            // own, like a WebSocket session, falls through to a dedicated request span.
+            if (httpContext.Features.Get<ExecuteHttpRequestSpan>() is { } httpRequestSpan)
             {
-                return EmptyScope;
+                httpContextActivity = httpRequestSpan.Activity;
             }
         }
 

src/HotChocolate/Diagnostics/src/Diagnostics/Listeners/ActivityServerDiagnosticListener.cs

@@ -1,5 +1,7 @@
 using Microsoft.AspNetCore.Http;
 using HotChocolate.AspNetCore.Instrumentation;
+using HotChocolate.AspNetCore.Subscriptions;
+using HotChocolate.AspNetCore.Subscriptions.Protocols;
 using HotChocolate.Execution;
 using HotChocolate.Language;
 using static HotChocolate.Diagnostics.HotChocolateActivitySource;
@@ -103,6 +105,11 @@ public override void ParserErrors(HttpContext context, IReadOnlyList<IError> err
         }
     }
 
+    public override void WebSocketConnectionInitialized(
+        ISocketSession session,
+        IOperationMessagePayload connectionInitMessage)
+        => enricher.OnWebSocketConnectionInitialized(session, connectionInitMessage);
+
     public override IDisposable FormatHttpResponse(HttpContext context, OperationResult result)
     {
         if (options.SkipFormatHttpResponse)