Add initial common infrastructure for wiring up distributed tracing

sdk/core/Azure.Core/src/Azure.Core.csproj

@@ -20,6 +20,7 @@
   <ItemGroup>
     <PackageReference Include="Microsoft.Bcl.AsyncInterfaces" />
     <PackageReference Include="System.Buffers" />
+    <PackageReference Include="System.Diagnostics.DiagnosticSource" />
     <PackageReference Include="System.Memory" />
     <PackageReference Include="System.Numerics.Vectors" />
     <PackageReference Include="System.Runtime.CompilerServices.Unsafe" />

sdk/core/Azure.Core/src/Pipeline/AzureOperationScope.cs

@@ -0,0 +1,86 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Diagnostics;
+
+namespace Azure.Core.Pipeline
+{
+    public struct DiagnosticScope: IDisposable
+    {
+        private Activity _activity;
+
+        private readonly string _name;
+
+        private readonly DiagnosticListener _source;
+
+        internal DiagnosticScope(string name, DiagnosticListener source)
+        {
+            _name = name;
+            _source = source;
+            _activity = _source.IsEnabled() ? new Activity(_name) : null;
+        }
+
+        public bool IsEnabled => _activity != null;
+
+        public void AddAttribute(string name, string value)
+        {
+            _activity?.AddTag(name, value);
+        }
+
+        public void AddAttribute<T>(string name, T value)
+        {
+            if (_activity != null)
+            {
+                AddAttribute(name, value.ToString());
+            }
+        }
+
+        public void AddAttribute<T>(string name, T value, Func<T, string> format)
+        {
+            if (_activity != null)
+            {
+                AddAttribute(name, format(value));
+            }
+        }
+
+        public void Start()
+        {
+            if (_activity != null && _source.IsEnabled(_name))
+            {
+                _source.StartActivity(_activity, null);
+            }
+        }
+
+        public void Dispose()
+        {
+            if (_activity == null)
+            {
+                return;
+            }
+
+            if (_source != null)
+            {
+                _source.StopActivity(_activity, null);
+            }
+            else
+            {
+                _activity?.Stop();
+            }
+        }
+
+        public void Failed(Exception e)
+        {
+            if (_activity == null)
+            {
+                return;
+            }
+
+            _source?.Write(_activity.OperationName + ".Exception", e);
+            _activity?.Stop();
+
+            _activity = null;
+
+        }
+    }
+}

sdk/core/Azure.Core/src/Pipeline/HttpPipeline.cs

@@ -15,11 +15,13 @@ public class HttpPipeline
         private readonly ResponseClassifier _responseClassifier;
         private readonly ReadOnlyMemory<HttpPipelinePolicy> _pipeline;
 
