Merge pull request #394 from Azure/dev

Merge dev branch to master for 3.0 preview release

src/Microsoft.Azure.ServiceBus/ClientEntity.cs

@@ -24,7 +24,7 @@ namespace Microsoft.Azure.ServiceBus
         {
             this.clientTypeName = clientTypeName;
             this.ClientId = GenerateClientId(clientTypeName, postfix);
-            this.RetryPolicy = retryPolicy ?? throw new ArgumentNullException(nameof(retryPolicy));
+            this.RetryPolicy = retryPolicy ?? RetryPolicy.Default;
             this.syncLock = new object();
         }
 

src/Microsoft.Azure.ServiceBus/Constants.cs

@@ -36,5 +36,14 @@ namespace Microsoft.Azure.ServiceBus
         public static readonly TimeSpan DefaultRetryDeltaBackoff = TimeSpan.FromSeconds(3);
 
         public static readonly TimeSpan NoMessageBackoffTimeSpan = TimeSpan.FromSeconds(5);
+
+        public const string SasTokenType = "servicebus.windows.net:sastoken";
+
+        public const string JsonWebTokenType = "jwt";
+
+        public const string AadServiceBusAudience = "https://servicebus.azure.net/";
+
+        /// Represents 00:00:00 UTC Thursday 1, January 1970.
+        public static readonly DateTime EpochTime = new DateTime(1970, 1, 1, 0, 0, 0, 0, DateTimeKind.Utc);
     }
 }

src/Microsoft.Azure.ServiceBus/Core/MessageReceiver.cs

@@ -5,6 +5,7 @@ namespace Microsoft.Azure.ServiceBus.Core
 {
     using System;
     using System.Collections.Generic;
+    using System.Diagnostics;
     using System.Linq;
     using System.Threading;
     using System.Threading.Tasks;
@@ -48,6 +49,7 @@ namespace Microsoft.Azure.ServiceBus.Core
         readonly object messageReceivePumpSyncLock;
         readonly bool ownsConnection;
         readonly ActiveClientLinkManager clientLinkManager;
+        readonly ServiceBusDiagnosticSource diagnosticSource;
 
         int prefetchCount;
         long lastPeekedSequenceNumber;
@@ -132,7 +134,7 @@ namespace Microsoft.Azure.ServiceBus.Core
             this.PrefetchCount = prefetchCount;
             this.messageReceivePumpSyncLock = new object();
             this.clientLinkManager = new ActiveClientLinkManager(this.ClientId, this.CbsTokenProvider);
-
+            this.diagnosticSource = new ServiceBusDiagnosticSource(entityPath, serviceBusConnection.Endpoint);
             MessagingEventSource.Log.MessageReceiverCreateStop(serviceBusConnection.Endpoint.Authority, entityPath, this.ClientId);
         }
 
@@ -298,21 +300,36 @@ namespace Microsoft.Azure.ServiceBus.Core
 
             MessagingEventSource.Log.MessageReceiveStart(this.ClientId, maxMessageCount);
 
+            bool isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+            Activity activity = isDiagnosticSourceEnabled ? this.diagnosticSource.ReceiveStart(maxMessageCount) : null;
+            Task receiveTask = null;
+
             IList<Message> unprocessedMessageList = null;
             try
             {
-                await this.RetryPolicy.RunOperation(
+                receiveTask = this.RetryPolicy.RunOperation(
                     async () =>
                     {
-                        unprocessedMessageList = await this.OnReceiveAsync(maxMessageCount, operationTimeout).ConfigureAwait(false);
-                    }, operationTimeout)
-                    .ConfigureAwait(false);
+                        unprocessedMessageList = await this.OnReceiveAsync(maxMessageCount, operationTimeout)
+                            .ConfigureAwait(false);
+                    }, operationTimeout);
+                await receiveTask.ConfigureAwait(false);
+
             }
             catch (Exception exception)
             {
+                if (isDiagnosticSourceEnabled)
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
+
                 MessagingEventSource.Log.MessageReceiveException(this.ClientId, exception);
                 throw;
             }
+            finally
+            {
+                this.diagnosticSource.ReceiveStop(activity, maxMessageCount, receiveTask?.Status, unprocessedMessageList);
+            }
 
             MessagingEventSource.Log.MessageReceiveStop(this.ClientId, unprocessedMessageList?.Count ?? 0);
 
@@ -366,22 +383,34 @@ namespace Microsoft.Azure.ServiceBus.Core
 
             MessagingEventSource.Log.MessageReceiveDeferredMessageStart(this.ClientId, sequenceNumberList.Length, sequenceNumberList);
 
+            bool isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+            Activity activity = isDiagnosticSourceEnabled ? this.diagnosticSource.ReceiveDeferredStart(sequenceNumberList) : null;
+            Task receiveTask = null;
+
             IList<Message> messages = null;
             try
             {
-                await this.RetryPolicy.RunOperation(
+                receiveTask = this.RetryPolicy.RunOperation(
                     async () =>
                     {
                         messages = await this.OnReceiveDeferredMessageAsync(sequenceNumberList).ConfigureAwait(false);
-                    }, this.OperationTimeout)
-                    .ConfigureAwait(false);
+                    }, this.OperationTimeout);
+                await receiveTask.ConfigureAwait(false);
             }
             catch (Exception exception)
             {
+                if (isDiagnosticSourceEnabled)
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
+
                 MessagingEventSource.Log.MessageReceiveDeferredMessageException(this.ClientId, exception);
                 throw;
             }
-
+            finally
+            {
+                this.diagnosticSource.ReceiveDeferredStop(activity, sequenceNumberList, receiveTask?.Status, messages);
+            }
             MessagingEventSource.Log.MessageReceiveDeferredMessageStop(this.ClientId, messages?.Count ?? 0);
 
             return messages;
@@ -425,17 +454,31 @@ namespace Microsoft.Azure.ServiceBus.Core
             }
 
             MessagingEventSource.Log.MessageCompleteStart(this.ClientId, lockTokenList.Count, lockTokenList);
+            bool isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+            Activity activity = isDiagnosticSourceEnabled ? this.diagnosticSource.CompleteStart(lockTokenList) : null;
+            Task completeTask = null;
 
             try
             {
-                await this.RetryPolicy.RunOperation(() => this.OnCompleteAsync(lockTokenList), this.OperationTimeout)
-                    .ConfigureAwait(false);
+                completeTask =
+                    this.RetryPolicy.RunOperation(() => this.OnCompleteAsync(lockTokenList), this.OperationTimeout);
+                await completeTask.ConfigureAwait(false);
             }
             catch (Exception exception)
             {
+                if (isDiagnosticSourceEnabled)
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
+
                 MessagingEventSource.Log.MessageCompleteException(this.ClientId, exception);
+
                 throw;
             }
+            finally
+            {
+                this.diagnosticSource.CompleteStop(activity, lockTokenList, completeTask?.Status);
+            }
 
             MessagingEventSource.Log.MessageCompleteStop(this.ClientId);
         }
@@ -456,16 +499,31 @@ namespace Microsoft.Azure.ServiceBus.Core
             this.ThrowIfNotPeekLockMode();
 
             MessagingEventSource.Log.MessageAbandonStart(this.ClientId, 1, lockToken);
+            bool isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+            Activity activity = isDiagnosticSourceEnabled ? this.diagnosticSource.DisposeStart("Abandon", lockToken) : null;
+            Task abandonTask = null;
+
             try
             {
-                await this.RetryPolicy.RunOperation(() => this.OnAbandonAsync(lockToken, propertiesToModify), this.OperationTimeout)
-                    .ConfigureAwait(false);
+                abandonTask = this.RetryPolicy.RunOperation(() => this.OnAbandonAsync(lockToken, propertiesToModify),
+                    this.OperationTimeout);
+                await abandonTask.ConfigureAwait(false);
             }
             catch (Exception exception)
             {
+                if (isDiagnosticSourceEnabled)
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
+
                 MessagingEventSource.Log.MessageAbandonException(this.ClientId, exception);
                 throw;
             }
+            finally
+            {
+                this.diagnosticSource.DisposeStop(activity, lockToken, abandonTask?.Status);
+            }
+
 
             MessagingEventSource.Log.MessageAbandonStop(this.ClientId);
         }
@@ -487,18 +545,30 @@ namespace Microsoft.Azure.ServiceBus.Core
             this.ThrowIfNotPeekLockMode();
 
             MessagingEventSource.Log.MessageDeferStart(this.ClientId, 1, lockToken);
+            bool isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+            Activity activity = isDiagnosticSourceEnabled ? this.diagnosticSource.DisposeStart("Defer", lockToken) : null;
+            Task deferTask = null;
 
             try
             {
-                await this.RetryPolicy.RunOperation(() => this.OnDeferAsync(lockToken, propertiesToModify), this.OperationTimeout)
-                    .ConfigureAwait(false);
+                deferTask = this.RetryPolicy.RunOperation(() => this.OnDeferAsync(lockToken, propertiesToModify),
+                    this.OperationTimeout);
+                await deferTask.ConfigureAwait(false);
             }
             catch (Exception exception)
             {
+                if (isDiagnosticSourceEnabled)
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
+
                 MessagingEventSource.Log.MessageDeferException(this.ClientId, exception);
                 throw;
             }
-
+            finally
+            {
+                this.diagnosticSource.DisposeStop(activity, lockToken, deferTask?.Status);
+            }
             MessagingEventSource.Log.MessageDeferStop(this.ClientId);
         }
 
@@ -520,18 +590,30 @@ namespace Microsoft.Azure.ServiceBus.Core
             this.ThrowIfNotPeekLockMode();
 
             MessagingEventSource.Log.MessageDeadLetterStart(this.ClientId, 1, lockToken);
+            bool isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+            Activity activity = isDiagnosticSourceEnabled ? this.diagnosticSource.DisposeStart("DeadLetter", lockToken) : null;
+            Task deadLetterTask = null;
 
             try
             {
-                await this.RetryPolicy.RunOperation(() => this.OnDeadLetterAsync(lockToken, propertiesToModify), this.OperationTimeout)
-                    .ConfigureAwait(false);
+                deadLetterTask = this.RetryPolicy.RunOperation(() => this.OnDeadLetterAsync(lockToken, propertiesToModify),
+                    this.OperationTimeout);
+                await deadLetterTask.ConfigureAwait(false);
             }
             catch (Exception exception)
             {
+                if (isDiagnosticSourceEnabled)
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
+
                 MessagingEventSource.Log.MessageDeadLetterException(this.ClientId, exception);
                 throw;
             }
-
+            finally
+            {
+                this.diagnosticSource.DisposeStop(activity, lockToken, deadLetterTask?.Status);
+            }
             MessagingEventSource.Log.MessageDeadLetterStop(this.ClientId);
         }
 
@@ -554,17 +636,32 @@ namespace Microsoft.Azure.ServiceBus.Core
             this.ThrowIfNotPeekLockMode();
 
             MessagingEventSource.Log.MessageDeadLetterStart(this.ClientId, 1, lockToken);
+            bool isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+            Activity activity = isDiagnosticSourceEnabled ? this.diagnosticSource.DisposeStart("DeadLetter", lockToken) : null;
+            Task deadLetterTask = null;
 
             try
             {
-                await this.RetryPolicy.RunOperation(() => this.OnDeadLetterAsync(lockToken, null, deadLetterReason, deadLetterErrorDescription), this.OperationTimeout)
-                    .ConfigureAwait(false);
+                deadLetterTask =
+                    this.RetryPolicy.RunOperation(
+                        () => this.OnDeadLetterAsync(lockToken, null, deadLetterReason, deadLetterErrorDescription),
+                        this.OperationTimeout);
+                await deadLetterTask.ConfigureAwait(false);
             }
             catch (Exception exception)
             {
+                if (isDiagnosticSourceEnabled)
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
+
                 MessagingEventSource.Log.MessageDeadLetterException(this.ClientId, exception);
                 throw;
             }
+            finally
+            {
+                this.diagnosticSource.DisposeStop(activity, lockToken, deadLetterTask?.Status);
+            }
 
             MessagingEventSource.Log.MessageDeadLetterStop(this.ClientId);
         }
@@ -600,21 +697,33 @@ namespace Microsoft.Azure.ServiceBus.Core
             this.ThrowIfNotPeekLockMode();
 
             MessagingEventSource.Log.MessageRenewLockStart(this.ClientId, 1, lockToken);
+            bool isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+            Activity activity = isDiagnosticSourceEnabled ? this.diagnosticSource.RenewLockStart(lockToken) : null;
+            Task renewTask = null;
 
             var lockedUntilUtc = DateTime.MinValue;
 
             try
             {
-                await this.RetryPolicy.RunOperation(
-                    async () => lockedUntilUtc = await this.OnRenewLockAsync(lockToken).ConfigureAwait(false), this.OperationTimeout)
-                    .ConfigureAwait(false);
+                renewTask = this.RetryPolicy.RunOperation(
+                    async () => lockedUntilUtc = await this.OnRenewLockAsync(lockToken).ConfigureAwait(false),
+                    this.OperationTimeout);
+                await renewTask.ConfigureAwait(false);
             }
             catch (Exception exception)
             {
+                if (isDiagnosticSourceEnabled)
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
+
                 MessagingEventSource.Log.MessageRenewLockException(this.ClientId, exception);
                 throw;
             }
-
+            finally
+            {
+                this.diagnosticSource.RenewLockStop(activity, lockToken, renewTask?.Status, lockedUntilUtc);
+            }
             MessagingEventSource.Log.MessageRenewLockStop(this.ClientId);
 
             return lockedUntilUtc;
@@ -671,20 +780,34 @@ namespace Microsoft.Azure.ServiceBus.Core
             IList<Message> messages = null;
 
             MessagingEventSource.Log.MessagePeekStart(this.ClientId, fromSequenceNumber, messageCount);
+            bool isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+            Activity activity = isDiagnosticSourceEnabled ? this.diagnosticSource.PeekStart(fromSequenceNumber, messageCount) : null;
+            Task peekTask = null;
+
             try
             {
-                await this.RetryPolicy.RunOperation(
+                peekTask = this.RetryPolicy.RunOperation(
                     async () =>
                     {
                         messages = await this.OnPeekAsync(fromSequenceNumber, messageCount).ConfigureAwait(false);
-                    }, this.OperationTimeout)
-                    .ConfigureAwait(false);
+                    }, this.OperationTimeout);
+
+                await peekTask.ConfigureAwait(false);
             }
             catch (Exception exception)
             {
+                if (isDiagnosticSourceEnabled)
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
+
                 MessagingEventSource.Log.MessagePeekException(this.ClientId, exception);
                 throw;
             }
+            finally
+            {
+                this.diagnosticSource.PeekStop(activity, fromSequenceNumber, messageCount, peekTask?.Status, messages);
+            }
 
             MessagingEventSource.Log.MessagePeekStop(this.ClientId, messages?.Count ?? 0);
             return messages;
@@ -1017,7 +1140,7 @@ namespace Microsoft.Azure.ServiceBus.Core
             {
                 throw new ArgumentOutOfRangeException(nameof(deadLetterErrorDescription), $"Maximum permitted length is {Constants.MaxDeadLetterReasonLength}");
             }
-            
+
             var lockTokens = new[] { new Guid(lockToken) };
             if (lockTokens.Any(lt => this.requestResponseLockedMessages.Contains(lt)))
             {
@@ -1071,7 +1194,7 @@ namespace Microsoft.Azure.ServiceBus.Core
                 }
 
                 this.receivePumpCancellationTokenSource = new CancellationTokenSource();
-                this.receivePump = new MessageReceivePump(this, registerHandlerOptions, callback, this.ServiceBusConnection.Endpoint.Authority, this.receivePumpCancellationTokenSource.Token);
+                this.receivePump = new MessageReceivePump(this, registerHandlerOptions, callback, this.ServiceBusConnection.Endpoint, this.receivePumpCancellationTokenSource.Token);
             }
 
             try
@@ -1246,7 +1369,8 @@ namespace Microsoft.Azure.ServiceBus.Core
                         }
                         else
                         {
-                            throw new NotSupportedException(Resources.InvalidAmqpMessageProperty.FormatForUser(pair.Key.GetType()));
+                            throw new NotSupportedException(
+                                Resources.InvalidAmqpMessageProperty.FormatForUser(pair.Key.GetType()));
                         }
                     }
 

src/Microsoft.Azure.ServiceBus/Core/MessageSender.cs

@@ -5,6 +5,7 @@ namespace Microsoft.Azure.ServiceBus.Core
 {
     using System;
     using System.Collections.Generic;
+    using System.Diagnostics;
     using System.Linq;
     using System.Threading;
     using System.Threading.Tasks;
@@ -37,6 +38,7 @@ namespace Microsoft.Azure.ServiceBus.Core
         int deliveryCount;
         readonly bool ownsConnection;
         readonly ActiveClientLinkManager clientLinkManager;
+        readonly ServiceBusDiagnosticSource diagnosticSource;
 
         /// <summary>
         /// Creates a new AMQP MessageSender.
@@ -96,6 +98,7 @@ namespace Microsoft.Azure.ServiceBus.Core
             this.SendLinkManager = new FaultTolerantAmqpObject<SendingAmqpLink>(this.CreateLinkAsync, CloseSession);
             this.RequestResponseLinkManager = new FaultTolerantAmqpObject<RequestResponseAmqpLink>(this.CreateRequestResponseLinkAsync, CloseRequestResponseSession);
             this.clientLinkManager = new ActiveClientLinkManager(this.ClientId, this.CbsTokenProvider);
+            this.diagnosticSource = new ServiceBusDiagnosticSource(entityPath, serviceBusConnection.Endpoint);
 
             MessagingEventSource.Log.MessageSenderCreateStop(serviceBusConnection.Endpoint.Authority, entityPath, this.ClientId);
         }
@@ -144,21 +147,35 @@ namespace Microsoft.Azure.ServiceBus.Core
         public async Task SendAsync(IList<Message> messageList)
         {
             this.ThrowIfClosed();
+
             var count = MessageSender.ValidateMessages(messageList);
             MessagingEventSource.Log.MessageSendStart(this.ClientId, count);
 
-            var processedMessages = await this.ProcessMessages(messageList).ConfigureAwait(false);
+            bool isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+            Activity activity = isDiagnosticSourceEnabled ? this.diagnosticSource.SendStart(messageList) : null;
+            Task sendTask = null;
 
             try
             {
-                await this.RetryPolicy.RunOperation(() => this.OnSendAsync(processedMessages), this.OperationTimeout)
-                    .ConfigureAwait(false);
+                var processedMessages = await this.ProcessMessages(messageList).ConfigureAwait(false);
+
+                sendTask = this.RetryPolicy.RunOperation(() => this.OnSendAsync(processedMessages), this.OperationTimeout);
+                await sendTask.ConfigureAwait(false);
             }
             catch (Exception exception)
             {
+                if (isDiagnosticSourceEnabled)
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
+
                 MessagingEventSource.Log.MessageSendException(this.ClientId, exception);
                 throw;
             }
+            finally
+            {
+                this.diagnosticSource.SendStop(activity, messageList, sendTask?.Status);
+            }
 
             MessagingEventSource.Log.MessageSendStop(this.ClientId);
         }
@@ -190,22 +207,35 @@ namespace Microsoft.Azure.ServiceBus.Core
             MessagingEventSource.Log.ScheduleMessageStart(this.ClientId, scheduleEnqueueTimeUtc);
             long result = 0;
 
-            var processedMessage = await this.ProcessMessage(message).ConfigureAwait(false);
+            bool isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+            Activity activity = isDiagnosticSourceEnabled ? this.diagnosticSource.ScheduleStart(message, scheduleEnqueueTimeUtc) : null;
+            Task scheduleTask = null;
 
             try
             {
-                await this.RetryPolicy.RunOperation(
+                var processedMessage = await this.ProcessMessage(message).ConfigureAwait(false);
+
+                scheduleTask = this.RetryPolicy.RunOperation(
                     async () =>
                     {
                         result = await this.OnScheduleMessageAsync(processedMessage).ConfigureAwait(false);
-                    }, this.OperationTimeout)
-                    .ConfigureAwait(false);
+                    }, this.OperationTimeout);
+                await scheduleTask.ConfigureAwait(false);
             }
             catch (Exception exception)
             {
+                if (isDiagnosticSourceEnabled)
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
+
                 MessagingEventSource.Log.ScheduleMessageException(this.ClientId, exception);
                 throw;
             }
+            finally
+            {
+                this.diagnosticSource.ScheduleStop(activity, message, scheduleEnqueueTimeUtc, scheduleTask?.Status, result);
+            }
 
             MessagingEventSource.Log.ScheduleMessageStop(this.ClientId);
             return result;
@@ -220,17 +250,30 @@ namespace Microsoft.Azure.ServiceBus.Core
             this.ThrowIfClosed();
             MessagingEventSource.Log.CancelScheduledMessageStart(this.ClientId, sequenceNumber);
 
+            bool isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+            Activity activity = isDiagnosticSourceEnabled ? this.diagnosticSource.CancelStart(sequenceNumber) : null;
+            Task cancelTask = null;
+
             try
             {
-                await this.RetryPolicy.RunOperation(() => this.OnCancelScheduledMessageAsync(sequenceNumber), this.OperationTimeout)
-                    .ConfigureAwait(false);
+                cancelTask = this.RetryPolicy.RunOperation(() => this.OnCancelScheduledMessageAsync(sequenceNumber),
+                    this.OperationTimeout);
+                await cancelTask.ConfigureAwait(false);
             }
             catch (Exception exception)
             {
+                if (isDiagnosticSourceEnabled)
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
+
                 MessagingEventSource.Log.CancelScheduledMessageException(this.ClientId, exception);
                 throw;
             }
-
+            finally
+            {
+                this.diagnosticSource.CancelStop(activity, sequenceNumber, cancelTask?.Status);
+            }
             MessagingEventSource.Log.CancelScheduledMessageStop(this.ClientId);
         }
 
@@ -403,6 +446,7 @@ namespace Microsoft.Azure.ServiceBus.Core
                     }
 
                     var outcome = await amqpLink.SendMessageAsync(amqpMessage, this.GetNextDeliveryTag(), AmqpConstants.NullBinary, timeoutHelper.RemainingTime()).ConfigureAwait(false);
+
                     if (outcome.DescriptorCode != Accepted.Code)
                     {
                         var rejected = (Rejected)outcome;

src/Microsoft.Azure.ServiceBus/Extensions/MessageDiagnosticsExtensions.cs

@@ -0,0 +1,153 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.Azure.ServiceBus.Diagnostics
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Diagnostics;
+
+    public static class MessageExtensions
+    {
+        /// <summary>
+        /// Creates <see cref="Activity"/> based on the tracing context stored in the <see cref="Message"/>
+        /// <param name="activityName">Optional Activity name</param>
+        /// <returns>New <see cref="Activity"/> with tracing context</returns>
+        /// </summary>
+        /// <remarks>
+        /// Tracing context is used to correlate telemetry between producer and consumer and 
+        /// represented by 'Diagnostic-Id' and 'Correlation-Context' properties in <see cref="Message.UserProperties"/>.
+        /// 
+        /// .NET SDK automatically injects context when sending message to the ServiceBus (if diagnostics is enabled by tracing system).
+        /// 
+        /// <para>
+        /// 'Diagnostic-Id' uniquely identifies operation that enqueued message
+        /// </para>
+        /// <para>
+        /// 'Correlation-Context' is comma separated list of sting key value pairs represeting optional context for the operation.
+        /// </para>
+        /// 
+        /// If there is no tracing context in the message, this method returns <see cref="Activity"/> without parent.
+        /// 
+        /// Returned <see cref="Activity"/> needs to be started before it can be used (see example below)
+        /// </remarks>
+        /// <example>
+        /// <code>
+        /// async Task ProcessAsync()
+        /// {
+        ///    var message = await messageReceiver.ReceiveAsync();
+        ///    var activity = message.ExtractActivity();
+        ///    activity.Start();
+        ///    Logger.LogInformation($"Message received, Id = {Activity.Current.Id}")
+        ///    try 
+        ///    {
+        ///       // process message
+        ///    }
+        ///    catch (Exception ex)
+        ///    {
+        ///         Logger.LogError($"Exception {ex}, Id = {Activity.Current.Id}")
+        ///    }
+        ///    finally 
+        ///    {
+        ///         activity.Stop();
+        ///         // Activity is stopped, we no longer have it in Activity.Current, let's user activity now
+        ///         Logger.LogInformation($"Message processed, Id = {activity.Id}, Duration = {activity.Duration}")
+        ///    }
+        /// }
+        /// </code>
+        /// 
+        /// Note that every log is stamped with <see cref="Activity.Current"/>.Id, that could be used within 
+        /// any nested method call (sync or async) - <see cref="Activity.Current"/> is an ambient context that flows with async method calls.
+        /// 
+        /// </example>
+
+        public static Activity ExtractActivity(this Message message, string activityName = null)
+        {
+            if (message == null)
+            {
+                throw new ArgumentNullException(nameof(message));
+            }
+
+            if (activityName == null)
+            {
+                activityName = ServiceBusDiagnosticSource.ProcessActivityName;
+            }
+
+            var activity = new Activity(activityName);
+
+            if (TryExtractId(message, out string id))
+            {
+                activity.SetParentId(id);
+
+                if (message.TryExtractContext(out IList<KeyValuePair<string, string>> ctx))
+                {
+                    foreach (var kvp in ctx)
+                    {
+                        activity.AddBaggage(kvp.Key, kvp.Value);
+                    }
+                }
+            }
+
+            return activity;
+        }
+
+        internal static bool TryExtractId(this Message message, out string id)
+        {
+            id = null;
+            if (message.UserProperties.TryGetValue(ServiceBusDiagnosticSource.ActivityIdPropertyName,
+                out object requestId))
+            {
+                var tmp = requestId as string;
+                if (tmp != null && tmp.Trim().Length > 0)
+                {
+                    id = tmp;
+                    return true;
+                }
+            }
+
+            return false;
+        }
+
+        internal static bool TryExtractContext(this Message message, out IList<KeyValuePair<string, string>> context)
+        {
+            context = null;
+            try
+            {
+                if (message.UserProperties.TryGetValue(ServiceBusDiagnosticSource.CorrelationContextPropertyName,
+                    out object ctxObj))
+                {
+                    string ctxStr = ctxObj as string;
+                    if (string.IsNullOrEmpty(ctxStr))
+                    {
+                        return false;
+                    }
+
+                    var ctxList = ctxStr.Split(',');
+                    if (ctxList.Length == 0)
+                    {
+                        return false;
+                    }
+
+                    context = new List<KeyValuePair<string, string>>();
+                    foreach (string item in ctxList)
+                    {
+                        var kvp = item.Split('=');
+                        if (kvp.Length == 2)
+                        {
+                            context.Add(new KeyValuePair<string, string>(kvp[0], kvp[1]));
+                        }
+                    }
+
+                    return true;
+                }
+            }
+            catch (Exception)
+            {
+                // ignored, if context is invalid, there nothing we can do:
+                // invalid context was created by consumer, but if we throw here, it will break message processing on producer
+                // and producer does not control which context it receives
+            }
+            return false;
+        }
+    }
+}

src/Microsoft.Azure.ServiceBus/Message.cs

@@ -9,8 +9,11 @@ namespace Microsoft.Azure.ServiceBus
     using Primitives;
 
     /// <summary>
-    /// The object used to communicate and transfer data with Service Bus.
+    /// The message object used to communicate and transfer data with Service Bus.
     /// </summary>
+    /// <remarks>
+    /// The message structure is discussed in detail in the <a href="https://docs.microsoft.com/azure/service-bus-messaging/service-bus-messages-payloads">product documentation.</a>
+    /// </remarks>
     public class Message
     {
         /// <summary>
@@ -63,8 +66,14 @@ namespace Microsoft.Azure.ServiceBus
         /// <summary>
         /// Gets or sets the MessageId to identify the message.
         /// </summary>
-        /// <remarks>A value set by the user to identify the message. In case message deduplication is enabled on the entity, this value will be used for deduplication.
-        /// Max MessageId size is 128 chars.</remarks>
+        /// <remarks>
+        ///    The message identifier is an application-defined value that uniquely identifies the 
+        ///    message and its payload. The identifier is a free-form string and can reflect a GUID 
+        ///    or an identifier derived from the application context. If enabled, the 
+        ///    <a href="https://docs.microsoft.com/azure/service-bus-messaging/duplicate-detection">duplicate detection</a> 
+        ///    feature identifies and removes second and further submissions of messages with the 
+        ///    same MessageId.
+        /// </remarks>
         public string MessageId
         {
             get => this.messageId;
@@ -76,10 +85,14 @@ namespace Microsoft.Azure.ServiceBus
             }
         }
 
-        /// <summary>Gets or sets a partition key for sending a transactional message to a queue or topic that is not session-aware.</summary>
-        /// <value>The partition key for sending a transactional message.</value>
-        /// <remarks>Transactions are not currently supported with this library. Messages with same partitionKey are sent to the same partition.
-        /// Max PartitionKey size is 128 chars.</remarks>
+        /// <summary>Gets or sets a partition key for sending a message to a partitioned entity.</summary>
+        /// <value>The partition key. Maximum length is 128 characters.</value>
+        /// <remarks>
+        ///    For <a href="https://docs.microsoft.com/azure/service-bus-messaging/service-bus-partitioning">partitioned entities</a>, 
+        ///    setting this value enables assigning related messages to the same internal partition, so that submission sequence 
+        ///    order is correctly recorded. The partition is chosen by a hash function over this value and cannot be chosen 
+        ///    directly. For session-aware entities, the <see cref="SessionId"/> property overrides this value.
+        /// </remarks>
         public string PartitionKey
         {
             get => this.partitionKey;
@@ -91,9 +104,14 @@ namespace Microsoft.Azure.ServiceBus
             }
         }
 
-        /// <summary>Gets or sets a partition key for sending a transactional message via a transfer queue.</summary>
-        /// <value>The partition key value when a transaction is to be used to send messages via a transfer queue.</value>
-        /// <remarks>Max size of ViaPartitionKey is 128 chars.</remarks>
+        /// <summary>Gets or sets a partition key for sending a message into an entity via a partitioned transfer queue.</summary>
+        /// <value>The partition key. Maximum length is 128 characters. </value>
+        /// <remarks>
+        ///    If a message is sent via a transfer queue in the scope of a transaction, this value selects the 
+        ///    transfer queue partition: This is functionally equivalent to <see cref="PartitionKey"/> and ensures that
+        ///    messages are kept together and in order as they are transferred.
+        ///    See <a href="https://docs.microsoft.com/azure/service-bus-messaging/service-bus-transactions#transfers-and-send-via">Transfers and Send Via</a>.
+        /// </remarks>
         public string ViaPartitionKey
         {
             get => this.viaPartitionKey;
@@ -105,9 +123,15 @@ namespace Microsoft.Azure.ServiceBus
             }
         }
 
-        /// <summary>Gets or sets a sessionId. A message with sessionId set can only be received using a <see cref="IMessageSession"/> object.</summary>
-        /// <value>The identifier of the session.</value>
-        /// <remarks>Max size of sessionId is 128 chars.</remarks>
+        /// <summary>Gets or sets the session identifier for a session-aware entity.</summary>
+        /// <value>The session identifier. Maximum length is 128 characters.</value>
+        /// <remarks>
+        ///    For session-aware entities, this application-defined value specifies the session 
+        ///    affiliation of the message. Messages with the same session identifier are subject 
+        ///    to summary locking and enable exact in-order processing and demultiplexing. 
+        ///    For session-unaware entities, this value is ignored.
+        ///    See <a href="https://docs.microsoft.com/azure/service-bus-messaging/message-sessions">Message Sessions</a>.
+        /// </remarks>
         public string SessionId
         {
             get => this.sessionId;
@@ -119,9 +143,12 @@ namespace Microsoft.Azure.ServiceBus
             }
         }
 
-        /// <summary>Gets or sets the session identifier to reply to.</summary>
-        /// <value>The session identifier to reply to.</value>
-        /// <remarks>Max size of ReplyToSessionId is 128.</remarks>
+        /// <summary>Gets or sets a session identifier augmenting the <see cref="ReplyTo"/> address.</summary>
+        /// <value>Session identifier. Maximum length is 128 characters.</value>
+        /// <remarks>
+        ///    This value augments the ReplyTo information and specifies which SessionId should be set 
+        ///    for the reply when sent to the reply entity. See <a href="https://docs.microsoft.com/azure/service-bus-messaging/service-bus-messages-payloads?#message-routing-and-correlation">Message Routing and Correlation</a>
+        /// </remarks>
         public string ReplyToSessionId
         {
             get => this.replyToSessionId;
@@ -134,9 +161,12 @@ namespace Microsoft.Azure.ServiceBus
         }
 
         /// <summary>Gets the date and time in UTC at which the message is set to expire.</summary>
-        /// <value>The message expiration time in UTC.</value>
+        /// <value>The message expiration time in UTC. This property is read-only.</value>
         /// <exception cref="System.InvalidOperationException">If the message has not been received. For example if a new message was created but not yet sent and received.</exception>
-        /// <remarks>Unless specifically set for a message, this value is controlled by the 'DefaultMessageTimeToLive' property set while creating the entity.</remarks>
+        /// <remarks>
+        ///  The UTC instant at which the message is marked for removal and no longer available for retrieval 
+        ///  from the entity due to expiration. Expiry is controlled by the <see cref="TimeToLive"/> property 
+        ///  and this property is computed from <see cref="SystemPropertiesCollection.EnqueuedTimeUtc"/>+<see cref="TimeToLive"/></remarks>
         public DateTime ExpiresAtUtc
         {
             get
@@ -151,14 +181,17 @@ namespace Microsoft.Azure.ServiceBus
         }
 
         /// <summary>
-        /// Gets or sets the message’s time to live value. This is the duration after which the message expires, starting from when the message is sent to the Service Bus.
-        /// Messages older than their TimeToLive value will expire and no longer be retained in the message store. Expired messages cannot be received.
-        /// TimeToLive is the maximum lifetime that a message can be received, but its value cannot exceed the entity specified value on the destination queue or subscription.
-        /// If a lower TimeToLive value is specified, it will be applied to the individual message. However, a larger value specified on the message will be overridden by the
-        /// entity’s DefaultMessageTimeToLive value.
+        /// Gets or sets the message’s "time to live" value. 
         /// </summary>
         /// <value>The message’s time to live value.</value>
