Add OTel support for Service Bus (Azure#34492)

* Add OTel support for Service Bus

* Update extension

* ref docs

* Don't include new shared source in Azure.Core.csproj

* Fix ref

* Fix tests

* PR fb

* Add batch count test cases

* Core updates to support setting tracestate

* PR fb

* PR fb

* Use destination.name for everything other than receive

* receive or process

Author: JoshLove-msft

sdk/core/Azure.Core.TestFramework/src/TestActivitySourceListener.cs

@@ -5,6 +5,7 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Linq;
+using NUnit.Framework;
 
 namespace Azure.Core.Tests
 {
@@ -16,11 +17,11 @@ public class TestActivitySourceListener: IDisposable
         public Queue<Activity> Activities { get; } =
             new Queue<Activity>();
 
-        public TestActivitySourceListener(string name) : this(source => source.Name == name)
+        public TestActivitySourceListener(string name, Action<Activity> activityStartedCallback = default) : this(source => source.Name.StartsWith(name), activityStartedCallback)
         {
         }
 
-        public TestActivitySourceListener(Func<ActivitySource, bool> sourceSelector)
+        public TestActivitySourceListener(Func<ActivitySource, bool> sourceSelector, Action<Activity> activityStartedCallback = default)
         {
             _listener = new ActivityListener
             {
@@ -31,6 +32,7 @@ public TestActivitySourceListener(Func<ActivitySource, bool> sourceSelector)
                     {
                         Activities.Enqueue(activity);
                     }
+                    activityStartedCallback?.Invoke(activity);
                 },
                 Sample = (ref ActivityCreationOptions<ActivityContext> options) =>
                 {
@@ -46,6 +48,13 @@ public TestActivitySourceListener(Func<ActivitySource, bool> sourceSelector)
             ActivitySource.AddActivityListener(_listener);
         }
 
+        public Activity AssertAndRemoveActivity(string name)
+        {
+            var activity = Activities.Dequeue();
+            Assert.AreEqual(name, activity.OperationName);
+            return activity;
+        }
+
         public void Dispose()
         {
             _listener?.Dispose();

sdk/core/Azure.Core/src/Shared/DiagnosticScope.cs

@@ -107,7 +107,7 @@ public void AddAttribute<T>(string name, T value, Func<T, string> format)
         /// <param name="traceparent">The traceparent for the link.</param>
         /// <param name="tracestate">The tracestate for the link.</param>
         /// <param name="attributes">Optional attributes to associate with the link.</param>
-        public void AddLink(string traceparent, string tracestate, IDictionary<string, string>? attributes = null)
+        public void AddLink(string traceparent, string? tracestate, IDictionary<string, string>? attributes = null)
         {
             _activityAdapter?.AddLink(traceparent, tracestate, attributes);
         }
@@ -124,12 +124,13 @@ public void SetStartTime(DateTime dateTime)
         }
 
         /// <summary>
-        /// Sets the trace parent for the current scope.
+        /// Sets the trace context for the current scope.
         /// </summary>
         /// <param name="traceparent">The trace parent to set for the current scope.</param>
-        public void SetTraceparent(string traceparent)
+        /// <param name="tracestate">The trace state to set for the current scope.</param>
+        public void SetTraceContext(string traceparent, string? tracestate = default)
         {
-            _activityAdapter?.SetTraceparent(traceparent);
+            _activityAdapter?.SetTraceContext(traceparent, tracestate);
         }
 
         public void Dispose()
@@ -199,6 +200,7 @@ private class ActivityAdapter : IDisposable
             private DateTimeOffset _startTime;
             private List<Activity>? _links;
             private string? _traceparent;
+            private string? _tracestate;
 
             public ActivityAdapter(object? activitySource, DiagnosticSource diagnosticSource, string activityName, ActivityKind kind, object? diagnosticSourceArgs)
             {
@@ -258,7 +260,7 @@ public void AddTag(string name, object value)
                 return linkCollection;
             }
 
-            public void AddLink(string traceparent, string tracestate, IDictionary<string, string>? attributes)
+            public void AddLink(string traceparent, string? tracestate, IDictionary<string, string>? attributes)
             {
                 var linkedActivity = new Activity("LinkedActivity");
                 linkedActivity.SetW3CFormat();
@@ -334,6 +336,11 @@ public void AddLink(string traceparent, string tracestate, IDictionary<string, s
                         _currentActivity.SetParentId(_traceparent);
                     }
 
+                    if (_tracestate != null)
+                    {
+                        _currentActivity.SetTraceState(_tracestate);
+                    }
+
                     _currentActivity.Start();
                 }
 
@@ -351,7 +358,8 @@ public void AddLink(string traceparent, string tracestate, IDictionary<string, s
                     startTime: _startTime,
                     tags: _tagCollection,
                     links: GetActivitySourceLinkCollection(),
-                    parentId: _traceparent);
+                    traceparent: _traceparent,
+                    tracestate: _tracestate);
             }
 
             public void SetStartTime(DateTime startTime)
@@ -365,13 +373,14 @@ public void MarkFailed(Exception exception)
                 _diagnosticSource?.Write(_activityName + ".Exception", exception);
             }
 
