Propagate trace parent to SignalR hub invocations (#57049)

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,24 @@ private void RecordRequestStartMetrics(HttpContext httpContext)
         hasDiagnosticListener = false;
 
         var headers = httpContext.Request.Headers;
-        _propagator.ExtractTraceIdAndState(headers,
+        var activity = ActivityCreator.CreateFromRemote(
+            _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,
+            ActivityKind.Server,
+            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,123 @@
+// 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? CreateFromRemote(
+        ActivitySource activitySource,
+        DistributedContextPropagator propagator,
+        object distributedContextCarrier,
+        DistributedContextPropagator.PropagatorGetterCallback propagatorGetter,
+        string activityName,
+        ActivityKind kind,
+        IEnumerable<KeyValuePair<string, object?>>? tags,
+        IEnumerable<ActivityLink>? links,
+        bool diagnosticsOrLoggingEnabled)
+    {
+        Activity? activity = null;
+        string? requestId = null;
+        string? traceState = null;
+
+        if (activitySource.HasListeners())
+        {
+            propagator.ExtractTraceIdAndState(
+                distributedContextCarrier,
+                propagatorGetter,
+                out requestId,
+                out traceState);
+
+            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, kind, context, tags: tags, links: links);
+            }
+            else
+            {
+                // Pass in the ID we got from the headers if there was one.
+                activity = activitySource.CreateActivity(activityName, kind, 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)
+            {
+                // Note that there is a very small chance that propagator has already been called.
+                // Requires that the activity source had listened, but it didn't create an activity.
+                // Can only happen if there is a race between HasListeners and CreateActivity calls,
+                // and someone removing the listener.
+                //
+                // The only negative of calling the propagator twice is a small performance hit.
+                // It's small and unlikely so it's not worth trying to optimize.
+                propagator.ExtractTraceIdAndState(
+                    distributedContextCarrier,
+                    propagatorGetter,
+                    out requestId,
+                    out traceState);
+
+                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
+        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/Shared/SignalR/InProcessTestServer.cs

@@ -27,6 +27,8 @@ public abstract class InProcessTestServer : IAsyncDisposable
 
     public abstract string Url { get; }
 
+    public abstract IServiceProvider Services { get; }
+
     public abstract ValueTask DisposeAsync();
 }
 
@@ -54,6 +56,8 @@ internal override event Action<LogRecord> ServerLogged
 
     public override string Url => _url;
 
+    public override IServiceProvider Services => _host.Services;
+
     public static async Task<InProcessTestServer<TStartup>> StartServer(ILoggerFactory loggerFactory, Action<KestrelServerOptions> configureKestrelServerOptions = null, IDisposable disposable = null)
     {
         var server = new InProcessTestServer<TStartup>(loggerFactory, configureKestrelServerOptions, disposable);

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,25 @@ private async Task InvokeStreamCore(ConnectionState connectionState, string meth
         }
     }
 
+    private static void InjectHeaders(HubInvocationMessage invocationMessage)
+    {
+        // TODO: Change when SignalR client has an activity.
+        // This sends info about the current activity, regardless of the activity source, to the SignalR server.
+        // When SignalR client supports client activities this logic should be updated to only send headers
+        // if the SignalR client activity is created. The goal is to match the behavior of distributed tracing in HttpClient.
+        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();