-        /// <remarks>If the TTL set on a message by the sender exceeds the destination's TTL, then the message's TTL will be overwritten by the later one.</remarks>
+        /// <remarks>
+        ///     This value is the relative duration after which the message expires, starting from the instant 
+        ///      the message has been accepted and stored by the broker, as captured in <see cref="SystemPropertiesCollection.EnqueuedTimeUtc"/>. 
+        ///      When not set explicitly, the assumed value is the DefaultTimeToLive for the respective queue or topic. 
+        ///      A message-level <see cref="TimeToLive"/> value cannot be longer than the entity's DefaultTimeToLive 
+        ///      setting and it is silently adjusted if it does. 
+        ///      See <a href="https://docs.microsoft.com/azure/service-bus-messaging/message-expiration">Expiration</a>
+        /// </remarks>
         public TimeSpan TimeToLive
         {
             get
@@ -179,25 +212,48 @@ namespace Microsoft.Azure.ServiceBus
         }
 
         /// <summary>Gets or sets the a correlation identifier.</summary>
-        /// <value>The identifier of the correlation.</value>
-        /// <remarks>Its a custom property that can be used to either transfer a correlation Id to the destination or be used in <see cref="CorrelationFilter"/></remarks>
+        /// <value>Correlation identifier.</value>
+        /// <remarks>
+        ///    Allows an application to specify a context for the message for the purposes of correlation, 
+        ///    for example reflecting the MessageId of a message that is being replied to.
+        ///    See <a href="https://docs.microsoft.com/azure/service-bus-messaging/service-bus-messages-payloads?#message-routing-and-correlation">Message Routing and Correlation</a>.
+        /// </remarks>
         public string CorrelationId { get; set; }
 
-        /// <summary>Gets or sets the application specific label.</summary>
-        /// <value>The application specific label.</value>
+        /// <summary>Gets or sets an application specific label.</summary>
+        /// <value>The application specific label</value>
+        /// <remarks>
+        ///   This property enables the application to indicate the purpose of the message to the receiver in a standardized 
+        ///   fashion, similar to an email subject line. The mapped AMQP property is "subject".
+        /// </remarks>
         public string Label { get; set; }
 
-        /// <summary>Gets or sets the send to address.</summary>
-        /// <value>The send to address.</value>
+        /// <summary>Gets or sets the "to" address.</summary>
+        /// <value>The "to" address.</value>
+        /// <remarks>
+        ///    This property is reserved for future use in routing scenarios and presently ignored by the broker itself. 
+        ///     Applications can use this value in rule-driven 
+        ///     <a href="https://docs.microsoft.com/azure/service-bus-messaging/service-bus-auto-forwarding">auto-forward chaining</a> scenarios to indicate the 
+        ///     intended logical destination of the message.
+        /// </remarks>
         public string To { get; set; }
 
-        /// <summary>Gets or sets the type of the content.</summary>
-        /// <value>The type of the content of the message body. This is a
-        /// content type identifier utilized by the sender and receiver for application specific logic.</value>
+        /// <summary>Gets or sets the content tpye descriptor.</summary>
+        /// <value>RFC2045 Content-Type descriptor.</value>
+        /// <remarks>
+        ///   Optionally describes the payload of the message, with a descriptor following the format of 
+        ///   RFC2045, Section 5, for example "application/json".
+        /// </remarks>
         public string ContentType { get; set; }
 
-        /// <summary>Gets or sets the address of the queue to reply to.</summary>
-        /// <value>The reply to queue address.</value>
+        /// <summary>Gets or sets the address of an entity to send replies to.</summary>
+        /// <value>The reply entity address.</value>
+        /// <remarks>
+        ///    This optional and application-defined value is a standard way to express a reply path 
+        ///    to the receiver of the message. When a sender expects a reply, it sets the value to the 
+        ///    absolute or relative path of the queue or topic it expects the reply to be sent to.
+        ///    See <a href="https://docs.microsoft.com/azure/service-bus-messaging/service-bus-messages-payloads?#message-routing-and-correlation">Message Routing and Correlation</a>.
+        /// </remarks>
         public string ReplyTo { get; set; }
 
         /// <summary>Gets or sets the date and time in UTC at which the message will be enqueued. This
@@ -215,7 +271,7 @@ namespace Microsoft.Azure.ServiceBus
         public long Size => Body.Length;
 
         /// <summary>
-        /// Gets the user property bag, which can be used for custom message properties.
+        /// Gets the "user properties" bag, which can be used for custom message metadata.
         /// </summary>
         /// <remarks>
         /// Only following value types are supported:
@@ -237,8 +293,10 @@ namespace Microsoft.Azure.ServiceBus
             return string.Format(CultureInfo.CurrentCulture, $"{{MessageId:{this.MessageId}}}");
         }
 
-        /// <summary>Clones a message, so that it is possible to send a clone of a message as a new message.</summary>
-        /// <returns>The <see cref="Message" /> that contains the cloned message.</returns>
+        /// <summary>Clones a message, so that it is possible to send a clone of an already received 
+        /// message as a new message. The system properties of original message 
+        /// are not copied.</summary>
+        /// <returns>A cloned <see cref="Message" />.</returns>
         public Message Clone()
         {
             var clone = (Message)this.MemberwiseClone();
@@ -304,16 +362,25 @@ namespace Microsoft.Azure.ServiceBus
             /// <summary>
             /// Gets the lock token for the current message.
             /// </summary>
-            /// <remarks>A lock token will only be specified if the message was received using <see cref="ReceiveMode.PeekLock"/></remarks>
+            /// <remarks>
+            ///   The lock token is a reference to the lock that is being held by the broker in <see cref="ReceiveMode.PeekLock"/> mode. 
+            ///   Locks are used to explicitly settle messages as explained in the <a href="https://docs.microsoft.com/azure/service-bus-messaging/message-transfers-locks-settlement">product documentation in more detail</a>.
+            ///   The token can also be used to pin the lock permanently through the <a href="https://docs.microsoft.com/azure/service-bus-messaging/message-deferral">Deferral API</a> and, with that, take the message out of the 
+            ///   regular delivery state flow. This property is read-only.
+            /// </remarks>
             public string LockToken => this.LockTokenGuid.ToString();
 
-            /// <summary>Specifies if message is a received message or not.</summary>
+            /// <summary>Specifies if the message has been obtained from the broker.</summary>
             public bool IsReceived => this.sequenceNumber > -1;
 
             /// <summary>
             /// Get the current delivery count.
             /// </summary>
             /// <value>This value starts at 1.</value>
+            /// <remarks>
+            ///    Number of deliveries that have been attempted for this message. The count is incremented when a message lock expires, 
+            ///    or the message is explicitly abandoned by the receiver. This property is read-only.
+            /// </remarks>
             public int DeliveryCount
             {
                 get
@@ -327,6 +394,11 @@ namespace Microsoft.Azure.ServiceBus
 
             /// <summary>Gets the date and time in UTC until which the message will be locked in the queue/subscription.</summary>
             /// <value>The date and time until which the message will be locked in the queue/subscription.</value>
+            /// <remarks>
+            /// 	For messages retrieved under a lock (peek-lock receive mode, not pre-settled) this property reflects the UTC 
+            ///     instant until which the message is held locked in the queue/subscription. When the lock expires, the <see cref="DeliveryCount"/> 
+            ///     is incremented and the message is again available for retrieval. This property is read-only.
+            /// </remarks>
             public DateTime LockedUntilUtc
             {
                 get
@@ -338,7 +410,13 @@ namespace Microsoft.Azure.ServiceBus
                 internal set => this.lockedUntilUtc = value;
             }
 
-            /// <summary>Gets the unique number assigned to a message by Service Bus, for this entity.</summary>
+            /// <summary>Gets the unique number assigned to a message by Service Bus.</summary>
+            /// <remarks>
+            ///     The sequence number is a unique 64-bit integer assigned to a message as it is accepted 
+            ///     and stored by the broker and functions as its true identifier. For partitioned entities, 
+            ///     the topmost 16 bits reflect the partition identifier. Sequence numbers monotonically increase. 
+            ///     They roll over to 0 when the 48-64 bit range is exhausted. This property is read-only.
+            /// </remarks>
             public long SequenceNumber
             {
                 get
@@ -353,6 +431,10 @@ namespace Microsoft.Azure.ServiceBus
             /// <summary>
             /// Gets the name of the queue or subscription that this message was enqueued on, before it was deadlettered.
             /// </summary>
+            /// <remarks>
+            /// 	Only set in messages that have been dead-lettered and subsequently auto-forwarded from the dead-letter queue 
+            ///     to another entity. Indicates the entity in which the message was dead-lettered. This property is read-only.
+            /// </remarks>
             public string DeadLetterSource
             {
                 get
@@ -375,10 +457,12 @@ namespace Microsoft.Azure.ServiceBus
                 set => this.partitionId = value;
             }
 
-            /// <summary>Gets or sets the enqueued sequence number of the message.</summary>
+            /// <summary>Gets or sets the original sequence number of the message.</summary>
             /// <value>The enqueued sequence number of the message.</value>
-            /// <remarks>In scenarios of Topic-Subscription or ForwardTo, the message is initially enqueued on a different entity as compared to the
-            /// entity from where the message is received. This returns the sequence number of the message in the initial entity.</remarks>
+            /// <remarks>
+            /// For messages that have been auto-forwarded, this property reflects the sequence number 
+            /// that had first been assigned to the message at its original point of submission. This property is read-only.
+            /// </remarks>
             public long EnqueuedSequenceNumber
             {
                 get
@@ -391,7 +475,12 @@ namespace Microsoft.Azure.ServiceBus
             }
 
             /// <summary>Gets or sets the date and time of the sent time in UTC.</summary>
-            /// <value>The enqueue time in UTC. This value represents the actual time of enqueuing the message.</value>
+            /// <value>The enqueue time in UTC. </value>
+            /// <remarks>
+            ///    The UTC instant at which the message has been accepted and stored in the entity. 
+            ///    This value can be used as an authoritative and neutral arrival time indicator when 
+            ///    the receiver does not want to trust the sender's clock. This property is read-only.
+            /// </remarks>
             public DateTime EnqueuedTimeUtc
             {
                 get

src/Microsoft.Azure.ServiceBus/MessageLockLostException.cs

@@ -10,12 +10,12 @@ namespace Microsoft.Azure.ServiceBus
     /// </summary>
     public sealed class MessageLockLostException : ServiceBusException
     {
-        internal MessageLockLostException(string message)
+        public MessageLockLostException(string message)
             : this(message, null)
         {
         }
 
-        internal MessageLockLostException(string message, Exception innerException)
+        public MessageLockLostException(string message, Exception innerException)
             : base(false, message, innerException)
         {
         }

src/Microsoft.Azure.ServiceBus/MessageReceivePump.cs

@@ -4,6 +4,7 @@
 namespace Microsoft.Azure.ServiceBus
 {
     using System;
+    using System.Diagnostics;
     using System.Threading;
     using System.Threading.Tasks;
     using Core;
@@ -17,19 +18,21 @@ namespace Microsoft.Azure.ServiceBus
         readonly IMessageReceiver messageReceiver;
         readonly CancellationToken pumpCancellationToken;
         readonly SemaphoreSlim maxConcurrentCallsSemaphoreSlim;
+        readonly ServiceBusDiagnosticSource diagnosticSource;
 
         public MessageReceivePump(IMessageReceiver messageReceiver,
             MessageHandlerOptions registerHandlerOptions,
             Func<Message, CancellationToken, Task> callback,
-            string endpoint,
+            Uri endpoint,
             CancellationToken pumpCancellationToken)
         {
             this.messageReceiver = messageReceiver ?? throw new ArgumentNullException(nameof(messageReceiver));
             this.registerHandlerOptions = registerHandlerOptions;
             this.onMessageCallback = callback;
-            this.endpoint = endpoint;
+            this.endpoint = endpoint.Authority;
             this.pumpCancellationToken = pumpCancellationToken;
             this.maxConcurrentCallsSemaphoreSlim = new SemaphoreSlim(this.registerHandlerOptions.MaxConcurrentCalls);
+            this.diagnosticSource = new ServiceBusDiagnosticSource(messageReceiver.Path, endpoint);
         }
 
         public void StartPump()
@@ -63,13 +66,28 @@ namespace Microsoft.Azure.ServiceBus
                     if (message != null)
                     {
                         MessagingEventSource.Log.MessageReceiverPumpTaskStart(this.messageReceiver.ClientId, message, this.maxConcurrentCallsSemaphoreSlim.CurrentCount);
-                        TaskExtensionHelper.Schedule(() => this.MessageDispatchTask(message));
+
+                        TaskExtensionHelper.Schedule(() =>
+                        {
+                            if (ServiceBusDiagnosticSource.IsEnabled())
+                            {
+                                return this.MessageDispatchTaskInstrumented(message);
+                            }
+                            else
+                            {
+                                return this.MessageDispatchTask(message);
+                            }
+                        });
                     }
                 }
                 catch (Exception exception)
                 {
-                    MessagingEventSource.Log.MessageReceivePumpTaskException(this.messageReceiver.ClientId, string.Empty, exception);
-                    await this.RaiseExceptionReceived(exception, ExceptionReceivedEventArgsAction.Receive).ConfigureAwait(false);
+                    // Not reporting an ObjectDisposedException as we're stopping the pump
+                    if (!(exception is ObjectDisposedException && this.pumpCancellationToken.IsCancellationRequested))
+                    {
+                        MessagingEventSource.Log.MessageReceivePumpTaskException(this.messageReceiver.ClientId, string.Empty, exception);
+                        await this.RaiseExceptionReceived(exception, ExceptionReceivedEventArgsAction.Receive).ConfigureAwait(false); 
+                    }
                 }
                 finally
                 {
@@ -83,6 +101,26 @@ namespace Microsoft.Azure.ServiceBus
             }
         }
 
+        async Task MessageDispatchTaskInstrumented(Message message)
+        {
+            Activity activity = this.diagnosticSource.ProcessStart(message);
+            Task processTask = null;
+            try
+            {
+                processTask = MessageDispatchTask(message);
+                await processTask.ConfigureAwait(false);
+            }
+            catch (Exception e)
+            {
+                this.diagnosticSource.ReportException(e);
+                throw;
+            }
+            finally
+            {
+                this.diagnosticSource.ProcessStop(activity, message, processTask?.Status);
+            }
+        }
+
         async Task MessageDispatchTask(Message message)
         {
             CancellationTokenSource renewLockCancellationTokenSource = null;
@@ -103,6 +141,7 @@ namespace Microsoft.Azure.ServiceBus
             {
                 MessagingEventSource.Log.MessageReceiverPumpUserCallbackStart(this.messageReceiver.ClientId, message);
                 await this.onMessageCallback(message, this.pumpCancellationToken).ConfigureAwait(false);
+
                 MessagingEventSource.Log.MessageReceiverPumpUserCallbackStop(this.messageReceiver.ClientId, message);
             }
             catch (Exception exception)
@@ -116,8 +155,13 @@ namespace Microsoft.Azure.ServiceBus
                     await this.AbandonMessageIfNeededAsync(message).ConfigureAwait(false);
                 }
 
+                if (ServiceBusDiagnosticSource.IsEnabled())
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
                 // AbandonMessageIfNeededAsync should take care of not throwing exception
                 this.maxConcurrentCallsSemaphoreSlim.Release();
+
                 return;
             }
             finally
@@ -204,9 +248,10 @@ namespace Microsoft.Azure.ServiceBus
                 {
                     MessagingEventSource.Log.MessageReceiverPumpRenewMessageException(this.messageReceiver.ClientId, message, exception);
 
-                    // TaskCancelled is expected here as renewTasks will be cancelled after the Complete call is made.
-                    // Lets not bother user with this exception.
-                    if (!(exception is TaskCanceledException))
+                    // TaskCanceled is expected here as renewTasks will be cancelled after the Complete call is made.
+                    // ObjectDisposedException should only happen here because the CancellationToken was disposed at which point 
+                    // this renew exception is not relevant anymore. Lets not bother user with this exception.
+                    if (!(exception is TaskCanceledException) && !(exception is ObjectDisposedException))
                     {
                         await this.RaiseExceptionReceived(exception, ExceptionReceivedEventArgsAction.RenewLock).ConfigureAwait(false);
                     }

src/Microsoft.Azure.ServiceBus/MessageSession.cs

@@ -4,6 +4,7 @@
 namespace Microsoft.Azure.ServiceBus
 {
     using System;
+    using System.Diagnostics;
     using System.Threading;
     using System.Threading.Tasks;
     using Amqp;
@@ -13,6 +14,8 @@ namespace Microsoft.Azure.ServiceBus
 
     internal class MessageSession : MessageReceiver, IMessageSession
     {
+        private readonly ServiceBusDiagnosticSource diagnosticSource;
+
         public MessageSession(
             string entityPath,
             MessagingEntityType? entityType,
@@ -25,6 +28,7 @@ namespace Microsoft.Azure.ServiceBus
             bool isSessionReceiver = false)
             : base(entityPath, entityType, receiveMode, serviceBusConnection, cbsTokenProvider, retryPolicy, prefetchCount, sessionId, isSessionReceiver)
         {
+            this.diagnosticSource = new ServiceBusDiagnosticSource(entityPath, serviceBusConnection.Endpoint);
         }
 
         /// <summary>
@@ -44,19 +48,19 @@ namespace Microsoft.Azure.ServiceBus
         public Task<byte[]> GetStateAsync()
         {
             this.ThrowIfClosed();
-            return this.OnGetStateAsync();
+            return ServiceBusDiagnosticSource.IsEnabled() ? this.OnGetStateInstrumentedAsync() : this.OnGetStateAsync();
         }
 
         public Task SetStateAsync(byte[] sessionState)
         {
             this.ThrowIfClosed();
-            return this.OnSetStateAsync(sessionState);
+            return ServiceBusDiagnosticSource.IsEnabled() ? this.OnSetStateInstrumentedAsync(sessionState) : this.OnSetStateAsync(sessionState);
         }
 
         public Task RenewSessionLockAsync()
         {
             this.ThrowIfClosed();
-            return this.OnRenewSessionLockAsync();
+            return ServiceBusDiagnosticSource.IsEnabled() ? this.OnRenewSessionLockInstrumentedAsync() : this.OnRenewSessionLockAsync();
         }
 
         protected override void OnMessageHandler(MessageHandlerOptions registerHandlerOptions, Func<Message, CancellationToken, Task> callback)
@@ -162,5 +166,71 @@ namespace Microsoft.Azure.ServiceBus
                 throw new ObjectDisposedException($"MessageSession with Id '{this.ClientId}' has already been closed. Please accept a new MessageSession.");
             }
         }
+
+        private async Task<byte[]> OnGetStateInstrumentedAsync()
+        {
+            Activity activity = this.diagnosticSource.GetSessionStateStart(this.SessionId);
+            Task<byte[]> getStateTask = null;
+            byte[] state = null;
+
+            try
+            {
+                getStateTask = this.OnGetStateAsync();
+                state = await getStateTask.ConfigureAwait(false);
+                return state;
+            }
+            catch (Exception ex)
+            {
+                this.diagnosticSource.ReportException(ex);
+                throw;
+            }
+            finally
+            {
+                this.diagnosticSource.GetSessionStateStop(activity, this.SessionId, state, getStateTask?.Status);
+            }
+        }
+
+        private async Task OnSetStateInstrumentedAsync(byte[] sessionState)
+        {
+            Activity activity = this.diagnosticSource.SetSessionStateStart(this.SessionId, sessionState);
+            Task setStateTask = null;
+
+            try
+            {
+                setStateTask = this.OnSetStateAsync(sessionState);
+                await setStateTask.ConfigureAwait(false);
+            }
+            catch (Exception ex)
+            {
+                this.diagnosticSource.ReportException(ex);
+                throw;
+            }
+            finally
+            {
+                this.diagnosticSource.SetSessionStateStop(activity, sessionState, this.SessionId, setStateTask?.Status);
+            }
+        }
+
+        private async Task OnRenewSessionLockInstrumentedAsync()
+        {
+            Activity activity = this.diagnosticSource.RenewSessionLockStart(this.SessionId);
+            Task renewTask = null;
+
+            try
+            {
+                renewTask = this.OnRenewSessionLockAsync();
+                await renewTask.ConfigureAwait(false);
+            }
+            catch (Exception ex)
+            {
+                this.diagnosticSource.ReportException(ex);
+                throw;
+            }
+            finally
+            {
+                this.diagnosticSource.RenewSessionLockStop(activity, this.SessionId, renewTask?.Status);
+            }
+        }
+
     }
 }

src/Microsoft.Azure.ServiceBus/MessagingEntityDisabledException.cs

@@ -10,12 +10,12 @@ namespace Microsoft.Azure.ServiceBus
     /// </summary>
     public sealed class MessagingEntityDisabledException : ServiceBusException
     {
-        internal MessagingEntityDisabledException(string message)
+        public MessagingEntityDisabledException(string message)
             : this(message, null)
         {
         }
 
-        internal MessagingEntityDisabledException(string message, Exception innerException)
+        public MessagingEntityDisabledException(string message, Exception innerException)
             : base(false, message, innerException)
         {
         }

src/Microsoft.Azure.ServiceBus/MessagingEntityNotFoundException.cs

@@ -10,12 +10,12 @@ namespace Microsoft.Azure.ServiceBus
     /// </summary>
     public sealed class MessagingEntityNotFoundException : ServiceBusException
     {
-        internal MessagingEntityNotFoundException(string message)
+        public MessagingEntityNotFoundException(string message)
             : this(message, null)
         {
         }
 
-        internal MessagingEntityNotFoundException(string message, Exception innerException)
+        public MessagingEntityNotFoundException(string message, Exception innerException)
             : base(false, message, innerException)
         {
         }

src/Microsoft.Azure.ServiceBus/Microsoft.Azure.ServiceBus.csproj

@@ -2,7 +2,7 @@
 
   <PropertyGroup>
     <Description>This is the next generation Azure Service Bus .NET Standard client library that focuses on queues &amp; topics. For more information about Service Bus, see https://azure.microsoft.com/en-us/services/service-bus/</Description>
-    <VersionPrefix>2.0.0</VersionPrefix>
+    <VersionPrefix>3.0.0-preview-01</VersionPrefix>
     <Authors>Microsoft</Authors>
     <TargetFrameworks>net461;netstandard2.0</TargetFrameworks>
     <AssemblyOriginatorKeyFile>../../build/keyfile.snk</AssemblyOriginatorKeyFile>
@@ -31,6 +31,10 @@
 
   <ItemGroup>
     <PackageReference Include="Microsoft.Azure.Amqp" Version="2.1.2" />
+    <PackageReference Include="Microsoft.Azure.Services.AppAuthentication" Version="1.1.0-preview" />
+    <PackageReference Include="Microsoft.IdentityModel.Clients.ActiveDirectory" Version="3.17.2" />
+    <PackageReference Include="System.Diagnostics.DiagnosticSource" Version="4.4.0" />
+    <PackageReference Include="System.IdentityModel.Tokens.Jwt" Version="5.2.0-preview2-41113220915" />
     <PackageReference Include="System.Runtime.Serialization.Xml" Version="4.3.0" />
   </ItemGroup>
 

src/Microsoft.Azure.ServiceBus/Primitives/AzureActiveDirectoryTokenProvider.cs

@@ -0,0 +1,96 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.Azure.ServiceBus.Primitives
+{
+    using System;
+    using System.Threading.Tasks;
+    using Microsoft.IdentityModel.Clients.ActiveDirectory;
+
+    /// <summary>
+    /// Represents the Azure Active Directory token provider for the Service Bus.
+    /// </summary>
+    public class AzureActiveDirectoryTokenProvider : TokenProvider
+    {
+        readonly AuthenticationContext authContext;
+        readonly ClientCredential clientCredential;
+#if !UAP10_0
+        readonly ClientAssertionCertificate clientAssertionCertificate;
+#endif
+        readonly string clientId;
+        readonly Uri redirectUri;
+        readonly IPlatformParameters platformParameters;
+        readonly UserIdentifier userIdentifier;
+
+        enum AuthType
+        {
+            ClientCredential,
+            UserPasswordCredential,
+            ClientAssertionCertificate,
+            InteractiveUserLogin
+        }
+
+        readonly AuthType authType;
+
+        internal AzureActiveDirectoryTokenProvider(AuthenticationContext authContext, ClientCredential credential)
+        {
+            this.clientCredential = credential;
+            this.authContext = authContext;
+            this.authType = AuthType.ClientCredential;
+            this.clientId = clientCredential.ClientId;
+        }
+
+#if !UAP10_0
+        internal AzureActiveDirectoryTokenProvider(AuthenticationContext authContext, ClientAssertionCertificate clientAssertionCertificate)
+        {
+            this.clientAssertionCertificate = clientAssertionCertificate;
+            this.authContext = authContext;
+            this.authType = AuthType.ClientAssertionCertificate;
+            this.clientId = clientAssertionCertificate.ClientId;
+        }
+#endif
+
+        internal AzureActiveDirectoryTokenProvider(AuthenticationContext authContext, string clientId, Uri redirectUri, IPlatformParameters platformParameters, UserIdentifier userIdentifier)
+        {
+            this.authContext = authContext;
+            this.clientId = clientId;
+            this.redirectUri = redirectUri;
+            this.platformParameters = platformParameters;
+            this.userIdentifier = userIdentifier;
+            this.authType = AuthType.InteractiveUserLogin;
+        }
+
+        /// <summary>
+        /// Gets a <see cref="SecurityToken"/> for the given audience and duration.
+        /// </summary>
+        /// <param name="appliesTo">The URI which the access token applies to</param>
+        /// <param name="timeout">The time span that specifies the timeout value for the message that gets the security token</param>
+        /// <returns><see cref="SecurityToken"/></returns>
+        public override async Task<SecurityToken> GetTokenAsync(string appliesTo, TimeSpan timeout)
+        {
+            AuthenticationResult authResult;
+
+            switch (this.authType)
+            {
+                case AuthType.ClientCredential:
+                    authResult = await this.authContext.AcquireTokenAsync(Constants.AadServiceBusAudience, this.clientCredential);
+                    break;
+
+#if !UAP10_0
+                case AuthType.ClientAssertionCertificate:
+                    authResult = await this.authContext.AcquireTokenAsync(Constants.AadServiceBusAudience, this.clientAssertionCertificate);
+                    break;
+#endif
+
+                case AuthType.InteractiveUserLogin:
+                    authResult = await this.authContext.AcquireTokenAsync(Constants.AadServiceBusAudience, this.clientId, this.redirectUri, this.platformParameters, this.userIdentifier);
+                    break;
+
+                default:
+                    throw new NotSupportedException();
+            }
+
+            return new JsonSecurityToken(authResult.AccessToken, appliesTo);
+        }
+    }
+}

src/Microsoft.Azure.ServiceBus/Primitives/ITokenProvider.cs

@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.Azure.ServiceBus.Primitives
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Linq;
+    using System.Text;
+    using System.Threading.Tasks;
+
+    /// <summary>
+    /// Provides interface definition of a token provider.
+    /// </summary>
+    public interface ITokenProvider
+    {
+        /// <summary>
+        /// Gets a <see cref="SecurityToken"/>.
+        /// </summary>
+        /// <param name="appliesTo">The URI which the access token applies to</param>
+        /// <param name="timeout">The time span that specifies the timeout value for the message that gets the security token</param>
+        /// <returns><see cref="SecurityToken"/></returns>
+        Task<SecurityToken> GetTokenAsync(string appliesTo, TimeSpan timeout);
+    }
+}

src/Microsoft.Azure.ServiceBus/Primitives/JsonSecurityToken.cs

@@ -0,0 +1,32 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.Azure.ServiceBus.Primitives
+{
+    using System;
+    using System.Collections.ObjectModel;
+    using System.IdentityModel.Tokens;
+    using System.IdentityModel.Tokens.Jwt;
+
+    /// <summary>
+    /// Extends SecurityToken for JWT specific properties
+    /// </summary>
+    public class JsonSecurityToken : SecurityToken
+    {
+        /// <summary>
+        /// Creates a new instance of the <see cref="JsonSecurityToken"/> class.
+        /// </summary>
+        /// <param name="rawToken">Raw JSON Web Token string</param>
+        /// <param name="audience">The audience</param>
+        public JsonSecurityToken(string rawToken, string audience)
+            : base(rawToken, GetExpirationDateTimeUtcFromToken(rawToken), audience, Constants.JsonWebTokenType)
+        {
+        }
+
+        static DateTime GetExpirationDateTimeUtcFromToken(string token)
+        {
+            var jwtSecurityToken = new JwtSecurityToken(token);
+            return jwtSecurityToken.ValidTo;
+        }
+    }
+}

src/Microsoft.Azure.ServiceBus/Primitives/ManagedServiceIdentityTokenProvider.cs

@@ -0,0 +1,29 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.Azure.ServiceBus.Primitives
+{
+    using System;
+    using System.Threading.Tasks;
+    using Azure.Services.AppAuthentication;
+
+    /// <summary>
+    /// Represents the Azure Active Directory token provider for Azure Managed Service Identity integration.
+    /// </summary>
+    public class ManagedServiceIdentityTokenProvider : TokenProvider
+    {
+        static AzureServiceTokenProvider azureServiceTokenProvider = new AzureServiceTokenProvider();
+
+        /// <summary>
+        /// Gets a <see cref="SecurityToken"/> for the given audience and duration.
+        /// </summary>
+        /// <param name="appliesTo">The URI which the access token applies to</param>
+        /// <param name="timeout">The time span that specifies the timeout value for the message that gets the security token</param>
+        /// <returns><see cref="SecurityToken"/></returns>
+        public async override Task<SecurityToken> GetTokenAsync(string appliesTo, TimeSpan timeout)
+        {
+            string accessToken = await azureServiceTokenProvider.GetAccessTokenAsync(Constants.AadServiceBusAudience);
+            return new JsonSecurityToken(accessToken, appliesTo);
+        }
+    }
+}

src/Microsoft.Azure.ServiceBus/Primitives/SecurityToken.cs

@@ -4,79 +4,55 @@
 namespace Microsoft.Azure.ServiceBus.Primitives
 {
     using System;
-    using System.Collections.Generic;
-    using System.Globalization;
-    using System.Net;
 
     /// <summary>
     /// Provides information about a security token such as audience, expiry time, and the string token value.
     /// </summary>
-    internal class SecurityToken
+    public class SecurityToken
     {
-        // per Simple Web Token draft specification
-        private const string TokenAudience = "Audience";
-        private const string TokenExpiresOn = "ExpiresOn";
-        private const string TokenIssuer = "Issuer";
-        private const string TokenDigest256 = "HMACSHA256";
+        /// <summary>
+        /// Token literal
+        /// </summary>
+        string token;
 
-        const string InternalExpiresOnFieldName = "ExpiresOn";
-        const string InternalAudienceFieldName = TokenAudience;
-        const string InternalKeyValueSeparator = "=";
-        const string InternalPairSeparator = "&";
-        static readonly Func<string, string> Decoder = WebUtility.UrlDecode;
-        static readonly DateTime EpochTime = new DateTime(1970, 1, 1, 0, 0, 0, 0, DateTimeKind.Utc);
-        readonly string token;
-        readonly DateTime expiresAtUtc;
-        readonly string audience;
+        /// <summary>
+        /// Expiry date-time
+        /// </summary>
+        DateTime expiresAtUtc;
+
+        /// <summary>
+        /// Token audience
+        /// </summary>
+        string audience;
+
+        /// <summary>
+        /// Token type
+        /// </summary>
+        string tokenType;
 
         /// <summary>
         /// Creates a new instance of the <see cref="SecurityToken"/> class.
         /// </summary>
+        /// <param name="tokenString">The token</param>
         /// <param name="expiresAtUtc">The expiration time</param>
-        public SecurityToken(string tokenString, DateTime expiresAtUtc, string audience)
+        /// <param name="audience">The audience</param>
+        /// <param name="tokenType">The type of the token</param>
+        public SecurityToken(string tokenString, DateTime expiresAtUtc, string audience, string tokenType)
         {
-            if (string.IsNullOrWhiteSpace(tokenString))
+            if (string.IsNullOrEmpty(tokenString))
             {
-                throw Fx.Exception.ArgumentNull(nameof(tokenString));
+                throw Fx.Exception.ArgumentNullOrWhiteSpace(nameof(tokenString));
             }
-            if (string.IsNullOrWhiteSpace(audience))
+
+            if (string.IsNullOrEmpty(audience))
             {
-                throw Fx.Exception.ArgumentNull(nameof(audience));
+                throw Fx.Exception.ArgumentNullOrWhiteSpace(nameof(audience));
             }
 
             this.token = tokenString;
             this.expiresAtUtc = expiresAtUtc;
             this.audience = audience;
-        }
-
-        /// <summary>
-        /// Creates a new instance of the <see cref="SecurityToken"/> class.
-        /// </summary>
-        /// <param name="expiresAtUtc">The expiration time</param>
-        public SecurityToken(string tokenString, DateTime expiresAtUtc)
-        {
-            if (string.IsNullOrWhiteSpace(tokenString))
-            {
-                throw Fx.Exception.ArgumentNull(nameof(tokenString));
-            }
-
-            this.token = tokenString;
-            this.expiresAtUtc = expiresAtUtc;
-            this.audience = this.GetAudienceFromToken(tokenString);
-        }
-
-        /// <summary>
-        /// Creates a new instance of the <see cref="SecurityToken"/> class.
-        /// </summary>
-        public SecurityToken(string tokenString)
-        {
-            if (string.IsNullOrWhiteSpace(tokenString))
-            {
-                throw Fx.Exception.ArgumentNull(nameof(tokenString));
-            }
-
-            this.token = tokenString;
-            this.GetExpirationDateAndAudienceFromToken(tokenString, out this.expiresAtUtc, out this.audience);
+            this.tokenType = tokenType;
         }
 
         /// <summary>
@@ -92,59 +68,11 @@ namespace Microsoft.Azure.ServiceBus.Primitives
         /// <summary>
         /// Gets the actual token.
         /// </summary>
-        public object TokenValue => this.token;
+        public virtual string TokenValue => this.token;
 
-        protected virtual string ExpiresOnFieldName => InternalExpiresOnFieldName;
-
-        protected virtual string AudienceFieldName => InternalAudienceFieldName;
-
-        protected virtual string KeyValueSeparator => InternalKeyValueSeparator;
-
-        protected virtual string PairSeparator => InternalPairSeparator;
-
-        static IDictionary<string, string> Decode(string encodedString, Func<string, string> keyDecoder, Func<string, string> valueDecoder, string keyValueSeparator, string pairSeparator)
-        {
-            var dictionary = new Dictionary<string, string>();
-            IEnumerable<string> valueEncodedPairs = encodedString.Split(new[] { pairSeparator }, StringSplitOptions.None);
-            foreach (var valueEncodedPairAsString in valueEncodedPairs)
-            {
-                var pair = valueEncodedPairAsString.Split(new[] { keyValueSeparator }, StringSplitOptions.None);
-                if (pair.Length != 2)
-                {
-                    throw new FormatException(Resources.InvalidEncoding);
-                }
-
-                dictionary.Add(keyDecoder(pair[0]), valueDecoder(pair[1]));
-            }
-
-            return dictionary;
-        }
-
-        string GetAudienceFromToken(string token)
-        {
-            var decodedToken = Decode(token, Decoder, Decoder, this.KeyValueSeparator, this.PairSeparator);
-            if (!decodedToken.TryGetValue(this.AudienceFieldName, out var audience))
-            {
-                throw new FormatException(Resources.TokenMissingAudience);
-            }
-
-            return audience;
-        }
-
-        void GetExpirationDateAndAudienceFromToken(string token, out DateTime expiresOn, out string audience)
-        {
-            IDictionary<string, string> decodedToken = Decode(token, Decoder, Decoder, this.KeyValueSeparator, this.PairSeparator);
-            if (!decodedToken.TryGetValue(this.ExpiresOnFieldName, out var expiresIn))
-            {
-                throw new FormatException(Resources.TokenMissingExpiresOn);
-            }
-
-            if (!decodedToken.TryGetValue(this.AudienceFieldName, out audience))
-            {
-                throw new FormatException(Resources.TokenMissingAudience);
-            }
-
-            expiresOn = (EpochTime + TimeSpan.FromSeconds(double.Parse(expiresIn, CultureInfo.InvariantCulture)));
-        }
+        /// <summary>
+        /// Gets the token type.
+        /// </summary>
+        public virtual string TokenType => this.tokenType;
     }
-}
+}

src/Microsoft.Azure.ServiceBus/Primitives/ServiceBusNamespaceConnection.cs

@@ -28,5 +28,22 @@ namespace Microsoft.Azure.ServiceBus.Primitives
 
             this.InitializeConnection(serviceBusConnectionStringBuilder);
         }
+
+        public ServiceBusNamespaceConnection(string endpoint, TransportType transportType, RetryPolicy retryPolicy)
+            : base(Constants.DefaultOperationTimeout, retryPolicy)
+        {
+            if (string.IsNullOrWhiteSpace(endpoint))
+            {
+                throw Fx.Exception.ArgumentNullOrWhiteSpace(nameof(endpoint));
+            }
+
+            var serviceBusConnectionStringBuilder = new ServiceBusConnectionStringBuilder()
+            {
+                Endpoint = endpoint,
+                TransportType = transportType
+            };
+
+            this.InitializeConnection(serviceBusConnectionStringBuilder);
+        }
     }
 }