-        public HttpPipeline(HttpPipelineTransport transport, HttpPipelinePolicy[] policies = null, ResponseClassifier responseClassifier = null)
+        public HttpPipeline(HttpPipelineTransport transport, HttpPipelinePolicy[] policies = null, ResponseClassifier responseClassifier = null, ClientDiagnostics clientDiagnostics = null)
         {
             _transport = transport ?? throw new ArgumentNullException(nameof(transport));
             _responseClassifier = responseClassifier ?? new ResponseClassifier();
 
+            Diagnostics = clientDiagnostics ?? new ClientDiagnostics(true);
+
             policies = policies ?? Array.Empty<HttpPipelinePolicy>();
 
             var all = new HttpPipelinePolicy[policies.Length + 1];
@@ -32,6 +34,8 @@ public HttpPipeline(HttpPipelineTransport transport, HttpPipelinePolicy[] polici
         public Request CreateRequest()
             => _transport.CreateRequest();
 
+        public ClientDiagnostics Diagnostics { get; }
+
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public async Task<Response> SendRequestAsync(Request request, CancellationToken cancellationToken)
         {

sdk/core/Azure.Core/src/Pipeline/HttpPipelineBuilder.cs

@@ -40,9 +40,11 @@ public static HttpPipeline Build(ClientOptions options, bool bufferResponse = tr
                 policies.Add(BufferResponsePolicy.Shared);
             }
 
+            policies.Add(new RequestActivityPolicy());
+
             policies.RemoveAll(policy => policy == null);
 
-            return new HttpPipeline(options.Transport, policies.ToArray(), options.ResponseClassifier);
+            return new HttpPipeline(options.Transport, policies.ToArray(), options.ResponseClassifier, new ClientDiagnostics(options.Diagnostics.IsLoggingEnabled));
         }
 
         // internal for testing

sdk/core/Azure.Core/src/Pipeline/HttpPipelineDiagnostics.cs

@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System.Diagnostics;
+
+namespace Azure.Core.Pipeline
+{
+    public sealed class ClientDiagnostics
+    {
+        private readonly bool _isActivityEnabled;
+
+        public ClientDiagnostics(bool isActivityEnabled)
+        {
+            _isActivityEnabled = isActivityEnabled;
+        }
+
+        private static readonly DiagnosticListener s_source = new DiagnosticListener("Azure");
+
+        public DiagnosticScope CreateScope(string name)
+        {
+            if (!_isActivityEnabled)
+            {
+                return default;
+            }
+            return new DiagnosticScope(name, s_source);
+        }
+    }
+}

sdk/core/Azure.Core/src/Pipeline/Policies/RequestActivityPolicy.cs

@@ -0,0 +1,83 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Diagnostics;
+using System.Globalization;
+using System.Threading.Tasks;
+
+namespace Azure.Core.Pipeline.Policies
+{
+    internal class RequestActivityPolicy : HttpPipelinePolicy
+    {
+        public static RequestActivityPolicy Shared { get; } = new RequestActivityPolicy();
+
+        private static readonly DiagnosticListener s_diagnosticSource = new DiagnosticListener("Azure.Pipeline");
+
+        public override Task ProcessAsync(HttpPipelineMessage message, ReadOnlyMemory<HttpPipelinePolicy> pipeline)
+        {
+            return ProcessAsync(message, pipeline, true);
+        }
+
+        public override void Process(HttpPipelineMessage message, ReadOnlyMemory<HttpPipelinePolicy> pipeline)
+        {
+            ProcessAsync(message, pipeline, false).EnsureCompleted();
+        }
+
+        private static async Task ProcessAsync(HttpPipelineMessage message, ReadOnlyMemory<HttpPipelinePolicy> pipeline, bool isAsync)
+        {
+            if (!s_diagnosticSource.IsEnabled())
+            {
+                if (isAsync)
+                {
+                    await ProcessNextAsync(message, pipeline).ConfigureAwait(false);
+                }
+                else
+                {
+                    ProcessNext(message, pipeline);
+                }
+
+                return;
+            }
+
+            var activity = new Activity("Azure.Core.Http.Request");
+            activity.AddTag("http.method", message.Request.Method.Method);
+            activity.AddTag("http.url", message.Request.UriBuilder.ToString());
+            if (message.Request.Headers.TryGetValue("User-Agent", out string userAgent))
+            {
+                activity.AddTag("http.user_agent", userAgent);
+            }
+
+            var diagnosticSourceActivityEnabled = s_diagnosticSource.IsEnabled(activity.OperationName);
+
+            if (diagnosticSourceActivityEnabled)
+            {
+                s_diagnosticSource.StartActivity(activity, null);
+            }
+            else
+            {
+                activity.Start();
+            }
+
+            if (isAsync)
+            {
+                await ProcessNextAsync(message, pipeline).ConfigureAwait(false);
+            }
+            else
+            {
+                ProcessNext(message, pipeline);
+            }
+
+            activity.AddTag("http.status_code", message.Response.Status.ToString(CultureInfo.InvariantCulture));
+
+            if (diagnosticSourceActivityEnabled)
+            {
+                s_diagnosticSource.StopActivity(activity, null);
+            }
+            else
+            {
+                activity.Stop();
+            }
+        }
+    }
+}

sdk/core/Azure.Core/tests/ClientDiagnosticsTests.cs

@@ -0,0 +1,91 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics;
+using Azure.Core.Pipeline;
+using NUnit.Framework;
+
+namespace Azure.Core.Tests
+{
+    public class ClientDiagnosticsTests
+    {
+        [Test]
+        public void CreatesActivityWithNameAndTags()
+        {
+
+            using var testListener = new TestDiagnosticListener("Azure");
+            ClientDiagnostics clientDiagnostics = new ClientDiagnostics(true);
+
+            DiagnosticScope scope = clientDiagnostics.CreateScope("ActivityName");
+
+            scope.AddAttribute("Attribute1", "Value1");
+            scope.AddAttribute("Attribute2", 2, i => i.ToString());
+            scope.AddAttribute("Attribute3", 3);
+
+            scope.Start();
+
+            KeyValuePair<string, object> startEvent = testListener.Events.Dequeue();
+
+            Activity activity = Activity.Current;
+
+            scope.Dispose();
+
+            KeyValuePair<string, object> stopEvent = testListener.Events.Dequeue();
+
+            Assert.Null(Activity.Current);
+            Assert.AreEqual("ActivityName.Start", startEvent.Key);
+            Assert.AreEqual("ActivityName.Stop", stopEvent.Key);
+
+            CollectionAssert.Contains(activity.Tags, new KeyValuePair<string, string>("Attribute1", "Value1"));
+            CollectionAssert.Contains(activity.Tags, new KeyValuePair<string, string>("Attribute2", "2"));
+            CollectionAssert.Contains(activity.Tags, new KeyValuePair<string, string>("Attribute3", "3"));
+        }
+
+        [Test]
+        public void FailedStopsActivityAndWritesExceptionEvent()
+        {
+            using var testListener = new TestDiagnosticListener("Azure");
+            ClientDiagnostics clientDiagnostics = new ClientDiagnostics(true);
+
+            DiagnosticScope scope = clientDiagnostics.CreateScope("ActivityName");
+
+            scope.AddAttribute("Attribute1", "Value1");
+            scope.AddAttribute("Attribute2", 2, i => i.ToString());
+
+            scope.Start();
+
+            KeyValuePair<string, object> startEvent = testListener.Events.Dequeue();
+
+            Activity activity = Activity.Current;
+
+            var exception = new Exception();
+            scope.Failed(exception);
+            scope.Dispose();
+
+            KeyValuePair<string, object> stopEvent = testListener.Events.Dequeue();
+
+            Assert.Null(Activity.Current);
+            Assert.AreEqual("ActivityName.Start", startEvent.Key);
+            Assert.AreEqual("ActivityName.Exception", stopEvent.Key);
+            Assert.AreEqual(exception, stopEvent.Value);
+            Assert.AreEqual(0, testListener.Events.Count);
+
+            CollectionAssert.Contains(activity.Tags, new KeyValuePair<string, string>("Attribute1", "Value1"));
+            CollectionAssert.Contains(activity.Tags, new KeyValuePair<string, string>("Attribute2", "2"));
+        }
+
+        [Test]
+        public void NoopsWhenDisabled()
+        {
+            ClientDiagnostics clientDiagnostics = new ClientDiagnostics(false);
+            DiagnosticScope scope = clientDiagnostics.CreateScope("");
+
+            scope.AddAttribute("Attribute1", "Value1");
+            scope.AddAttribute("Attribute2", 2, i => i.ToString());
+            scope.Failed(new Exception());
+            scope.Dispose();
+        }
+    }
+}

sdk/core/Azure.Core/tests/RequestActivityPolicyTests.cs

@@ -0,0 +1,58 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics;
+using System.Threading.Tasks;
+using Azure.Core.Http;
+using Azure.Core.Pipeline;
+using Azure.Core.Pipeline.Policies;
+using Azure.Core.Testing;
+using NUnit.Framework;
+
+namespace Azure.Core.Tests
+{
+    public class RequestActivityPolicyTests : SyncAsyncPolicyTestBase
+    {
+        public RequestActivityPolicyTests(bool isAsync) : base(isAsync)
+        {
+        }
+
+        [Test]
+        [NonParallelizable]
+        public async Task ActivityIsCreatedForRequest()
+        {
+            Activity activity = null;
+            KeyValuePair<string, object> startEvent = default;
+            using var testListener = new TestDiagnosticListener("Azure.Pipeline");
+
+            MockTransport mockTransport = CreateMockTransport(_ =>
+            {
+                activity = Activity.Current;
+                startEvent = testListener.Events.Dequeue();
+                return new MockResponse(201);
+            });
+
+            using Request request = mockTransport.CreateRequest();
+            request.Method = RequestMethod.Get;
+            request.UriBuilder.Uri = new Uri("http://example.com");
+            request.Headers.Add("User-Agent", "agent");
+
+            Task<Response> requestTask = SendRequestAsync(mockTransport, request, RequestActivityPolicy.Shared);
+
+            await requestTask;
+
+            KeyValuePair<string, object> stopEvent = testListener.Events.Dequeue();
+
+            Assert.AreEqual("Azure.Core.Http.Request.Start", startEvent.Key);
+            Assert.AreEqual("Azure.Core.Http.Request.Stop", stopEvent.Key);
+
+            Assert.AreEqual("Azure.Core.Http.Request", activity.OperationName);
+            CollectionAssert.Contains(activity.Tags, new KeyValuePair<string, string>("http.status_code", "201"));
+            CollectionAssert.Contains(activity.Tags, new KeyValuePair<string, string>("http.url", "http://example.com/"));
+            CollectionAssert.Contains(activity.Tags, new KeyValuePair<string, string>("http.method", "GET"));
+            CollectionAssert.Contains(activity.Tags, new KeyValuePair<string, string>("http.user_agent", "agent"));
+        }
+    }
+}