Propagate trace parent to SignalR hub invocations

Author: JamesNK

src/Hosting/Hosting/src/Internal/HostingApplicationDiagnostics.cs

@@ -9,6 +9,7 @@
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Http.Metadata;
+using Microsoft.AspNetCore.Shared;
 using Microsoft.Extensions.Logging;
 
 namespace Microsoft.AspNetCore.Hosting;
@@ -389,80 +390,23 @@ private void RecordRequestStartMetrics(HttpContext httpContext)
         hasDiagnosticListener = false;
 
         var headers = httpContext.Request.Headers;
-        _propagator.ExtractTraceIdAndState(headers,
+        var activity = ActivityCreator.CreateActivity(
+            _activitySource,
+            _propagator,
+            headers,
             static (object? carrier, string fieldName, out string? fieldValue, out IEnumerable<string>? fieldValues) =>
             {
                 fieldValues = default;
                 var headers = (IHeaderDictionary)carrier!;
                 fieldValue = headers[fieldName];
             },
-            out var requestId,
-            out var traceState);
-
-        Activity? activity = null;
-        if (_activitySource.HasListeners())
-        {
-            if (ActivityContext.TryParse(requestId, traceState, isRemote: true, out ActivityContext context))
-            {
-                // The requestId used the W3C ID format. Unfortunately, the ActivitySource.CreateActivity overload that
-                // takes a string parentId never sets HasRemoteParent to true. We work around that by calling the
-                // ActivityContext overload instead which sets HasRemoteParent to parentContext.IsRemote.
-                // https://github.com/dotnet/aspnetcore/pull/41568#discussion_r868733305
-                activity = _activitySource.CreateActivity(ActivityName, ActivityKind.Server, context);
-            }
-            else
-            {
-                // Pass in the ID we got from the headers if there was one.
-                activity = _activitySource.CreateActivity(ActivityName, ActivityKind.Server, string.IsNullOrEmpty(requestId) ? null! : requestId);
-            }
-        }
-
+            ActivityName,
+            tags: null,
+            links: null,
+            loggingEnabled || diagnosticListenerActivityCreationEnabled);
         if (activity is null)
         {
-            // CreateActivity didn't create an Activity (this is an optimization for the
-            // case when there are no listeners). Let's create it here if needed.
-            if (loggingEnabled || diagnosticListenerActivityCreationEnabled)
-            {
-                activity = new Activity(ActivityName);
-                if (!string.IsNullOrEmpty(requestId))
-                {
-                    activity.SetParentId(requestId);
-                }
-            }
-            else
-            {
-                return null;
-            }
-        }
-
-        // The trace id was successfully extracted, so we can set the trace state
-        // https://www.w3.org/TR/trace-context/#tracestate-header
-        if (!string.IsNullOrEmpty(requestId))
-        {
-            if (!string.IsNullOrEmpty(traceState))
-            {
-                activity.TraceStateString = traceState;
-            }
-        }
-
-        // Baggage can be used regardless of whether a distributed trace id was present on the inbound request.
-        // https://www.w3.org/TR/baggage/#abstract
-        var baggage = _propagator.ExtractBaggage(headers, static (object? carrier, string fieldName, out string? fieldValue, out IEnumerable<string>? fieldValues) =>
-        {
-            fieldValues = default;
-            var headers = (IHeaderDictionary)carrier!;
-            fieldValue = headers[fieldName];
-        });
-
-        // AddBaggage adds items at the beginning  of the list, so we need to add them in reverse to keep the same order as the client
-        // By contract, the propagator has already reversed the order of items so we need not reverse it again
-        // Order could be important if baggage has two items with the same key (that is allowed by the contract)
-        if (baggage is not null)
-        {
-            foreach (var baggageItem in baggage)
-            {
-                activity.AddBaggage(baggageItem.Key, baggageItem.Value);
-            }
+            return null;
         }
 
         _diagnosticListener.OnActivityImport(activity, httpContext);

src/Hosting/Hosting/src/Microsoft.AspNetCore.Hosting.csproj

@@ -18,6 +18,7 @@
     <Compile Include="$(SharedSourceRoot)StaticWebAssets\**\*.cs" LinkBase="StaticWebAssets" />
     <Compile Include="$(SharedSourceRoot)Metrics\MetricsExtensions.cs" />
     <Compile Include="$(SharedSourceRoot)Metrics\MetricsConstants.cs" />
+    <Compile Include="$(SharedSourceRoot)Diagnostics\ActivityCreator.cs" />
   </ItemGroup>
 
   <ItemGroup>

src/Shared/Diagnostics/ActivityCreator.cs