src/Microsoft.Azure.ServiceBus/Primitives/SharedAccessSignatureToken.cs

@@ -0,0 +1,150 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.Azure.ServiceBus.Primitives
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Globalization;
+    using System.Net;
+
+    /// <summary>
+    /// A WCF SecurityToken that wraps a Shared Access Signature
+    /// </summary>
+    class SharedAccessSignatureToken : SecurityToken
+    {
+        internal const string SharedAccessSignature = "SharedAccessSignature";
+        internal const string SignedResource = "sr";
+        internal const string Signature = "sig";
+        internal const string SignedKeyName = "skn";
+        internal const string SignedExpiry = "se";
+        internal const int MaxKeyNameLength = 256;
+        internal const int MaxKeyLength = 256;
+
+        const string SignedResourceFullFieldName = SharedAccessSignature + " " + SignedResource;
+        const string SasPairSeparator = "&";
+        const string SasKeyValueSeparator = "=";
+
+        static readonly Func<string, string> Decoder = WebUtility.UrlDecode;
+
+        /// <summary>
+        /// Creates a new instance of the <see cref="SharedAccessSignatureToken"/> class.
+        /// </summary>
+        /// <param name="tokenString">The token</param>
+        public SharedAccessSignatureToken(string tokenString)
+            : base(tokenString, GetExpirationDateTimeUtcFromToken(tokenString), GetAudienceFromToken(tokenString), Constants.SasTokenType)
+        {
+        }
+
+        internal static void Validate(string sharedAccessSignature)
+        {
+            if (string.IsNullOrEmpty(sharedAccessSignature))
+            {
+                throw new ArgumentNullException(nameof(sharedAccessSignature));
+            }
+
+            IDictionary<string, string> parsedFields = ExtractFieldValues(sharedAccessSignature);
+
+            string signature;
+            if (!parsedFields.TryGetValue(Signature, out signature))
+            {
+                throw new ArgumentNullException(Signature);
+            }
+
+            string expiry;
+            if (!parsedFields.TryGetValue(SignedExpiry, out expiry))
+            {
+                throw new ArgumentNullException(SignedExpiry);
+            }
+
+            string keyName;
+            if (!parsedFields.TryGetValue(SignedKeyName, out keyName))
+            {
+                throw new ArgumentNullException(SignedKeyName);
+            }
+
+            string encodedAudience;
+            if (!parsedFields.TryGetValue(SignedResource, out encodedAudience))
+            {
+                throw new ArgumentNullException(SignedResource);
+            }
+        }
+
+        static IDictionary<string, string> ExtractFieldValues(string sharedAccessSignature)
+        {
+            string[] tokenLines = sharedAccessSignature.Split();
+
+            if (!string.Equals(tokenLines[0].Trim(), SharedAccessSignature, StringComparison.OrdinalIgnoreCase) || tokenLines.Length != 2)
+            {
+                throw new ArgumentNullException(nameof(sharedAccessSignature));
+            }
+
+            IDictionary<string, string> parsedFields = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
+            string[] tokenFields = tokenLines[1].Trim().Split(new[] { SasPairSeparator }, StringSplitOptions.None);
+
+            foreach (string tokenField in tokenFields)
+            {
+                if (tokenField != string.Empty)
+                {
+                    string[] fieldParts = tokenField.Split(new[] { SasKeyValueSeparator }, StringSplitOptions.None);
+                    if (string.Equals(fieldParts[0], SignedResource, StringComparison.OrdinalIgnoreCase))
+                    {
+                        // We need to preserve the casing of the escape characters in the audience,
+                        // so defer decoding the URL until later.
+                        parsedFields.Add(fieldParts[0], fieldParts[1]);
+                    }
+                    else
+                    {
+                        parsedFields.Add(fieldParts[0], WebUtility.UrlDecode(fieldParts[1]));
+                    }
+                }
+            }
+
+            return parsedFields;
+        }
+
+        static string GetAudienceFromToken(string token)
+        {
+            string audience;
+            IDictionary<string, string> decodedToken = Decode(token, Decoder, Decoder, SasKeyValueSeparator, SasPairSeparator);
+            if (!decodedToken.TryGetValue(SignedResourceFullFieldName, out audience))
+            {
+                throw new FormatException(Resources.TokenMissingAudience);
+            }
+
+            return audience;
+        }
+
+        static DateTime GetExpirationDateTimeUtcFromToken(string token)
+        {
+            string expiresIn;
+            IDictionary<string, string> decodedToken = Decode(token, Decoder, Decoder, SasKeyValueSeparator, SasPairSeparator);
+            if (!decodedToken.TryGetValue(SignedExpiry, out expiresIn))
+            {
+                throw new FormatException(Resources.TokenMissingExpiresOn);
+            }
+
+            var expiresOn = (Constants.EpochTime + TimeSpan.FromSeconds(double.Parse(expiresIn, CultureInfo.InvariantCulture)));
+
+            return expiresOn;
+        }
+
+        static IDictionary<string, string> Decode(string encodedString, Func<string, string> keyDecoder, Func<string, string> valueDecoder, string keyValueSeparator, string pairSeparator)
+        {
+            IDictionary<string, string> dictionary = new Dictionary<string, string>();
+            IEnumerable<string> valueEncodedPairs = encodedString.Split(new[] { pairSeparator }, StringSplitOptions.None);
+            foreach (string valueEncodedPair in valueEncodedPairs)
+            {
+                string[] pair = valueEncodedPair.Split(new[] { keyValueSeparator }, StringSplitOptions.None);
+                if (pair.Length != 2)
+                {
+                    throw new FormatException(Resources.InvalidEncoding);
+                }
+
+                dictionary.Add(keyDecoder(pair[0]), valueDecoder(pair[1]));
+            }
+
+            return dictionary;
+        }
+    }
+}

src/Microsoft.Azure.ServiceBus/Primitives/SharedAccessSignatureTokenProvider.cs

@@ -15,36 +15,42 @@ namespace Microsoft.Azure.ServiceBus.Primitives
     /// <summary>
     /// The SharedAccessSignatureTokenProvider generates tokens using a shared access key or existing signature.
     /// </summary>
-    internal class SharedAccessSignatureTokenProvider : TokenProvider
+    public class SharedAccessSignatureTokenProvider : TokenProvider
     {
-        /// <summary>
-        /// Represents 00:00:00 UTC Thursday 1, January 1970.
-        /// </summary>
-        public static readonly DateTime EpochTime = new DateTime(1970, 1, 1, 0, 0, 0, 0, DateTimeKind.Utc);
+        const TokenScope DefaultTokenScope = TokenScope.Entity;
+
+        internal static readonly TimeSpan DefaultTokenTimeout = TimeSpan.FromMinutes(60);
+
         readonly byte[] encodedSharedAccessKey;
         readonly string keyName;
         readonly TimeSpan tokenTimeToLive;
+        readonly TokenScope tokenScope;
         readonly string sharedAccessSignature;
+        internal static readonly Func<string, byte[]> MessagingTokenProviderKeyEncoder = Encoding.UTF8.GetBytes;
 
         internal SharedAccessSignatureTokenProvider(string sharedAccessSignature)
-            : base(TokenScope.Entity)
         {
             SharedAccessSignatureToken.Validate(sharedAccessSignature);
             this.sharedAccessSignature = sharedAccessSignature;
         }
 
-        internal SharedAccessSignatureTokenProvider(string keyName, string sharedAccessKey, TimeSpan tokenTimeToLive)
-            : this(keyName, sharedAccessKey, tokenTimeToLive, TokenScope.Entity)
+        internal SharedAccessSignatureTokenProvider(string keyName, string sharedAccessKey, TokenScope tokenScope = TokenScope.Entity)
+            : this(keyName, sharedAccessKey, MessagingTokenProviderKeyEncoder, DefaultTokenTimeout, tokenScope)
         {
         }
 
-        internal SharedAccessSignatureTokenProvider(string keyName, string sharedAccessKey, TimeSpan tokenTimeToLive, TokenScope tokenScope)
-            : this(keyName, sharedAccessKey, TokenProvider.MessagingTokenProviderKeyEncoder, tokenTimeToLive, tokenScope)
+        internal SharedAccessSignatureTokenProvider(string keyName, string sharedAccessKey, TimeSpan tokenTimeToLive, TokenScope tokenScope = TokenScope.Entity)
+            : this(keyName, sharedAccessKey, MessagingTokenProviderKeyEncoder, tokenTimeToLive, tokenScope)
         {
         }
 
+        /// <summary></summary>
+        /// <param name="keyName"></param>
+        /// <param name="sharedAccessKey"></param>
+        /// <param name="customKeyEncoder"></param>
+        /// <param name="tokenTimeToLive"></param>
+        /// <param name="tokenScope"></param>
         protected SharedAccessSignatureTokenProvider(string keyName, string sharedAccessKey, Func<string, byte[]> customKeyEncoder, TimeSpan tokenTimeToLive, TokenScope tokenScope)
-            : base(tokenScope)
         {
             if (string.IsNullOrEmpty(keyName))
             {
@@ -74,16 +80,28 @@ namespace Microsoft.Azure.ServiceBus.Primitives
             this.tokenTimeToLive = tokenTimeToLive;
             this.encodedSharedAccessKey = customKeyEncoder != null ?
                 customKeyEncoder(sharedAccessKey) :
-                TokenProvider.MessagingTokenProviderKeyEncoder(sharedAccessKey);
+                MessagingTokenProviderKeyEncoder(sharedAccessKey);
+            this.tokenScope = tokenScope;
         }
 
-        protected override Task<SecurityToken> OnGetTokenAsync(string appliesTo, string action, TimeSpan timeout)
+        /// <summary>
+        /// Gets a <see cref="SecurityToken"/> for the given audience and duration.
+        /// </summary>
+        /// <param name="appliesTo">The URI which the access token applies to</param>
+        /// <param name="timeout">The time span that specifies the timeout value for the message that gets the security token</param>
+        /// <returns><see cref="SecurityToken"/></returns>
+        public override Task<SecurityToken> GetTokenAsync(string appliesTo, TimeSpan timeout)
         {
-            var tokenString = this.BuildSignature(appliesTo);
-            var sharedAccessSignatureToken = new SharedAccessSignatureToken(tokenString);
-            return Task.FromResult<SecurityToken>(sharedAccessSignatureToken);
+            TimeoutHelper.ThrowIfNegativeArgument(timeout);
+            appliesTo = NormalizeAppliesTo(appliesTo);
+            string tokenString = this.BuildSignature(appliesTo);
+            var securityToken = new SharedAccessSignatureToken(tokenString);
+            return Task.FromResult<SecurityToken>(securityToken);
         }
 
+        /// <summary></summary>
+        /// <param name="targetUri"></param>
+        /// <returns></returns>
         protected virtual string BuildSignature(string targetUri)
         {
             return string.IsNullOrWhiteSpace(this.sharedAccessSignature)
@@ -95,6 +113,11 @@ namespace Microsoft.Azure.ServiceBus.Primitives
                 : this.sharedAccessSignature;
         }
 
+        string NormalizeAppliesTo(string appliesTo)
+        {
+            return ServiceBusUriHelper.NormalizeUri(appliesTo, "http", true, stripPath: this.tokenScope == TokenScope.Namespace, ensureTrailingSlash: true);
+        }
+
         static class SharedAccessSignatureBuilder
         {
             [SuppressMessage("Microsoft.Globalization", "CA1308:NormalizeStringsToUppercase", Justification = "Uris are normalized to lowercase")]
@@ -104,39 +127,34 @@ namespace Microsoft.Azure.ServiceBus.Primitives
                 string targetUri,
                 TimeSpan timeToLive)
             {
-                // Note that target URI is not normalized because in IoT scenario it
+                // Note that target URI is not normalized because in IoT scenario it 
                 // is case sensitive.
                 string expiresOn = BuildExpiresOn(timeToLive);
                 string audienceUri = WebUtility.UrlEncode(targetUri);
-                var fields = new List<string> { audienceUri, expiresOn };
+                List<string> fields = new List<string> { audienceUri, expiresOn };
 
                 // Example string to be signed:
                 // http://mynamespace.servicebus.windows.net/a/b/c?myvalue1=a
                 // <Value for ExpiresOn>
-                var signature = Sign(string.Join("\n", fields), encodedSharedAccessKey);
+                string signature = Sign(string.Join("\n", fields), encodedSharedAccessKey);
 
                 // Example returned string:
                 // SharedAccessKeySignature
                 // sr=ENCODED(http://mynamespace.servicebus.windows.net/a/b/c?myvalue1=a)&sig=<Signature>&se=<ExpiresOnValue>&skn=<KeyName>
-                return string.Format(
-                    CultureInfo.InvariantCulture,
-                    "{0} {1}={2}&{3}={4}&{5}={6}&{7}={8}",
+
+                return string.Format(CultureInfo.InvariantCulture, "{0} {1}={2}&{3}={4}&{5}={6}&{7}={8}",
                     SharedAccessSignatureToken.SharedAccessSignature,
-                    SharedAccessSignatureToken.SignedResource,
-                    audienceUri,
-                    SharedAccessSignatureToken.Signature,
-                    WebUtility.UrlEncode(signature),
-                    SharedAccessSignatureToken.SignedExpiry,
-                    WebUtility.UrlEncode(expiresOn),
-                    SharedAccessSignatureToken.SignedKeyName,
-                    WebUtility.UrlEncode(keyName));
+                    SharedAccessSignatureToken.SignedResource, audienceUri,
+                    SharedAccessSignatureToken.Signature, WebUtility.UrlEncode(signature),
+                    SharedAccessSignatureToken.SignedExpiry, WebUtility.UrlEncode(expiresOn),
+                    SharedAccessSignatureToken.SignedKeyName, WebUtility.UrlEncode(keyName));
             }
 
             static string BuildExpiresOn(TimeSpan timeToLive)
             {
-                var expiresOn = DateTime.UtcNow.Add(timeToLive);
-                var secondsFromBaseTime = expiresOn.Subtract(EpochTime);
-                var seconds = Convert.ToInt64(secondsFromBaseTime.TotalSeconds, CultureInfo.InvariantCulture);
+                DateTime expiresOn = DateTime.UtcNow.Add(timeToLive);
+                TimeSpan secondsFromBaseTime = expiresOn.Subtract(Constants.EpochTime);
+                long seconds = Convert.ToInt64(secondsFromBaseTime.TotalSeconds, CultureInfo.InvariantCulture);
                 return Convert.ToString(seconds, CultureInfo.InvariantCulture);
             }
 
@@ -148,98 +166,5 @@ namespace Microsoft.Azure.ServiceBus.Primitives
                 }
             }
         }
-
-        /// <summary>
-        /// A WCF SecurityToken that wraps a Shared Access Signature
-        /// </summary>
-        class SharedAccessSignatureToken : SecurityToken
-        {
-            public const int MaxKeyNameLength = 256;
-            public const int MaxKeyLength = 256;
-            public const string SharedAccessSignature = "SharedAccessSignature";
-            public const string SignedResource = "sr";
-            public const string Signature = "sig";
-            public const string SignedKeyName = "skn";
-            public const string SignedExpiry = "se";
-            public const string SignedResourceFullFieldName = SharedAccessSignature + " " + SignedResource;
-            public const string SasKeyValueSeparator = "=";
-            public const string SasPairSeparator = "&";
-
-            public SharedAccessSignatureToken(string tokenString)
-                : base(tokenString)
-            {
-            }
-
-            protected override string AudienceFieldName => SignedResourceFullFieldName;
-
-            protected override string ExpiresOnFieldName => SignedExpiry;
-
-            protected override string KeyValueSeparator => SasKeyValueSeparator;
-
-            protected override string PairSeparator => SasPairSeparator;
-
-            internal static void Validate(string sharedAccessSignature)
-            {
-                if (string.IsNullOrEmpty(sharedAccessSignature))
-                {
-                    throw new ArgumentNullException(nameof(sharedAccessSignature));
-                }
-
-                IDictionary<string, string> parsedFields = ExtractFieldValues(sharedAccessSignature);
-
-                if (!parsedFields.TryGetValue(Signature, out _))
-                {
-                    throw new ArgumentNullException(Signature);
-                }
-
-                if (!parsedFields.TryGetValue(SignedExpiry, out _))
-                {
-                    throw new ArgumentNullException(SignedExpiry);
-                }
-
-                if (!parsedFields.TryGetValue(SignedKeyName, out _))
-                {
-                    throw new ArgumentNullException(SignedKeyName);
-                }
-
-                if (!parsedFields.TryGetValue(SignedResource, out _))
-                {
-                    throw new ArgumentNullException(SignedResource);
-                }
-            }
-
-            static IDictionary<string, string> ExtractFieldValues(string sharedAccessSignature)
-            {
-                var tokenLines = sharedAccessSignature.Split();
-
-                if (!string.Equals(tokenLines[0].Trim(), SharedAccessSignature, StringComparison.OrdinalIgnoreCase) || tokenLines.Length != 2)
-                {
-                    throw new ArgumentNullException(nameof(sharedAccessSignature));
-                }
-
-                var parsedFields = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
-                var tokenFields = tokenLines[1].Trim().Split(new[] { SasPairSeparator }, StringSplitOptions.None);
-
-                foreach (var tokenField in tokenFields)
-                {
-                    if (tokenField != string.Empty)
-                    {
-                        var fieldParts = tokenField.Split(new[] { SasKeyValueSeparator }, StringSplitOptions.None);
-                        if (string.Equals(fieldParts[0], SignedResource, StringComparison.OrdinalIgnoreCase))
-                        {
-                            // We need to preserve the casing of the escape characters in the audience,
-                            // so defer decoding the URL until later.
-                            parsedFields.Add(fieldParts[0], fieldParts[1]);
-                        }
-                        else
-                        {
-                            parsedFields.Add(fieldParts[0], WebUtility.UrlDecode(fieldParts[1]));
-                        }
-                    }
-                }
-
-                return parsedFields;
-            }
-        }
     }
-}
+}

src/Microsoft.Azure.ServiceBus/Primitives/TokenProvider.cs

@@ -4,36 +4,14 @@
 namespace Microsoft.Azure.ServiceBus.Primitives
 {
     using System;
-    using System.Text;
     using System.Threading.Tasks;
+    using Microsoft.IdentityModel.Clients.ActiveDirectory;
 
     /// <summary>
     /// This abstract base class can be extended to implement additional token providers.
     /// </summary>
-    internal abstract class TokenProvider
+    public abstract class TokenProvider : ITokenProvider
     {
-        internal static readonly TimeSpan DefaultTokenTimeout = TimeSpan.FromMinutes(60);
-        internal static readonly Func<string, byte[]> MessagingTokenProviderKeyEncoder = Encoding.UTF8.GetBytes;
-        const TokenScope DefaultTokenScope = TokenScope.Entity;
-
-        protected TokenProvider()
-            : this(TokenProvider.DefaultTokenScope)
-        {
-        }
-
-        protected TokenProvider(TokenScope tokenScope)
-        {
-            this.TokenScope = tokenScope;
-            this.ThisLock = new object();
-        }
-
-        /// <summary>
-        /// Gets the scope or permissions associated with the token.
-        /// </summary>
-        public TokenScope TokenScope { get; }
-
-        protected object ThisLock { get; }
-
         /// <summary>
         /// Construct a TokenProvider based on a sharedAccessSignature.
         /// </summary>
@@ -52,19 +30,20 @@ namespace Microsoft.Azure.ServiceBus.Primitives
         /// <returns>A TokenProvider initialized with the provided RuleId and Password</returns>
         public static TokenProvider CreateSharedAccessSignatureTokenProvider(string keyName, string sharedAccessKey)
         {
-            return new SharedAccessSignatureTokenProvider(keyName, sharedAccessKey, DefaultTokenTimeout);
+            return new SharedAccessSignatureTokenProvider(keyName, sharedAccessKey);
         }
 
-        ////internal static TokenProvider CreateIoTTokenProvider(string keyName, string sharedAccessKey)
-        ////{
-        ////    return new IoTTokenProvider(keyName, sharedAccessKey, DefaultTokenTimeout);
-        ////}
+        //internal static TokenProvider CreateIoTTokenProvider(string keyName, string sharedAccessKey)
+        //{
+        //    return new IoTTokenProvider(keyName, sharedAccessKey, DefaultTokenTimeout);
+        //}
 
         /// <summary>
         /// Construct a TokenProvider based on the provided Key Name and Shared Access Key.
         /// </summary>
         /// <param name="keyName">The key name of the corresponding SharedAccessKeyAuthorizationRule.</param>
         /// <param name="sharedAccessKey">The key associated with the SharedAccessKeyAuthorizationRule</param>
+        /// <param name="tokenTimeToLive">The token time to live</param>
         /// <returns>A TokenProvider initialized with the provided RuleId and Password</returns>
         public static TokenProvider CreateSharedAccessSignatureTokenProvider(string keyName, string sharedAccessKey, TimeSpan tokenTimeToLive)
         {
@@ -80,7 +59,7 @@ namespace Microsoft.Azure.ServiceBus.Primitives
         /// <returns>A TokenProvider initialized with the provided RuleId and Password</returns>
         public static TokenProvider CreateSharedAccessSignatureTokenProvider(string keyName, string sharedAccessKey, TokenScope tokenScope)
         {
-            return new SharedAccessSignatureTokenProvider(keyName, sharedAccessKey, DefaultTokenTimeout, tokenScope);
+            return new SharedAccessSignatureTokenProvider(keyName, sharedAccessKey, tokenScope);
         }
 
         /// <summary>
@@ -88,6 +67,7 @@ namespace Microsoft.Azure.ServiceBus.Primitives
         /// </summary>
         /// <param name="keyName">The key name of the corresponding SharedAccessKeyAuthorizationRule.</param>
         /// <param name="sharedAccessKey">The key associated with the SharedAccessKeyAuthorizationRule</param>
+        /// <param name="tokenTimeToLive">The token time to live</param> 
         /// <param name="tokenScope">The tokenScope of tokens to request.</param>
         /// <returns>A TokenProvider initialized with the provided RuleId and Password</returns>
         public static TokenProvider CreateSharedAccessSignatureTokenProvider(string keyName, string sharedAccessKey, TimeSpan tokenTimeToLive, TokenScope tokenScope)
@@ -95,23 +75,96 @@ namespace Microsoft.Azure.ServiceBus.Primitives
             return new SharedAccessSignatureTokenProvider(keyName, sharedAccessKey, tokenTimeToLive, tokenScope);
         }
 
+        /// <summary>Creates an Azure Active Directory token provider.</summary>
+        /// <param name="authContext">AuthenticationContext for AAD.</param>
+        /// <param name="clientCredential">The app credential.</param>
+        /// <returns>The <see cref="TokenProvider" /> for returning Json web token.</returns>
+        public static TokenProvider CreateAadTokenProvider(AuthenticationContext authContext, ClientCredential clientCredential)
+        {
+            if (authContext == null)
+            {
+                throw new ArgumentNullException(nameof(authContext));
+            }
+
+            if (clientCredential == null)
+            {
+                throw new ArgumentNullException(nameof(clientCredential));
+            }
+
+            return new AzureActiveDirectoryTokenProvider(authContext, clientCredential);
+        }
+
+        /// <summary>Creates an Azure Active Directory token provider.</summary>
+        /// <param name="authContext">AuthenticationContext for AAD.</param>
+        /// <param name="clientId">ClientId for AAD.</param>
+        /// <param name="redirectUri">The redirectUri on Client App.</param>
+        /// <param name="platformParameters">Platform parameters</param>
+        /// <param name="userIdentifier">User Identifier</param>
+        /// <returns>The <see cref="TokenProvider" /> for returning Json web token.</returns>
+        public static TokenProvider CreateAadTokenProvider(
+            AuthenticationContext authContext,
+            string clientId,
+            Uri redirectUri,
+            IPlatformParameters platformParameters,
+            UserIdentifier userIdentifier = null)
+        {
+            if (authContext == null)
+            {
+                throw new ArgumentNullException(nameof(authContext));
+            }
+
+            if (string.IsNullOrEmpty(clientId))
+            {
+                throw new ArgumentNullException(nameof(clientId));
+            }
+
+            if (redirectUri == null)
+            {
+                throw new ArgumentNullException(nameof(redirectUri));
+            }
+
+            if (platformParameters == null)
+            {
+                throw new ArgumentNullException(nameof(platformParameters));
+            }
+
+            return new AzureActiveDirectoryTokenProvider(authContext, clientId, redirectUri, platformParameters, userIdentifier);
+        }
+
+#if !UAP10_0
+        /// <summary>Creates an Azure Active Directory token provider.</summary>
+        /// <param name="authContext">AuthenticationContext for AAD.</param>
+        /// <param name="clientAssertionCertificate">The client assertion certificate credential.</param>
+        /// <returns>The <see cref="TokenProvider" /> for returning Json web token.</returns>
+        public static TokenProvider CreateAadTokenProvider(AuthenticationContext authContext, ClientAssertionCertificate clientAssertionCertificate)
+        {
+            if (authContext == null)
+            {
+                throw new ArgumentNullException(nameof(authContext));
+            }
+
+            if (clientAssertionCertificate == null)
+            {
+                throw new ArgumentNullException(nameof(clientAssertionCertificate));
+            }
+
+            return new AzureActiveDirectoryTokenProvider(authContext, clientAssertionCertificate);
+        }
+#endif
+
+        /// <summary>Creates Azure Managed Service Identity token provider.</summary>
+        /// <returns>The <see cref="TokenProvider" /> for returning Json web token.</returns>
+        public static TokenProvider CreateManagedServiceIdentityTokenProvider()
+        {
+            return new ManagedServiceIdentityTokenProvider();
+        }
+
         /// <summary>
         /// Gets a <see cref="SecurityToken"/> for the given audience and duration.
         /// </summary>
         /// <param name="appliesTo">The URI which the access token applies to</param>
         /// <param name="timeout">The time span that specifies the timeout value for the message that gets the security token</param>
-        public Task<SecurityToken> GetTokenAsync(string appliesTo, string action, TimeSpan timeout)
-        {
-            TimeoutHelper.ThrowIfNegativeArgument(timeout);
-            appliesTo = this.NormalizeAppliesTo(appliesTo);
-            return this.OnGetTokenAsync(appliesTo, action, timeout);
-        }
-
-        protected abstract Task<SecurityToken> OnGetTokenAsync(string appliesTo, string action, TimeSpan timeout);
-
-        protected virtual string NormalizeAppliesTo(string appliesTo)
-        {
-            return ServiceBusUriHelper.NormalizeUri(appliesTo, "http", true, stripPath: this.TokenScope == TokenScope.Namespace, ensureTrailingSlash: true);
-        }
+        /// <returns></returns>
+        public abstract Task<SecurityToken> GetTokenAsync(string appliesTo, TimeSpan timeout);
     }
-}
+}

src/Microsoft.Azure.ServiceBus/Primitives/TokenProviderAdapter.cs

@@ -14,10 +14,10 @@ namespace Microsoft.Azure.ServiceBus.Primitives
     /// </summary>
     sealed class TokenProviderAdapter : ICbsTokenProvider
     {
-        readonly TokenProvider tokenProvider;
+        readonly ITokenProvider tokenProvider;
         readonly TimeSpan operationTimeout;
 
-        public TokenProviderAdapter(TokenProvider tokenProvider, TimeSpan operationTimeout)
+        public TokenProviderAdapter(ITokenProvider tokenProvider, TimeSpan operationTimeout)
         {
             Debug.Assert(tokenProvider != null, "tokenProvider cannot be null");
             this.tokenProvider = tokenProvider;
@@ -27,7 +27,7 @@ namespace Microsoft.Azure.ServiceBus.Primitives
         public async Task<CbsToken> GetTokenAsync(Uri namespaceAddress, string appliesTo, string[] requiredClaims)
         {
             var claim = requiredClaims?.FirstOrDefault();
-            var securityToken = await this.tokenProvider.GetTokenAsync(appliesTo, claim, this.operationTimeout).ConfigureAwait(false);
+            var securityToken = await this.tokenProvider.GetTokenAsync(appliesTo, this.operationTimeout).ConfigureAwait(false);
             return new CbsToken(securityToken.TokenValue, CbsConstants.ServiceBusSasTokenType, securityToken.ExpiresAtUtc);
         }
     }

src/Microsoft.Azure.ServiceBus/Primitives/TokenScope.cs

@@ -6,7 +6,7 @@ namespace Microsoft.Azure.ServiceBus.Primitives
     /// <summary>
     /// A enum representing the scope of the <see cref="SecurityToken"/>.
     /// </summary>
-    internal enum TokenScope
+    public enum TokenScope
     {
         /// <summary>
         /// The namespace.

src/Microsoft.Azure.ServiceBus/QueueClient.cs

@@ -7,9 +7,9 @@ namespace Microsoft.Azure.ServiceBus
     using System.Collections.Generic;
     using System.Threading;
     using System.Threading.Tasks;
-    using Core;
-    using Azure.Amqp;
-    using Primitives;
+    using Microsoft.Azure.Amqp;
+    using Microsoft.Azure.ServiceBus.Core;
+    using Microsoft.Azure.ServiceBus.Primitives;
 
     /// <summary>
     /// QueueClient can be used for all basic interactions with a Service Bus Queue.
@@ -56,6 +56,7 @@ namespace Microsoft.Azure.ServiceBus
     {
         readonly bool ownsConnection;
         readonly object syncLock;
+
         int prefetchCount;
         MessageSender innerSender;
         MessageReceiver innerReceiver;
@@ -94,6 +95,37 @@ namespace Microsoft.Azure.ServiceBus
                 throw Fx.Exception.ArgumentNullOrWhiteSpace(entityPath);
             }
 
+            this.InternalTokenProvider = this.ServiceBusConnection.CreateTokenProvider();
+            this.CbsTokenProvider = new TokenProviderAdapter(this.InternalTokenProvider, this.ServiceBusConnection.OperationTimeout);
+            this.ownsConnection = true;
+        }
+
+        /// <summary>
+        /// Creates a new instance of the Queue client using the specified endpoint, entity path, and token provider.
+        /// </summary>
+        /// <param name="endpoint">Fully qualified domain name for Service Bus. Most likely, {yournamespace}.servicebus.windows.net</param>
+        /// <param name="entityPath">Queue path.</param>
+        /// <param name="tokenProvider">Token provider which will generate security tokens for authorization.</param>
+        /// <param name="transportType">Transport type.</param>
+        /// <param name="receiveMode">Mode of receive of messages. Defaults to <see cref="ReceiveMode"/>.PeekLock.</param>
+        /// <param name="retryPolicy">Retry policy for queue operations. Defaults to <see cref="RetryPolicy.Default"/></param>
+        /// <returns></returns>
+        public QueueClient(
+            string endpoint,
+            string entityPath,
+            ITokenProvider tokenProvider,
+            TransportType transportType = TransportType.Amqp,
+            ReceiveMode receiveMode = ReceiveMode.PeekLock,
+            RetryPolicy retryPolicy = null)
+            : this(new ServiceBusNamespaceConnection(endpoint, transportType, retryPolicy), entityPath, receiveMode, retryPolicy)
+        {
+            if (tokenProvider == null)
+            {
+                throw Fx.Exception.ArgumentNull(nameof(tokenProvider));
+            }
+
+            this.InternalTokenProvider = tokenProvider;
+            this.CbsTokenProvider = new TokenProviderAdapter(this.InternalTokenProvider, this.ServiceBusConnection.OperationTimeout);
             this.ownsConnection = true;
         }
 
@@ -107,8 +139,6 @@ namespace Microsoft.Azure.ServiceBus
             this.syncLock = new object();
             this.QueueName = entityPath;
             this.ReceiveMode = receiveMode;
-            this.TokenProvider = this.ServiceBusConnection.CreateTokenProvider();
-            this.CbsTokenProvider = new TokenProviderAdapter(this.TokenProvider, serviceBusConnection.OperationTimeout);
 
             MessagingEventSource.Log.QueueClientCreateStop(serviceBusConnection.Endpoint.Authority, entityPath, this.ClientId);
         }
