Flow trace header when listener is attached (Azure#7058)

szimmer-dev · Jul 31, 2019 · 2ed4fbc · 2ed4fbc
1 parent 9560df8
commit 2ed4fbc
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 2 deletions.
diff --git a/sdk/core/Azure.Core/src/Pipeline/RequestActivityPolicy.cs b/sdk/core/Azure.Core/src/Pipeline/RequestActivityPolicy.cs
@@ -59,11 +59,11 @@ private static async Task ProcessAsync(HttpPipelineMessage message, ReadOnlyMemo
 
             if (isAsync)
             {
-                await ProcessNextAsync(message, pipeline).ConfigureAwait(false);
+                await ProcessNextAsync(message, pipeline, true).ConfigureAwait(false);
             }
             else
             {
-                ProcessNext(message, pipeline);
+                ProcessNextAsync(message, pipeline, false).EnsureCompleted();
             }
 
             activity.AddTag("http.status_code", message.Response.Status.ToString(CultureInfo.InvariantCulture));

diff --git a/sdk/core/Azure.Core/tests/RequestActivityPolicyTests.cs b/sdk/core/Azure.Core/tests/RequestActivityPolicyTests.cs
@@ -55,6 +55,42 @@ public async Task ActivityIsCreatedForRequest()
             CollectionAssert.Contains(activity.Tags, new KeyValuePair<string, string>("http.user_agent", "agent"));
         }
 
+
+        [Test]
+        [NonParallelizable]
+        public async Task ActivityIdIsStampedOnRequest()
+        {
+            using var testListener = new TestDiagnosticListener("Azure.Pipeline");
+
+            var previousFormat = Activity.DefaultIdFormat;
+            Activity.DefaultIdFormat = ActivityIdFormat.W3C;
+            try
+            {
+                Activity activity = null;
+
+                MockTransport mockTransport = CreateMockTransport(_ =>
+                {
+                    activity = Activity.Current;
+                    return new MockResponse(201);
+                });
+
+                using Request request = mockTransport.CreateRequest();
+                request.Method = RequestMethod.Get;
+                request.UriBuilder.Uri = new Uri("http://example.com");
+
+                Task<Response> requestTask = SendRequestAsync(mockTransport, request, RequestActivityPolicy.Shared);
+
+                await requestTask;
+
+                Assert.True(mockTransport.SingleRequest.TryGetHeader("traceparent", out string requestId));
+                Assert.AreEqual(activity.Id, requestId);
+            }
+            finally
+            {
+                Activity.DefaultIdFormat = previousFormat;
+            }
+        }
+
         [Test]
         [NonParallelizable]
         public async Task CurrentActivityIsInjectedIntoRequest()