@@ -0,0 +1,114 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Diagnostics;
+
+namespace Microsoft.AspNetCore.Shared;
+
+internal static class ActivityCreator
+{
+    /// <summary>
+    /// Create an activity with details received from a remote source.
+    /// </summary>
+    public static Activity? CreateActivity(
+        ActivitySource activitySource,
+        DistributedContextPropagator propagator,
+        object? distributedContextCarrier,
+        DistributedContextPropagator.PropagatorGetterCallback propagatorGetter,
+        string activityName,
+        IEnumerable<KeyValuePair<string, object?>>? tags,
+        IEnumerable<ActivityLink>? links,
+        bool diagnosticsOrLoggingEnabled)
+    {
+        string? requestId = null;
+        string? traceState = null;
+        if (distributedContextCarrier != null)
+        {
+            propagator.ExtractTraceIdAndState(
+                distributedContextCarrier,
+                propagatorGetter,
+                out requestId,
+                out traceState);
+        }
+
+        Activity? activity = null;
+        if (activitySource.HasListeners())
+        {
+            if (ActivityContext.TryParse(requestId, traceState, isRemote: true, out ActivityContext context))
+            {
+                // The requestId used the W3C ID format. Unfortunately, the ActivitySource.CreateActivity overload that
+                // takes a string parentId never sets HasRemoteParent to true. We work around that by calling the
+                // ActivityContext overload instead which sets HasRemoteParent to parentContext.IsRemote.
+                // https://github.com/dotnet/aspnetcore/pull/41568#discussion_r868733305
+                activity = activitySource.CreateActivity(activityName, ActivityKind.Server, context, tags: tags, links: links);
+            }
+            else
+            {
+                // Pass in the ID we got from the headers if there was one.
+                activity = activitySource.CreateActivity(activityName, ActivityKind.Server, string.IsNullOrEmpty(requestId) ? null! : requestId, tags: tags, links: links);
+            }
+        }
+
+        if (activity is null)
+        {
+            // CreateActivity didn't create an Activity (this is an optimization for the
+            // case when there are no listeners). Let's create it here if needed.
+            if (diagnosticsOrLoggingEnabled)
+            {
+                activity = new Activity(activityName);
+                if (!string.IsNullOrEmpty(requestId))
+                {
+                    activity.SetParentId(requestId);
+                }
+                if (tags != null)
+                {
+                    foreach (var tag in tags)
+                    {
+                        activity.AddTag(tag.Key, tag.Value);
+                    }
+                }
+                if (links != null)
+                {
+                    foreach (var link in links)
+                    {
+                        activity.AddLink(link);
+                    }
+                }
+            }
+            else
+            {
+                return null;
+            }
+        }
+
+        // The trace id was successfully extracted, so we can set the trace state
+        // https://www.w3.org/TR/trace-context/#tracestate-header
+        if (!string.IsNullOrEmpty(requestId))
+        {
+            if (!string.IsNullOrEmpty(traceState))
+            {
+                activity.TraceStateString = traceState;
+            }
+        }
+
+        // Baggage can be used regardless of whether a distributed trace id was present on the inbound request.
+        // https://www.w3.org/TR/baggage/#abstract
+        if (distributedContextCarrier != null)
+        {
+            var baggage = propagator.ExtractBaggage(distributedContextCarrier, propagatorGetter);
+
+            // AddBaggage adds items at the beginning  of the list, so we need to add them in reverse to keep the same order as the client
+            // By contract, the propagator has already reversed the order of items so we need not reverse it again
+            // Order could be important if baggage has two items with the same key (that is allowed by the contract)
+            if (baggage is not null)
+            {
+                foreach (var baggageItem in baggage)
+                {
+                    activity.AddBaggage(baggageItem.Key, baggageItem.Value);
+                }
+            }
+        }
+
+        return activity;
+    }
+}

src/SignalR/clients/csharp/Client.Core/src/HubConnection.cs