@@ -276,7 +306,7 @@ namespace Microsoft.Azure.ServiceBus
                                 this.ClientId,
                                 this.ReceiveMode,
                                 this.SessionClient,
-                                this.ServiceBusConnection.Endpoint.Authority);
+                                this.ServiceBusConnection.Endpoint);
                         }
                     }
                 }
@@ -289,7 +319,7 @@ namespace Microsoft.Azure.ServiceBus
 
         ICbsTokenProvider CbsTokenProvider { get; }
 
-        TokenProvider TokenProvider { get; }
+        ITokenProvider InternalTokenProvider { get; }
 
         /// <summary>
         /// Sends a message to Service Bus.
@@ -305,6 +335,7 @@ namespace Microsoft.Azure.ServiceBus
         public Task SendAsync(IList<Message> messageList)
         {
             this.ThrowIfClosed();
+
             return this.InnerSender.SendAsync(messageList);
         }
 

src/Microsoft.Azure.ServiceBus/QuotaExceededException.cs

@@ -11,12 +11,12 @@ namespace Microsoft.Azure.ServiceBus
     /// </summary>
     public sealed class QuotaExceededException : ServiceBusException
     {
-        internal QuotaExceededException(string message)
+        public QuotaExceededException(string message)
             : this(message, null)
         {
         }
 
-        internal QuotaExceededException(string message, Exception innerException)
+        public QuotaExceededException(string message, Exception innerException)
             : base(false, message, innerException)
         {
         }

src/Microsoft.Azure.ServiceBus/ServerBusyException.cs

@@ -10,12 +10,12 @@ namespace Microsoft.Azure.ServiceBus
     /// </summary>
     public sealed class ServerBusyException : ServiceBusException
     {
-        internal ServerBusyException(string message)
+        public ServerBusyException(string message)
             : this(message, null)
         {
         }
 
-        internal ServerBusyException(string message, Exception innerException)
+        public ServerBusyException(string message, Exception innerException)
             : base(true, message, innerException)
         {
         }

src/Microsoft.Azure.ServiceBus/ServiceBusCommunicationException.cs

@@ -10,12 +10,12 @@ namespace Microsoft.Azure.ServiceBus
     /// </summary>
     public class ServiceBusCommunicationException : ServiceBusException
     {
-        protected internal ServiceBusCommunicationException(string message)
+        public ServiceBusCommunicationException(string message)
             : this(message, null)
         {
         }
 
-        protected internal ServiceBusCommunicationException(string message, Exception innerException)
+        public ServiceBusCommunicationException(string message, Exception innerException)
             : base(true, message, innerException)
         {
         }

src/Microsoft.Azure.ServiceBus/ServiceBusDiagnosticsSource.cs

@@ -0,0 +1,773 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.Azure.ServiceBus
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Diagnostics;
+    using System.Linq;
+    using System.Threading.Tasks;
+
+    using Microsoft.Azure.ServiceBus.Diagnostics;
+
+    internal class ServiceBusDiagnosticSource
+    {
+        public const string DiagnosticListenerName = "Microsoft.Azure.ServiceBus";
+        public const string BaseActivityName = "Microsoft.Azure.ServiceBus.";
+
+        public const string ExceptionEventName = BaseActivityName + "Exception";
+        public const string ProcessActivityName =  BaseActivityName + "Process";
+
+        public const string ActivityIdPropertyName = "Diagnostic-Id";
+        public const string CorrelationContextPropertyName = "Correlation-Context";
+        public const string RelatedToTag = "RelatedTo";
+        public const string MessageIdTag = "MessageId";
+        public const string SessionIdTag = "SessionId";
+
+        private static readonly DiagnosticListener DiagnosticListener = new DiagnosticListener(DiagnosticListenerName);
+        private readonly string entityPath;
+        private readonly Uri endpoint;
+
+        public ServiceBusDiagnosticSource(string entityPath, Uri endpoint)
+        {
+            this.entityPath = entityPath;
+            this.endpoint = endpoint;
+        }
+
+        public static bool IsEnabled()
+        {
+            return DiagnosticListener.IsEnabled();
+        }
+
+
+        #region Send
+
+        internal Activity SendStart(IList<Message> messageList)
+        {
+            Activity activity = Start("Send", () => new
+                {
+                    Messages = messageList,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint
+                },
+                a => SetTags(a, messageList)
+            );
+
+            Inject(messageList);
+
+            return activity;
+        }
+
+        internal void SendStop(Activity activity, IList<Message> messageList, TaskStatus? status)
+        {
+            if (activity != null)
+            {
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    Messages = messageList,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    Status = status ?? TaskStatus.Faulted
+                });
+            }
+        }
+
+        #endregion
+
+
+        #region Process
+
+        internal Activity ProcessStart(Message message)
+        {
+            return ProcessStart("Process", message, () => new
+                {
+                    Message = message,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint
+                },
+                a => SetTags(a, message));
+        }
+
+        internal void ProcessStop(Activity activity, Message message, TaskStatus? status)
+        {
+            if (activity != null)
+            {
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    Message = message,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    Status = status ?? TaskStatus.Faulted
+                });
+            }
+        }
+
+        #endregion
+
+
+        #region ProcessSession
+
+        internal Activity ProcessSessionStart(IMessageSession session, Message message)
+        {
+            return ProcessStart("ProcessSession", message, () => new
+                {
+                    Session = session,
+                    Message = message,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint
+                },
+                a => SetTags(a, message));
+        }
+
+        internal void ProcessSessionStop(Activity activity, IMessageSession session, Message message, TaskStatus? status)
+        {
+            if (activity != null)
+            {
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    Session = session,
+                    Message = message,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    Status = status ?? TaskStatus.Faulted
+                });
+            }
+        }
+
+        #endregion
+
+
+        #region Schedule
+
+        internal Activity ScheduleStart(Message message, DateTimeOffset scheduleEnqueueTimeUtc)
+        {
+            Activity activity = Start("Schedule", () => new
+                {
+                    Message = message,
+                    ScheduleEnqueueTimeUtc = scheduleEnqueueTimeUtc,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint
+                },
+                a => SetTags(a, message));
+
+            Inject(message);
+
+            return activity;
+        }
+
+        internal void ScheduleStop(Activity activity, Message message, DateTimeOffset scheduleEnqueueTimeUtc, TaskStatus? status, long sequenceNumber)
+        {
+            if (activity != null)
+            {
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    Message = message,
+                    ScheduleEnqueueTimeUtc = scheduleEnqueueTimeUtc,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    SequenceNumber = sequenceNumber,
+                    Status = status ?? TaskStatus.Faulted
+                });
+            }
+        }
+
+        #endregion
+
+
+        #region Cancel
+
+        internal Activity CancelStart(long sequenceNumber)
+        {
+            return Start("Cancel", () => new
+            {
+                SequenceNumber = sequenceNumber,
+                Entity = this.entityPath,
+                Endpoint = this.endpoint
+            },
+            null);
+        }
+
+        internal void CancelStop(Activity activity, long sequenceNumber, TaskStatus? status)
+        {
+            if (activity != null)
+            {
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    SequenceNumber = sequenceNumber,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    Status = status ?? TaskStatus.Faulted
+                });
+            }
+        }
+
+        #endregion
+
+
+        #region Receive
+
+        internal Activity ReceiveStart(int messageCount)
+        {
+            return Start("Receive", () => new
+            {
+                RequestedMessageCount = messageCount,
+                Entity = this.entityPath,
+                Endpoint = this.endpoint
+            },
+            null);
+        }
+
+        internal void ReceiveStop(Activity activity, int messageCount, TaskStatus? status, IList<Message> messageList)
+        {
+            if (activity != null)
+            {
+                SetRelatedOperations(activity, messageList);
+                SetTags(activity, messageList);
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    RequestedMessageCount = messageCount,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    Status = status ?? TaskStatus.Faulted,
+                    Messages = messageList
+                });
+            }
+        }
+
+        #endregion
+
+
+        #region Peek
+
+        internal Activity PeekStart(long fromSequenceNumber, int messageCount)
+        {
+            return Start("Peek", () => new
+            {
+                FromSequenceNumber = fromSequenceNumber,
+                RequestedMessageCount = messageCount,
+                Entity = this.entityPath,
+                Endpoint = this.endpoint
+            },
+            null);
+        }
+
+        internal void PeekStop(Activity activity, long fromSequenceNumber, int messageCount, TaskStatus? status, IList<Message> messageList)
+        {
+            if (activity != null)
+            {
+                SetRelatedOperations(activity, messageList);
+                SetTags(activity, messageList);
+
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    FromSequenceNumber = fromSequenceNumber,
+                    RequestedMessageCount = messageCount,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    Status = status ?? TaskStatus.Faulted,
+                    Messages = messageList
+                });
+            }
+        }
+
+        #endregion
+
+
+        #region ReceiveDeferred
+
+        internal Activity ReceiveDeferredStart(IEnumerable<long> sequenceNumbers)
+        {
+            return Start("ReceiveDeferred", () => new
+            {
+                SequenceNumbers = sequenceNumbers,
+                Entity = this.entityPath,
+                Endpoint = this.endpoint
+            },
+            null);
+        }
+
+        internal void ReceiveDeferredStop(Activity activity, IEnumerable<long> sequenceNumbers, TaskStatus? status, IList<Message> messageList)
+        {
+            if (activity != null)
+            {
+                SetRelatedOperations(activity, messageList);
+                SetTags(activity, messageList);
+
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    SequenceNumbers = sequenceNumbers,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    Messages = messageList,
+                    Status = status ?? TaskStatus.Faulted
+                });
+            }
+        }
+
+        #endregion
+
+
+        #region  Complete
+
+        internal Activity CompleteStart(IList<string> lockTokens)
+        {
+            return Start("Complete", () => new
+            {
+                LockTokens = lockTokens,
+                Entity = this.entityPath,
+                Endpoint = this.endpoint
+            },
+            null);
+        }
+
+        internal void CompleteStop(Activity activity, IList<string> lockTokens, TaskStatus? status)
+        {
+            if (activity != null)
+            {
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    LockTokens = lockTokens,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    Status = status ?? TaskStatus.Faulted
+                });
+            }
+        }
+
+        #endregion
+
+
+        #region Dispose
+
+        internal Activity DisposeStart(string operationName, string lockToken)
+        {
+            return Start(operationName, () => new
+            {
+                LockToken = lockToken,
+                Entity = this.entityPath,
+                Endpoint = this.endpoint
+            },
+            null);
+        }
+
+        internal void DisposeStop(Activity activity, string lockToken, TaskStatus? status)
+        {
+            if (activity != null)
+            {
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    LockToken = lockToken,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    Status = status ?? TaskStatus.Faulted
+                });
+            }
+        }
+
+        #endregion
+
+ 
+        #region RenewLock
+
+        internal Activity RenewLockStart(string lockToken)
+        {
+            return Start("RenewLock", () => new
+            {
+                LockToken = lockToken,
+                Entity = this.entityPath,
+                Endpoint = this.endpoint
+            },
+            null);
+        }
+
+        internal void RenewLockStop(Activity activity, string lockToken, TaskStatus? status, DateTime lockedUntilUtc)
+        {
+            if (activity != null)
+            {
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    LockToken = lockToken,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    Status = status ?? TaskStatus.Faulted,
+                    LockedUntilUtc = lockedUntilUtc
+                });
+            }
+        }
+
+        #endregion
+
+
+        #region AddRule
+
+        internal Activity AddRuleStart(RuleDescription description)
+        {
+            return Start("AddRule", () => new
+            {
+                Rule = description,
+                Entity = this.entityPath,
+                Endpoint = this.endpoint
+            },
+            null);
+        }
+
+        internal void AddRuleStop(Activity activity, RuleDescription description, TaskStatus? status)
+        {
+            if (activity != null)
+            {
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    Rule = description,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    Status = status ?? TaskStatus.Faulted
+                });
+            }
+        }
+
+        #endregion
+
+
+        #region RemoveRule
+
+        internal Activity RemoveRuleStart(string ruleName)
+        {
+            return Start("RemoveRule", () => new
+            {
+                RuleName = ruleName,
+                Entity = this.entityPath,
+                Endpoint = this.endpoint
+            },
+            null);
+        }
+
+        internal void RemoveRuleStop(Activity activity, string ruleName, TaskStatus? status)
+        {
+            if (activity != null)
+            {
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    RuleName = ruleName,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    Status = status ?? TaskStatus.Faulted
+                });
+            }
+        }
+
+        #endregion
+
+
+        #region GetRules
+
+        internal Activity GetRulesStart()
+        {
+            return Start("GetRules", () => new
+            {
+                Entity = this.entityPath,
+                Endpoint = this.endpoint
+            },
+            null);
+        }
+
+        internal void GetRulesStop(Activity activity, IEnumerable<RuleDescription> rules, TaskStatus? status)
+        {
+            if (activity != null)
+            {
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    Rules = rules,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    Status = status ?? TaskStatus.Faulted
+                });
+            }
+        }
+
+        #endregion
+
+
+        #region AcceptMessageSession
+
+        internal Activity AcceptMessageSessionStart(string sessionId)
+        {
+            return Start("AcceptMessageSession", () => new
+                {
+                    SessionId = sessionId,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint
+                },
+                a => SetSessionTag(a, sessionId)
+            );
+        }
+
+        internal void AcceptMessageSessionStop(Activity activity, string sessionId, TaskStatus? status)
+        {
+            if (activity != null)
+            {
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    SessionId = sessionId,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    Status = status ?? TaskStatus.Faulted
+                });
+            }
+        }
+
+        #endregion
+
+
+        #region GetSessionStateAsync
+
+        internal Activity GetSessionStateStart(string sessionId)
+        {
+            return Start("GetSessionState", () => new
+                {
+                    SessionId = sessionId,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint
+                },
+                a => SetSessionTag(a, sessionId));
+        }
+
+        internal void GetSessionStateStop(Activity activity, string sessionId, byte[] state, TaskStatus? status)
+        {
+            if (activity != null)
+            {
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    SessionId = sessionId,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    Status = status ?? TaskStatus.Faulted,
+                    State = state
+                });
+            }
+        }
+
+        #endregion
+
+
+        #region SetSessionState
+
+        internal Activity SetSessionStateStart(string sessionId, byte[] state)
+        {
+            return Start("SetSessionState", () => new
+                {
+                    State = state,
+                    SessionId = sessionId,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint
+                },
+                a => SetSessionTag(a, sessionId));
+        }
+
+        internal void SetSessionStateStop(Activity activity, byte[] state, string sessionId, TaskStatus? status)
+        {
+            if (activity != null)
+            {
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    State = state,
+                    SessionId = sessionId,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    Status = status ?? TaskStatus.Faulted
+                });
+            }
+        }
+
+        #endregion
+
+
+        #region RenewSessionLock
+
+        internal Activity RenewSessionLockStart(string sessionId)
+        {
+            return Start("RenewSessionLock", () => new
+            {
+                SessionId = sessionId,
+                Entity = this.entityPath,
+                Endpoint = this.endpoint
+            },
+            a => SetSessionTag(a, sessionId));
+        }
+
+        internal void RenewSessionLockStop(Activity activity, string sessionId, TaskStatus? status)
+        {
+            if (activity != null)
+            {
+                DiagnosticListener.StopActivity(activity, new
+                {
+                    SessionId = sessionId,
+                    Entity = this.entityPath,
+                    Endpoint = this.endpoint,
+                    Status = status ?? TaskStatus.Faulted
+                });
+            }
+        }
+
+        #endregion
+
+        internal void ReportException(Exception ex)
+        {
+            if (DiagnosticListener.IsEnabled(ExceptionEventName))
+            {
+                DiagnosticListener.Write(ExceptionEventName,
+                    new
+                    {
+                        Exception = ex,
+                        Entity = this.entityPath,
+                        Endpoint = this.endpoint
+                    });
+            }
+        }
+
+        private Activity Start(string operationName, Func<object> getPayload, Action<Activity> setTags)
+        {
+            Activity activity = null;
+            string activityName = BaseActivityName + operationName;
+            if (DiagnosticListener.IsEnabled(activityName, this.entityPath))
+            {
+                activity = new Activity(activityName);
+                setTags?.Invoke(activity);
+
+                if (DiagnosticListener.IsEnabled(activityName + ".Start"))
+                {
+                    DiagnosticListener.StartActivity(activity, getPayload());
+                }
+                else
+                {
+                    activity.Start();
+                }
+            }
+
+            return activity;
+        }
+
+        private void Inject(IList<Message> messageList)
+        {
+            var currentActivity = Activity.Current;
+            if (currentActivity != null)
+            {
+                var correlationContext = SerializeCorrelationContext(currentActivity.Baggage.ToList());
+
+                foreach (var message in messageList)
+                {
+                    Inject(message, currentActivity.Id, correlationContext);
+                }
+            }
+        }
+
+        private void Inject(Message message)
+        {
+            var currentActivity = Activity.Current;
+            if (currentActivity != null)
+            {
+                Inject(message, currentActivity.Id, SerializeCorrelationContext(currentActivity.Baggage.ToList()));
+            }
+        }
+
+        private void Inject(Message message, string id, string correlationContext)
+        {
+            if (!message.UserProperties.ContainsKey(ActivityIdPropertyName))
+            {
+                message.UserProperties[ActivityIdPropertyName] = id;
+                if (correlationContext != null)
+                {
+                    message.UserProperties[CorrelationContextPropertyName] = correlationContext;
+                }
+            }
+        }
+
+        private string SerializeCorrelationContext(IList<KeyValuePair<string,string>> baggage)
+        {
+            if (baggage.Any())
+            {
+                return string.Join(",", baggage.Select(kvp => kvp.Key + "=" + kvp.Value));
+            }
+            return null;
+        }
+
+        private void SetRelatedOperations(Activity activity, IList<Message> messageList)
+        {
+            if (messageList != null && messageList.Count > 0)
+            {
+                var relatedTo = new List<string>();
+                foreach (var message in messageList)
+                {
+                    if (message.TryExtractId(out string id))
+                    {
+                        relatedTo.Add(id);
+                    }
+                }
+
+                if (relatedTo.Count > 0)
+                {
+                    activity.AddTag(RelatedToTag, string.Join(",", relatedTo.Distinct()));
+                }
+            }
+        }
+
+        private Activity ProcessStart(string operationName, Message message, Func<object> getPayload, Action<Activity> setTags)
+        {
+            Activity activity = null;
+            string activityName = BaseActivityName + operationName;
+
+            if (DiagnosticListener.IsEnabled(activityName, entityPath))
+            {
+                var tmpActivity = message.ExtractActivity(activityName);
+                setTags?.Invoke(tmpActivity);
+                
+                if (DiagnosticListener.IsEnabled(activityName, entityPath, tmpActivity))
+                {
+                    activity = tmpActivity;
+                    if (DiagnosticListener.IsEnabled(activityName + ".Start"))
+                    {
+                        DiagnosticListener.StartActivity(activity, getPayload());
+                    }
+                    else
+                    {
+                        activity.Start();
+                    }
+                }
+            }
+            return activity;
+        }
+
+        private void SetTags(Activity activity, IList<Message> messageList)
+        {
+            var messageIds = messageList.Where(m => m.MessageId != null).Select(m => m.MessageId).ToArray();
+            if (messageIds.Any())
+            {
+                activity.AddTag(MessageIdTag, string.Join(",", messageIds));
+            }
+
+            var sessionIds = messageList.Where(m => m.SessionId != null).Select(m => m.SessionId).Distinct().ToArray();
+            if (sessionIds.Any())
+            {
+                activity.AddTag(SessionIdTag, string.Join(",", sessionIds));
+            }
+        }
+
+        private void SetTags(Activity activity, Message message)
+        {
+            if (message.MessageId != null)
+            {
+                activity.AddTag(MessageIdTag, message.MessageId);
+            }
+
+            SetSessionTag(activity, message.SessionId);
+        }
+
+        private void SetSessionTag(Activity activity, string sessionId)
+        {
+            if (sessionId != null)
+            {
+                activity.AddTag(SessionIdTag, sessionId);
+            }
+        }
+    }
+}

src/Microsoft.Azure.ServiceBus/ServiceBusTimeoutException.cs

@@ -10,11 +10,11 @@ namespace Microsoft.Azure.ServiceBus
     /// </summary>
     public class ServiceBusTimeoutException : ServiceBusException
     {
-        internal ServiceBusTimeoutException(string message) : this(message, null)
+        public ServiceBusTimeoutException(string message) : this(message, null)
         {
         }
 
-        internal ServiceBusTimeoutException(string message, Exception innerException) : base(true, message, innerException)
+        public ServiceBusTimeoutException(string message, Exception innerException) : base(true, message, innerException)
         {
         }
     }

src/Microsoft.Azure.ServiceBus/SessionClient.cs

@@ -5,6 +5,7 @@ namespace Microsoft.Azure.ServiceBus
 {
     using System;
     using System.Collections.Generic;
+    using System.Diagnostics;
     using System.Linq;
     using System.Threading.Tasks;
     using Amqp;
@@ -45,6 +46,7 @@ namespace Microsoft.Azure.ServiceBus
     {
         const int DefaultPrefetchCount = 0;
         readonly bool ownsConnection;
+        readonly ServiceBusDiagnosticSource diagnosticSource;
 
         /// <summary>
         /// Creates a new SessionClient from a <see cref="ServiceBusConnectionStringBuilder"/>
@@ -123,6 +125,7 @@ namespace Microsoft.Azure.ServiceBus
             this.ReceiveMode = receiveMode;
             this.PrefetchCount = prefetchCount;
             this.CbsTokenProvider = cbsTokenProvider;
+            this.diagnosticSource = new ServiceBusDiagnosticSource(entityPath, serviceBusConnection.Endpoint);
 
             // Register plugins on the message session.
             if (registeredPlugins != null)
@@ -213,6 +216,10 @@ namespace Microsoft.Azure.ServiceBus
                 this.PrefetchCount,
                 sessionId);
 
+            bool isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+            Activity activity = isDiagnosticSourceEnabled ? this.diagnosticSource.AcceptMessageSessionStart(sessionId) : null;
+            Task acceptMessageSessionTask = null;
+
             var session = new MessageSession(
                 this.EntityPath,
                 this.EntityType,
@@ -226,11 +233,18 @@ namespace Microsoft.Azure.ServiceBus
 
             try
             {
-                await this.RetryPolicy.RunOperation(() => session.GetSessionReceiverLinkAsync(serverWaitTime), serverWaitTime)
-                    .ConfigureAwait(false);
+                acceptMessageSessionTask = this.RetryPolicy.RunOperation(
+                    () => session.GetSessionReceiverLinkAsync(serverWaitTime),
+                    serverWaitTime);
+                await acceptMessageSessionTask.ConfigureAwait(false);
             }
             catch (Exception exception)
             {
+                if (isDiagnosticSourceEnabled)
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
+
                 MessagingEventSource.Log.AmqpSessionClientAcceptMessageSessionException(
                     this.ClientId,
                     this.EntityPath,
@@ -239,6 +253,10 @@ namespace Microsoft.Azure.ServiceBus
                 await session.CloseAsync().ConfigureAwait(false);
                 throw AmqpExceptionHelper.GetClientException(exception);
             }
+            finally
+            {
+                this.diagnosticSource.AcceptMessageSessionStop(activity, session.SessionId, acceptMessageSessionTask?.Status);
+            }
 
             MessagingEventSource.Log.AmqpSessionClientAcceptMessageSessionStop(
                 this.ClientId,

src/Microsoft.Azure.ServiceBus/SessionLockLostException.cs

@@ -10,12 +10,12 @@ namespace Microsoft.Azure.ServiceBus
     /// </summary>
     public sealed class SessionLockLostException : ServiceBusException
     {
-        internal SessionLockLostException(string message)
+        public SessionLockLostException(string message)
             : this(message, null)
         {
         }
 
-        internal SessionLockLostException(string message, Exception innerException)
+        public SessionLockLostException(string message, Exception innerException)
             : base(false, message, innerException)
         {
         }

src/Microsoft.Azure.ServiceBus/SessionPumpHost.cs

@@ -12,9 +12,9 @@ namespace Microsoft.Azure.ServiceBus
         readonly object syncLock;
         SessionReceivePump sessionReceivePump;
         CancellationTokenSource sessionPumpCancellationTokenSource;
-        readonly string endpoint;
+        readonly Uri endpoint;
 
-        public SessionPumpHost(string clientId, ReceiveMode receiveMode, ISessionClient sessionClient, string endpoint)
+        public SessionPumpHost(string clientId, ReceiveMode receiveMode, ISessionClient sessionClient, Uri endpoint)
         {
             this.syncLock = new object();
             this.ClientId = clientId;

src/Microsoft.Azure.ServiceBus/SessionReceivePump.cs

@@ -4,6 +4,7 @@
 namespace Microsoft.Azure.ServiceBus
 {
     using System;
+    using System.Diagnostics;
     using System.Threading;
     using System.Threading.Tasks;
     using Primitives;
@@ -19,13 +20,14 @@ namespace Microsoft.Azure.ServiceBus
         readonly CancellationToken pumpCancellationToken;
         readonly SemaphoreSlim maxConcurrentSessionsSemaphoreSlim;
         readonly SemaphoreSlim maxPendingAcceptSessionsSemaphoreSlim;
+        private readonly ServiceBusDiagnosticSource diagnosticSource;
 
         public SessionReceivePump(string clientId,
             ISessionClient client,
             ReceiveMode receiveMode,
             SessionHandlerOptions sessionHandlerOptions,
             Func<IMessageSession, Message, CancellationToken, Task> callback,
-            string endpoint,
+            Uri endpoint,
             CancellationToken token)
         {
             this.client = client ?? throw new ArgumentException(nameof(client));
@@ -33,11 +35,12 @@ namespace Microsoft.Azure.ServiceBus
             this.ReceiveMode = receiveMode;
             this.sessionHandlerOptions = sessionHandlerOptions;
             this.userOnSessionCallback = callback;
-            this.endpoint = endpoint;
+            this.endpoint = endpoint.Authority;
             this.entityPath = client.EntityPath;
             this.pumpCancellationToken = token;
             this.maxConcurrentSessionsSemaphoreSlim = new SemaphoreSlim(this.sessionHandlerOptions.MaxConcurrentSessions);
             this.maxPendingAcceptSessionsSemaphoreSlim = new SemaphoreSlim(this.sessionHandlerOptions.MaxConcurrentAcceptSessionCalls);
+            this.diagnosticSource = new ServiceBusDiagnosticSource(client.EntityPath, endpoint);
         }
 
         ReceiveMode ReceiveMode { get; }
@@ -146,7 +149,10 @@ namespace Microsoft.Azure.ServiceBus
                     }
                     else
                     {
-                        await this.RaiseExceptionReceived(exception, ExceptionReceivedEventArgsAction.AcceptMessageSession).ConfigureAwait(false);
+                        if (!(exception is ObjectDisposedException && this.pumpCancellationToken.IsCancellationRequested))
+                        {
+                            await this.RaiseExceptionReceived(exception, ExceptionReceivedEventArgsAction.AcceptMessageSession).ConfigureAwait(false); 
+                        }
                         if (!MessagingUtilities.ShouldRetry(exception))
                         {
                             break;
@@ -197,7 +203,10 @@ namespace Microsoft.Azure.ServiceBus
                             continue;
                         }
 
-                        await this.RaiseExceptionReceived(exception, ExceptionReceivedEventArgsAction.Receive).ConfigureAwait(false);
+                        if (!(exception is ObjectDisposedException && this.pumpCancellationToken.IsCancellationRequested))
+                        {
+                            await this.RaiseExceptionReceived(exception, ExceptionReceivedEventArgsAction.Receive).ConfigureAwait(false); 
+                        }
                         break;
                     }
 
@@ -207,36 +216,54 @@ namespace Microsoft.Azure.ServiceBus
                         break;
                     }
 
-                    // Set the timer
-                    userCallbackTimer.Change(this.sessionHandlerOptions.MaxAutoRenewDuration, TimeSpan.FromMilliseconds(-1));
-                    var callbackExceptionOccurred = false;
+                    bool isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+                    Activity activity = isDiagnosticSourceEnabled ? this.diagnosticSource.ProcessSessionStart(session, message) : null;
+                    Task processTask = null;
+
                     try
                     {
-                        await this.userOnSessionCallback(session, message, this.pumpCancellationToken).ConfigureAwait(false);
-                    }
-                    catch (Exception exception)
-                    {
-                        MessagingEventSource.Log.MessageReceivePumpTaskException(this.clientId, session.SessionId, exception);
-                        await this.RaiseExceptionReceived(exception, ExceptionReceivedEventArgsAction.UserCallback).ConfigureAwait(false);
-                        callbackExceptionOccurred = true;
-                        if (!(exception is MessageLockLostException || exception is SessionLockLostException))
+                        // Set the timer
+                        userCallbackTimer.Change(this.sessionHandlerOptions.MaxAutoRenewDuration,
+                            TimeSpan.FromMilliseconds(-1));
+                        var callbackExceptionOccurred = false;
+                        try
                         {
-                            await this.AbandonMessageIfNeededAsync(session, message).ConfigureAwait(false);
+                            processTask = this.userOnSessionCallback(session, message, this.pumpCancellationToken);
+                            await processTask.ConfigureAwait(false);
+                        }
+                        catch (Exception exception)
+                        {
+                            if (isDiagnosticSourceEnabled)
+                            {
+                                this.diagnosticSource.ReportException(exception);
+                            }
+
+                            MessagingEventSource.Log.MessageReceivePumpTaskException(this.clientId, session.SessionId, exception);
+                            await this.RaiseExceptionReceived(exception, ExceptionReceivedEventArgsAction.UserCallback).ConfigureAwait(false);
+                            callbackExceptionOccurred = true;
+                            if (!(exception is MessageLockLostException || exception is SessionLockLostException))
+                            {
+                                await this.AbandonMessageIfNeededAsync(session, message).ConfigureAwait(false);
+                            }
+                        }
+                        finally
+                        {
+                            userCallbackTimer.Change(Timeout.Infinite, Timeout.Infinite);
+                        }
+
+                        if (!callbackExceptionOccurred)
+                        {
+                            await this.CompleteMessageIfNeededAsync(session, message).ConfigureAwait(false);
+                        }
+                        else if (session.IsClosedOrClosing)
+                        {
+                            // If User closed the session as part of the callback, break out of the loop
+                            break;
                         }
                     }
                     finally
                     {
-                        userCallbackTimer.Change(Timeout.Infinite, Timeout.Infinite);
-                    }
-
-                    if (!callbackExceptionOccurred)
-                    {
-                        await this.CompleteMessageIfNeededAsync(session, message).ConfigureAwait(false);
-                    }
-                    else if (session.IsClosedOrClosing)
-                    {
-                        // If User closed the session as part of the callback, break out of the loop
-                        break;
+                        this.diagnosticSource.ProcessSessionStop(activity, session, message, processTask?.Status);
                     }
                 }
             }
@@ -293,9 +320,10 @@ namespace Microsoft.Azure.ServiceBus
                 {
                     MessagingEventSource.Log.SessionReceivePumpSessionRenewLockException(this.clientId, session.SessionId, exception);
 
-                    // TaskCancelled is expected here as renewTasks will be cancelled after the Complete call is made.
-                    // Lets not bother user with this exception.
-                    if (!(exception is TaskCanceledException))
+                    // TaskCanceled is expected here as renewTasks will be cancelled after the Complete call is made.
+                    // ObjectDisposedException should only happen here because the CancellationToken was disposed at which point 
+                    // this renew exception is not relevant anymore. Lets not bother user with this exception.
+                    if (!(exception is TaskCanceledException) && !(exception is ObjectDisposedException))
                     {
                         await this.RaiseExceptionReceived(exception, ExceptionReceivedEventArgsAction.RenewLock).ConfigureAwait(false);
                     }

src/Microsoft.Azure.ServiceBus/SubscriptionClient.cs

@@ -5,6 +5,7 @@ namespace Microsoft.Azure.ServiceBus
 {
     using System;
     using System.Collections.Generic;
+    using System.Diagnostics;
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Amqp;
@@ -52,6 +53,8 @@ namespace Microsoft.Azure.ServiceBus
         int prefetchCount;
         readonly object syncLock;
         readonly bool ownsConnection;
+        readonly ServiceBusDiagnosticSource diagnosticSource;
+
         IInnerSubscriptionClient innerSubscriptionClient;
         SessionClient sessionClient;
         SessionPumpHost sessionPumpHost;
@@ -91,6 +94,39 @@ namespace Microsoft.Azure.ServiceBus
                 throw Fx.Exception.ArgumentNullOrWhiteSpace(subscriptionName);
             }
 
+            this.InternalTokenProvider = this.ServiceBusConnection.CreateTokenProvider();
+            this.CbsTokenProvider = new TokenProviderAdapter(this.InternalTokenProvider, this.ServiceBusConnection.OperationTimeout);
+            this.ownsConnection = true;
+        }
+
+        /// <summary>
+        /// Creates a new instance of the Subscription client using the specified endpoint, entity path, and token provider.
+        /// </summary>
+        /// <param name="endpoint">Fully qualified domain name for Service Bus. Most likely, {yournamespace}.servicebus.windows.net</param>
+        /// <param name="topicPath">Topic path.</param>
+        /// <param name="subscriptionName">Subscription name.</param>
+        /// <param name="tokenProvider">Token provider which will generate security tokens for authorization.</param>
+        /// <param name="transportType">Transport type.</param>
+        /// <param name="receiveMode">Mode of receive of messages. Defaults to <see cref="ReceiveMode"/>.PeekLock.</param>
+        /// <param name="retryPolicy">Retry policy for subscription operations. Defaults to <see cref="RetryPolicy.Default"/></param>
+        /// <returns></returns>
+        public SubscriptionClient(
+            string endpoint,
+            string topicPath,
+            string subscriptionName,
+            ITokenProvider tokenProvider,
+            TransportType transportType = TransportType.Amqp,
+            ReceiveMode receiveMode = ReceiveMode.PeekLock,
+            RetryPolicy retryPolicy = null)
+            : this(new ServiceBusNamespaceConnection(endpoint, transportType, retryPolicy), topicPath, subscriptionName, receiveMode, retryPolicy)
+        {
+            if (tokenProvider == null)
+            {
+                throw Fx.Exception.ArgumentNull(nameof(tokenProvider));
+            }
+
+            this.InternalTokenProvider = tokenProvider;
+            this.CbsTokenProvider = new TokenProviderAdapter(this.InternalTokenProvider, this.ServiceBusConnection.OperationTimeout);
             this.ownsConnection = true;
         }
 
@@ -106,8 +142,7 @@ namespace Microsoft.Azure.ServiceBus
             this.SubscriptionName = subscriptionName;
             this.Path = EntityNameHelper.FormatSubscriptionPath(this.TopicPath, this.SubscriptionName);
             this.ReceiveMode = receiveMode;
-            this.TokenProvider = this.ServiceBusConnection.CreateTokenProvider();
-            this.CbsTokenProvider = new TokenProviderAdapter(this.TokenProvider, serviceBusConnection.OperationTimeout);
+            this.diagnosticSource = new ServiceBusDiagnosticSource(this.Path, serviceBusConnection.Endpoint);
 
             MessagingEventSource.Log.SubscriptionClientCreateStop(serviceBusConnection.Endpoint.Authority, topicPath, subscriptionName, this.ClientId);
         }
@@ -246,7 +281,7 @@ namespace Microsoft.Azure.ServiceBus
                                 this.ClientId,
                                 this.ReceiveMode,
                                 this.SessionClient,
-                                this.ServiceBusConnection.Endpoint.Authority);
+                                this.ServiceBusConnection.Endpoint);
                         }
                     }
                 }
