Merge pull request #15 from nenoNaninu/improve_filter

improve HubInstrumentationFilter
nenoNaninu · Feb 3, 2024 · a5a9ef9 · a5a9ef9
2 parents 172d193 + ee1c685
commit a5a9ef9
Showing 1 changed file with 47 additions and 11 deletions.
diff --git a/src/AspNetCore.SignalR.OpenTelemetry/HubInstrumentationFilter.cs b/src/AspNetCore.SignalR.OpenTelemetry/HubInstrumentationFilter.cs
@@ -49,32 +49,68 @@ public HubInstrumentationFilter(ILoggerFactory loggerFactory)
         }
     }
 
-    public Task OnConnectedAsync(HubLifetimeContext context, Func<HubLifetimeContext, Task> next)
+    public async Task OnConnectedAsync(HubLifetimeContext context, Func<HubLifetimeContext, Task> next)
     {
         var hubName = context.Hub.GetType().Name;
 
-        var transport = context.Context.Features.Get<IHttpTransportFeature>();
-        HubLogger.LogOnConnected(_logger, hubName, transport?.TransportType ?? HttpTransportType.None);
+        using var scope = HubLogger.BeginHubMethodInvocationScope(_logger, hubName, nameof(OnConnectedAsync));
+        using var activity = HubActivitySource.StartInvocationActivity(hubName, nameof(OnConnectedAsync));
 
-        return next(context);
+        try
+        {
+            var transport = context.Context.Features.Get<IHttpTransportFeature>();
+            HubLogger.LogOnConnected(_logger, hubName, transport?.TransportType ?? HttpTransportType.None);
+
+            var stopwatch = ValueStopwatch.StartNew();
+
+            await next(context);
+
+            var duration = stopwatch.GetElapsedTime();
+
+            HubLogger.LogHubMethodInvocationDuration(_logger, duration.TotalMilliseconds);
+            HubActivitySource.StopInvocationActivityOk(activity);
+        }
+        catch (Exception exception)
+        {
+            HubActivitySource.StopInvocationActivityError(activity, exception);
+            throw;
+        }
     }
 
-    public Task OnDisconnectedAsync(
+    public async Task OnDisconnectedAsync(
         HubLifetimeContext context,
         Exception? exception,
         Func<HubLifetimeContext, Exception?, Task> next)
     {
         var hubName = context.Hub.GetType().Name;
 
-        if (exception is null)
+        using var scope = HubLogger.BeginHubMethodInvocationScope(_logger, hubName, nameof(OnDisconnectedAsync));
+        using var activity = HubActivitySource.StartInvocationActivity(hubName, nameof(OnDisconnectedAsync));
+
+        try
         {
-            HubLogger.LogOnDisconnected(_logger, hubName);
+            if (exception is null)
+            {
+                HubLogger.LogOnDisconnected(_logger, hubName);
+            }
+            else
+            {
+                HubLogger.LogOnDisconnectedWithError(_logger, hubName, exception);
+            }
+
+            var stopwatch = ValueStopwatch.StartNew();
+
+            await next(context, exception);
+
+            var duration = stopwatch.GetElapsedTime();
+
+            HubLogger.LogHubMethodInvocationDuration(_logger, duration.TotalMilliseconds);
+            HubActivitySource.StopInvocationActivityOk(activity);
         }
-        else
+        catch (Exception ex)
         {
-            HubLogger.LogOnDisconnectedWithError(_logger, hubName, exception);
+            HubActivitySource.StopInvocationActivityError(activity, ex);
+            throw;
         }
-
-        return next(context, exception);
     }
 }