-            public void SetTraceparent(string traceparent)
+            public void SetTraceContext(string traceparent, string? tracestate)
             {
                 if (_currentActivity != null)
                 {
                     throw new InvalidOperationException("Traceparent can not be set after the activity is started.");
                 }
                 _traceparent = traceparent;
+                _tracestate = tracestate;
             }
 
             public void Dispose()
@@ -418,13 +427,14 @@ static ActivityExtensions()
         private static Action<Activity, string?>? SetTraceStateStringMethod;
         private static Func<Activity, int>? GetIdFormatMethod;
         private static Action<Activity, string, object?>? ActivityAddTagMethod;
-        private static Func<object, string, int, string?, ICollection<KeyValuePair<string, object>>?, IList?, DateTimeOffset, Activity?>? ActivitySourceStartActivityMethod;
+        private static Func<object, string, int, object?, ICollection<KeyValuePair<string, object>>?, IList?, DateTimeOffset, Activity?>? ActivitySourceStartActivityMethod;
         private static Func<object, bool>? ActivitySourceHasListenersMethod;
         private static Func<string, string?, ICollection<KeyValuePair<string, object>>?, object?>? CreateActivityLinkMethod;
         private static Func<ICollection<KeyValuePair<string,object>>?>? CreateTagsCollectionMethod;
         private static Func<Activity, string, object?>? GetCustomPropertyMethod;
         private static Action<Activity, string, object>? SetCustomPropertyMethod;
         private static readonly ParameterExpression ActivityParameter = Expression.Parameter(typeof(Activity));
+        private static MethodInfo? ParseActivityContextMethod;
 
         public static object? GetCustomProperty(this Activity activity, string propertyName)
         {
@@ -687,13 +697,23 @@ public static bool ActivitySourceHasListeners(object? activitySource)
             return ActivitySourceHasListenersMethod.Invoke(activitySource);
         }
 