@@ -259,7 +294,7 @@ namespace Microsoft.Azure.ServiceBus
 
         ICbsTokenProvider CbsTokenProvider { get; }
 
-        TokenProvider TokenProvider { get; }
+        ITokenProvider InternalTokenProvider { get; }
 
         /// <summary>
         /// Completes a <see cref="Message"/> using its lock token. This will delete the message from the subscription.
@@ -299,7 +334,7 @@ namespace Microsoft.Azure.ServiceBus
         /// <remarks>
         /// A lock token can be found in <see cref="Message.SystemPropertiesCollection.LockToken"/>,
         /// only when <see cref="ReceiveMode"/> is set to <see cref="ServiceBus.ReceiveMode.PeekLock"/>.
-        /// In order to receive a message from the deadletter sub-queue, you will need a new <see cref="IMessageReceiver"/> or <see cref="IQueueClient"/>, with the corresponding path.
+        /// In order to receive a message from the deadletter sub-queue, you will need a new <see cref="IMessageReceiver"/> or <see cref="ISubscriptionClient"/>, with the corresponding path.
         /// You can use <see cref="EntityNameHelper.FormatDeadLetterPath(string)"/> to help with this.
         /// This operation can only be performed on messages that were received by this client.
         /// </remarks>
@@ -444,15 +479,29 @@ namespace Microsoft.Azure.ServiceBus
             description.ValidateDescriptionName();
             MessagingEventSource.Log.AddRuleStart(this.ClientId, description.Name);
 
+            bool isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+            Activity activity = isDiagnosticSourceEnabled ? this.diagnosticSource.AddRuleStart(description) : null;
+            Task addRuleTask = null;
+
             try
             {
-                await this.InnerSubscriptionClient.OnAddRuleAsync(description).ConfigureAwait(false);
+                addRuleTask = this.InnerSubscriptionClient.OnAddRuleAsync(description);
+                await addRuleTask.ConfigureAwait(false);
             }
             catch (Exception exception)
             {
+                if (isDiagnosticSourceEnabled)
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
+
                 MessagingEventSource.Log.AddRuleException(this.ClientId, exception);
                 throw;
             }
+            finally
+            {
+                this.diagnosticSource.AddRuleStop(activity, description, addRuleTask?.Status);
+            }
 
             MessagingEventSource.Log.AddRuleStop(this.ClientId);
         }
@@ -471,16 +520,30 @@ namespace Microsoft.Azure.ServiceBus
             }
 
             MessagingEventSource.Log.RemoveRuleStart(this.ClientId, ruleName);
+            bool isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+            Activity activity = isDiagnosticSourceEnabled ? this.diagnosticSource.RemoveRuleStart(ruleName) : null;
+            Task removeRuleTask = null;
 
             try
             {
-                await this.InnerSubscriptionClient.OnRemoveRuleAsync(ruleName).ConfigureAwait(false);
+                removeRuleTask = this.InnerSubscriptionClient.OnRemoveRuleAsync(ruleName);
+                await removeRuleTask.ConfigureAwait(false);
             }
             catch (Exception exception)
             {
+                if (isDiagnosticSourceEnabled)
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
+
                 MessagingEventSource.Log.RemoveRuleException(this.ClientId, exception);
+
                 throw;
             }
+            finally
+            {
+                this.diagnosticSource.RemoveRuleStop(activity, ruleName, removeRuleTask?.Status);
+            }
 
             MessagingEventSource.Log.RemoveRuleStop(this.ClientId);
         }
@@ -493,19 +556,34 @@ namespace Microsoft.Azure.ServiceBus
             this.ThrowIfClosed();
 
             MessagingEventSource.Log.GetRulesStart(this.ClientId);
+            bool isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+            Activity activity = isDiagnosticSourceEnabled ? this.diagnosticSource.GetRulesStart() : null;
+            Task<IEnumerable<RuleDescription>> getRulesTask = null;
+
             var skip = 0;
             var top = int.MaxValue;
-            IEnumerable<RuleDescription> rules;
+            IEnumerable<RuleDescription> rules = null;
 
             try
             {
-                rules = await this.InnerSubscriptionClient.OnGetRulesAsync(top, skip);
+                getRulesTask = this.InnerSubscriptionClient.OnGetRulesAsync(top, skip);
+                rules = await getRulesTask.ConfigureAwait(false);
             }
             catch (Exception exception)
             {
+                if (isDiagnosticSourceEnabled)
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
+
                 MessagingEventSource.Log.GetRulesException(this.ClientId, exception);
+
                 throw;
             }
+            finally
+            {
+                this.diagnosticSource.GetRulesStop(activity, rules, getRulesTask?.Status);
+            }
 
             MessagingEventSource.Log.GetRulesStop(this.ClientId);
             return rules;

src/Microsoft.Azure.ServiceBus/TopicClient.cs

@@ -65,6 +65,35 @@ namespace Microsoft.Azure.ServiceBus
                 throw Fx.Exception.ArgumentNullOrWhiteSpace(entityPath);
             }
 
+            this.InternalTokenProvider = this.ServiceBusConnection.CreateTokenProvider();
+            this.CbsTokenProvider = new TokenProviderAdapter(this.InternalTokenProvider, this.ServiceBusConnection.OperationTimeout);
+            this.ownsConnection = true;
+        }
+
+        /// <summary>
+        /// Creates a new instance of the Topic client using the specified endpoint, entity path, and token provider.
+        /// </summary>
+        /// <param name="endpoint">Fully qualified domain name for Service Bus. Most likely, {yournamespace}.servicebus.windows.net</param>
+        /// <param name="entityPath">Topic path.</param>
+        /// <param name="tokenProvider">Token provider which will generate security tokens for authorization.</param>
+        /// <param name="transportType">Transport type.</param>
+        /// <param name="retryPolicy">Retry policy for topic operations. Defaults to <see cref="RetryPolicy.Default"/></param>
+        /// <returns></returns>
+        public TopicClient(
+            string endpoint,
+            string entityPath,
+            ITokenProvider tokenProvider,
+            TransportType transportType = TransportType.Amqp,
+            RetryPolicy retryPolicy = null)
+            : this(new ServiceBusNamespaceConnection(endpoint, transportType, retryPolicy), entityPath, retryPolicy)
+        {
+            if (tokenProvider == null)
+            {
+                throw Fx.Exception.ArgumentNull(nameof(tokenProvider));
+            }
+
+            this.InternalTokenProvider = tokenProvider;
+            this.CbsTokenProvider = new TokenProviderAdapter(this.InternalTokenProvider, this.ServiceBusConnection.OperationTimeout);
             this.ownsConnection = true;
         }
 
@@ -77,8 +106,6 @@ namespace Microsoft.Azure.ServiceBus
             this.OperationTimeout = this.ServiceBusConnection.OperationTimeout;
             this.syncLock = new object();
             this.TopicName = entityPath;
-            this.TokenProvider = this.ServiceBusConnection.CreateTokenProvider();
-            this.CbsTokenProvider = new TokenProviderAdapter(this.TokenProvider, serviceBusConnection.OperationTimeout);
 
             MessagingEventSource.Log.TopicClientCreateStop(serviceBusConnection?.Endpoint.Authority, entityPath, this.ClientId);
         }
@@ -130,7 +157,7 @@ namespace Microsoft.Azure.ServiceBus
 
         ICbsTokenProvider CbsTokenProvider { get; }
 
-        TokenProvider TokenProvider { get; }
+        ITokenProvider InternalTokenProvider { get; }
 
         /// <summary>
         /// Sends a message to Service Bus.

test/Microsoft.Azure.ServiceBus.UnitTests/API/APIApprovals.cs

@@ -1,22 +1,25 @@
-#if NET461
-namespace Microsoft.Azure.ServiceBus.UnitTests.API
+namespace Microsoft.Azure.ServiceBus.UnitTests.API
 {
-    using System.Runtime.CompilerServices;
+    using System;
+    using System.Linq;
     using ApprovalTests;
-    using ApprovalTests.Reporters;
-    using PublicApiGenerator;
     using Xunit;
 
     public class ApiApprovals
     {
         [Fact]
-        [MethodImpl(MethodImplOptions.NoInlining)]
-        [UseReporter(typeof(DiffReporter), typeof(ClipboardReporter))]
         public void ApproveAzureServiceBus()
         {
-            var publicApi = ApiGenerator.GeneratePublicApi(typeof(Message).Assembly);
+            var assembly = typeof(Message).Assembly;
+            var publicApi = Filter(PublicApiGenerator.ApiGenerator.GeneratePublicApi(assembly, whitelistedNamespacePrefixes: new[] { "Microsoft.Azure.ServiceBus." }));
             Approvals.Verify(publicApi);
         }
+
+        string Filter(string text)
+        {
+            return string.Join(Environment.NewLine, text.Split(new[] {Environment.NewLine}, StringSplitOptions.RemoveEmptyEntries)
+                .Where(l => !l.StartsWith("[assembly: System.Runtime.Versioning.TargetFrameworkAttribute"))
+            );
+        }
     }
-}
-#endif
+}

test/Microsoft.Azure.ServiceBus.UnitTests/API/ApiApprovals.ApproveAzureServiceBus.approved.txt

@@ -1,8 +1,6 @@
 [assembly: System.Runtime.CompilerServices.InternalsVisibleToAttribute(@"Microsoft.Azure.ServiceBus.UnitTests,PublicKey=0024000004800000940000000602000000240000525341310004000001000100fdf4acac3b2244dd8a96737e5385b31414369dc3e42f371172127252856a0650793e1f5673a16d5d78e2ac852a104bc51e6f018dca44fdd26a219c27cb2b263956a80620223c8e9c2f8913c3c903e1e453e9e4e84098afdad5f4badb8c1ebe0a7b0a4b57a08454646a65886afe3e290a791ff3260099ce0edf0bdbccafadfeb6")]
 [assembly: System.Runtime.InteropServices.ComVisibleAttribute(false)]
 [assembly: System.Runtime.InteropServices.GuidAttribute("a042adf0-ef65-4f87-b634-322a409f3d61")]
-[assembly: System.Runtime.Versioning.TargetFrameworkAttribute(".NETFramework,Version=v4.6.1", FrameworkDisplayName=".NET Framework 4.6.1")]
-
 namespace Microsoft.Azure.ServiceBus
 {
     
@@ -168,9 +166,21 @@ namespace Microsoft.Azure.ServiceBus
         public System.TimeSpan MaxAutoRenewDuration { get; set; }
         public int MaxConcurrentCalls { get; set; }
     }
