Add client span to SignalR .NET client (#57101)

Author: JamesNK

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

@@ -59,6 +59,8 @@ public partial class HubConnection : IAsyncDisposable
     // Default amount of bytes we'll buffer when using Stateful Reconnect until applying backpressure to sends from the client.
     internal const long DefaultStatefulReconnectBufferSize = 100_000;
 
+    internal const string ActivityName = "Microsoft.AspNetCore.SignalR.Client.InvocationOut";
+
     // The receive loop has a single reader and single writer at a time so optimize the channel for that
     private static readonly UnboundedChannelOptions _receiveLoopOptions = new UnboundedChannelOptions
     {
@@ -73,11 +75,13 @@ public partial class HubConnection : IAsyncDisposable
     private readonly ILoggerFactory _loggerFactory;
     private readonly ILogger _logger;
     private readonly ConnectionLogScope _logScope;
+    private readonly ActivitySource _activitySource;
     private readonly IHubProtocol _protocol;
     private readonly IServiceProvider _serviceProvider;
     private readonly IConnectionFactory _connectionFactory;
     private readonly IRetryPolicy? _reconnectPolicy;
     private readonly EndPoint _endPoint;
+    private readonly string? _serviceName;
     private readonly ConcurrentDictionary<string, InvocationHandlerList> _handlers = new ConcurrentDictionary<string, InvocationHandlerList>(StringComparer.Ordinal);
 
     // Holds all mutable state other than user-defined handlers and settable properties.
@@ -235,6 +239,10 @@ public HubConnection(IConnectionFactory connectionFactory,
 
         _logScope = new ConnectionLogScope();
 
+        // ActivitySource can be resolved from the service provider when unit testing.
+        _activitySource = (serviceProvider.GetService<SignalRClientActivitySource>() ?? SignalRClientActivitySource.Instance).ActivitySource;
+        _serviceName = (_endPoint is UriEndPoint e) ? e.Uri.AbsolutePath.Trim('/') : null;
+
         var options = serviceProvider.GetService<IOptions<HubConnectionOptions>>();
 
         ServerTimeout = options?.Value.ServerTimeout ?? DefaultServerTimeout;
@@ -720,7 +728,8 @@ async Task OnStreamCanceled(InvocationRequest irq)
         var readers = default(Dictionary<string, object>);
 
         CheckDisposed();
-        var connectionState = await _state.WaitForActiveConnectionAsync(nameof(StreamAsChannelCoreAsync), token: cancellationToken).ConfigureAwait(false);
+
+        var (connectionState, activity) = await WaitForActiveConnectionWithActivityAsync(nameof(StreamAsChannelCoreAsync), methodName, token: cancellationToken).ConfigureAwait(false);
 
         ChannelReader<object?> channel;
         try
@@ -731,7 +740,7 @@ async Task OnStreamCanceled(InvocationRequest irq)
             readers = PackageStreamingParams(connectionState, ref args, out var streamIds);
 
             // I just want an excuse to use 'irq' as a variable name...
-            var irq = InvocationRequest.Stream(cancellationToken, returnType, connectionState.GetNextId(), _loggerFactory, this, out channel);
+            var irq = InvocationRequest.Stream(cancellationToken, returnType, connectionState.GetNextId(), _loggerFactory, this, activity, out channel);
             await InvokeStreamCore(connectionState, methodName, irq, args, streamIds?.ToArray(), cancellationToken).ConfigureAwait(false);
 
             if (cancellationToken.CanBeCanceled)
@@ -1003,12 +1012,71 @@ private async Task CommonStreaming(ConnectionState connectionState, string strea
         }
     }
 
+    private async Task<(ConnectionState, Activity?)> WaitForActiveConnectionWithActivityAsync(string sendingMethodName, string invokedMethodName, CancellationToken token)
+    {
+        // Start the activity before waiting on the connection.
+        // Starting the activity here means time to connect or reconnect is included in the invoke.
+        var activity = CreateActivity(invokedMethodName);
+
+        try
+        {
+            ConnectionState connectionState;
+            var connectionStateTask = _state.WaitForActiveConnectionAsync(sendingMethodName, token);
+            if (connectionStateTask.Status == TaskStatus.RanToCompletion)
+            {
+                // Attempt to get already connected connection and set server tags using it.
+                connectionState = connectionStateTask.Result;
+                SetServerTags(activity, connectionState.ConnectionUrl);
+                activity?.Start();
+            }
+            else
+            {
+                // Fallback to using configured endpoint.
+                var initialUri = (_endPoint as UriEndPoint)?.Uri;
+                SetServerTags(activity, initialUri);
+                activity?.Start();
+
+                connectionState = await connectionStateTask.ConfigureAwait(false);
+
+                // After connection is returned, check if URL is different. If so, update activity server tags.
+                if (connectionState.ConnectionUrl != null && connectionState.ConnectionUrl != initialUri)
+                {
+                    SetServerTags(activity, connectionState.ConnectionUrl);
+                }
+            }
+
+            return (connectionState, activity);
+        }
+        catch (Exception ex)
+        {
+            // If there is an error getting an active connection then the invocation has failed.
+            if (activity is not null)
+            {
+                activity.SetStatus(ActivityStatusCode.Error);
+                activity.SetTag("error.type", ex.GetType().FullName);
+                activity.Stop();
+            }
+
+            throw;
+        }
+
+        static void SetServerTags(Activity? activity, Uri? uri)
+        {
+            if (activity != null && uri != null)
+            {
+                activity.SetTag("server.address", uri.Host);
+                activity.SetTag("server.port", uri.Port);
+            }
+        }
+    }
+
     private async Task<object?> InvokeCoreAsyncCore(string methodName, Type returnType, object?[] args, CancellationToken cancellationToken)
     {
         var readers = default(Dictionary<string, object>);
 
         CheckDisposed();
-        var connectionState = await _state.WaitForActiveConnectionAsync(nameof(InvokeCoreAsync), token: cancellationToken).ConfigureAwait(false);
+
+        var (connectionState, activity) = await WaitForActiveConnectionWithActivityAsync(nameof(InvokeCoreAsync), methodName, token: cancellationToken).ConfigureAwait(false);
 
         Task<object?> invocationTask;
         try
@@ -1017,7 +1085,7 @@ private async Task CommonStreaming(ConnectionState connectionState, string strea
 
             readers = PackageStreamingParams(connectionState, ref args, out var streamIds);
 
-            var irq = InvocationRequest.Invoke(cancellationToken, returnType, connectionState.GetNextId(), _loggerFactory, this, out invocationTask);
+            var irq = InvocationRequest.Invoke(cancellationToken, returnType, connectionState.GetNextId(), _loggerFactory, this, activity, out invocationTask);
             await InvokeCore(connectionState, methodName, irq, args, streamIds?.ToArray(), cancellationToken).ConfigureAwait(false);
 
             LaunchStreams(connectionState, readers, cancellationToken);
@@ -1031,13 +1099,43 @@ private async Task CommonStreaming(ConnectionState connectionState, string strea
         return await invocationTask.ConfigureAwait(false);
     }
 
+    private Activity? CreateActivity(string methodName)
+    {
+        var activity = _activitySource.CreateActivity(ActivityName, ActivityKind.Client);
+        if (activity is null && Activity.Current is not null && _logger.IsEnabled(LogLevel.Critical))
+        {
+            activity = new Activity(ActivityName);
+        }
+
+        if (activity is not null)
+        {
+            if (!string.IsNullOrEmpty(_serviceName))
+            {
+                activity.DisplayName = $"{_serviceName}/{methodName}";
+                activity.SetTag("rpc.service", _serviceName);
+            }
+            else
+            {
+                activity.DisplayName = methodName;
+            }
+
+            activity.SetTag("rpc.system", "signalr");
+            activity.SetTag("rpc.method", methodName);
+        }
+
+        return activity;
+    }
+
     private async Task InvokeCore(ConnectionState connectionState, string methodName, InvocationRequest irq, object?[] args, string[]? streams, CancellationToken cancellationToken)
     {
         Log.PreparingBlockingInvocation(_logger, irq.InvocationId, methodName, irq.ResultType.FullName!, args.Length);
 
         // Client invocations are always blocking
         var invocationMessage = new InvocationMessage(irq.InvocationId, methodName, args, streams);
-        InjectHeaders(invocationMessage);
+        if (irq.Activity is not null)
+        {
+            InjectHeaders(irq.Activity, invocationMessage);
+        }
 
         Log.RegisteringInvocation(_logger, irq.InvocationId);
         connectionState.AddInvocation(irq);
@@ -1064,7 +1162,10 @@ 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);
+        if (irq.Activity is not null)
+        {
+            InjectHeaders(irq.Activity, invocationMessage);
+        }
 
         Log.RegisteringInvocation(_logger, irq.InvocationId);
 
@@ -1085,23 +1186,16 @@ private async Task InvokeStreamCore(ConnectionState connectionState, string meth
         }
     }
 
-    private static void InjectHeaders(HubInvocationMessage invocationMessage)
+    private static void InjectHeaders(Activity currentActivity, 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) =>
         {
-            DistributedContextPropagator.Current.Inject(currentActivity, invocationMessage, static (carrier, key, value) =>
+            if (carrier is HubInvocationMessage invocationMessage)
             {
-                if (carrier is HubInvocationMessage invocationMessage)
-                {
-                    invocationMessage.Headers ??= new Dictionary<string, string>();
-                    invocationMessage.Headers[key] = value;
-                }
-            });
-        }
+                invocationMessage.Headers ??= new Dictionary<string, string>();
+                invocationMessage.Headers[key] = value;
+            }
+        });
     }
 
     private async Task SendHubMessage(ConnectionState connectionState, HubMessage hubMessage, CancellationToken cancellationToken = default)