@@ -1037,6 +1037,7 @@ private async Task InvokeCore(ConnectionState connectionState, string methodName
 
         // Client invocations are always blocking
         var invocationMessage = new InvocationMessage(irq.InvocationId, methodName, args, streams);
+        InjectHeaders(invocationMessage);
 
         Log.RegisteringInvocation(_logger, irq.InvocationId);
         connectionState.AddInvocation(irq);
@@ -1063,6 +1064,7 @@ private async Task InvokeStreamCore(ConnectionState connectionState, string meth
         Log.PreparingStreamingInvocation(_logger, irq.InvocationId, methodName, irq.ResultType.FullName!, args.Length);
 
         var invocationMessage = new StreamInvocationMessage(irq.InvocationId, methodName, args, streams);
+        InjectHeaders(invocationMessage);
 
         Log.RegisteringInvocation(_logger, irq.InvocationId);
 
@@ -1083,6 +1085,21 @@ private async Task InvokeStreamCore(ConnectionState connectionState, string meth
         }
     }
 
+    private static void InjectHeaders(HubInvocationMessage invocationMessage)
+    {
+        if (Activity.Current is { } currentActivity)
+        {
+            DistributedContextPropagator.Current.Inject(currentActivity, invocationMessage, static (carrier, key, value) =>
+            {
+                if (carrier is HubInvocationMessage invocationMessage)
+                {
+                    invocationMessage.Headers ??= new Dictionary<string, string>();
+                    invocationMessage.Headers[key] = value;
+                }
+            });
+        }
+    }
+
     private async Task SendHubMessage(ConnectionState connectionState, HubMessage hubMessage, CancellationToken cancellationToken = default)
     {
         _state.AssertConnectionValid();

src/SignalR/common/testassets/Tests.Utils/TestClient.cs

@@ -8,6 +8,8 @@
 using Microsoft.AspNetCore.Connections.Features;
 using Microsoft.AspNetCore.Internal;
 using Microsoft.AspNetCore.SignalR.Protocol;
+using System.Reflection.PortableExecutable;
+
 #if TESTUTILS
 using Microsoft.AspNetCore.InternalTesting;
 #else
@@ -99,6 +101,12 @@ public async Task<IList<HubMessage>> StreamAsync(string methodName, string[] str
         return await ListenAllAsync(invocationId);
     }
 
+    public async Task<IList<HubMessage>> StreamAsync(string methodName, string[] streamIds, IDictionary<string, string> headers, params object[] args)
+    {
+        var invocationId = await SendStreamInvocationAsync(methodName, streamIds, headers, args);
+        return await ListenAllAsync(invocationId);
+    }
+
     public async Task<IList<HubMessage>> ListenAllAsync(string invocationId)
     {
         var result = new List<HubMessage>();
@@ -185,10 +193,20 @@ public Task<string> SendInvocationAsync(string methodName, params object[] args)
         return SendInvocationAsync(methodName, nonBlocking: false, args: args);
     }
 
+    public Task<string> SendInvocationAsync(string methodName, IDictionary<string, string> headers, params object[] args)
+    {
+        return SendInvocationAsync(methodName, nonBlocking: false, headers: headers, args: args);
+    }
+
     public Task<string> SendInvocationAsync(string methodName, bool nonBlocking, params object[] args)
+    {
+        return SendInvocationAsync(methodName, nonBlocking: nonBlocking, headers: null, args: args);
+    }
+
+    public Task<string> SendInvocationAsync(string methodName, bool nonBlocking, IDictionary<string, string> headers, params object[] args)
     {
         var invocationId = nonBlocking ? null : GetInvocationId();
-        return SendHubMessageAsync(new InvocationMessage(invocationId, methodName, args));
+        return SendHubMessageAsync(new InvocationMessage(invocationId, methodName, args) { Headers = headers });
     }
 
     public Task<string> SendStreamInvocationAsync(string methodName, params object[] args)
@@ -197,9 +215,14 @@ public Task<string> SendStreamInvocationAsync(string methodName, params object[]
     }
 
     public Task<string> SendStreamInvocationAsync(string methodName, string[] streamIds, params object[] args)
+    {
+        return SendStreamInvocationAsync(methodName, streamIds: streamIds, headers: null, args);
+    }
+
+    public Task<string> SendStreamInvocationAsync(string methodName, string[] streamIds, IDictionary<string, string> headers, params object[] args)
     {
         var invocationId = GetInvocationId();
-        return SendHubMessageAsync(new StreamInvocationMessage(invocationId, methodName, args, streamIds));
+        return SendHubMessageAsync(new StreamInvocationMessage(invocationId, methodName, args, streamIds) { Headers = headers });
     }
 
     public Task<string> BeginUploadStreamAsync(string invocationId, string methodName, string[] streamIds, params object[] args)

src/SignalR/server/Core/src/HubConnectionHandler.cs

@@ -136,10 +136,6 @@ public override async Task OnConnectedAsync(ConnectionContext connection)
             OriginalActivity = Activity.Current,
         };
 
-        // Get off the parent span.
-        // This is likely the Http Request span and we want Hub method invocations to not be collected under a long running span.
-        Activity.Current = null;
-
         var resolvedSupportedProtocols = (supportedProtocols as IReadOnlyList<string>) ?? supportedProtocols.ToList();
         if (!await connectionContext.HandshakeAsync(handshakeTimeout, resolvedSupportedProtocols, _protocolResolver, _userIdProvider, _enableDetailedErrors))
         {