-    public sealed class MessageLockLostException : Microsoft.Azure.ServiceBus.ServiceBusException { }
-    public sealed class MessagingEntityDisabledException : Microsoft.Azure.ServiceBus.ServiceBusException { }
-    public sealed class MessagingEntityNotFoundException : Microsoft.Azure.ServiceBus.ServiceBusException { }
+    public sealed class MessageLockLostException : Microsoft.Azure.ServiceBus.ServiceBusException
+    {
+        public MessageLockLostException(string message) { }
+        public MessageLockLostException(string message, System.Exception innerException) { }
+    }
+    public sealed class MessagingEntityDisabledException : Microsoft.Azure.ServiceBus.ServiceBusException
+    {
+        public MessagingEntityDisabledException(string message) { }
+        public MessagingEntityDisabledException(string message, System.Exception innerException) { }
+    }
+    public sealed class MessagingEntityNotFoundException : Microsoft.Azure.ServiceBus.ServiceBusException
+    {
+        public MessagingEntityNotFoundException(string message) { }
+        public MessagingEntityNotFoundException(string message, System.Exception innerException) { }
+    }
     public sealed class NoRetry : Microsoft.Azure.ServiceBus.RetryPolicy
     {
         public NoRetry() { }
@@ -180,6 +190,7 @@ namespace Microsoft.Azure.ServiceBus
     {
         public QueueClient(Microsoft.Azure.ServiceBus.ServiceBusConnectionStringBuilder connectionStringBuilder, Microsoft.Azure.ServiceBus.ReceiveMode receiveMode = 0, Microsoft.Azure.ServiceBus.RetryPolicy retryPolicy = null) { }
         public QueueClient(string connectionString, string entityPath, Microsoft.Azure.ServiceBus.ReceiveMode receiveMode = 0, Microsoft.Azure.ServiceBus.RetryPolicy retryPolicy = null) { }
+        public QueueClient(string endpoint, string entityPath, Microsoft.Azure.ServiceBus.Primitives.ITokenProvider tokenProvider, Microsoft.Azure.ServiceBus.TransportType transportType = 0, Microsoft.Azure.ServiceBus.ReceiveMode receiveMode = 0, Microsoft.Azure.ServiceBus.RetryPolicy retryPolicy = null) { }
         public override System.TimeSpan OperationTimeout { get; set; }
         public string Path { get; }
         public int PrefetchCount { get; set; }
@@ -202,7 +213,11 @@ namespace Microsoft.Azure.ServiceBus
         public System.Threading.Tasks.Task SendAsync(System.Collections.Generic.IList<Microsoft.Azure.ServiceBus.Message> messageList) { }
         public override void UnregisterPlugin(string serviceBusPluginName) { }
     }
-    public sealed class QuotaExceededException : Microsoft.Azure.ServiceBus.ServiceBusException { }
+    public sealed class QuotaExceededException : Microsoft.Azure.ServiceBus.ServiceBusException
+    {
+        public QuotaExceededException(string message) { }
+        public QuotaExceededException(string message, System.Exception innerException) { }
+    }
     public enum ReceiveMode
     {
         PeekLock = 0,
@@ -239,11 +254,15 @@ namespace Microsoft.Azure.ServiceBus
         public Microsoft.Azure.ServiceBus.Filter Filter { get; set; }
         public string Name { get; set; }
     }
-    public sealed class ServerBusyException : Microsoft.Azure.ServiceBus.ServiceBusException { }
+    public sealed class ServerBusyException : Microsoft.Azure.ServiceBus.ServiceBusException
+    {
+        public ServerBusyException(string message) { }
+        public ServerBusyException(string message, System.Exception innerException) { }
+    }
     public class ServiceBusCommunicationException : Microsoft.Azure.ServiceBus.ServiceBusException
     {
-        protected internal ServiceBusCommunicationException(string message) { }
-        protected internal ServiceBusCommunicationException(string message, System.Exception innerException) { }
+        public ServiceBusCommunicationException(string message) { }
+        public ServiceBusCommunicationException(string message, System.Exception innerException) { }
     }
     public class ServiceBusConnectionStringBuilder
     {
@@ -273,7 +292,11 @@ namespace Microsoft.Azure.ServiceBus
         public override string Message { get; }
         public string ServiceBusNamespace { get; }
     }
-    public class ServiceBusTimeoutException : Microsoft.Azure.ServiceBus.ServiceBusException { }
+    public class ServiceBusTimeoutException : Microsoft.Azure.ServiceBus.ServiceBusException
+    {
+        public ServiceBusTimeoutException(string message) { }
+        public ServiceBusTimeoutException(string message, System.Exception innerException) { }
+    }
     public sealed class SessionClient : Microsoft.Azure.ServiceBus.ClientEntity, Microsoft.Azure.ServiceBus.IClientEntity, Microsoft.Azure.ServiceBus.ISessionClient
     {
         public SessionClient(Microsoft.Azure.ServiceBus.ServiceBusConnectionStringBuilder connectionStringBuilder, Microsoft.Azure.ServiceBus.ReceiveMode receiveMode = 0, Microsoft.Azure.ServiceBus.RetryPolicy retryPolicy = null, int prefetchCount = 0) { }
@@ -298,7 +321,11 @@ namespace Microsoft.Azure.ServiceBus
         public int MaxConcurrentSessions { get; set; }
         public System.TimeSpan MessageWaitTimeout { get; set; }
     }
-    public sealed class SessionLockLostException : Microsoft.Azure.ServiceBus.ServiceBusException { }
+    public sealed class SessionLockLostException : Microsoft.Azure.ServiceBus.ServiceBusException
+    {
+        public SessionLockLostException(string message) { }
+        public SessionLockLostException(string message, System.Exception innerException) { }
+    }
     public class SqlFilter : Microsoft.Azure.ServiceBus.Filter
     {
         public SqlFilter(string sqlExpression) { }
@@ -317,6 +344,7 @@ namespace Microsoft.Azure.ServiceBus
     {
         public SubscriptionClient(Microsoft.Azure.ServiceBus.ServiceBusConnectionStringBuilder connectionStringBuilder, string subscriptionName, Microsoft.Azure.ServiceBus.ReceiveMode receiveMode = 0, Microsoft.Azure.ServiceBus.RetryPolicy retryPolicy = null) { }
         public SubscriptionClient(string connectionString, string topicPath, string subscriptionName, Microsoft.Azure.ServiceBus.ReceiveMode receiveMode = 0, Microsoft.Azure.ServiceBus.RetryPolicy retryPolicy = null) { }
+        public SubscriptionClient(string endpoint, string topicPath, string subscriptionName, Microsoft.Azure.ServiceBus.Primitives.ITokenProvider tokenProvider, Microsoft.Azure.ServiceBus.TransportType transportType = 0, Microsoft.Azure.ServiceBus.ReceiveMode receiveMode = 0, Microsoft.Azure.ServiceBus.RetryPolicy retryPolicy = null) { }
         public override System.TimeSpan OperationTimeout { get; set; }
         public string Path { get; }
         public int PrefetchCount { get; set; }
@@ -345,6 +373,7 @@ namespace Microsoft.Azure.ServiceBus
     {
         public TopicClient(Microsoft.Azure.ServiceBus.ServiceBusConnectionStringBuilder connectionStringBuilder, Microsoft.Azure.ServiceBus.RetryPolicy retryPolicy = null) { }
         public TopicClient(string connectionString, string entityPath, Microsoft.Azure.ServiceBus.RetryPolicy retryPolicy = null) { }
+        public TopicClient(string endpoint, string entityPath, Microsoft.Azure.ServiceBus.Primitives.ITokenProvider tokenProvider, Microsoft.Azure.ServiceBus.TransportType transportType = 0, Microsoft.Azure.ServiceBus.RetryPolicy retryPolicy = null) { }
         public override System.TimeSpan OperationTimeout { get; set; }
         public string Path { get; }
         public override System.Collections.Generic.IList<Microsoft.Azure.ServiceBus.Core.ServiceBusPlugin> RegisteredPlugins { get; }
@@ -476,6 +505,14 @@ namespace Microsoft.Azure.ServiceBus.Core
         public virtual System.Threading.Tasks.Task<Microsoft.Azure.ServiceBus.Message> BeforeMessageSend(Microsoft.Azure.ServiceBus.Message message) { }
     }
 }
+namespace Microsoft.Azure.ServiceBus.Diagnostics
+{
+    
+    public class static MessageExtensions
+    {
+        public static System.Diagnostics.Activity ExtractActivity(this Microsoft.Azure.ServiceBus.Message message, string activityName = null) { }
+    }
+}
 namespace Microsoft.Azure.ServiceBus.InteropExtensions
 {
     
@@ -488,4 +525,58 @@ namespace Microsoft.Azure.ServiceBus.InteropExtensions
     {
         public static T GetBody<T>(this Microsoft.Azure.ServiceBus.Message message, System.Runtime.Serialization.XmlObjectSerializer serializer = null) { }
     }
+}
+namespace Microsoft.Azure.ServiceBus.Primitives
+{
+    
+    public class AzureActiveDirectoryTokenProvider : Microsoft.Azure.ServiceBus.Primitives.TokenProvider
+    {
+        public override System.Threading.Tasks.Task<Microsoft.Azure.ServiceBus.Primitives.SecurityToken> GetTokenAsync(string appliesTo, System.TimeSpan timeout) { }
+    }
+    public interface ITokenProvider
+    {
+        System.Threading.Tasks.Task<Microsoft.Azure.ServiceBus.Primitives.SecurityToken> GetTokenAsync(string appliesTo, System.TimeSpan timeout);
+    }
+    public class JsonSecurityToken : Microsoft.Azure.ServiceBus.Primitives.SecurityToken
+    {
+        public JsonSecurityToken(string rawToken, string audience) { }
+    }
+    public class ManagedServiceIdentityTokenProvider : Microsoft.Azure.ServiceBus.Primitives.TokenProvider
+    {
+        public ManagedServiceIdentityTokenProvider() { }
+        public override System.Threading.Tasks.Task<Microsoft.Azure.ServiceBus.Primitives.SecurityToken> GetTokenAsync(string appliesTo, System.TimeSpan timeout) { }
+    }
+    public class SecurityToken
+    {
+        public SecurityToken(string tokenString, System.DateTime expiresAtUtc, string audience, string tokenType) { }
+        public string Audience { get; }
+        public System.DateTime ExpiresAtUtc { get; }
+        public virtual string TokenType { get; }
+        public virtual string TokenValue { get; }
+    }
+    public class SharedAccessSignatureTokenProvider : Microsoft.Azure.ServiceBus.Primitives.TokenProvider
+    {
+        protected SharedAccessSignatureTokenProvider(string keyName, string sharedAccessKey, System.Func<string, byte[]> customKeyEncoder, System.TimeSpan tokenTimeToLive, Microsoft.Azure.ServiceBus.Primitives.TokenScope tokenScope) { }
+        protected virtual string BuildSignature(string targetUri) { }
+        public override System.Threading.Tasks.Task<Microsoft.Azure.ServiceBus.Primitives.SecurityToken> GetTokenAsync(string appliesTo, System.TimeSpan timeout) { }
+    }
+    public abstract class TokenProvider : Microsoft.Azure.ServiceBus.Primitives.ITokenProvider
+    {
+        protected TokenProvider() { }
+        public static Microsoft.Azure.ServiceBus.Primitives.TokenProvider CreateAadTokenProvider(Microsoft.IdentityModel.Clients.ActiveDirectory.AuthenticationContext authContext, Microsoft.IdentityModel.Clients.ActiveDirectory.ClientCredential clientCredential) { }
+        public static Microsoft.Azure.ServiceBus.Primitives.TokenProvider CreateAadTokenProvider(Microsoft.IdentityModel.Clients.ActiveDirectory.AuthenticationContext authContext, string clientId, System.Uri redirectUri, Microsoft.IdentityModel.Clients.ActiveDirectory.IPlatformParameters platformParameters, Microsoft.IdentityModel.Clients.ActiveDirectory.UserIdentifier userIdentifier = null) { }
+        public static Microsoft.Azure.ServiceBus.Primitives.TokenProvider CreateAadTokenProvider(Microsoft.IdentityModel.Clients.ActiveDirectory.AuthenticationContext authContext, Microsoft.IdentityModel.Clients.ActiveDirectory.ClientAssertionCertificate clientAssertionCertificate) { }
+        public static Microsoft.Azure.ServiceBus.Primitives.TokenProvider CreateManagedServiceIdentityTokenProvider() { }
+        public static Microsoft.Azure.ServiceBus.Primitives.TokenProvider CreateSharedAccessSignatureTokenProvider(string sharedAccessSignature) { }
+        public static Microsoft.Azure.ServiceBus.Primitives.TokenProvider CreateSharedAccessSignatureTokenProvider(string keyName, string sharedAccessKey) { }
+        public static Microsoft.Azure.ServiceBus.Primitives.TokenProvider CreateSharedAccessSignatureTokenProvider(string keyName, string sharedAccessKey, System.TimeSpan tokenTimeToLive) { }
+        public static Microsoft.Azure.ServiceBus.Primitives.TokenProvider CreateSharedAccessSignatureTokenProvider(string keyName, string sharedAccessKey, Microsoft.Azure.ServiceBus.Primitives.TokenScope tokenScope) { }
+        public static Microsoft.Azure.ServiceBus.Primitives.TokenProvider CreateSharedAccessSignatureTokenProvider(string keyName, string sharedAccessKey, System.TimeSpan tokenTimeToLive, Microsoft.Azure.ServiceBus.Primitives.TokenScope tokenScope) { }
+        public abstract System.Threading.Tasks.Task<Microsoft.Azure.ServiceBus.Primitives.SecurityToken> GetTokenAsync(string appliesTo, System.TimeSpan timeout);
+    }
+    public enum TokenScope
+    {
+        Namespace = 0,
+        Entity = 1,
+    }
 }

test/Microsoft.Azure.ServiceBus.UnitTests/API/ApiApprover.6.0.0/PublicApiApprover.cs

@@ -1,34 +0,0 @@
-#if NET461
-
-using System.IO;
-using System.Reflection;
-using ApprovalTests;
-using ApprovalTests.Namers;
-
-namespace ApiApprover
-{
-    public static class PublicApiApprover
-    {
-        public static void ApprovePublicApi(Assembly assembly)
-        {
-            var publicApi = PublicApiGenerator.ApiGenerator.GeneratePublicApi(assembly);
-            var writer = new ApprovalTextWriter(publicApi, "cs");
-            var approvalNamer = new AssemblyPathNamer(assembly.Location);
-            Approvals.Verify(writer, approvalNamer, Approvals.GetReporter());
-        }
-
-        private class AssemblyPathNamer : UnitTestFrameworkNamer
-        {
-            private readonly string name;
-
-            public AssemblyPathNamer(string assemblyPath)
-            {
-                name = Path.GetFileNameWithoutExtension(assemblyPath);
-            }
-
-            public override string Name => name;
-        }
-    }
-}
-
-#endif

test/Microsoft.Azure.ServiceBus.UnitTests/API/ApprovalTestConfig.cs

@@ -0,0 +1,7 @@
+using ApprovalTests.Reporters;
+
+#if NET461
+[assembly: UseReporter(typeof(XUnit2Reporter), typeof(AllFailingTestsClipboardReporter))]
+#else
+[assembly: UseReporter(typeof(XUnit2Reporter))]
+#endif

test/Microsoft.Azure.ServiceBus.UnitTests/API/PublicApiGenerator.6.0.0/ApiGenerator.cs

@@ -1,801 +0,0 @@
-#if NET461
-
-using System;
-using System.CodeDom;
-using System.CodeDom.Compiler;
-using System.Collections.Generic;
-using System.IO;
-using System.Linq;
-using System.Reflection;
-using System.Text;
-using System.Text.RegularExpressions;
-using Microsoft.CSharp;
-using Mono.Cecil;
-using Mono.Cecil.Rocks;
-using ICustomAttributeProvider = Mono.Cecil.ICustomAttributeProvider;
-using TypeAttributes = System.Reflection.TypeAttributes;
-using System.Globalization;
-
-// ReSharper disable BitwiseOperatorOnEnumWithoutFlags
-namespace PublicApiGenerator
-{
-    public static class ApiGenerator
-    {
-        public static string GeneratePublicApi(Assembly assembly, Type[] includeTypes = null, bool shouldIncludeAssemblyAttributes = true)
-        {
-            var assemblyResolver = new DefaultAssemblyResolver();
-            var assemblyPath = assembly.Location;
-            assemblyResolver.AddSearchDirectory(Path.GetDirectoryName(assemblyPath));
-
-            var readSymbols = File.Exists(Path.ChangeExtension(assemblyPath, ".pdb"));
-            var asm = AssemblyDefinition.ReadAssembly(assemblyPath, new ReaderParameters(ReadingMode.Deferred)
-            {
-                ReadSymbols = readSymbols,
-                AssemblyResolver = assemblyResolver,
-            });
-
-            return CreatePublicApiForAssembly(asm, tr => includeTypes == null || includeTypes.Any(t => t.FullName == tr.FullName && t.Assembly.FullName == tr.Module.Assembly.FullName), shouldIncludeAssemblyAttributes);
-        }
-
-        // TODO: Assembly references?
-        // TODO: Better handle namespaces - using statements? - requires non-qualified type names
-        static string CreatePublicApiForAssembly(AssemblyDefinition assembly, Func<TypeDefinition, bool> shouldIncludeType, bool shouldIncludeAssemblyAttributes)
-        {
-            var publicApiBuilder = new StringBuilder();
-            var cgo = new CodeGeneratorOptions
-            {
-                BracingStyle = "C",
-                BlankLinesBetweenMembers = false,
-                VerbatimOrder = false
-            };
-
-            using (var provider = new CSharpCodeProvider())
-            {
-                var compileUnit = new CodeCompileUnit();
-                if (shouldIncludeAssemblyAttributes && assembly.HasCustomAttributes)
-                {
-                    PopulateCustomAttributes(assembly, compileUnit.AssemblyCustomAttributes);
-                }
-
-                var publicTypes = assembly.Modules.SelectMany(m => m.GetTypes())
-                    .Where(t => !t.IsNested && ShouldIncludeType(t) && shouldIncludeType(t))
-                    .OrderBy(t => t.FullName);
-                foreach (var publicType in publicTypes)
-                {
-                    var @namespace = compileUnit.Namespaces.Cast<CodeNamespace>()
-                        .FirstOrDefault(n => n.Name == publicType.Namespace);
-                    if (@namespace == null)
-                    {
-                        @namespace = new CodeNamespace(publicType.Namespace);
-                        compileUnit.Namespaces.Add(@namespace);
-                    }
-
-                    var typeDeclaration = CreateTypeDeclaration(publicType);
-                    @namespace.Types.Add(typeDeclaration);
-                }
-
-                using (var writer = new StringWriter())
-                {
-                    provider.GenerateCodeFromCompileUnit(compileUnit, writer, cgo);
-                    var typeDeclarationText = NormaliseGeneratedCode(writer);
-                    publicApiBuilder.AppendLine(typeDeclarationText);
-                }
-            }
-            return NormaliseLineEndings(publicApiBuilder.ToString().Trim());
-        }
-
-        static string NormaliseLineEndings(string value)
-        {
-            return Regex.Replace(value, @"\r\n|\n\r|\r|\n", Environment.NewLine);
-        }
-
-        static bool IsDelegate(TypeDefinition publicType)
-        {
-            return publicType.BaseType != null && publicType.BaseType.FullName == "System.MulticastDelegate";
-        }
-
-        static bool ShouldIncludeType(TypeDefinition t)
-        {
-            return (t.IsPublic || t.IsNestedPublic || t.IsNestedFamily) && !IsCompilerGenerated(t);
-        }
-
-        static bool ShouldIncludeMember(IMemberDefinition m)
-        {
-            return !IsCompilerGenerated(m) && !IsDotNetTypeMember(m) && !(m is FieldDefinition);
-        }
-
-        static bool IsCompilerGenerated(IMemberDefinition m)
-        {
-            return m.CustomAttributes.Any(a => a.AttributeType.FullName == "System.Runtime.CompilerServices.CompilerGeneratedAttribute");
-        }
-
-        static bool IsDotNetTypeMember(IMemberDefinition m)
-        {
-            if (m.DeclaringType?.FullName == null)
-                return false;
-            return m.DeclaringType.FullName.StartsWith("System") || m.DeclaringType.FullName.StartsWith("Microsoft") && !m.DeclaringType.FullName.StartsWith("Microsoft.Azure.ServiceBus");
-        }
-
-        static void AddMemberToTypeDeclaration(CodeTypeDeclaration typeDeclaration, IMemberDefinition memberInfo)
-        {
-            if (memberInfo is MethodDefinition methodDefinition)
-            {
-                if (methodDefinition.IsConstructor)
-                    AddCtorToTypeDeclaration(typeDeclaration, methodDefinition);
-                else
-                    AddMethodToTypeDeclaration(typeDeclaration, methodDefinition);
-            }
-            else if (memberInfo is PropertyDefinition)
-            {
-                AddPropertyToTypeDeclaration(typeDeclaration, (PropertyDefinition) memberInfo);
-            }
-            else if (memberInfo is EventDefinition)
-            {
-                typeDeclaration.Members.Add(GenerateEvent((EventDefinition)memberInfo));
-            }
-            else if (memberInfo is FieldDefinition)
-            {
-                AddFieldToTypeDeclaration(typeDeclaration, (FieldDefinition) memberInfo);
-            }
-        }
-
-        static string NormaliseGeneratedCode(StringWriter writer)
-        {
-            var gennedClass = writer.ToString();
-            const string autoGeneratedHeader = @"^//-+\s*$.*^//-+\s*$";
-            const string emptyGetSet = @"\s+{\s+get\s+{\s+}\s+set\s+{\s+}\s+}";
-            const string emptyGet = @"\s+{\s+get\s+{\s+}\s+}";
-            const string emptySet = @"\s+{\s+set\s+{\s+}\s+}";
-            const string getSet = @"\s+{\s+get;\s+set;\s+}";
-            const string get = @"\s+{\s+get;\s+}";
-            const string set = @"\s+{\s+set;\s+}";
-            gennedClass = Regex.Replace(gennedClass, autoGeneratedHeader, string.Empty,
-                RegexOptions.IgnorePatternWhitespace | RegexOptions.Multiline | RegexOptions.Singleline);
-            gennedClass = Regex.Replace(gennedClass, emptyGetSet, " { get; set; }", RegexOptions.IgnorePatternWhitespace);
-            gennedClass = Regex.Replace(gennedClass, getSet, " { get; set; }", RegexOptions.IgnorePatternWhitespace);
-            gennedClass = Regex.Replace(gennedClass, emptyGet, " { get; }", RegexOptions.IgnorePatternWhitespace);
-            gennedClass = Regex.Replace(gennedClass, emptySet, " { set; }", RegexOptions.IgnorePatternWhitespace);
-            gennedClass = Regex.Replace(gennedClass, get, " { get; }", RegexOptions.IgnorePatternWhitespace);
-            gennedClass = Regex.Replace(gennedClass, set, " { set; }", RegexOptions.IgnorePatternWhitespace);
-            gennedClass = Regex.Replace(gennedClass, @"\s+{\s+}", " { }", RegexOptions.IgnorePatternWhitespace);
-            gennedClass = Regex.Replace(gennedClass, @"\)\s+;", ");", RegexOptions.IgnorePatternWhitespace);
-            return gennedClass;
-        }
-
-        static CodeTypeDeclaration CreateTypeDeclaration(TypeDefinition publicType)
-        {
-            if (IsDelegate(publicType))
-                return CreateDelegateDeclaration(publicType);
-
-            var @static = false;
-            TypeAttributes attributes = 0;
-            if (publicType.IsPublic || publicType.IsNestedPublic)
-                attributes |= TypeAttributes.Public;
-            if (publicType.IsNestedFamily)
-                attributes |= TypeAttributes.NestedFamily;
-            if (publicType.IsSealed && !publicType.IsAbstract)
-                attributes |= TypeAttributes.Sealed;
-            else if (!publicType.IsSealed && publicType.IsAbstract && !publicType.IsInterface)
-                attributes |= TypeAttributes.Abstract;
-            else if (publicType.IsSealed && publicType.IsAbstract)
-                @static = true;
-
-            // Static support is a hack. CodeDOM does support it, and this isn't
-            // correct C#, but it's good enough for our API outline
-            var name = publicType.Name;
-
-            var index = name.IndexOf('`');
-            if (index != -1)
-                name = name.Substring(0, index);
-            var declaration = new CodeTypeDeclaration(@static ? "static " + name : name)
-            {
-                CustomAttributes = CreateCustomAttributes(publicType),
-                // TypeAttributes must be specified before the IsXXX as they manipulate TypeAttributes!
-                TypeAttributes = attributes,
-                IsClass = publicType.IsClass,
-                IsEnum = publicType.IsEnum,
-                IsInterface = publicType.IsInterface,
-                IsStruct = publicType.IsValueType && !publicType.IsPrimitive && !publicType.IsEnum,
-            };
-
-            if (declaration.IsInterface && publicType.BaseType != null)
-                throw new NotImplementedException("Base types for interfaces needs testing");
-
-            PopulateGenericParameters(publicType, declaration.TypeParameters);
-
-            if (publicType.BaseType != null && ShouldOutputBaseType(publicType))
-            {
-                if (publicType.BaseType.FullName == "System.Enum")
-                {
-                    var underlyingType = publicType.GetEnumUnderlyingType();
-                    if (underlyingType.FullName != "System.Int32")
-                        declaration.BaseTypes.Add(CreateCodeTypeReference(underlyingType));
-                }
-                else
-                    declaration.BaseTypes.Add(CreateCodeTypeReference(publicType.BaseType));
-            }
-            foreach(var @interface in publicType.Interfaces.OrderBy(i => i.FullName)
-                .Select(t => new { Reference = t, Definition = t.Resolve() })
-                .Where(t => ShouldIncludeType(t.Definition))
-                .Select(t => t.Reference))
-                declaration.BaseTypes.Add(CreateCodeTypeReference(@interface));
-
-            foreach (var memberInfo in publicType.GetMembers().Where(ShouldIncludeMember).OrderBy(m => m.Name))
-                AddMemberToTypeDeclaration(declaration, memberInfo);
-
-            // Fields should be in defined order for an enum
-            var fields = !publicType.IsEnum
-                ? publicType.Fields.OrderBy(f => f.Name)
-                : (IEnumerable<FieldDefinition>)publicType.Fields;
-            foreach (var field in fields)
-                AddMemberToTypeDeclaration(declaration, field);
-
-            foreach (var nestedType in publicType.NestedTypes.Where(ShouldIncludeType).OrderBy(t => t.FullName))
-            {
-                var nestedTypeDeclaration = CreateTypeDeclaration(nestedType);
-                declaration.Members.Add(nestedTypeDeclaration);
-            }
-
-            return declaration;
-        }
-
-        static CodeTypeDeclaration CreateDelegateDeclaration(TypeDefinition publicType)
-        {
-            var invokeMethod = publicType.Methods.Single(m => m.Name == "Invoke");
-            var name = publicType.Name;
-            var index = name.IndexOf('`');
-            if (index != -1)
-                name = name.Substring(0, index);
-            var declaration = new CodeTypeDelegate(name)
-            {
-                Attributes = MemberAttributes.Public,
-                CustomAttributes = CreateCustomAttributes(publicType),
-                ReturnType = CreateCodeTypeReference(invokeMethod.ReturnType),
-            };
-
-            // CodeDOM. No support. Return type attributes.
-            PopulateCustomAttributes(invokeMethod.MethodReturnType, declaration.CustomAttributes, type => ModifyCodeTypeReference(type, "return:"));
-            PopulateGenericParameters(publicType, declaration.TypeParameters);
-            PopulateMethodParameters(invokeMethod, declaration.Parameters);
-
-            // Of course, CodeDOM doesn't support generic type parameters for delegates. Of course.
-            if (declaration.TypeParameters.Count > 0)
-            {
-                var parameterNames = from parameterType in declaration.TypeParameters.Cast<CodeTypeParameter>()
-                    select parameterType.Name;
-                declaration.Name = string.Format(CultureInfo.InvariantCulture, "{0}<{1}>", declaration.Name, string.Join(", ", parameterNames));
-            }
-
-            return declaration;
-        }
-
-        static bool ShouldOutputBaseType(TypeDefinition publicType)
-        {
-            return publicType.BaseType.FullName != "System.Object" && publicType.BaseType.FullName != "System.ValueType";
-        }
-
-        static void PopulateGenericParameters(IGenericParameterProvider publicType, CodeTypeParameterCollection parameters)
-        {
-            foreach (var parameter in publicType.GenericParameters)
-            {
-                if (parameter.HasCustomAttributes)
-                    throw new NotImplementedException("Attributes on type parameters is not supported. And weird");
-
-                // A little hacky. Means we get "in" and "out" prefixed on any constraints, but it's either that
-                // or add it as a custom attribute, which looks even weirder
-                var name = parameter.Name;
-                if (parameter.IsCovariant)
-                    name = "out " + name;
-                if (parameter.IsContravariant)
-                    name = "in " + name;
-
-                var typeParameter = new CodeTypeParameter(name)
-                {
-                    HasConstructorConstraint =
-                        parameter.HasDefaultConstructorConstraint && !parameter.HasNotNullableValueTypeConstraint
-                };
-                if (parameter.HasNotNullableValueTypeConstraint)
-                    typeParameter.Constraints.Add(" struct"); // Extra space is a hack!
-                if (parameter.HasReferenceTypeConstraint)
-                    typeParameter.Constraints.Add(" class");
-                foreach (var constraint in parameter.Constraints.Where(t => t.FullName != "System.ValueType"))
-                {
-                    typeParameter.Constraints.Add(CreateCodeTypeReference(constraint.GetElementType()));
-                }
-                parameters.Add(typeParameter);
-            }
-        }
-
-        static CodeAttributeDeclarationCollection CreateCustomAttributes(ICustomAttributeProvider type)
-        {
-            var attributes = new CodeAttributeDeclarationCollection();
-            PopulateCustomAttributes(type, attributes);
-            return attributes;
-        }
-
-        static void PopulateCustomAttributes(ICustomAttributeProvider type,
-            CodeAttributeDeclarationCollection attributes)
-        {
-            PopulateCustomAttributes(type, attributes, ctr => ctr);
-        }
-
-        static void PopulateCustomAttributes(ICustomAttributeProvider type,
-            CodeAttributeDeclarationCollection attributes, Func<CodeTypeReference, CodeTypeReference> codeTypeModifier)
-        {
-            foreach (var customAttribute in type.CustomAttributes.Where(ShouldIncludeAttribute).OrderBy(a => a.AttributeType.FullName).ThenBy(a => ConvertAttributeToCode(codeTypeModifier, a)))
-            {
-                var attribute = GenerateCodeAttributeDeclaration(codeTypeModifier, customAttribute);
-                attributes.Add(attribute);
-            }
-        }
-
-        static CodeAttributeDeclaration GenerateCodeAttributeDeclaration(Func<CodeTypeReference, CodeTypeReference> codeTypeModifier, CustomAttribute customAttribute)
-        {
-            var attribute = new CodeAttributeDeclaration(codeTypeModifier(CreateCodeTypeReference(customAttribute.AttributeType)));
-            foreach (var arg in customAttribute.ConstructorArguments)
-            {
-                attribute.Arguments.Add(new CodeAttributeArgument(CreateInitialiserExpression(arg)));
-            }
-            foreach (var field in customAttribute.Fields.OrderBy(f => f.Name))
-            {
-                attribute.Arguments.Add(new CodeAttributeArgument(field.Name, CreateInitialiserExpression(field.Argument)));
-            }
-            foreach (var property in customAttribute.Properties.OrderBy(p => p.Name))
-            {
-                attribute.Arguments.Add(new CodeAttributeArgument(property.Name, CreateInitialiserExpression(property.Argument)));
-            }
-            return attribute;
-        }
-
-        // Litee: This method is used for additional sorting of custom attributes when multiple values are allowed
-        static object ConvertAttributeToCode(Func<CodeTypeReference, CodeTypeReference> codeTypeModifier, CustomAttribute customAttribute)
-        {
-            using (var provider = new CSharpCodeProvider())
-            {
-                var cgo = new CodeGeneratorOptions
-                {
-                    BracingStyle = "C",
-                    BlankLinesBetweenMembers = false,
-                    VerbatimOrder = false
-                };
-                var attribute = GenerateCodeAttributeDeclaration(codeTypeModifier, customAttribute);
-                var declaration = new CodeTypeDeclaration("DummyClass")
-                {
-                    CustomAttributes = new CodeAttributeDeclarationCollection(new[] { attribute }),
-                };
-                using (var writer = new StringWriter())
-                {
-                    provider.GenerateCodeFromType(declaration, writer, cgo);
-                    return writer.ToString();
-                }
-            }
-        }
-
-        static readonly HashSet<string> SkipAttributeNames = new HashSet<string>
-        {
-            "System.CodeDom.Compiler.GeneratedCodeAttribute",
-            "System.ComponentModel.EditorBrowsableAttribute",
-            "System.Runtime.CompilerServices.AsyncStateMachineAttribute",
-            "System.Runtime.CompilerServices.CompilerGeneratedAttribute",
-            "System.Runtime.CompilerServices.CompilationRelaxationsAttribute",
-            "System.Runtime.CompilerServices.ExtensionAttribute",
-            "System.Runtime.CompilerServices.RuntimeCompatibilityAttribute",
-            "System.Reflection.DefaultMemberAttribute",
-            "System.Diagnostics.DebuggableAttribute",
-            "System.Diagnostics.DebuggerNonUserCodeAttribute",
-            "System.Diagnostics.DebuggerStepThroughAttribute",
-            "System.Reflection.AssemblyCompanyAttribute",
-            "System.Reflection.AssemblyConfigurationAttribute",
-            "System.Reflection.AssemblyCopyrightAttribute",
-            "System.Reflection.AssemblyDescriptionAttribute",
-            "System.Reflection.AssemblyFileVersionAttribute",
-            "System.Reflection.AssemblyInformationalVersionAttribute",
-            "System.Reflection.AssemblyProductAttribute",
-            "System.Reflection.AssemblyTitleAttribute",
-            "System.Reflection.AssemblyTrademarkAttribute"
-        };
-
-        static bool ShouldIncludeAttribute(CustomAttribute attribute)
-        {
-            var attributeTypeDefinition = attribute.AttributeType.Resolve();
-            return !SkipAttributeNames.Contains(attribute.AttributeType.FullName) && attributeTypeDefinition.IsPublic;
-        }
-
-        static CodeExpression CreateInitialiserExpression(CustomAttributeArgument attributeArgument)
-        {
-            if (attributeArgument.Value is CustomAttributeArgument)
-            {
-                return CreateInitialiserExpression((CustomAttributeArgument) attributeArgument.Value);
-            }
-
-            if (attributeArgument.Value is CustomAttributeArgument[])
-            {
-                var initialisers = from argument in (CustomAttributeArgument[]) attributeArgument.Value
-                    select CreateInitialiserExpression(argument);
-                return new CodeArrayCreateExpression(CreateCodeTypeReference(attributeArgument.Type), initialisers.ToArray());
-            }
-
-            var type = attributeArgument.Type.Resolve();
-            var value = attributeArgument.Value;
-            if (type.BaseType != null && type.BaseType.FullName == "System.Enum")
-            {
-                var originalValue = Convert.ToInt64(value);
-                if (type.CustomAttributes.Any(a => a.AttributeType.FullName == "System.FlagsAttribute"))
-                {
-                    //var allFlags = from f in type.Fields
-                    //    where f.Constant != null
-                    //    let v = Convert.ToInt64(f.Constant)
-                    //    where v == 0 || (originalValue & v) != 0
-                    //    select (CodeExpression)new CodeFieldReferenceExpression(typeExpression, f.Name);
-                    //return allFlags.Aggregate((current, next) => new CodeBinaryOperatorExpression(current, CodeBinaryOperatorType.BitwiseOr, next));
-
-                    // I'd rather use the above, as it's just using the CodeDOM, but it puts
-                    // brackets around each CodeBinaryOperatorExpression
-                    var flags = from f in type.Fields
-                                where f.Constant != null
-                                let v = Convert.ToInt64(f.Constant)
-                                where v == 0 || (originalValue & v) != 0
-                                select type.FullName + "." + f.Name;
-                    return new CodeSnippetExpression(flags.Aggregate((current, next) => current + " | " + next));
-                }
-
-                var allFlags = from f in type.Fields
-                               where f.Constant != null
-                               let v = Convert.ToInt64(f.Constant)
-                               where v == originalValue
-                               select new CodeFieldReferenceExpression(new CodeTypeReferenceExpression(CreateCodeTypeReference(type)), f.Name);
-                return allFlags.FirstOrDefault();
-            }
-
-            if (type.FullName == "System.Type" && value is TypeReference)
-            {
-                return new CodeTypeOfExpression(CreateCodeTypeReference((TypeReference)value));
-            }
-
-            if (value is string)
-            {
-                // CodeDOM outputs a verbatim string. Any string with \n is treated as such, so normalise
-                // it to make it easier for comparisons
-                value = Regex.Replace((string)value, @"\n", "\\n");
-                value = Regex.Replace((string)value, @"\r\n|\r\\n", "\\r\\n");
-            }
-
-            return new CodePrimitiveExpression(value);
-        }
-
-        static void AddCtorToTypeDeclaration(CodeTypeDeclaration typeDeclaration, MethodDefinition member)
-        {
-            if (member.IsAssembly || member.IsPrivate)
-                return;
-
-            var method = new CodeConstructor
-            {
-                CustomAttributes = CreateCustomAttributes(member),
-                Name = member.Name,
-                Attributes = GetMethodAttributes(member)
-            };
-            PopulateMethodParameters(member, method.Parameters);
-
-            typeDeclaration.Members.Add(method);
-        }
-
-        static void AddMethodToTypeDeclaration(CodeTypeDeclaration typeDeclaration, MethodDefinition member)
-        {
-            if (member.IsAssembly || member.IsPrivate || member.IsSpecialName)
-                return;
-
-            var returnType = CreateCodeTypeReference(member.ReturnType);
-
-            var method = new CodeMemberMethod
-            {
-                Name = member.Name,
-                Attributes = GetMethodAttributes(member),
-                CustomAttributes = CreateCustomAttributes(member),
-                ReturnType = returnType,
-            };
-            PopulateCustomAttributes(member.MethodReturnType, method.ReturnTypeCustomAttributes);
-            PopulateGenericParameters(member, method.TypeParameters);
-            PopulateMethodParameters(member, method.Parameters, IsExtensionMethod(member));
-
-            typeDeclaration.Members.Add(method);
-        }
-
-        static bool IsExtensionMethod(ICustomAttributeProvider method)
-        {
-            return method.CustomAttributes.Any(a => a.AttributeType.FullName == "System.Runtime.CompilerServices.ExtensionAttribute");
-        }
-
-        static void PopulateMethodParameters(IMethodSignature member,
-            CodeParameterDeclarationExpressionCollection parameters, bool isExtension = false)
-        {
-            foreach (var parameter in member.Parameters)
-            {
-                FieldDirection direction = 0;
-                if (parameter.IsOut)
-                    direction |= FieldDirection.Out;
-                else if (parameter.ParameterType.IsByReference)
-                    direction |= FieldDirection.Ref;
-
-                var parameterType = parameter.ParameterType.IsByReference
-                    ? parameter.ParameterType.GetElementType()
-                    : parameter.ParameterType;
-
-                var type = CreateCodeTypeReference(parameterType);
-
-                if (isExtension)
-                {
-                    type = ModifyCodeTypeReference(type, "this");
-                    isExtension = false;
-                }
-
-                var name = parameter.HasConstant
-                    ? string.Format(CultureInfo.InvariantCulture, "{0} = {1}", parameter.Name, FormatParameterConstant(parameter))
-                    : parameter.Name;
-                var expression = new CodeParameterDeclarationExpression(type, name)
-                {
-                    Direction = direction,
-                    CustomAttributes = CreateCustomAttributes(parameter)
-                };
-                parameters.Add(expression);
-            }
-        }
-
-        static object FormatParameterConstant(IConstantProvider parameter)
-        {
-            return parameter.Constant is string ? string.Format(CultureInfo.InvariantCulture, "\"{0}\"", parameter.Constant) : (parameter.Constant ?? "null");
-        }
-
-        static MemberAttributes GetMethodAttributes(MethodDefinition method)
-        {
-            MemberAttributes access = 0;
-            if (method.IsFamily)
-                access = MemberAttributes.Family;
-            if (method.IsPublic)
-                access = MemberAttributes.Public;
-            if (method.IsAssembly)
-                access = MemberAttributes.Assembly;
-            if (method.IsFamilyAndAssembly)
-                access = MemberAttributes.FamilyAndAssembly;
-            if (method.IsFamilyOrAssembly)
-                access = MemberAttributes.FamilyOrAssembly;
-
-            MemberAttributes scope = 0;
-            if (method.IsStatic)
-                scope |= MemberAttributes.Static;
-            if (method.IsFinal || !method.IsVirtual)
-                scope |= MemberAttributes.Final;
-            if (method.IsAbstract)
-                scope |= MemberAttributes.Abstract;
-            if (method.IsVirtual && !method.IsNewSlot)
-                scope |= MemberAttributes.Override;
-
-            MemberAttributes vtable = 0;
-            if (IsHidingMethod(method))
-                vtable = MemberAttributes.New;
-
-            return access | scope | vtable;
-        }
-
-        static bool IsHidingMethod(MethodDefinition method)
-        {
-            var typeDefinition = method.DeclaringType;
-
-            // If we're an interface, just try and find any method with the same signature
-            // in any of the interfaces that we implement
-            if (typeDefinition.IsInterface)
-            {
-                var interfaceMethods = from @interfaceReference in typeDefinition.Interfaces
-                    let interfaceDefinition = @interfaceReference.Resolve()
-                    where interfaceDefinition != null
-                    select interfaceDefinition.Methods;
-
-                return interfaceMethods.Any(ms => MetadataResolver.GetMethod(ms, method) != null);
-            }
-
-            // If we're not an interface, find a base method that isn't virtual
-            return !method.IsVirtual && GetBaseTypes(typeDefinition).Any(d => MetadataResolver.GetMethod(d.Methods, method) != null);
-        }
-
-        static IEnumerable<TypeDefinition> GetBaseTypes(TypeDefinition type)
-        {
-            var baseType = type.BaseType;
-            while (baseType != null)
-            {
-                var definition = baseType.Resolve();
-                if (definition == null)
-                    yield break;
-                yield return definition;
-
-                baseType = baseType.DeclaringType;
-            }
-        }
-
-        static void AddPropertyToTypeDeclaration(CodeTypeDeclaration typeDeclaration, PropertyDefinition member)
-        {
-            var getterAttributes = member.GetMethod != null ? GetMethodAttributes(member.GetMethod) : 0;
-            var setterAttributes = member.SetMethod != null ? GetMethodAttributes(member.SetMethod) : 0;
-
-            if (!HasVisiblePropertyMethod(getterAttributes) && !HasVisiblePropertyMethod(setterAttributes))
-                return;
-
-            var propertyAttributes = GetPropertyAttributes(getterAttributes, setterAttributes);
-
-            var propertyType = member.PropertyType.IsGenericParameter
-                ? new CodeTypeReference(member.PropertyType.Name)
-                : CreateCodeTypeReference(member.PropertyType);
-
-            var property = new CodeMemberProperty
-            {
-                Name = member.Name,
-                Type = propertyType,
-                Attributes = propertyAttributes,
-                CustomAttributes = CreateCustomAttributes(member),
-                HasGet = member.GetMethod != null && HasVisiblePropertyMethod(getterAttributes),
-                HasSet = member.SetMethod != null && HasVisiblePropertyMethod(setterAttributes)
-            };
-
-            // Here's a nice hack, because hey, guess what, the CodeDOM doesn't support
-            // attributes on getters or setters
-            if (member.GetMethod != null && member.GetMethod.HasCustomAttributes)
-            {
-                PopulateCustomAttributes(member.GetMethod, property.CustomAttributes, type => ModifyCodeTypeReference(type, "get:"));
-            }
-            if (member.SetMethod != null && member.SetMethod.HasCustomAttributes)
-            {
-                PopulateCustomAttributes(member.SetMethod, property.CustomAttributes, type => ModifyCodeTypeReference(type, "set:"));
-            }
-
-            foreach (var parameter in member.Parameters)
-            {
-                property.Parameters.Add(
-                    new CodeParameterDeclarationExpression(CreateCodeTypeReference(parameter.ParameterType),
-                        parameter.Name));
-            }
-
-            // TODO: CodeDOM has no support for different access modifiers for getters and setters
-            // TODO: CodeDOM has no support for attributes on setters or getters - promote to property?
-
-            typeDeclaration.Members.Add(property);
-        }
-
-        static MemberAttributes GetPropertyAttributes(MemberAttributes getterAttributes, MemberAttributes setterAttributes)
-        {
-            MemberAttributes access = 0;
-            var getterAccess = getterAttributes & MemberAttributes.AccessMask;
-            var setterAccess = setterAttributes & MemberAttributes.AccessMask;
-            if (getterAccess == MemberAttributes.Public || setterAccess == MemberAttributes.Public)
-                access = MemberAttributes.Public;
-            else if (getterAccess == MemberAttributes.Family || setterAccess == MemberAttributes.Family)
-                access = MemberAttributes.Family;
-            else if (getterAccess == MemberAttributes.FamilyAndAssembly || setterAccess == MemberAttributes.FamilyAndAssembly)
-                access = MemberAttributes.FamilyAndAssembly;
-            else if (getterAccess == MemberAttributes.FamilyOrAssembly || setterAccess == MemberAttributes.FamilyOrAssembly)
-                access = MemberAttributes.FamilyOrAssembly;
-            else if (getterAccess == MemberAttributes.Assembly || setterAccess == MemberAttributes.Assembly)
-                access = MemberAttributes.Assembly;
-            else if (getterAccess == MemberAttributes.Private || setterAccess == MemberAttributes.Private)
-                access = MemberAttributes.Private;
-
-            // Scope should be the same for getter and setter. If one isn't specified, it'll be 0
-            var getterScope = getterAttributes & MemberAttributes.ScopeMask;
-            var setterScope = setterAttributes & MemberAttributes.ScopeMask;
-            var scope = (MemberAttributes) Math.Max((int) getterScope, (int) setterScope);
-
-            // Vtable should be the same for getter and setter. If one isn't specified, it'll be 0
-            var getterVtable = getterAttributes & MemberAttributes.VTableMask;
-            var setterVtable = setterAttributes & MemberAttributes.VTableMask;
-            var vtable = (MemberAttributes) Math.Max((int) getterVtable, (int) setterVtable);
-
-            return access | scope | vtable;
-        }
-
-        static bool HasVisiblePropertyMethod(MemberAttributes attributes)
-        {
-            var access = attributes & MemberAttributes.AccessMask;
-            return access == MemberAttributes.Public || access == MemberAttributes.Family ||
-                   access == MemberAttributes.FamilyOrAssembly;
-        }
-
-        static CodeTypeMember GenerateEvent(EventDefinition eventDefinition)
-        {
-            var @event = new CodeMemberEvent
-            {
-                Name = eventDefinition.Name,
-                Attributes = MemberAttributes.Public | MemberAttributes.Final,
-                CustomAttributes = CreateCustomAttributes(eventDefinition),
-                Type = CreateCodeTypeReference(eventDefinition.EventType)
-            };
-
-            return @event;
-        }
-
-        static void AddFieldToTypeDeclaration(CodeTypeDeclaration typeDeclaration, FieldDefinition memberInfo)
-        {
-            if (memberInfo.IsPrivate || memberInfo.IsAssembly || memberInfo.IsSpecialName)
-                return;
-
-            MemberAttributes attributes = 0;
-            if (memberInfo.HasConstant)
-                attributes |= MemberAttributes.Const;
-            if (memberInfo.IsFamily)
-                attributes |= MemberAttributes.Family;
-            if (memberInfo.IsPublic)
-                attributes |= MemberAttributes.Public;
-            if (memberInfo.IsStatic && !memberInfo.HasConstant)
-                attributes |= MemberAttributes.Static;
-
-            // TODO: Values for readonly fields are set in the ctor
-            var codeTypeReference = CreateCodeTypeReference(memberInfo.FieldType);
-            if (memberInfo.IsInitOnly)
-                codeTypeReference = MakeReadonly(codeTypeReference);
-            var field = new CodeMemberField(codeTypeReference, memberInfo.Name)
-            {
-                Attributes = attributes,
-                CustomAttributes = CreateCustomAttributes(memberInfo)
-            };
-
-            if (memberInfo.HasConstant)
-                field.InitExpression = new CodePrimitiveExpression(memberInfo.Constant);
-
-            typeDeclaration.Members.Add(field);
-        }
-
-        static CodeTypeReference MakeReadonly(CodeTypeReference typeReference)
-        {
-            return ModifyCodeTypeReference(typeReference, "readonly");
-        }
-
-        static CodeTypeReference ModifyCodeTypeReference(CodeTypeReference typeReference, string modifier)
-        {
-            using (var provider = new CSharpCodeProvider())
-                return new CodeTypeReference(modifier + " " + provider.GetTypeOutput(typeReference));
-        }
-
-        static CodeTypeReference CreateCodeTypeReference(TypeReference type)
-        {
-            var typeName = GetTypeName(type);
-            return new CodeTypeReference(typeName, CreateGenericArguments(type));
-        }
-
-        static string GetTypeName(TypeReference type)
-        {
-            if (type.IsGenericParameter)
-                return type.Name;
-
-            if (!type.IsNested)
-            {
-                return (!string.IsNullOrEmpty(type.Namespace) ? (type.Namespace + ".") : "") + type.Name;
-            }
-
-            return GetTypeName(type.DeclaringType) + "." + type.Name;
-        }
-
-        static CodeTypeReference[] CreateGenericArguments(TypeReference type)
-        {
-            if (!(type is IGenericInstance genericInstance))
-            {
-                return null;
-            }
-
-            var genericArguments = new List<CodeTypeReference>();
-            foreach (var argument in genericInstance.GenericArguments)
-            {
-                genericArguments.Add(CreateCodeTypeReference(argument));
-            }
-            return genericArguments.ToArray();
-        }
-    }
-
-    static class CecilEx
-    {
-        public static IEnumerable<IMemberDefinition> GetMembers(this TypeDefinition type)
-        {
-            return type.Fields.Cast<IMemberDefinition>()
-                .Concat(type.Methods)
-                .Concat(type.Properties)
-                .Concat(type.Events);
-        }
-    }
-}
-
-#endif

test/Microsoft.Azure.ServiceBus.UnitTests/Diagnostics/DiagnosticsTests.cs

@@ -0,0 +1,535 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.Azure.ServiceBus.UnitTests.Diagnostics
+{
+    using System;
+    using System.Collections.Concurrent;
+    using System.Collections.Generic;
+    using System.Diagnostics;
+    using System.Linq;
+    using System.Reflection;
+    using System.Threading.Tasks;
+    using Xunit;
+
+    public abstract class DiagnosticsTests : IDisposable
+    {
+        protected ConcurrentQueue<(string eventName, object payload, Activity activity)> events;
+        protected FakeDiagnosticListener listener;
+        protected IDisposable subscription;
+        protected const int maxWaitSec = 10;
+
+        protected abstract string EntityName { get; }
+        private bool disposed = false;
+
+        internal DiagnosticsTests()
+        {
+            this.events = new ConcurrentQueue<(string eventName, object payload, Activity activity)>();
+            this.listener = new FakeDiagnosticListener(kvp =>
+            {
+                TestUtility.Log($"Diagnostics event: {kvp.Key}, Activity Id: {Activity.Current?.Id}");
+                if (kvp.Key.Contains("Exception"))
+                {
+                    TestUtility.Log($"Exception {kvp.Value}");
+                }
+
+                this.events.Enqueue((kvp.Key, kvp.Value, Activity.Current));
+            });
+            this.subscription = DiagnosticListener.AllListeners.Subscribe(this.listener);
+        }
+
+        #region Send
+
+        protected void AssertSendStart(string name, object payload, Activity activity, Activity parentActivity, int messageCount = 1)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Send.Start", name);
+            AssertCommonPayloadProperties(payload);
+            var messages = GetPropertyValueFromAnonymousTypeInstance<IList<Message>>(payload, "Messages");
+            Assert.Equal(messageCount, messages.Count);
+
+            Assert.NotNull(activity);
+            Assert.Equal(parentActivity, activity.Parent);
+
+            AssertTags(messages, activity);
+        }
+
+        protected void AssertSendStop(string name, object payload, Activity activity, Activity sendActivity, int messageCount = 1)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Send.Stop", name);
+            AssertCommonStopPayloadProperties(payload);
+
+            if (sendActivity != null)
+            {
+                Assert.Equal(sendActivity, activity);
+            }
+
+            var messages = GetPropertyValueFromAnonymousTypeInstance<IList<Message>>(payload, "Messages");
+            Assert.Equal(messageCount, messages.Count);
+        }
+
+        #endregion
+
+        #region Complete
+
+        protected void AssertCompleteStart(string name, object payload, Activity activity, Activity parentActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Complete.Start", name);
+            AssertCommonPayloadProperties(payload);
+            GetPropertyValueFromAnonymousTypeInstance<IList<string>>(payload, "LockTokens");
+
+            Assert.NotNull(activity);
+            Assert.Equal(parentActivity, activity.Parent);
+        }
+
+        protected void AssertCompleteStop(string name, object payload, Activity activity, Activity completeActivity, Activity parentActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Complete.Stop", name);
+            AssertCommonStopPayloadProperties(payload);
+
+            if (completeActivity != null)
+            {
+                Assert.Equal(completeActivity, activity);
+            }
+
+            var tokens = GetPropertyValueFromAnonymousTypeInstance<IList<string>>(payload, "LockTokens");
+            Assert.Equal(1, tokens.Count);
+        }
+
+        #endregion
+
+
+        #region Abandon
+
+        protected void AssertAbandonStart(string name, object payload, Activity activity, Activity parentActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Abandon.Start", name);
+            AssertCommonPayloadProperties(payload);
+            GetPropertyValueFromAnonymousTypeInstance<string>(payload, "LockToken");
+
+            Assert.NotNull(activity);
+            Assert.Equal(parentActivity, activity.Parent);
+        }
+
+        protected void AssertAbandonStop(string name, object payload, Activity activity, Activity abandonActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Abandon.Stop", name);
+            AssertCommonStopPayloadProperties(payload);
+
+            if (abandonActivity != null)
+            {
+                Assert.Equal(abandonActivity, activity);
+            }
+
+            GetPropertyValueFromAnonymousTypeInstance<string>(payload, "LockToken");
+        }
+
+        #endregion
+
+
+        #region Defer
+
+        protected void AssertDeferStart(string name, object payload, Activity activity, Activity parentActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Defer.Start", name);
+            AssertCommonPayloadProperties(payload);
+            GetPropertyValueFromAnonymousTypeInstance<string>(payload, "LockToken");
+
+            Assert.NotNull(activity);
+            Assert.Equal(parentActivity, activity.Parent);
+        }
+
+        protected void AssertDeferStop(string name, object payload, Activity activity, Activity deferActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Defer.Stop", name);
+            AssertCommonStopPayloadProperties(payload);
+
+            if (deferActivity != null)
+            {
+                Assert.Equal(deferActivity, activity);
+            }
+
+            GetPropertyValueFromAnonymousTypeInstance<string>(payload, "LockToken");
+        }
+
+        #endregion
+
+        #region DeadLetter
+
+        protected void AssertDeadLetterStart(string name, object payload, Activity activity, Activity parentActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.DeadLetter.Start", name);
+            AssertCommonPayloadProperties(payload);
+            GetPropertyValueFromAnonymousTypeInstance<string>(payload, "LockToken");
+
+            Assert.NotNull(activity);
+            Assert.Equal(parentActivity, activity.Parent);
+        }
+
+        protected void AssertDeadLetterStop(string name, object payload, Activity activity, Activity deadLetterActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.DeadLetter.Stop", name);
+            AssertCommonStopPayloadProperties(payload);
+
+            if (deadLetterActivity != null)
+            {
+                Assert.Equal(deadLetterActivity, activity);
+            }
+
+            GetPropertyValueFromAnonymousTypeInstance<string>(payload, "LockToken");
+        }
+
+        #endregion
+
+        #region Schedule
+
+        protected void AssertScheduleStart(string name, object payload, Activity activity, Activity parentActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Schedule.Start", name);
+            AssertCommonPayloadProperties(payload);
+            var message = GetPropertyValueFromAnonymousTypeInstance<Message>(payload, "Message");
+            GetPropertyValueFromAnonymousTypeInstance<DateTimeOffset>(payload, "ScheduleEnqueueTimeUtc");
+
+            Assert.NotNull(activity);
+            Assert.Equal(parentActivity, activity.Parent);
+            AssertTags(message, activity);
+        }
+
+        protected void AssertScheduleStop(string name, object payload, Activity activity, Activity scheduleActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Schedule.Stop", name);
+            AssertCommonStopPayloadProperties(payload);
+
+            Assert.Equal(scheduleActivity, activity);
+            var message = GetPropertyValueFromAnonymousTypeInstance<Message>(payload, "Message");
+            GetPropertyValueFromAnonymousTypeInstance<DateTimeOffset>(payload, "ScheduleEnqueueTimeUtc");
+            GetPropertyValueFromAnonymousTypeInstance<long>(payload, "SequenceNumber");
+            Assert.NotNull(message);
+            Assert.Contains("Diagnostic-Id", message.UserProperties.Keys);
+            if (scheduleActivity.Baggage.Any())
+            {
+                Assert.Contains("Correlation-Context", message.UserProperties.Keys);
+            }
+        }
+
+        #endregion
+
+        #region Cancel
+
+        protected void AssertCancelStart(string name, object payload, Activity activity, Activity parentActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Cancel.Start", name);
+            AssertCommonPayloadProperties(payload);
+            GetPropertyValueFromAnonymousTypeInstance<long>(payload, "SequenceNumber");
+
+            Assert.NotNull(activity);
+            Assert.Equal(parentActivity, activity.Parent);
+        }
+
+        protected void AssertCancelStop(string name, object payload, Activity activity, Activity scheduleActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Cancel.Stop", name);
+            AssertCommonStopPayloadProperties(payload);
+
+            Assert.Equal(scheduleActivity, activity);
+            GetPropertyValueFromAnonymousTypeInstance<long>(payload, "SequenceNumber");
+        }
+
+        #endregion
+
+        #region Receive
+
+        protected int AssertReceiveStart(string name, object payload, Activity activity, int messagesCount = 1)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Receive.Start", name);
+            AssertCommonPayloadProperties(payload);
+
+            var count = GetPropertyValueFromAnonymousTypeInstance<int>(payload, "RequestedMessageCount");
+            if (messagesCount != -1)
+            {
+                Assert.Equal(messagesCount, count);
+            }
+            Assert.NotNull(activity);
+            Assert.Null(activity.Parent);
+
+            return count;
+        }
+
+        protected int AssertReceiveStop(string name, object payload, Activity activity, Activity receiveActivity, Activity sendActivity, int sentMessagesCount = 1, int receivedMessagesCount = 1)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Receive.Stop", name);
+            AssertCommonStopPayloadProperties(payload);
+
+            if (receiveActivity != null)
+            {
+                Assert.Equal(receiveActivity, activity);
+            }
+
+            Assert.Equal(sentMessagesCount, GetPropertyValueFromAnonymousTypeInstance<int>(payload, "RequestedMessageCount"));
+            var messages = GetPropertyValueFromAnonymousTypeInstance<IList<Message>>(payload, "Messages");
+
+            if (receivedMessagesCount != -1)
+            {
+                Assert.Equal(receivedMessagesCount, messages.Count);
+            }
+
+            if (sendActivity != null)
+            {
+                Assert.Contains(sendActivity.Id, activity.Tags.Single(t => t.Key == "RelatedTo").Value);
+            }
+
+            AssertTags(messages, activity);
+            return messages.Count;
+        }
+
+        #endregion
+
+        #region ReceiveDeferred
+
+        protected void AssertReceiveDeferredStart(string name, object payload, Activity activity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.ReceiveDeferred.Start", name);
+            AssertCommonPayloadProperties(payload);
+
+            Assert.Single(GetPropertyValueFromAnonymousTypeInstance<IEnumerable<long>>(payload, "SequenceNumbers"));
+
+            Assert.NotNull(activity);
+            Assert.Null(activity.Parent);
+        }
+
+        protected void AssertReceiveDeferredStop(string name, object payload, Activity activity, Activity receiveActivity, Activity sendActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.ReceiveDeferred.Stop", name);
+            AssertCommonStopPayloadProperties(payload);
+
+            if (receiveActivity != null)
+            {
+                Assert.Equal(receiveActivity, activity);
+            }
+
+            var messages = GetPropertyValueFromAnonymousTypeInstance<IList<Message>>(payload, "Messages");
+            Assert.Equal(1, messages.Count);
+            Assert.Single(GetPropertyValueFromAnonymousTypeInstance<IEnumerable<long>>(payload, "SequenceNumbers"));
+
+            Assert.Equal(sendActivity.Id, activity.Tags.Single(t => t.Key == "RelatedTo").Value);
+
+            AssertTags(messages, activity);
+        }
+
+        #endregion
+
+
+        #region Process
+
+        protected void AssertProcessStart(string name, object payload, Activity activity, Activity sendActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Process.Start", name);
+            AssertCommonPayloadProperties(payload);
+
+            var message = GetPropertyValueFromAnonymousTypeInstance<Message>(payload, "Message");
+
+            Assert.NotNull(activity);
+            Assert.Null(activity.Parent);
+            Assert.Equal(sendActivity.Id, activity.ParentId);
+            Assert.Equal(sendActivity.Baggage.OrderBy(p => p.Key), activity.Baggage.OrderBy(p => p.Key));
+
+            AssertTags(message, activity);
+        }
+
+        protected void AssertProcessStop(string name, object payload, Activity activity, Activity processActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Process.Stop", name);
+            AssertCommonStopPayloadProperties(payload);
+
+            if (processActivity != null)
+            {
+                Assert.Equal(processActivity, activity);
+            }
+        }
+
+        #endregion
+
+
+        #region Peek
+
+        protected void AssertPeekStart(string name, object payload, Activity activity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Peek.Start", name);
+            AssertCommonPayloadProperties(payload);
+
+            GetPropertyValueFromAnonymousTypeInstance<long>(payload, "FromSequenceNumber");
+            Assert.Equal(1, GetPropertyValueFromAnonymousTypeInstance<int>(payload, "RequestedMessageCount"));
+
+            Assert.NotNull(activity);
+            Assert.Null(activity.Parent);
+        }
+
+        protected void AssertPeekStop(string name, object payload, Activity activity, Activity peekActivity, Activity sendActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Peek.Stop", name);
+            AssertCommonStopPayloadProperties(payload);
+
+            if (peekActivity != null)
+            {
+                Assert.Equal(peekActivity, activity);
+            }
+
+            GetPropertyValueFromAnonymousTypeInstance<long>(payload, "FromSequenceNumber");
+            Assert.Equal(1, GetPropertyValueFromAnonymousTypeInstance<int>(payload, "RequestedMessageCount"));
+            var messages = GetPropertyValueFromAnonymousTypeInstance<IList<Message>>(payload, "Messages");
+            AssertTags(messages, activity);
+
+            Assert.Equal(sendActivity.Id, activity.Tags.Single(t => t.Key == "RelatedTo").Value);
+        }
+
+        #endregion
+
+
+        #region RenewLock
+
+        protected void AssertRenewLockStart(string name, object payload, Activity activity, Activity parentActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.RenewLock.Start", name);
+            AssertCommonPayloadProperties(payload);
+            GetPropertyValueFromAnonymousTypeInstance<string>(payload, "LockToken");
+
+            Assert.NotNull(activity);
+            Assert.Equal(parentActivity, activity.Parent);
+        }
+
+        protected void AssertRenewLockStop(string name, object payload, Activity activity, Activity renewLockActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.RenewLock.Stop", name);
+            AssertCommonStopPayloadProperties(payload);
+
+            if (renewLockActivity != null)
+            {
+                Assert.Equal(renewLockActivity, activity);
+            }
+
+            GetPropertyValueFromAnonymousTypeInstance<string>(payload, "LockToken");
+            GetPropertyValueFromAnonymousTypeInstance<DateTime>(payload, "LockedUntilUtc");
+        }
+
+        #endregion
+
+        protected void AssertTags(IList<Message> messageList, Activity activity)
+        {
+            var messagesWithId = messageList.Where(m => m.MessageId != null).ToArray();
+            if (messagesWithId.Any())
+            {
+                Assert.Contains("MessageId", activity.Tags.Select(t => t.Key));
+
+                string messageIdTag = activity.Tags.Single(t => t.Key == "MessageId").Value;
+                foreach (var m in messagesWithId)
+                {
+                    Assert.Contains(m.MessageId, messageIdTag);
+                }
+            }
+
+            var messagesWithSessionId = messageList.Where(m => m.SessionId != null).ToArray();
+            if (messagesWithSessionId.Any())
+            {
+                Assert.Contains("SessionId", activity.Tags.Select(t => t.Key));
+
+                string sessionIdTag = activity.Tags.Single(t => t.Key == "SessionId").Value;
+                foreach (var m in messagesWithSessionId)
+                {
+                    Assert.Contains(m.SessionId, sessionIdTag);
+                }
+            }
+        }
+
+        protected void AssertTags(Message message, Activity activity)
+        {
+            if (message.MessageId != null)
+            {
+                Assert.Contains("MessageId", activity.Tags.Select(t => t.Key));
+                Assert.Equal(message.MessageId, activity.Tags.Single(t => t.Key == "MessageId").Value);
+            }
+
+            if (message.SessionId != null)
+            {
+                Assert.Contains("SessionId", activity.Tags.Select(t => t.Key));
+                Assert.Equal(message.SessionId, activity.Tags.Single(t => t.Key == "SessionId").Value);
+            }
+        }
+
+        protected void AssertException(string name, object payload, Activity activity, Activity parentActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.Exception", name);
+            AssertCommonPayloadProperties(payload);
+
+            GetPropertyValueFromAnonymousTypeInstance<Exception>(payload, "Exception");
+
+            Assert.NotNull(activity);
+            if (parentActivity != null)
+            {
+                Assert.Equal(parentActivity, activity.Parent);
+            }
+        }
+
+        protected void AssertCommonPayloadProperties(object eventPayload)
+        {
+            var entity = GetPropertyValueFromAnonymousTypeInstance<string>(eventPayload, "Entity");
+            GetPropertyValueFromAnonymousTypeInstance<Uri>(eventPayload, "Endpoint");
+
+            Assert.Equal(this.EntityName, entity);
+        }
+
+        protected void AssertCommonStopPayloadProperties(object eventPayload)
+        {
+            AssertCommonPayloadProperties(eventPayload);
+            var status = GetPropertyValueFromAnonymousTypeInstance<TaskStatus>(eventPayload, "Status");
+            Assert.Equal(TaskStatus.RanToCompletion, status);
+        }
+
+        protected T GetPropertyValueFromAnonymousTypeInstance<T>(object obj, string propertyName)
+        {
+            Type t = obj.GetType();
+
+            PropertyInfo p = t.GetRuntimeProperty(propertyName);
+
+            object propertyValue = p.GetValue(obj);
+            Assert.NotNull(propertyValue);
+            Assert.IsAssignableFrom<T>(propertyValue);
+
+            return (T)propertyValue;
+        }
+
+        public void Dispose()
+        {
+            Dispose(true);
+            GC.SuppressFinalize(this);
+        }
+
+        protected virtual void Dispose(bool disposing)
+        {
+            if (this.disposed)
+                return;
+
+            if (disposing)
+            {
+                //Thread.Sleep(90000);
+                this.listener?.Disable();
+
+                while (this.events.TryDequeue(out var evnt))
+                {
+                }
+
+                while (Activity.Current != null)
+                {
+                    Activity.Current.Stop();
+                }
+
+                this.listener?.Dispose();
+                this.subscription?.Dispose();
+
+                this.events = null;
+                this.listener = null;
+                this.subscription = null;
+            }
+
+            this.disposed = true;
+        }
+    }
+}