-        public static Activity? ActivitySourceStartActivity(object? activitySource, string activityName, int kind, DateTimeOffset startTime, ICollection<KeyValuePair<string, object>>? tags, IList? links, string? parentId)
+        public static Activity? ActivitySourceStartActivity(
+            object? activitySource,
+            string activityName,
+            int kind,
+            DateTimeOffset startTime,
+            ICollection<KeyValuePair<string, object>>? tags,
+            IList? links,
+            string? traceparent,
+            string? tracestate)
         {
             if (activitySource == null)
             {
                 return null;
             }
 
+            object? activityContext = default;
+
             if (ActivitySourceStartActivityMethod == null)
             {
                 if (ActivityLinkType == null ||
@@ -709,7 +729,7 @@ public static bool ActivitySourceHasListeners(object? activitySource)
                     {
                         typeof(string),
                         ActivityKindType,
-                        typeof(string),
+                        ActivityContextType,
                         typeof(IEnumerable<KeyValuePair<string, object>>),
                         typeof(IEnumerable<>).MakeGenericType(ActivityLinkType),
                         typeof(DateTimeOffset)
@@ -724,27 +744,38 @@ public static bool ActivitySourceHasListeners(object? activitySource)
                         var sourceParameter = Expression.Parameter(typeof(object));
                         var nameParameter = Expression.Parameter(typeof(string));
                         var kindParameter = Expression.Parameter(typeof(int));
-                        var parentIdParameter = Expression.Parameter(typeof(string));
+                        var contextParameter = Expression.Parameter(typeof(object));
                         var startTimeParameter = Expression.Parameter(typeof(DateTimeOffset));
                         var tagsParameter = Expression.Parameter(typeof(ICollection<KeyValuePair<string, object>>));
                         var linksParameter = Expression.Parameter(typeof(IList));
                         var methodParameter = method.GetParameters();
-                        ActivitySourceStartActivityMethod = Expression.Lambda<Func<object, string, int, string?, ICollection<KeyValuePair<string, object>>?, IList?, DateTimeOffset, Activity?>>(
+                        ParseActivityContextMethod = ActivityContextType.GetMethod("Parse", BindingFlags.Static | BindingFlags.Public);
+
+                        ActivitySourceStartActivityMethod = Expression.Lambda<Func<object, string, int, object?, ICollection<KeyValuePair<string, object>>?, IList?, DateTimeOffset, Activity?>>(
                             Expression.Call(
                                 Expression.Convert(sourceParameter, method.DeclaringType!),
                                 method,
                                 nameParameter,
                                 Expression.Convert(kindParameter,  methodParameter[1].ParameterType),
-                                Expression.Convert(parentIdParameter, methodParameter[2].ParameterType),
+                                Expression.Convert(contextParameter, methodParameter[2].ParameterType),
                                 Expression.Convert(tagsParameter,  methodParameter[3].ParameterType),
                                 Expression.Convert(linksParameter,  methodParameter[4].ParameterType),
                                 Expression.Convert(startTimeParameter,  methodParameter[5].ParameterType)),
-                            sourceParameter, nameParameter, kindParameter, parentIdParameter, tagsParameter, linksParameter, startTimeParameter).Compile();
+                            sourceParameter, nameParameter, kindParameter, contextParameter, tagsParameter, linksParameter, startTimeParameter).Compile();
                     }
                 }
             }
 
-            return ActivitySourceStartActivityMethod.Invoke(activitySource, activityName, kind, parentId, tags, links, startTime);
+            if (ActivityContextType != null && ParseActivityContextMethod != null)
+            {
+                if (traceparent != null)
+                    activityContext = ParseActivityContextMethod.Invoke(null, new[] {traceparent, tracestate})!;
+                else
+                    // because ActivityContext is a struct, we need to create a default instance rather than allowing the argument to be null
+                    activityContext = Activator.CreateInstance(ActivityContextType);
+            }
+
+            return ActivitySourceStartActivityMethod.Invoke(activitySource, activityName, kind, activityContext, tags, links, startTime);
         }
 
         public static object? CreateActivitySource(string name)

sdk/core/Azure.Core/src/Shared/MessagingClientDiagnostics.cs