@@ -1131,7 +1225,8 @@ private async Task SendCoreAsyncCore(string methodName, object?[] args, Cancella
         var readers = default(Dictionary<string, object>);
 
         CheckDisposed();
-        var connectionState = await _state.WaitForActiveConnectionAsync(nameof(SendCoreAsync), token: cancellationToken).ConfigureAwait(false);
+
+        var (connectionState, activity) = await WaitForActiveConnectionWithActivityAsync(nameof(SendCoreAsync), methodName, token: cancellationToken).ConfigureAwait(false);
         try
         {
             CheckDisposed();
@@ -1140,12 +1235,27 @@ private async Task SendCoreAsyncCore(string methodName, object?[] args, Cancella
 
             Log.PreparingNonBlockingInvocation(_logger, methodName, args.Length);
             var invocationMessage = new InvocationMessage(null, methodName, args, streamIds?.ToArray());
+            if (activity is not null)
+            {
+                InjectHeaders(activity, invocationMessage);
+            }
             await SendHubMessage(connectionState, invocationMessage, cancellationToken).ConfigureAwait(false);
 
             LaunchStreams(connectionState, readers, cancellationToken);
         }
+        catch (Exception ex)
+        {
+            if (activity is not null)
+            {
+                activity.SetStatus(ActivityStatusCode.Error);
+                activity.SetTag("error.type", ex.GetType().FullName);
+                activity.Stop();
+            }
+            throw;
+        }
         finally
         {
+            activity?.Stop();
             _state.ReleaseConnectionLock();
         }
     }
@@ -2018,6 +2128,7 @@ private sealed class ConnectionState : IInvocationBinder
         private long _nextActivationSendPing;
 
         public ConnectionContext Connection { get; }
+        public Uri? ConnectionUrl { get; }
         public Task? ReceiveTask { get; set; }
         public Exception? CloseException { get; set; }
         public CancellationToken UploadStreamToken { get; set; }
@@ -2036,6 +2147,7 @@ public bool Stopping
         public ConnectionState(ConnectionContext connection, HubConnection hubConnection)
         {
             Connection = connection;
+            ConnectionUrl = (connection.RemoteEndPoint is UriEndPoint ep) ? ep.Uri : null;
 
             _hubConnection = hubConnection;
             _hubConnection._logScope.ConnectionId = connection.ConnectionId;