test/Microsoft.Azure.ServiceBus.UnitTests/Diagnostics/ExtractActivityTests.cs

@@ -0,0 +1,131 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.Azure.ServiceBus.UnitTests.Diagnostics
+{
+    using System.Linq;
+    using Microsoft.Azure.ServiceBus.Diagnostics;
+    using Xunit;
+
+    public class ExtractActivityTests
+    {
+        [Fact]
+        [DisplayTestMethodName]
+        void ValidIdAndContextAreExtracted()
+        {
+            var message = new Message();
+            message.UserProperties["Diagnostic-Id"] = "diagnostic-id";
+            message.UserProperties["Correlation-Context"] = "k1=v1";
+
+            var activity = message.ExtractActivity();
+
+            Assert.Equal("diagnostic-id", activity.ParentId);
+            Assert.Equal("diagnostic-id", activity.RootId);
+            
+            Assert.Null(activity.Id);
+
+            var baggage = activity.Baggage.ToDictionary(kvp => kvp.Key, kvp => kvp.Value);
+            Assert.Equal(1, baggage.Count);
+            Assert.Contains("k1", baggage.Keys);
+            Assert.Equal("v1", baggage["k1"]);
+        }
+
+        [Fact]
+        [DisplayTestMethodName]
+        void ValidIdAndMultipleContextAreExtracted()
+        {
+            var message = new Message();
+            message.UserProperties["Diagnostic-Id"] = "diagnostic-id";
+            message.UserProperties["Correlation-Context"] = "k1=v1,k2=v2,k3=v3";
+
+            var activity = message.ExtractActivity();
+
+            Assert.Equal("diagnostic-id", activity.ParentId);
+            Assert.Equal("diagnostic-id", activity.RootId);
+
+            Assert.Null(activity.Id);
+
+            var baggage = activity.Baggage.ToDictionary(kvp => kvp.Key, kvp => kvp.Value);
+            Assert.Equal(3, baggage.Count);
+            Assert.Contains("k1", baggage.Keys);
+            Assert.Contains("k2", baggage.Keys);
+            Assert.Contains("k3", baggage.Keys);
+            Assert.Equal("v1", baggage["k1"]);
+            Assert.Equal("v2", baggage["k2"]);
+            Assert.Equal("v3", baggage["k3"]);
+        }
+
+        [Fact]
+        [DisplayTestMethodName]
+        void ActivityNameCouldBeChanged()
+        {
+            var message = new Message();
+            message.UserProperties["Diagnostic-Id"] = "diagnostic-id";
+
+            var activity = message.ExtractActivity("My activity");
+
+            Assert.Equal("My activity", activity.OperationName);
+        }
+
+        [Fact]
+        [DisplayTestMethodName]
+        void ValidIdAndNoContextAreExtracted()
+        {
+            var message = new Message();
+            message.UserProperties["Diagnostic-Id"] = "diagnostic-id";
+
+            var activity = message.ExtractActivity();
+
+            Assert.Equal("diagnostic-id", activity.ParentId);
+            Assert.Equal("diagnostic-id", activity.RootId);
+            Assert.Null(activity.Id);
+
+            var baggage = activity.Baggage.ToDictionary(kvp => kvp.Key, kvp => kvp.Value);
+            Assert.Equal(0, baggage.Count);
+        }
+
+
+        [Theory]
+        [InlineData(null)]
+        [InlineData("")]
+        [InlineData("not valid context")]
+        [InlineData("not,valid,context")]
+        [DisplayTestMethodName]
+        void ValidIdAndInvalidContextAreExtracted(string context)
+        {
+            var message = new Message();
+            message.UserProperties["Diagnostic-Id"] = "diagnostic-id";
+            message.UserProperties["Correlation-Context"] = context;
+
+            var activity = message.ExtractActivity();
+
+            Assert.Equal("diagnostic-id", activity.ParentId);
+            Assert.Equal("diagnostic-id", activity.RootId);
+            Assert.Null(activity.Id);
+
+            var baggage = activity.Baggage.ToDictionary(kvp => kvp.Key, kvp => kvp.Value);
+            Assert.Equal(0, baggage.Count);
+        }
+
+        [Theory]
+        [InlineData(null)]
+        [InlineData("")]
+        [DisplayTestMethodName]
+        void EmptyIdResultsInActivityWithoutParent(string id)
+        {
+            var message = new Message();
+            message.UserProperties["Diagnostic-Id"] = id;
+            message.UserProperties["Correlation-Context"] = "k1=v1,k2=v2";
+
+            var activity = message.ExtractActivity();
+
+            Assert.Null(activity.ParentId);
+            Assert.Null(activity.RootId);
+            Assert.Null(activity.Id);
+
+            var baggage = activity.Baggage.ToDictionary(kvp => kvp.Key, kvp => kvp.Value);
+            // baggage is ignored in absence of Id
+            Assert.Equal(0, baggage.Count);
+        }
+    }
+}

test/Microsoft.Azure.ServiceBus.UnitTests/Diagnostics/FakeDiagnosticListener.cs

@@ -0,0 +1,92 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.Azure.ServiceBus.UnitTests.Diagnostics
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Diagnostics;
+
+    public sealed class FakeDiagnosticListener : IObserver<DiagnosticListener>, IDisposable
+    {
+        private  IDisposable subscription;
+        private class FakeDiagnosticSourceWriteObserver : IObserver<KeyValuePair<string, object>>
+        {
+            private readonly Action<KeyValuePair<string, object>> writeCallback;
+
+            public FakeDiagnosticSourceWriteObserver(Action<KeyValuePair<string, object>> writeCallback)
+            {
+                this.writeCallback = writeCallback;
+            }
+
+            public void OnCompleted()
+            {
+            }
+
+            public void OnError(Exception error)
+            {
+            }
+
+            public void OnNext(KeyValuePair<string, object> value)
+            {
+                this.writeCallback(value);
+            }
+        }
+
+        private readonly Action<KeyValuePair<string, object>> writeCallback;
+
+        private Func<string, object, object, bool> writeObserverEnabled = (name, arg1, arg2) => false;
+
+        public FakeDiagnosticListener(Action<KeyValuePair<string, object>> writeCallback)
+        {
+            this.writeCallback = writeCallback;
+        }
+
+        public void OnCompleted()
+        {
+        }
+
+        public void OnError(Exception error)
+        {
+        }
+
+        public void OnNext(DiagnosticListener value)
+        {
+            if (value.Name.Equals("Microsoft.Azure.ServiceBus"))
+            {
+                this.subscription = value.Subscribe(new FakeDiagnosticSourceWriteObserver(this.writeCallback), this.IsEnabled);
+            }
+        }
+
+        public void Enable()
+        {
+            this.writeObserverEnabled = (name, arg1, arg2) => true;
+        }
+
+        public void Enable(Func<string, bool> writeObserverEnabled)
+        {
+            this.writeObserverEnabled = (name, arg1, arg2) => writeObserverEnabled(name);
+        }
+
+        public void Enable(Func<string, object, object, bool> writeObserverEnabled)
+        {
+            this.writeObserverEnabled = writeObserverEnabled;
+        }
+
+        public void Disable()
+        {
+            this.writeObserverEnabled = (name, arg1, arg2) => false;
+        }
+
+        private bool IsEnabled(string s, object arg1, object arg2)
+        {
+            return this.writeObserverEnabled.Invoke(s, arg1, arg2);
+        }
+
+        public void Dispose()
+        {
+            this.Disable();
+            this.subscription?.Dispose();
+        }
+    }
+}

test/Microsoft.Azure.ServiceBus.UnitTests/Diagnostics/QueueClientDiagnosticsTests.cs

@@ -0,0 +1,498 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.Azure.ServiceBus.UnitTests.Diagnostics
+{
+    using System;
+    using System.Diagnostics;
+    using System.Linq;
+    using System.Threading;
+    using System.Threading.Tasks;
+    using Xunit;
+
+    public sealed class QueueClientDiagnosticsTests : DiagnosticsTests
+    {
+        protected override string EntityName => TestConstants.NonPartitionedQueueName;
+        private QueueClient queueClient;
+        private bool disposed = false;
+
+        [Fact]
+        [DisplayTestMethodName]
+        async Task EventsAreNotFiredWhenDiagnosticsIsDisabled()
+        {
+            this.queueClient = new QueueClient(TestUtility.NamespaceConnectionString, TestConstants.NonPartitionedQueueName, ReceiveMode.ReceiveAndDelete);
+
+            Activity processActivity = null;
+
+            ManualResetEvent processingDone = new ManualResetEvent(false);
+            this.listener.Disable();
+
+            await TestUtility.SendMessagesAsync(this.queueClient.InnerSender, 1);
+            this.queueClient.RegisterMessageHandler((msg, ct) => {
+                    processActivity = Activity.Current;
+                    processingDone.Set();
+                    return Task.CompletedTask;
+                },
+                exArgs => Task.CompletedTask);
+
+            processingDone.WaitOne(TimeSpan.FromSeconds(maxWaitSec));
+            Assert.True(this.events.IsEmpty);
+            Assert.Null(processActivity);
+        }
+
+        [Fact]
+        [DisplayTestMethodName]
+        async Task EventsAreNotFiredWhenDiagnosticsIsDisabledForQueue()
+        {
+            this.queueClient = new QueueClient(TestUtility.NamespaceConnectionString, TestConstants.NonPartitionedQueueName,
+                ReceiveMode.ReceiveAndDelete);
+
+            Activity processActivity = null;
+
+            ManualResetEvent processingDone = new ManualResetEvent(false);
+            this.listener.Enable((name, queueName, arg) =>
+                queueName == null || queueName.ToString() != TestConstants.NonPartitionedQueueName);
+
+            await TestUtility.SendMessagesAsync(this.queueClient.InnerSender, 1);
+            this.queueClient.RegisterMessageHandler((msg, ct) =>
+                {
+                    processActivity = Activity.Current;
+                    processingDone.Set();
+                    return Task.CompletedTask;
+                },
+                exArgs => Task.CompletedTask);
+
+            processingDone.WaitOne(TimeSpan.FromSeconds(maxWaitSec));
+
+            Assert.True(this.events.IsEmpty);
+            Assert.Null(processActivity);
+        }
+
+        [Fact]
+        [DisplayTestMethodName]
+        async Task SendAndHandlerFireEvents()
+        {
+            this.queueClient = new QueueClient(TestUtility.NamespaceConnectionString, TestConstants.NonPartitionedQueueName,
+                ReceiveMode.PeekLock);
+
+            Activity parentActivity = new Activity("test").AddBaggage("k1", "v1").AddBaggage("k2", "v2");
+            Activity processActivity = null;
+            bool exceptionCalled = false;
+            ManualResetEvent processingDone = new ManualResetEvent(false);
+            this.listener.Enable((name, queueName, arg) => !name.Contains("Receive") && !name.Contains("Exception"));
+
+            parentActivity.Start();
+
+            await TestUtility.SendSessionMessagesAsync(this.queueClient.InnerSender, 1, 1);
+            parentActivity.Stop();
+
+            this.queueClient.RegisterMessageHandler((msg, ct) =>
+                {
+                    processActivity = Activity.Current;
+                    processingDone.Set();
+                    return Task.CompletedTask;
+                },
+                exArgs =>
+                {
+                    exceptionCalled = true;
+                    return Task.CompletedTask;
+                });
+
+            processingDone.WaitOne(TimeSpan.FromSeconds(maxWaitSec));
+
+            Assert.True(this.events.TryDequeue(out var sendStart));
+            AssertSendStart(sendStart.eventName, sendStart.payload, sendStart.activity, parentActivity);
+
+            Assert.True(this.events.TryDequeue(out var sendStop));
+            AssertSendStop(sendStop.eventName, sendStop.payload, sendStop.activity, sendStart.activity);
+
+            Assert.True(this.events.TryDequeue(out var processStart));
+            AssertProcessStart(processStart.eventName, processStart.payload, processStart.activity,
+                sendStart.activity);
+
+            // message is processed, but complete happens after that
+            // let's wat until Complete starts and ends and Process ends
+            int wait = 0;
+            while (wait++ < maxWaitSec && this.events.Count < 3)
+            {
+                await Task.Delay(TimeSpan.FromSeconds(1));
+            }
+
+            Assert.True(this.events.TryDequeue(out var completeStart));
+            AssertCompleteStart(completeStart.eventName, completeStart.payload, completeStart.activity,
+                processStart.activity);
+
+            Assert.True(this.events.TryDequeue(out var completeStop));
+            AssertCompleteStop(completeStop.eventName, completeStop.payload, completeStop.activity,
+                completeStart.activity, processStart.activity);
+
+            Assert.True(this.events.TryDequeue(out var processStop));
+            AssertProcessStop(processStop.eventName, processStop.payload, processStop.activity,
+                processStart.activity);
+
+            Assert.False(this.events.TryDequeue(out var evnt));
+
+            Assert.Equal(processStop.activity, processActivity);
+            Assert.False(exceptionCalled);
+        }
+
+        [Fact]
+        [DisplayTestMethodName]
+        async Task SendAndHandlerFireExceptionEvents()
+        {
+            this.queueClient = new QueueClient(TestUtility.NamespaceConnectionString, TestConstants.NonPartitionedQueueName,
+                ReceiveMode.PeekLock);
+
+            bool exceptionCalled = false;
+
+            ManualResetEvent processingDone = new ManualResetEvent(false);
+            await TestUtility.SendMessagesAsync(this.queueClient.InnerSender, 1);
+
+            this.listener.Enable((name, queueName, arg) => !name.EndsWith(".Start") && !name.Contains("Receive") );
+
+            int count = 0;
+            this.queueClient.RegisterMessageHandler((msg, ct) =>
+                {
+                    if (count++ == 0)
+                    {
+                        throw new Exception("123");
+                    }
+
+                    processingDone.Set();
+                    return Task.CompletedTask;
+                },
+                exArgs =>
+                {
+                    exceptionCalled = true;
+                    return Task.CompletedTask;
+                });
+            processingDone.WaitOne(TimeSpan.FromSeconds(maxWaitSec));
+            Assert.True(exceptionCalled);
+
+            // message is processed, but abandon happens after that
+            // let's spin until Complete call starts and ends
+            int wait = 0;
+            while (wait++ < maxWaitSec && this.events.Count < 3)
+            {
+                await Task.Delay(TimeSpan.FromSeconds(1));
+            }
+
+            Assert.True(this.events.TryDequeue(out var abandonStop));
+            AssertAbandonStop(abandonStop.eventName, abandonStop.payload, abandonStop.activity, null);
+
+            Assert.True(this.events.TryDequeue(out var exception));
+            AssertException(exception.eventName, exception.payload, exception.activity, null);
+
+            Assert.True(this.events.TryDequeue(out var processStop));
+            AssertProcessStop(processStop.eventName, processStop.payload, processStop.activity, null);
+
+            Assert.Equal(processStop.activity, abandonStop.activity.Parent);
+            Assert.Equal(processStop.activity, exception.activity);
+
+            // message will be processed and compelted again
+            wait = 0;
+            while (wait++ < maxWaitSec && this.events.Count < 2 )
+            {
+                await Task.Delay(TimeSpan.FromSeconds(1));
+            }
+
+            Assert.True(this.events.TryDequeue(out var completeStop));
+            AssertCompleteStop(completeStop.eventName, completeStop.payload, completeStop.activity, null, null);
+
+            Assert.True(this.events.TryDequeue(out processStop));
+            AssertProcessStop(processStop.eventName, processStop.payload, processStop.activity, null);
+
+            Assert.True(this.events.IsEmpty);
+        }
+
+        [Fact]
+        [DisplayTestMethodName]
+        async Task AbandonCompleteFireEvents()
+        {
+            this.queueClient = new QueueClient(TestUtility.NamespaceConnectionString, TestConstants.NonPartitionedQueueName,
+                ReceiveMode.PeekLock);
+
+            await TestUtility.SendMessagesAsync(this.queueClient.InnerSender, 1);
+            var messages = await TestUtility.ReceiveMessagesAsync(this.queueClient.InnerReceiver, 1);
+
+            this.listener.Enable((name, queueName, arg) => name.Contains("Abandon") || name.Contains("Complete"));
+            await TestUtility.AbandonMessagesAsync(this.queueClient.InnerReceiver, messages);
+
+            messages = await TestUtility.ReceiveMessagesAsync(this.queueClient.InnerReceiver, 1);
+
+            await TestUtility.CompleteMessagesAsync(this.queueClient.InnerReceiver, messages);
+
+            Assert.True(this.events.TryDequeue(out var abandonStart));
+            AssertAbandonStart(abandonStart.eventName, abandonStart.payload, abandonStart.activity, null);
+
+            Assert.True(this.events.TryDequeue(out var abandonStop));
+            AssertAbandonStop(abandonStop.eventName, abandonStop.payload, abandonStop.activity,
+                abandonStart.activity);
+
+            Assert.True(this.events.TryDequeue(out var completeStart));
+            AssertCompleteStart(completeStart.eventName, completeStart.payload, completeStart.activity, null);
+
+            Assert.True(this.events.TryDequeue(out var completeStop));
+            AssertCompleteStop(completeStop.eventName, completeStop.payload, completeStop.activity,
+                completeStart.activity, null);
+
+            Assert.True(this.events.IsEmpty);
+        }
+
+        [Fact]
+        [DisplayTestMethodName]
+        async Task BatchSendReceiveFireEvents()
+        {
+            this.queueClient = new QueueClient(TestUtility.NamespaceConnectionString, TestConstants.NonPartitionedQueueName,
+                ReceiveMode.ReceiveAndDelete);
+
+            this.listener.Enable( (name, queuName, arg) => name.Contains("Send") || name.Contains("Receive") );
+            await TestUtility.SendMessagesAsync(this.queueClient.InnerSender, 2);
+            await TestUtility.SendMessagesAsync(this.queueClient.InnerSender, 3);
+            var messages = await TestUtility.ReceiveMessagesAsync(this.queueClient.InnerReceiver, 5);
+
+            Assert.True(this.events.TryDequeue(out var sendStart1));
+            AssertSendStart(sendStart1.eventName, sendStart1.payload, sendStart1.activity, null, 2);
+
+            Assert.True(this.events.TryDequeue(out var sendStop1));
+            AssertSendStop(sendStop1.eventName, sendStop1.payload, sendStop1.activity, sendStop1.activity, 2);
+
+            Assert.True(this.events.TryDequeue(out var sendStart2));
+            AssertSendStart(sendStart2.eventName, sendStart2.payload, sendStart2.activity, null, 3);
+
+            Assert.True(this.events.TryDequeue(out var sendStop2));
+            AssertSendStop(sendStop2.eventName, sendStop2.payload, sendStop2.activity, sendStop2.activity, 3);
+
+            int receivedStopCount = 0;
+            string relatedTo = "";
+            while (this.events.TryDequeue(out var receiveStart))
+            {
+                var startCount = AssertReceiveStart(receiveStart.eventName, receiveStart.payload, receiveStart.activity,
+                    -1);
+
+                Assert.True(this.events.TryDequeue(out var receiveStop));
+                receivedStopCount += AssertReceiveStop(receiveStop.eventName, receiveStop.payload, receiveStop.activity,
+                    receiveStart.activity, null, startCount, -1);
+                relatedTo += receiveStop.activity.Tags.Single(t => t.Key == "RelatedTo").Value;
+            }
+
+            Assert.Equal(5, receivedStopCount);
+            Assert.Contains(sendStart1.activity.Id, relatedTo);
+            Assert.Contains(sendStart2.activity.Id, relatedTo);
+
+            Assert.True(this.events.IsEmpty);
+        }
+
+        [Fact]
+        [DisplayTestMethodName]
+        async Task PeekFireEvents()
+        {
+            this.queueClient = new QueueClient(TestUtility.NamespaceConnectionString, TestConstants.NonPartitionedQueueName, ReceiveMode.PeekLock);
+            this.listener.Enable((name, queuName, arg) => name.Contains("Send") || name.Contains("Peek"));
+
+            await TestUtility.SendMessagesAsync(this.queueClient.InnerSender, 1);
+            await TestUtility.PeekMessageAsync(this.queueClient.InnerReceiver);
+
+            this.listener.Disable();
+
+            var messages = await TestUtility.ReceiveMessagesAsync(this.queueClient.InnerReceiver, 1);
+            await TestUtility.CompleteMessagesAsync(this.queueClient.InnerReceiver, messages);
+
+            Assert.True(this.events.TryDequeue(out var sendStart));
+            AssertSendStart(sendStart.eventName, sendStart.payload, sendStart.activity, null);
+
+            Assert.True(this.events.TryDequeue(out var sendStop));
+            AssertSendStop(sendStop.eventName, sendStop.payload, sendStop.activity, sendStart.activity);
+
+            Assert.True(this.events.TryDequeue(out var peekStart));
+            AssertPeekStart(peekStart.eventName, peekStart.payload, peekStart.activity);
+
+            Assert.True(this.events.TryDequeue(out var peekStop));
+            AssertPeekStop(peekStop.eventName, peekStop.payload, peekStop.activity, peekStart.activity, sendStart.activity);
+
+            Assert.True(this.events.IsEmpty);
+        }
+
+        [Fact]
+        [DisplayTestMethodName]
+        async Task DeadLetterFireEvents()
+        {
+            this.queueClient = new QueueClient(TestUtility.NamespaceConnectionString, TestConstants.NonPartitionedQueueName,
+                ReceiveMode.PeekLock);
+            await TestUtility.SendMessagesAsync(this.queueClient.InnerSender, 1);
+            var messages = await TestUtility.ReceiveMessagesAsync(this.queueClient.InnerReceiver, 1);
+
+            this.listener.Enable((name, queuName, arg) => name.Contains("DeadLetter"));
+            await TestUtility.DeadLetterMessagesAsync(this.queueClient.InnerReceiver, messages);
+            this.listener.Disable();
+
+            QueueClient deadLetterQueueClient = null;
+            try
+            {
+                deadLetterQueueClient = new QueueClient(TestUtility.NamespaceConnectionString,
+                    EntityNameHelper.FormatDeadLetterPath(this.queueClient.QueueName), ReceiveMode.ReceiveAndDelete);
+                await TestUtility.ReceiveMessagesAsync(deadLetterQueueClient.InnerReceiver, 1);
+            }
+            finally
+            {
+                deadLetterQueueClient?.CloseAsync().Wait(TimeSpan.FromSeconds(maxWaitSec));
+            }
+
+            Assert.True(this.events.TryDequeue(out var deadLetterStart));
+            AssertDeadLetterStart(deadLetterStart.eventName, deadLetterStart.payload, deadLetterStart.activity, null);
+
+            Assert.True(this.events.TryDequeue(out var deadLetterStop));
+            AssertDeadLetterStop(deadLetterStop.eventName, deadLetterStop.payload, deadLetterStop.activity,
+                deadLetterStart.activity);
+
+            Assert.True(this.events.IsEmpty);
+        }
+
+        [Fact]
+        [DisplayTestMethodName]
+        async Task RenewLockFireEvents()
+        {
+            this.queueClient = new QueueClient(TestUtility.NamespaceConnectionString, TestConstants.NonPartitionedQueueName,
+                ReceiveMode.PeekLock);
+
+            await TestUtility.SendMessagesAsync(this.queueClient.InnerSender, 1);
+            var messages = await TestUtility.ReceiveMessagesAsync(this.queueClient.InnerReceiver, 1);
+
+            this.listener.Enable((name, queuName, arg) => name.Contains("RenewLock"));
+            await this.queueClient.InnerReceiver.RenewLockAsync(messages[0]);
+            this.listener.Disable();
+
+            await TestUtility.CompleteMessagesAsync(this.queueClient.InnerReceiver, messages);
+
+            Assert.True(this.events.TryDequeue(out var renewStart));
+            AssertRenewLockStart(renewStart.eventName, renewStart.payload, renewStart.activity, null);
+
+            Assert.True(this.events.TryDequeue(out var renewStop));
+            AssertRenewLockStop(renewStop.eventName, renewStop.payload, renewStop.activity, renewStart.activity);
+
+            Assert.True(this.events.IsEmpty);
+        }
+
+        [Fact]
+        [DisplayTestMethodName]
+        async Task DeferReceiveDeferredFireEvents()
+        {
+            this.queueClient = new QueueClient(TestUtility.NamespaceConnectionString, TestConstants.NonPartitionedQueueName,
+                ReceiveMode.PeekLock);
+
+            this.listener.Enable((name, queuName, arg) => name.Contains("Send") || name.Contains("Defer") || name.Contains("Receive"));
+
+            await TestUtility.SendMessagesAsync(this.queueClient.InnerSender, 1);
+            var messages = await TestUtility.ReceiveMessagesAsync(this.queueClient.InnerReceiver, 1);
+            await TestUtility.DeferMessagesAsync(this.queueClient.InnerReceiver, messages);
+            var message = await this.queueClient.InnerReceiver.ReceiveDeferredMessageAsync(messages[0]
+                .SystemProperties
+                .SequenceNumber);
+
+            this.listener.Disable();
+            await TestUtility.CompleteMessagesAsync(this.queueClient.InnerReceiver, new[] {message});
+
+            Assert.True(this.events.TryDequeue(out var sendStart));
+            Assert.True(this.events.TryDequeue(out var sendStop));
+            Assert.True(this.events.TryDequeue(out var receiveStart));
+            Assert.True(this.events.TryDequeue(out var receiveStop));
+
+            Assert.True(this.events.TryDequeue(out var deferStart));
+            AssertDeferStart(deferStart.eventName, deferStart.payload, deferStart.activity, null);
+
+            Assert.True(this.events.TryDequeue(out var deferStop));
+            AssertDeferStop(deferStop.eventName, deferStop.payload, deferStop.activity, deferStart.activity);
+
+            Assert.True(this.events.TryDequeue(out var receiveDeferredStart));
+            AssertReceiveDeferredStart(receiveDeferredStart.eventName, receiveDeferredStart.payload,
+                receiveDeferredStart.activity);
+
+            Assert.True(this.events.TryDequeue(out var receiveDeferredStop));
+            AssertReceiveDeferredStop(receiveDeferredStop.eventName, receiveDeferredStop.payload,
+                receiveDeferredStop.activity, receiveDeferredStart.activity, sendStart.activity);
+
+            Assert.True(this.events.IsEmpty);
+        }
+
+
+        [Fact]
+        [DisplayTestMethodName]
+        async Task SendAndHandlerFilterOutStartEvents()
+        {
+            this.queueClient = new QueueClient(TestUtility.NamespaceConnectionString, TestConstants.NonPartitionedQueueName, ReceiveMode.ReceiveAndDelete);
+
+            ManualResetEvent processingDone = new ManualResetEvent(false);
+            this.listener.Enable((name, queueName, arg) => !name.EndsWith("Start") && !name.Contains("Receive") && !name.Contains("Exception"));
+
+            await TestUtility.SendMessagesAsync(this.queueClient.InnerSender, 1);
+            this.queueClient.RegisterMessageHandler((msg, ct) =>
+                {
+                    processingDone.Set();
+                    return Task.CompletedTask;
+                },
+                exArgs => Task.CompletedTask);
+            processingDone.WaitOne(TimeSpan.FromSeconds(maxWaitSec));
+
+            Assert.True(this.events.TryDequeue(out var sendStop));
+            AssertSendStop(sendStop.eventName, sendStop.payload, sendStop.activity, null);
+
+            int wait = 0;
+            while (wait++ < maxWaitSec && this.events.Count < 1)
+            {
+                await Task.Delay(TimeSpan.FromSeconds(1));
+            }
+
+            Assert.True(this.events.TryDequeue(out var processStop));
+            AssertProcessStop(processStop.eventName, processStop.payload, processStop.activity, null);
+
+            Assert.True(this.events.IsEmpty);
+        }
+
+        [Fact]
+        [DisplayTestMethodName]
+        async Task ScheduleAndCancelFireEvents()
+        {
+            this.queueClient = new QueueClient(TestUtility.NamespaceConnectionString, TestConstants.NonPartitionedQueueName,
+                ReceiveMode.ReceiveAndDelete);
+
+            Activity parentActivity = new Activity("test");
+            this.listener.Enable((name, queuName, arg) => name.Contains("Schedule") || name.Contains("Cancel"));
+
+            parentActivity.Start();
+
+            var sequenceNumber = await this.queueClient.InnerSender.ScheduleMessageAsync(new Message(), DateTimeOffset.UtcNow.AddHours(1));
+            await this.queueClient.InnerSender.CancelScheduledMessageAsync(sequenceNumber);
+
+            Assert.True(this.events.TryDequeue(out var scheduleStart));
+            AssertScheduleStart(scheduleStart.eventName, scheduleStart.payload, scheduleStart.activity,
+                parentActivity);
+
+            Assert.True(this.events.TryDequeue(out var scheduleStop));
+            AssertScheduleStop(scheduleStop.eventName, scheduleStop.payload, scheduleStop.activity,
+                scheduleStart.activity);
+
+            Assert.True(this.events.TryDequeue(out var cancelStart));
+            AssertCancelStart(cancelStart.eventName, cancelStart.payload, cancelStart.activity, parentActivity);
+
+            Assert.True(this.events.TryDequeue(out var cancelStop));
+            AssertCancelStop(cancelStop.eventName, cancelStop.payload, cancelStop.activity, cancelStart.activity);
+
+            Assert.True(this.events.IsEmpty);
+        }
+
+        protected override void Dispose(bool disposing)
+        {
+            if (this.disposed)
+                return;
+
+            if (disposing)
+            {
+                this.queueClient?.CloseAsync().Wait(TimeSpan.FromSeconds(maxWaitSec));
+            }
+
+            this.disposed = true;
+
+            base.Dispose(disposing);
+        }
+    }
+}