@@ -0,0 +1,171 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System.Collections.Generic;
+using System.Diagnostics;
+using Azure.Core.Pipeline;
+
+#nullable enable
+
+namespace Azure.Core.Shared
+{
+    /// <summary>
+    /// Client Diagnostics support for messaging clients. Currently, this is only used for AMQP clients.
+    /// HTTP libraries should use the ClientDiagnostics type instead.
+    /// </summary>
+    internal class MessagingClientDiagnostics
+    {
+        private readonly string _fullyQualifiedNamespace;
+        private readonly string _entityPath;
+        private readonly string _messagingSystem;
+        private readonly DiagnosticScopeFactory _scopeFactory;
+
+        #region OTel-specific messaging attributes
+        public const string MessagingSystem = "messaging.system";
+        public const string DestinationName = "messaging.destination.name";
+        public const string SourceName = "messaging.source.name";
+        public const string MessagingOperation = "messaging.operation";
+        public const string NetPeerName = "net.peer.name";
+        public const string BatchCount = "messaging.batch.message_count";
+        public const string TraceParent = "traceparent";
+        public const string TraceState = "tracestate";
+        #endregion
+
+        #region legacy compat attributes
+        public const string MessageBusDestination = "message_bus.destination";
+        public const string PeerAddress = "peer.address";
+        public const string Component = "component";
+        #endregion
+
+        public const string DiagnosticIdAttribute = "Diagnostic-Id";
+
+        public MessagingClientDiagnostics(string clientNamespace, string? resourceProviderNamespace, string messagingSystem, string fullyQualifiedNamespace, string entityPath)
+        {
+            Argument.AssertNotNull(messagingSystem, nameof(messagingSystem));
+            Argument.AssertNotNull(fullyQualifiedNamespace, nameof(fullyQualifiedNamespace));
+            Argument.AssertNotNull(entityPath, nameof(entityPath));
+
+            _messagingSystem = messagingSystem;
+            _fullyQualifiedNamespace = fullyQualifiedNamespace;
+            _entityPath = entityPath;
+            _scopeFactory = new DiagnosticScopeFactory(clientNamespace, resourceProviderNamespace, true, false);
+        }
+
+        public DiagnosticScope CreateScope(
+            string activityName,
+            DiagnosticScope.ActivityKind kind,
+            MessagingDiagnosticOperation operation = default)
+        {
+            DiagnosticScope scope = _scopeFactory.CreateScope(activityName, kind);
+            if (ActivityExtensions.SupportsActivitySource())
+            {
+                scope.AddAttribute(MessagingSystem, _messagingSystem);
+                if (operation != default)
+                {
+                    scope.AddAttribute(MessagingOperation, operation.ToString());
+                }
+
+                scope.AddAttribute(NetPeerName, _fullyQualifiedNamespace);
+                scope.AddAttribute(operation == MessagingDiagnosticOperation.Receive || operation == MessagingDiagnosticOperation.Process ? SourceName : DestinationName, _entityPath);
+            }
+            else
+            {
+                scope.AddAttribute(Component, _messagingSystem);
+                scope.AddAttribute(PeerAddress, _fullyQualifiedNamespace);
+                scope.AddAttribute(MessageBusDestination, _entityPath);
+            }
+
+            return scope;
+        }
+
+        /// <summary>
+        ///   Attempts to extract the trace context from a message's properties.
+        /// </summary>
+        ///
+        /// <param name="properties">The properties holding the trace context.</param>
+        /// <param name="traceparent">The trace parent of the message.</param>
+        /// <param name="tracestate">The trace state of the message.</param>
+        /// <returns><c>true</c> if the message properties contained the diagnostic id; otherwise, <c>false</c>.</returns>
+        public static bool TryExtractTraceContext(IReadOnlyDictionary<string, object> properties, out string? traceparent, out string? tracestate)
+        {
+            traceparent = null;
+            tracestate = null;
+
+            if (properties.TryGetValue(TraceParent, out var traceParent) && traceParent is string traceParentString)
+            {
+                traceparent = traceParentString;
+                if (properties.TryGetValue(TraceState, out object state) && state is string stateString)
+                {
+                    tracestate = stateString;
+                }
+                return true;
+            }
+
+            // trace state is not valid without trace parent, so we don't need to check for it for the legacy attribute
+            if (properties.TryGetValue(DiagnosticIdAttribute, out var diagnosticId) && diagnosticId is string diagnosticIdString)
+            {
+                traceparent = diagnosticIdString;
+                return true;
+            }
+            return false;
+        }
+
+        /// <summary>
+        ///   Attempts to extract the trace context from a message's properties.
+        /// </summary>
+        ///
+        /// <param name="properties">The properties holding the trace context.</param>
+        /// <param name="traceparent">The trace parent of the message.</param>
+        /// <param name="tracestate">The trace state of the message.</param>
+        /// <returns><c>true</c> if the message properties contained the diagnostic id; otherwise, <c>false</c>.</returns>
+        public static bool TryExtractTraceContext(IDictionary<string, object> properties, out string? traceparent, out string? tracestate)
+        {
+            traceparent = null;
+            tracestate = null;
+
+            if (properties.TryGetValue(TraceParent, out var traceParent) && traceParent is string traceParentString)
+            {
+                traceparent = traceParentString;
+                if (properties.TryGetValue(TraceState, out object state) && state is string stateString)
+                {
+                    tracestate = stateString;
+                }
+                return true;
+            }
+
+            // trace state is not valid without trace parent, so we don't need to check for it for the legacy attribute
+            if (properties.TryGetValue(DiagnosticIdAttribute, out var diagnosticId) && diagnosticId is string diagnosticIdString)
+            {
+                traceparent = diagnosticIdString;
+                return true;
+            }
+            return false;
+        }
+
+        /// <summary>
+        /// Instrument the message properties for tracing. If tracing is enabled, a diagnostic id will be added to the message properties,
+        /// which alters the message size.
+        /// </summary>
+        /// <param name="properties">The dictionary of application message properties.</param>
+        /// <param name="activityName">The activity name to use for the diagnostic scope.</param>
+        public void InstrumentMessage(IDictionary<string, object> properties, string activityName)
+        {
+            if (!properties.ContainsKey(DiagnosticIdAttribute) && !properties.ContainsKey(TraceParent))
+            {
+                using DiagnosticScope messageScope = CreateScope(
+                    activityName,
+                    DiagnosticScope.ActivityKind.Producer);
+                messageScope.Start();
+
+                Activity activity = Activity.Current;
+                if (activity != null)
+                {
+                    properties[DiagnosticIdAttribute] = activity.Id;
+                    properties[TraceParent] = activity.Id;
+                    if (activity.TraceStateString != null)
+                        properties[TraceState] = activity.TraceStateString;
+                }
+            }
+        }
+    }
+}