Merge pull request #1404 from vivlimmsft/vivlim/activity-tracing-earlier

Move ApplyInboundActivity earlier so dispatch traces have activity IDs

Author: AArnott

src/StreamJsonRpc/JsonRpc.cs

@@ -1903,7 +1903,6 @@ protected virtual async ValueTask<JsonRpcMessage> DispatchRequestAsync(JsonRpcRe
         Requires.NotNull(targetMethod);
 
         object? result;
-        using IDisposable? activityTracingState = this.ActivityTracingStrategy?.ApplyInboundActivity(request);
 
         try
         {
@@ -2807,10 +2806,17 @@ private async Task HandleRpcAsync(JsonRpcMessage rpc)
     {
         Requires.NotNull(rpc, nameof(rpc));
         OutstandingCallData? data = null;
+        IDisposable? activityTracingState = null;
         try
         {
             if (rpc is JsonRpcRequest request)
             {
+                // Set up activity tracing before any dispatch-related traces so that
+                // RequestReceived, LocalInvocation, and other server-side trace events
+                // are emitted under the correct activity ID for cross-process correlation
+                // in tools like SvcTraceViewer.
+                activityTracingState = this.ActivityTracingStrategy?.ApplyInboundActivity(request);
+
                 if (this.TraceSource.Switch.ShouldTrace(TraceEventType.Information))
                 {
                     if (request.IsResponseExpected)
@@ -2971,6 +2977,10 @@ private async Task HandleRpcAsync(JsonRpcMessage rpc)
             // If we extracted this callback from the collection already, take care to complete it to avoid hanging our client.
             data?.CompletionHandler(null);
         }
+        finally
+        {
+            activityTracingState?.Dispose();
+        }
     }
 
     private void FaultPendingRequests()

test/StreamJsonRpc.Tests/CollectingTraceListener.cs

@@ -17,6 +17,8 @@ public class CollectingTraceListener : TraceListener
 
     private readonly ImmutableList<(Guid RelatedActivityId, Guid CurrentActivityId)>.Builder transfers = ImmutableList.CreateBuilder<(Guid RelatedActivityId, Guid CurrentActivityId)>();
 
+    private readonly ImmutableList<(JsonRpc.TraceEvents TraceEventId, Guid ActivityId)>.Builder eventsWithActivityIds = ImmutableList.CreateBuilder<(JsonRpc.TraceEvents TraceEventId, Guid ActivityId)>();
+
     public override bool IsThreadSafe => false;
 
     public IReadOnlyList<string> Messages
@@ -63,6 +65,17 @@ public ImmutableList<JsonRpc.TraceEvents> Ids
         }
     }
 
+    public ImmutableList<(JsonRpc.TraceEvents TraceEventId, Guid ActivityId)> EventsWithActivityIds
+    {
+        get
+        {
+            lock (this.eventsWithActivityIds)
+            {
+                return this.eventsWithActivityIds.ToImmutable();
+            }
+        }
+    }
+
     public AsyncAutoResetEvent MessageReceived { get; } = new AsyncAutoResetEvent();
 
     public override void TraceTransfer(TraceEventCache? eventCache, string source, int id, string? message, Guid relatedActivityId)
@@ -87,6 +100,11 @@ public override void TraceEvent(TraceEventCache? eventCache, string source, Trac
             this.events.Add((eventType, format is null ? null : string.Format(CultureInfo.InvariantCulture, format, args ?? Array.Empty<object?>())));
         }
 
+        lock (this.eventsWithActivityIds)
+        {
+            this.eventsWithActivityIds.Add(((JsonRpc.TraceEvents)id, Trace.CorrelationManager.ActivityId));
+        }
+
         base.TraceEvent(eventCache, source, eventType, id, format, args);
     }
 
@@ -102,6 +120,11 @@ public override void TraceEvent(TraceEventCache? eventCache, string source, Trac
             this.events.Add((eventType, message));
         }
 
+        lock (this.eventsWithActivityIds)
+        {
+            this.eventsWithActivityIds.Add(((JsonRpc.TraceEvents)id, Trace.CorrelationManager.ActivityId));
+        }
+
         base.TraceEvent(eventCache, source, eventType, id, message);
     }
 
@@ -117,6 +140,11 @@ public override void TraceEvent(TraceEventCache? eventCache, string source, Trac
             this.events.Add((eventType, null));
         }
 
+        lock (this.eventsWithActivityIds)
+        {
+            this.eventsWithActivityIds.Add(((JsonRpc.TraceEvents)id, Trace.CorrelationManager.ActivityId));
+        }
+
         base.TraceEvent(eventCache, source, eventType, id);
     }
 

test/StreamJsonRpc.Tests/JsonRpcTests.cs

@@ -2801,6 +2801,52 @@ public async Task CorrelationManagerActivitiesPropagate(string? traceState)
         }
     }
 
+    [Fact]
+    public async Task CorrelationManagerActivity_ServerDispatchTracesHaveActivityId()
+    {
+        // Verify that all server-side StreamJsonRpc traces emitted during RPC dispatch
+        // (RequestReceived, LocalInvocation) have a non-empty activity ID that matches
+        // the activity ID seen by the server method during execution.
+        var strategy = new CorrelationManagerTracingStrategy
+        {
+            TraceSource = new TraceSource("strategy", SourceLevels.ActivityTracing | SourceLevels.Information),
+        };
+        this.clientRpc.AllowModificationWhileListening = true;
+        this.clientRpc.ActivityTracingStrategy = strategy;
+        this.serverRpc.AllowModificationWhileListening = true;
+        this.serverRpc.ActivityTracingStrategy = strategy;
+
+        Guid clientActivityId = Guid.NewGuid();
+        Trace.CorrelationManager.ActivityId = clientActivityId;
+        Guid serverActivityId;
+        try
+        {
+            serverActivityId = await this.clientRpc.InvokeWithCancellationAsync<Guid>(nameof(Server.GetCorrelationManagerActivityId), cancellationToken: this.TimeoutToken);
+        }
+        finally
+        {
+            Trace.CorrelationManager.ActivityId = Guid.Empty;
+        }
+
+        // The server should have seen a non-empty, distinct activity ID.
+        Assert.NotEqual(Guid.Empty, serverActivityId);
+        Assert.NotEqual(clientActivityId, serverActivityId);
+
+        // Every dispatch-related trace on the server should have the same activity ID
+        // that the server method observed during execution.
+        var serverDispatchTraces = this.serverTraces.EventsWithActivityIds
+            .Where(e => e.TraceEventId == JsonRpc.TraceEvents.RequestReceived
+                     || e.TraceEventId == JsonRpc.TraceEvents.LocalInvocation)
+            .ToList();
+
+        Assert.NotEmpty(serverDispatchTraces);
+
+        foreach (var (traceEventId, activityId) in serverDispatchTraces)
+        {
+            Assert.Equal(serverActivityId, activityId);
+        }
+    }
+
     [Theory]
     [InlineData(null)]
     [InlineData("")]