test/Microsoft.Azure.ServiceBus.UnitTests/Diagnostics/SessionDiagnosticsTests.cs

@@ -0,0 +1,337 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.Azure.ServiceBus.UnitTests.Diagnostics
+{
+    using System;
+    using System.Diagnostics;
+    using System.Linq;
+    using System.Threading;
+    using System.Threading.Tasks;
+
+    using Microsoft.Azure.ServiceBus.Core;
+    using Xunit;
+
+    public class SessionDiagnosticsTests : DiagnosticsTests
+    {
+        protected override string EntityName => TestConstants.SessionNonPartitionedQueueName;
+        private IMessageSession messageSession;
+        private SessionClient sessionClient;
+        private MessageSender messageSender;
+        private QueueClient queueClient;
+        private bool disposed;
+
+        [Fact]
+        [DisplayTestMethodName]
+        async Task AcceptSetAndGetStateGetFireEvents()
+        {
+            this.messageSender = new MessageSender(TestUtility.NamespaceConnectionString,
+                TestConstants.SessionNonPartitionedQueueName);
+            this.sessionClient = new SessionClient(TestUtility.NamespaceConnectionString,
+                TestConstants.SessionNonPartitionedQueueName, ReceiveMode.ReceiveAndDelete);
+
+            this.listener.Enable();
+
+            var sessionId = Guid.NewGuid().ToString();
+            await this.messageSender.SendAsync(new Message
+            {
+                MessageId = "messageId",
+                SessionId = sessionId
+            });
+            this.messageSession = await this.sessionClient.AcceptMessageSessionAsync(sessionId);
+
+            await this.messageSession.SetStateAsync(new byte[] {1});
+            await this.messageSession.GetStateAsync();
+            await this.messageSession.SetStateAsync(new byte[] { });
+
+            await this.messageSession.ReceiveAsync();
+
+            Assert.True(events.TryDequeue(out var sendStart));
+            AssertSendStart(sendStart.eventName, sendStart.payload, sendStart.activity, null);
+
+            Assert.True(events.TryDequeue(out var sendStop));
+            AssertSendStop(sendStop.eventName, sendStop.payload, sendStop.activity, sendStart.activity);
+
+            Assert.True(events.TryDequeue(out var acceptStart));
+            AssertAcceptMessageSessionStart(acceptStart.eventName, acceptStart.payload, acceptStart.activity);
+
+            Assert.True(events.TryDequeue(out var acceptStop));
+            AssertAcceptMessageSessionStop(acceptStop.eventName, acceptStop.payload, acceptStop.activity,
+                acceptStart.activity);
+
+            Assert.True(events.TryDequeue(out var setStateStart));
+            AssertSetSessionStateStart(setStateStart.eventName, setStateStart.payload, setStateStart.activity);
+
+            Assert.True(events.TryDequeue(out var setStateStop));
+            AssertSetSessionStateStop(setStateStop.eventName, setStateStop.payload, setStateStop.activity,
+                setStateStart.activity);
+
+            Assert.True(events.TryDequeue(out var getStateStart));
+            AssertGetSessionStateStart(getStateStart.eventName, getStateStart.payload, getStateStart.activity);
+
+            Assert.True(events.TryDequeue(out var getStateStop));
+            AssertGetSessionStateStop(getStateStop.eventName, getStateStop.payload, getStateStop.activity,
+                getStateStop.activity);
+
+            Assert.True(events.TryDequeue(out var setStateStart2));
+            Assert.True(events.TryDequeue(out var setStateStop2));
+
+            Assert.True(events.TryDequeue(out var receiveStart));
+            AssertReceiveStart(receiveStart.eventName, receiveStart.payload, receiveStart.activity);
+
+            Assert.True(events.TryDequeue(out var receiveStop));
+            AssertReceiveStop(receiveStop.eventName, receiveStop.payload, receiveStop.activity, receiveStart.activity,
+                sendStart.activity);
+
+            Assert.True(events.IsEmpty);
+        }
+
+        [Fact]
+        [DisplayTestMethodName]
+        async Task EventsAreNotFiredWhenDiagnosticsIsDisabled()
+        {
+            this.messageSender = new MessageSender(TestUtility.NamespaceConnectionString,
+                TestConstants.SessionNonPartitionedQueueName);
+            this.sessionClient = new SessionClient(TestUtility.NamespaceConnectionString,
+                TestConstants.SessionNonPartitionedQueueName, ReceiveMode.ReceiveAndDelete);
+
+            this.listener.Disable();
+
+            var sessionId = Guid.NewGuid().ToString();
+            await this.messageSender.SendAsync(new Message
+            {
+                MessageId = "messageId",
+                SessionId = sessionId
+            });
+            this.messageSession = await sessionClient.AcceptMessageSessionAsync(sessionId);
+
+            await this.messageSession.SetStateAsync(new byte[] {1});
+            await this.messageSession.GetStateAsync();
+            await this.messageSession.SetStateAsync(new byte[] { });
+
+            await this.messageSession.ReceiveAsync();
+
+            Assert.True(events.IsEmpty);
+        }
+
+        [Fact]
+        [DisplayTestMethodName]
+        async Task SessionHandlerFireEvents()
+        {
+            TimeSpan timeout = TimeSpan.FromSeconds(5);
+            this.queueClient = new QueueClient(TestUtility.NamespaceConnectionString,
+                TestConstants.SessionNonPartitionedQueueName, ReceiveMode.ReceiveAndDelete, new NoRetry())
+            {
+                OperationTimeout = timeout
+            };
+
+            this.queueClient.ServiceBusConnection.OperationTimeout = timeout;
+            this.queueClient.SessionClient.OperationTimeout = timeout;
+
+            Stopwatch sw = Stopwatch.StartNew();
+            ManualResetEvent processingDone = new ManualResetEvent(false);
+            this.listener.Enable((name, queueName, arg) => !name.Contains("AcceptMessageSession") &&
+                                                      !name.Contains("Receive") &&
+                                                      !name.Contains("Exception"));
+            var sessionId = Guid.NewGuid().ToString();
+            var message = new Message
+            {
+                MessageId = "messageId",
+                SessionId = sessionId
+            };
+            await this.queueClient.SendAsync(message);
+
+            this.queueClient.RegisterSessionHandler((session, msg, ct) =>
+                {
+                    processingDone.Set();
+                    return Task.CompletedTask;
+                },
+                exArgs => Task.CompletedTask);
+
+            processingDone.WaitOne(TimeSpan.FromSeconds(maxWaitSec));
+
+            Assert.True(events.TryDequeue(out var sendStart));
+            AssertSendStart(sendStart.eventName, sendStart.payload, sendStart.activity, null);
+
+            Assert.True(events.TryDequeue(out var sendStop));
+            AssertSendStop(sendStop.eventName, sendStop.payload, sendStop.activity, sendStart.activity);
+
+            Assert.True(events.TryDequeue(out var processStart));
+            AssertProcessSessionStart(processStart.eventName, processStart.payload, processStart.activity,
+                sendStart.activity);
+
+            int wait = 0;
+            while (wait++ < maxWaitSec && events.Count < 1)
+            {
+                await Task.Delay(TimeSpan.FromSeconds(1));
+            }
+
+            Assert.True(events.TryDequeue(out var processStop));
+            AssertProcessSessionStop(processStop.eventName, processStop.payload, processStop.activity,
+                processStart.activity);
+
+            Assert.True(events.IsEmpty);
+
+            // workaround for https://github.com/Azure/azure-service-bus-dotnet/issues/372:
+            // SessionPumpTaskAsync calls AcceptMessageSessionAsync() without cancellation token.
+            // Even after SessionPump is stopped, this Task may still wait for session during operation timeout
+            // It may interferee with other tests by acception it's sessions and throwing exceptions.
+            // So, let's wait for timeout and a bit more to make sure all created tasks are completed
+            sw.Stop();
+
+            var timeToWait = (timeout - sw.Elapsed).TotalMilliseconds + 1000;
+            if (timeToWait > 0)
+            {
+                await Task.Delay((int)timeToWait);
+            }
+        }
+
+        protected override void Dispose(bool disposing)
+        {
+            if (this.disposed)
+                return;
+
+            if (disposing)
+            {
+                this.queueClient?.SessionPumpHost?.Close();
+                this.queueClient?.SessionClient.CloseAsync().Wait(TimeSpan.FromSeconds(maxWaitSec));
+                this.queueClient?.CloseAsync().Wait(TimeSpan.FromSeconds(maxWaitSec));
+                this.messageSession?.CloseAsync().Wait(TimeSpan.FromSeconds(maxWaitSec));
+                this.sessionClient?.CloseAsync().Wait(TimeSpan.FromSeconds(maxWaitSec));
+                this.messageSender?.CloseAsync().Wait(TimeSpan.FromSeconds(maxWaitSec));
+            }
+
+            this.disposed = true;
+
+            base.Dispose(disposing);
+        }
+
+        protected void AssertAcceptMessageSessionStart(string name, object payload, Activity activity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.AcceptMessageSession.Start", name);
+            this.AssertCommonPayloadProperties(payload);
+
+            var sessionId = this.GetPropertyValueFromAnonymousTypeInstance<string>(payload, "SessionId");
+            
+            Assert.NotNull(activity);
+            Assert.Null(activity.Parent);
+            Assert.Equal(sessionId, activity.Tags.Single(t => t.Key == "SessionId").Value);
+        }
+
+        protected void AssertAcceptMessageSessionStop(string name, object payload, Activity activity, Activity acceptActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.AcceptMessageSession.Stop", name);
+            this.AssertCommonStopPayloadProperties(payload);
+            this.GetPropertyValueFromAnonymousTypeInstance<string>(payload, "SessionId");
+
+            if (acceptActivity != null)
+            {
+                Assert.Equal(acceptActivity, activity);
+            }
+        }
+
+        protected void AssertGetSessionStateStart(string name, object payload, Activity activity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.GetSessionState.Start", name);
+            this.AssertCommonPayloadProperties(payload);
+
+            var sessionId = this.GetPropertyValueFromAnonymousTypeInstance<string>(payload, "SessionId");
+            
+            Assert.NotNull(activity);
+            Assert.Null(activity.Parent);
+            Assert.Equal(sessionId, activity.Tags.Single(t => t.Key == "SessionId").Value);
+        }
+
+        protected void AssertGetSessionStateStop(string name, object payload, Activity activity, Activity getStateActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.GetSessionState.Stop", name);
+            this.AssertCommonStopPayloadProperties(payload);
+            this.GetPropertyValueFromAnonymousTypeInstance<string>(payload, "SessionId");
+            this.GetPropertyValueFromAnonymousTypeInstance<byte[]>(payload, "State");
+
+            if (getStateActivity != null)
+            {
+                Assert.Equal(getStateActivity, activity);
+            }
+        }
+
+        protected void AssertSetSessionStateStart(string name, object payload, Activity activity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.SetSessionState.Start", name);
+            this.AssertCommonPayloadProperties(payload);
+            var sessionId = this.GetPropertyValueFromAnonymousTypeInstance<string>(payload, "SessionId");
+            this.GetPropertyValueFromAnonymousTypeInstance<byte[]>(payload, "State");
+
+            Assert.NotNull(activity);
+            Assert.Null(activity.Parent);
+            Assert.Equal(sessionId, activity.Tags.Single(t => t.Key == "SessionId").Value);
+        }
+
+        protected void AssertSetSessionStateStop(string name, object payload, Activity activity, Activity setStateActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.SetSessionState.Stop", name);
+
+            this.AssertCommonStopPayloadProperties(payload);
+            this.GetPropertyValueFromAnonymousTypeInstance<string>(payload, "SessionId");
+            this.GetPropertyValueFromAnonymousTypeInstance<byte[]>(payload, "State");
+
+            if (setStateActivity != null)
+            {
+                Assert.Equal(setStateActivity, activity);
+            }
+        }
+
+        protected void AssertRenewSessionLockStart(string name, object payload, Activity activity, Activity parentActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.RenewSessionLock.Start", name);
+            this.AssertCommonPayloadProperties(payload);
+            var sessionId= this.GetPropertyValueFromAnonymousTypeInstance<string>(payload, "SessionId");
+
+            Assert.NotNull(activity);
+            Assert.Null(activity.Parent);
+            Assert.Equal(sessionId, activity.Tags.Single(t => t.Key == "SessionId").Value);
+        }
+
+        protected void AssertRenewSessionLockStop(string name, object payload, Activity activity, Activity renewActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.RenewSessionLock.Stop", name);
+
+            this.AssertCommonStopPayloadProperties(payload);
+            this.GetPropertyValueFromAnonymousTypeInstance<string>(payload, "SessionId");
+
+            if (renewActivity != null)
+            {
+                Assert.Equal(renewActivity, activity);
+            }
+        }
+
+        protected void AssertProcessSessionStart(string name, object payload, Activity activity, Activity sendActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.ProcessSession.Start", name);
+            AssertCommonPayloadProperties(payload);
+
+            GetPropertyValueFromAnonymousTypeInstance<IMessageSession>(payload, "Session");
+            var message = GetPropertyValueFromAnonymousTypeInstance<Message>(payload, "Message");
+
+            Assert.NotNull(activity);
+            Assert.Null(activity.Parent);
+            Assert.Equal(sendActivity.Id, activity.ParentId);
+            Assert.Equal(sendActivity.Baggage.OrderBy(p => p.Key), activity.Baggage.OrderBy(p => p.Key));
+
+            AssertTags(message, activity);
+        }
+
+        protected void AssertProcessSessionStop(string name, object payload, Activity activity, Activity processActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.ProcessSession.Stop", name);
+            AssertCommonStopPayloadProperties(payload);
+            GetPropertyValueFromAnonymousTypeInstance<IMessageSession>(payload, "Session");
+            var message = GetPropertyValueFromAnonymousTypeInstance<Message>(payload, "Message");
+
+            if (processActivity != null)
+            {
+                Assert.Equal(processActivity, activity);
+            }
+        }
+    }
+}

test/Microsoft.Azure.ServiceBus.UnitTests/Diagnostics/SubscriptionClientDiagnosticsTests.cs

@@ -0,0 +1,155 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.Azure.ServiceBus.UnitTests.Diagnostics
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Diagnostics;
+    using System.Threading.Tasks;
+    using Xunit;
+
+    public class SubscriptionClientDiagnosticsTests : DiagnosticsTests
+    {
+        protected override string EntityName => $"{TestConstants.NonPartitionedTopicName}/Subscriptions/{TestConstants.SubscriptionName}";
+        private SubscriptionClient subscriptionClient;
+        private bool disposed;
+
+        [Fact]
+        [DisplayTestMethodName]
+        async Task AddRemoveGetFireEvents()
+        {
+            this.subscriptionClient = new SubscriptionClient(
+                TestUtility.NamespaceConnectionString,
+                TestConstants.NonPartitionedTopicName,
+                TestConstants.SubscriptionName,
+                ReceiveMode.ReceiveAndDelete);
+
+            this.listener.Enable((name, queueName, id) => name.Contains("Rule"));
+
+            var ruleName = Guid.NewGuid().ToString();
+            await this.subscriptionClient.AddRuleAsync(ruleName, new TrueFilter());
+            await this.subscriptionClient.GetRulesAsync();
+            await this.subscriptionClient.RemoveRuleAsync(ruleName);
+
+            Assert.True(this.events.TryDequeue(out var addRuleStart));
+            AssertAddRuleStart(addRuleStart.eventName, addRuleStart.payload, addRuleStart.activity);
+
+            Assert.True(this.events.TryDequeue(out var addRuleStop));
+            AssertAddRuleStop(addRuleStop.eventName, addRuleStop.payload, addRuleStop.activity, addRuleStart.activity);
+
+            Assert.True(this.events.TryDequeue(out var getRulesStart));
+            AssertGetRulesStart(getRulesStart.eventName, getRulesStart.payload, getRulesStart.activity);
+
+            Assert.True(this.events.TryDequeue(out var getRulesStop));
+            AssertGetRulesStop(getRulesStop.eventName, getRulesStop.payload, getRulesStop.activity, getRulesStart.activity);
+
+            Assert.True(this.events.TryDequeue(out var removeRuleStart));
+            AssertRemoveRuleStart(removeRuleStart.eventName, removeRuleStart.payload, removeRuleStart.activity);
+
+            Assert.True(this.events.TryDequeue(out var removeRuleStop));
+            AssertRemoveRuleStop(removeRuleStop.eventName, removeRuleStop.payload, removeRuleStop.activity, removeRuleStart.activity);
+
+            Assert.True(this.events.IsEmpty);
+        }
+
+        [Fact]
+        [DisplayTestMethodName]
+        async Task EventsAreNotFiredWhenDiagnosticsIsDisabled()
+        {
+            this.subscriptionClient = new SubscriptionClient(
+                TestUtility.NamespaceConnectionString,
+                TestConstants.NonPartitionedTopicName,
+                TestConstants.SubscriptionName,
+                ReceiveMode.ReceiveAndDelete);
+
+            this.listener.Disable();
+
+            var ruleName = Guid.NewGuid().ToString();
+            await this.subscriptionClient.AddRuleAsync(ruleName, new TrueFilter());
+            await this.subscriptionClient.GetRulesAsync();
+            await this.subscriptionClient.RemoveRuleAsync(ruleName);
+
+            Assert.True(this.events.IsEmpty);
+        }
+
+        protected override void Dispose(bool disposing)
+        {
+            if (this.disposed)
+                return;
+
+            if (disposing)
+            {
+                this.subscriptionClient?.CloseAsync().Wait(TimeSpan.FromSeconds(maxWaitSec));
+            }
+
+            this.disposed = true;
+
+            base.Dispose(disposing);
+        }
+
+        protected void AssertAddRuleStart(string name, object payload, Activity activity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.AddRule.Start", name);
+            AssertCommonPayloadProperties(payload);
+            GetPropertyValueFromAnonymousTypeInstance<RuleDescription>(payload, "Rule");
+            Assert.NotNull(activity);
+            Assert.Null(activity.Parent);
+        }
+
+        protected void AssertAddRuleStop(string name, object payload, Activity activity, Activity addRuleActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.AddRule.Stop", name);
+            AssertCommonStopPayloadProperties(payload);
+            GetPropertyValueFromAnonymousTypeInstance<RuleDescription>(payload, "Rule");
+
+            if (addRuleActivity != null)
+            {
+                Assert.Equal(addRuleActivity, activity);
+            }
+        }
+
+        protected void AssertGetRulesStart(string name, object payload, Activity activity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.GetRules.Start", name);
+            AssertCommonPayloadProperties(payload);
+            Assert.NotNull(activity);
+            Assert.Null(activity.Parent);
+        }
+
+        protected void AssertGetRulesStop(string name, object payload, Activity activity, Activity getRulesActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.GetRules.Stop", name);
+            
+            AssertCommonStopPayloadProperties(payload);
+            GetPropertyValueFromAnonymousTypeInstance<IEnumerable<RuleDescription>>(payload, "Rules");
+
+            if (getRulesActivity != null)
+            {
+                Assert.Equal(getRulesActivity, activity);
+            }
+        }
+
+        protected void AssertRemoveRuleStart(string name, object payload, Activity activity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.RemoveRule.Start", name);
+            AssertCommonPayloadProperties(payload);
+            GetPropertyValueFromAnonymousTypeInstance<string>(payload, "RuleName");
+            Assert.NotNull(activity);
+            Assert.Null(activity.Parent);
+        }
+
+        protected void AssertRemoveRuleStop(string name, object payload, Activity activity, Activity removeRuleActivity)
+        {
+            Assert.Equal("Microsoft.Azure.ServiceBus.RemoveRule.Stop", name);
+
+            AssertCommonStopPayloadProperties(payload);
+            GetPropertyValueFromAnonymousTypeInstance<string>(payload, "RuleName");
+
+            if (removeRuleActivity != null)
+            {
+                Assert.Equal(removeRuleActivity, activity);
+            }
+        }
+    }
+}

test/Microsoft.Azure.ServiceBus.UnitTests/Microsoft.Azure.ServiceBus.UnitTests.csproj

@@ -17,11 +17,31 @@
     <ProjectReference Include="..\..\src\Microsoft.Azure.ServiceBus\Microsoft.Azure.ServiceBus.csproj" />
   </ItemGroup>
 
+  <ItemGroup Condition="'$(TargetFramework)' == 'net461'">
+    <compile Remove="..\..\test\Microsoft.Azure.ServiceBus.UnitTests\ExpectedMessagingExceptionTests.cs" />
+    <compile Remove="..\..\test\Microsoft.Azure.ServiceBus.UnitTests\OnMessageQueueTests.cs" />
+    <compile Remove="..\..\test\Microsoft.Azure.ServiceBus.UnitTests\OnSessionQueueTests.cs" />
+    <compile Remove="..\..\test\Microsoft.Azure.ServiceBus.UnitTests\RetryTests.cs" />
+    <compile Remove="..\..\test\Microsoft.Azure.ServiceBus.UnitTests\Diagnostics\ExtractActivityTests.cs" />
+    <compile Remove="..\..\test\Microsoft.Azure.ServiceBus.UnitTests\Diagnostics\QueueClientDiagnosticsTests.cs" />
+    <compile Remove="..\..\test\Microsoft.Azure.ServiceBus.UnitTests\Diagnostics\SubscriptionClientDiagnosticsTests.cs" />
+    <compile Remove="..\..\test\Microsoft.Azure.ServiceBus.UnitTests\Diagnostics\SessionDiagnosticsTests.cs" />
+  </ItemGroup>
+
+  <ItemGroup Condition="'$(TargetFramework)' == 'netcoreapp2.0'">
+    <compile Remove="..\..\test\Microsoft.Azure.ServiceBus.UnitTests\OnMessageTopicSubscriptionTests.cs" />
+    <compile Remove="..\..\test\Microsoft.Azure.ServiceBus.UnitTests\OnSessionTopicSubscriptionTests.cs" />    
+  </ItemGroup>
+
   <ItemGroup>
-    <PackageReference Include="Microsoft.NET.Test.Sdk" Version="15.0.0" />
+    <PackageReference Include="Microsoft.NET.Test.Sdk" Version="15.3.0" />
     <PackageReference Include="System.Net.WebSockets.Client" Version="4.3.1" />
+    <PackageReference Include="System.ValueTuple" Version="4.4.0" />
     <PackageReference Include="xunit.runner.visualstudio" Version="2.2.0" />
     <PackageReference Include="xunit" Version="2.2.0" />
+    <PackageReference Include="PublicApiGenerator" Version="6.5.0" />
+    <PackageReference Include="ApprovalTests" Version="3.0.13" NoWarn="NU1701" />
+    <PackageReference Include="ApprovalUtilities" Version="3.0.13" NoWarn="NU1701" />
     <PackageReference Include="Microsoft.Extensions.PlatformAbstractions" Version="1.1.0" />
   </ItemGroup>
 
@@ -30,24 +50,8 @@
   </ItemGroup>
 
   <ItemGroup Condition="'$(TargetFramework)' == 'net461'">
-    <PackageReference Include="ApiApprover">
-      <Version>6.0.0</Version>
-    </PackageReference>
-    <PackageReference Include="ApprovalTests">
-      <Version>3.0.13</Version>
-    </PackageReference>
-    <PackageReference Include="ApprovalUtilities">
-      <Version>3.0.13</Version>
-    </PackageReference>
-    <PackageReference Include="Mono.Cecil">
-      <Version>0.9.6.4</Version>
-    </PackageReference>
-    <PackageReference Include="PublicApiGenerator">
-      <Version>6.0.0</Version>
-    </PackageReference>
-    <PackageReference Include="WindowsAzure.ServiceBus">
-      <Version>4.1.2</Version>
-    </PackageReference>
+    <PackageReference Include="WindowsAzure.ServiceBus" Version="4.1.2" />
   </ItemGroup>
+  
 
-</Project>
+</Project>

test/Microsoft.Azure.ServiceBus.UnitTests/RetryTests.cs

@@ -124,7 +124,7 @@ namespace Microsoft.Azure.ServiceBus.UnitTests
             Assert.True(retryExponential.IsServerBusy, "policy1.IsServerBusy is not true");
         }
 
-        [Fact]
+        [Fact(Skip = "Flaky Test in Appveyor, fix and enable back")]
         async void RunOperationShouldReturnImmediatelyIfRetryIntervalIsGreaterThanOperationTimeout()
         {
             var policy = RetryPolicy.Default;

test/Microsoft.Azure.ServiceBus.UnitTests/SenderReceiverClientTestBase.cs

@@ -166,6 +166,7 @@ namespace Microsoft.Azure.ServiceBus.UnitTests
         internal async Task ReceiveShouldReturnNoLaterThanServerWaitTimeTestCase(IMessageSender messageSender, IMessageReceiver messageReceiver, int messageCount)
         {
             var timer = Stopwatch.StartNew();
+            TestUtility.Log($"starting Receive");
             var message = await messageReceiver.ReceiveAsync(TimeSpan.FromSeconds(2));
             timer.Stop();
 
@@ -175,6 +176,7 @@ namespace Microsoft.Azure.ServiceBus.UnitTests
             // Ensuring total time taken is less than 60 seconds, which is the default timeout for receive.
             // Keeping the value of 40 to avoid flakiness in test infrastructure which may lead to extended time taken.
             // Todo: Change this value to a lower number once test infra is performant.
+            TestUtility.Log($"Elapsed Milliseconds: {timer.Elapsed.TotalMilliseconds}");
             Assert.True(timer.Elapsed.TotalSeconds < 40);
         }
 

test/Microsoft.Azure.ServiceBus.UnitTests/SenderReceiverTests.cs

@@ -6,6 +6,7 @@ namespace Microsoft.Azure.ServiceBus.UnitTests
     using System;
     using System.Text;
     using System.Collections.Generic;
+    using System.Threading;
     using System.Threading.Tasks;
     using Core;
     using Xunit;
@@ -184,7 +185,7 @@ namespace Microsoft.Azure.ServiceBus.UnitTests
             }
         }
 
-        [Fact]
+        [Fact(Skip = "Flaky Test in Appveyor, fix and enable back")]
         [DisplayTestMethodName]
         public async Task WaitingReceiveShouldReturnImmediatelyWhenReceiverIsClosed()
         {
@@ -215,8 +216,24 @@ namespace Microsoft.Azure.ServiceBus.UnitTests
             }
 
             TestUtility.Log("Waiting for maximum 10 Secs");
-            var waitingTask = Task.Delay(10000);
-            Assert.Equal(quickTask, await Task.WhenAny(quickTask, waitingTask));
+            bool receiverReturnedInTime = false;
+            using (var timeoutCancellationTokenSource = new CancellationTokenSource())
+            {
+
+                var completedTask = await Task.WhenAny(quickTask, Task.Delay(10000, timeoutCancellationTokenSource.Token));
+                if (completedTask == quickTask)
+                {
+                    timeoutCancellationTokenSource.Cancel();
+                    receiverReturnedInTime = true;
+                    TestUtility.Log("The Receiver closed in time.");
+                }
+                else
+                {
+                    TestUtility.Log("The Receiver did not close in time.");
+                }
+            }
+
+            Assert.True(receiverReturnedInTime);
         }
 
         [Fact]

test/Microsoft.Azure.ServiceBus.UnitTests/TokenProviderTests.cs

@@ -0,0 +1,44 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.Azure.ServiceBus.UnitTests
+{
+    using System;
+    using System.Threading.Tasks;
+    using Microsoft.Azure.ServiceBus.Primitives;
+    using Microsoft.IdentityModel.Clients.ActiveDirectory;
+    using Xunit;
+
+    public class TokenProviderTests : SenderReceiverClientTestBase
+    {
+        /// <summary>
+        /// This test is for manual only purpose. Fill in the tenant-id, app-id and app-secret before running.
+        /// </summary>
+        /// <returns></returns>
+        [Fact]
+        [DisplayTestMethodName]
+        async Task UseITokenProviderWithAad()
+        {
+            var tenantId = "";
+            var aadAppId = "";
+            var aadAppSecret = "";
+
+            if (string.IsNullOrEmpty(tenantId))
+            {
+                TestUtility.Log($"Skipping test during scheduled runs.");
+                return;
+            }
+
+            var authContext = new AuthenticationContext($"https://login.windows.net/{tenantId}");
+            var cc = new ClientCredential(aadAppId, aadAppSecret);
+            var tokenProvider = TokenProvider.CreateAadTokenProvider(authContext, cc);
+
+            // Create new client with updated connection string.
+            var csb = new ServiceBusConnectionStringBuilder(TestUtility.NamespaceConnectionString);
+            var queueClient = new QueueClient(csb.Endpoint, csb.EntityPath, tokenProvider);
+
+            // Send and receive messages.
+            await this.PeekLockTestCase(queueClient.InnerSender, queueClient.InnerReceiver, 10);
+        }
+    }
+}