[Tracer] Add Azure Service Bus batching support (#7553)

## Summary of changes

Follow-up to #7413

Add support for ServiceBusMessageBatch messages for Service Bus.

## Reason for change

Service Bus instrumentation completeness.

## Implementation details

## Test coverage

Integration tests.

## Other details
<!-- Fixes #{issue} -->


<!--  ⚠️ Note:

Where possible, please obtain 2 approvals prior to merging. Unless
CODEOWNERS specifies otherwise, for external teams it is typically best
to have one review from a team member, and one review from apm-dotnet.
Trivial changes do not require 2 reviews.

MergeQueue is NOT enabled in this repository. If you have write access
to the repo, the PR has 1-2 approvals (see above), and all of the
required checks have passed, you can use the Squash and Merge button to
merge the PR. If you don't have write access, or you need help, reach
out in the #apm-dotnet channel in Slack.
-->

Author: pablomartinezbernardo

tracer/build/supported_calltargets.g.json

@@ -3762,6 +3762,31 @@
     "IsAdoNetIntegration": false,
     "InstrumentationCategory": 1
   },
+  {
+    "IntegrationName": "AzureServiceBus",
+    "AssemblyName": "Azure.Messaging.ServiceBus",
+    "TargetTypeName": "Azure.Messaging.ServiceBus.ServiceBusSender",
+    "TargetMethodName": "SendMessagesAsync",
+    "TargetReturnType": "System.Threading.Tasks.Task",
+    "TargetParameterTypes": [
+      "Azure.Messaging.ServiceBus.ServiceBusMessageBatch",
+      "System.Threading.CancellationToken"
+    ],
+    "MinimumVersion": {
+      "Item1": 7,
+      "Item2": 14,
+      "Item3": 0
+    },
+    "MaximumVersion": {
+      "Item1": 7,
+      "Item2": 65535,
+      "Item3": 65535
+    },
+    "InstrumentationTypeName": "Datadog.Trace.ClrProfiler.AutoInstrumentation.Azure.ServiceBus.ServiceBusSenderSendMessageBatchAsyncIntegration",
+    "IntegrationKind": 0,
+    "IsAdoNetIntegration": false,
+    "InstrumentationCategory": 1
+  },
   {
     "IntegrationName": "AzureServiceBus",
     "AssemblyName": "Azure.Messaging.ServiceBus",

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/ServiceBus/AzureServiceBusCommon.cs

@@ -60,8 +60,59 @@ internal static long GetMessageSize<T>(T message)
             return size;
         }
 
-        internal static CallTargetState CreateSenderSpan<TTarget>(TTarget instance, string operationName = "azure_servicebus.send", IEnumerable? messages = null)
-            where TTarget : IServiceBusSender, IDuckType
+        internal static CallTargetState CreateSenderSpan(
+            IServiceBusSender instance,
+            string operationName,
+            IEnumerable? messages = null,
+            int? messageCount = null,
+            IEnumerable<SpanLink>? spanLinks = null)
+        {
+            var entityPath = instance.EntityPath;
+            var endpoint = instance.Connection?.ServiceEndpoint;
+            var networkDestinationName = endpoint?.Host;
+            // https://learn.microsoft.com/en-us/dotnet/api/system.uri.port#remarks
+            var networkDestinationPort = endpoint?.Port is null or -1 or 5671 ?
+                                            "5671" :
+                                            endpoint.Port.ToString();
+
+            return CreateSenderSpanInternal(
+                entityPath,
+                networkDestinationName,
+                networkDestinationPort,
+                operationName,
+                messages,
+                messageCount,
+                spanLinks);
+        }
+
+        internal static CallTargetState CreateSenderSpan(
+            IMessagingClientDiagnostics clientDiagnostics,
+            string operationName,
+            IEnumerable? messages = null,
+            int? messageCount = null,
+            IEnumerable<SpanLink>? spanLinks = null)
+        {
+            var entityPath = clientDiagnostics.EntityPath;
+            var networkDestinationName = clientDiagnostics.FullyQualifiedNamespace;
+
+            return CreateSenderSpanInternal(
+                entityPath,
+                networkDestinationName,
+                null,
+                operationName,
+                messages,
+                messageCount,
+                spanLinks);
+        }
+
+        private static CallTargetState CreateSenderSpanInternal(
+            string? entityPath,
+            string? networkDestinationName,
+            string? networkDestinationPort,
+            string operationName,
+            IEnumerable? messages,
+            int? messageCount,
+            IEnumerable<SpanLink>? spanLinks)
         {
             var tracer = Tracer.Instance;
             if (!tracer.Settings.IsIntegrationEnabled(IntegrationId.AzureServiceBus, false))
@@ -71,31 +122,31 @@ internal static CallTargetState CreateSenderSpan<TTarget>(TTarget instance, stri
 
             var tags = tracer.CurrentTraceSettings.Schema.Messaging.CreateAzureServiceBusTags(SpanKinds.Producer);
 
-            var entityPath = instance.EntityPath ?? "unknown";
             tags.MessagingDestinationName = entityPath;
-            tags.MessagingOperation = "send";
+            tags.MessagingOperation = operationName;
             tags.MessagingSystem = "servicebus";
             tags.InstrumentationName = "AzureServiceBus";
 
             string serviceName = tracer.CurrentTraceSettings.Schema.Messaging.GetServiceName("azureservicebus");
             var scope = tracer.StartActiveInternal(
-                operationName,
+                "azure_servicebus." + operationName,
                 tags: tags,
-                serviceName: serviceName);
+                serviceName: serviceName,
+                links: spanLinks);
             var span = scope.Span;
 
             span.Type = SpanTypes.Queue;
             span.ResourceName = entityPath;
 
-            var messageCount = messages is ICollection collection ? collection.Count : 0;
+            var actualMessageCount = messageCount ?? (messages is ICollection collection ? collection.Count : 0);
             string? singleMessageId = null;
 
-            if (messageCount > 1)
+            if (actualMessageCount > 1)
             {
-                span.SetTag(Tags.MessagingBatchMessageCount, messageCount.ToString());
+                span.SetTag(Tags.MessagingBatchMessageCount, actualMessageCount.ToString());
             }
 
-            if (messageCount == 1 && messages != null)
+            if (actualMessageCount == 1 && messages != null)
             {
                 foreach (var message in messages)
                 {
@@ -110,14 +161,14 @@ internal static CallTargetState CreateSenderSpan<TTarget>(TTarget instance, stri
                 }
             }
 
-            var endpoint = instance.Connection?.ServiceEndpoint;
-            if (endpoint != null)
+            if (!string.IsNullOrEmpty(networkDestinationName))
+            {
+                tags.NetworkDestinationName = networkDestinationName;
+            }
+
+            if (!string.IsNullOrEmpty(networkDestinationPort))
             {
-                tags.NetworkDestinationName = endpoint.Host;
-                // https://learn.microsoft.com/en-us/dotnet/api/system.uri.port?view=net-8.0#remarks
-                tags.NetworkDestinationPort = endpoint.Port is -1 or 5671 ?
-                                    "5671" :
-                                    endpoint.Port.ToString();
+                tags.NetworkDestinationPort = networkDestinationPort;
             }
 
             return new CallTargetState(scope);

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/ServiceBus/IMessagingClientDiagnostics.cs

@@ -0,0 +1,23 @@
+// <copyright file="IMessagingClientDiagnostics.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+#nullable enable
+
+using Datadog.Trace.DuckTyping;
+
+namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Azure.ServiceBus
+{
+    /// <summary>
+    /// Duck type interface for Azure.Core.Shared.MessagingClientDiagnostics
+    /// </summary>
+    internal interface IMessagingClientDiagnostics : IDuckType
+    {
+        [DuckField(Name = "_entityPath")]
+        string? EntityPath { get; }
+
+        [DuckField(Name = "_fullyQualifiedNamespace")]
+        string? FullyQualifiedNamespace { get; }
+    }
+}

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/ServiceBus/IServiceBusMessageBatch.cs

@@ -0,0 +1,22 @@
+// <copyright file="IServiceBusMessageBatch.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+#nullable enable
+
+using Datadog.Trace.DuckTyping;
+
+namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Azure.ServiceBus
+{
+    /// <summary>
+    /// Duck type interface for Azure.Messaging.ServiceBus.ServiceBusMessageBatch
+    /// </summary>
+    internal interface IServiceBusMessageBatch : IDuckType
+    {
+        int Count { get; }
+
+        [DuckField(Name = "_clientDiagnostics")]
+        IMessagingClientDiagnostics ClientDiagnostics { get; }
+    }
+}

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/ServiceBus/SendServiceBusMessageBatchIntegration.cs

@@ -5,13 +5,16 @@
 
 #nullable enable
 
+using System;
 using System.Collections;
 using System.ComponentModel;
 using System.Threading;
 using Datadog.Trace.ClrProfiler.CallTarget;
 using Datadog.Trace.Configuration;
 using Datadog.Trace.DataStreamsMonitoring;
 using Datadog.Trace.DuckTyping;
+using Datadog.Trace.Logging;
+using Datadog.Trace.Propagators;
 
 namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Azure.ServiceBus
 {
@@ -31,6 +34,8 @@ namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Azure.ServiceBus
     [EditorBrowsable(EditorBrowsableState.Never)]
     public class SendServiceBusMessageBatchIntegration
     {
+        private static readonly IDatadogLogger Log = DatadogLogging.GetLoggerFor(typeof(SendServiceBusMessageBatchIntegration));
+
         /// <summary>
         /// OnMethodBegin callback
         /// </summary>
@@ -40,8 +45,11 @@ public class SendServiceBusMessageBatchIntegration
         /// <param name="message">The message instance</param>
         /// <returns>Calltarget state value</returns>
         internal static CallTargetState OnMethodBegin<TTarget, TMessage>(TTarget instance, TMessage message)
+            where TTarget : IServiceBusMessageBatch, IDuckType
             where TMessage : IServiceBusMessage
         {
+            Scope? messageScope = null;
+
             if (Tracer.Instance.Settings.IsIntegrationEnabled(IntegrationId.AzureServiceBus)
                 && Tracer.Instance.TracerManager.DataStreamsManager.IsEnabled)
             {
@@ -51,7 +59,41 @@ internal static CallTargetState OnMethodBegin<TTarget, TMessage>(TTarget instanc
                 AzureServiceBusCommon.SetMessage(message.ApplicationProperties, message.Instance);
             }
 
-            return CallTargetState.GetDefault();
+            // Create TryAdd message spans for batch with links when enabled
+            if (Tracer.Instance.Settings.IsIntegrationEnabled(IntegrationId.AzureServiceBus, false) &&
+                Tracer.Instance.Settings.AzureServiceBusBatchLinksEnabled)
+            {
+                messageScope = CreateAddMessageSpan(instance, message);
+            }
+
+            return new CallTargetState(messageScope);
+        }
+
+        internal static CallTargetReturn<bool> OnMethodEnd<TTarget>(TTarget instance, bool returnValue, Exception? exception, in CallTargetState state)
+        {
+            if (state.Scope != null)
+            {
+                if (returnValue && instance != null)
+                {
+                    ServiceBusBatchSpanContext.AddMessageSpanContext(instance, state.Scope.Span.Context);
+                }
+
+                state.Scope.DisposeWithException(exception);
+            }
+
+            return new CallTargetReturn<bool>(returnValue);
+        }
+
+        private static Scope? CreateAddMessageSpan(IServiceBusMessageBatch batch, IServiceBusMessage message)
+        {
+            var messageEnumerable = new[] { message };
+            var state = AzureServiceBusCommon.CreateSenderSpan(
+                batch.ClientDiagnostics,
+                operationName: "create",
+                messages: messageEnumerable,
+                messageCount: 1);
+
+            return state.Scope;
         }
     }
 }

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/ServiceBus/ServiceBusBatchSpanContext.cs

@@ -0,0 +1,51 @@
+// <copyright file="ServiceBusBatchSpanContext.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+#nullable enable
+
+using System.Collections.Concurrent;
+using System.Runtime.CompilerServices;
+
+namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Azure.ServiceBus
+{
+    /// <summary>
+    /// Helper class to manage span contexts for Service Bus message batch operations.
+    /// Tracks individual message span contexts per batch for span linking purposes.
+    /// </summary>
+    internal static class ServiceBusBatchSpanContext
+    {
+        // Maps ServiceBusMessageBatch instances to their collection of message span contexts
+        private static readonly ConditionalWeakTable<object, ConcurrentBag<SpanContext>> BatchToMessageSpanContexts = new();
+
+        public static void AddMessageSpanContext(object batch, SpanContext spanContext)
+        {
+            if (batch == null || spanContext == null)
+            {
+                return;
+            }
+
+            var spanContexts = BatchToMessageSpanContexts.GetValue(batch, _ => new ConcurrentBag<SpanContext>());
+            spanContexts.Add(spanContext);
+        }
+
+        public static SpanContext[] ExtractMessageSpanContexts(object batch)
+        {
+            if (batch == null)
+            {
+                return [];
+            }
+
+            if (BatchToMessageSpanContexts.TryGetValue(batch, out var spanContexts))
+            {
+                var contexts = spanContexts.ToArray();
+                BatchToMessageSpanContexts.Remove(batch);
+
+                return contexts;
+            }
+
+            return [];
+        }
+    }
+}

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/ServiceBus/ServiceBusReceiverReceiveMessagesAsyncIntegration.cs

@@ -87,6 +87,11 @@ internal static CallTargetState OnMethodBegin<TTarget>(TTarget instance, int max
             return null;
         }
 
+        if (!tracer.Settings.AzureServiceBusBatchLinksEnabled)
+        {
+            return null;
+        }
+
         var extractedContexts = new HashSet<SpanContext>(new SpanContextComparer());
 
         try

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/ServiceBus/ServiceBusSenderScheduleMessagesAsyncIntegration.cs

@@ -30,7 +30,7 @@ namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Azure.ServiceBus;
 [EditorBrowsable(EditorBrowsableState.Never)]
 public class ServiceBusSenderScheduleMessagesAsyncIntegration
 {
-    private const string OperationName = "azure_servicebus.send";
+    private const string OperationName = "send";
 
     internal static CallTargetState OnMethodBegin<TTarget>(TTarget instance, ref IEnumerable messages, ref DateTimeOffset scheduledEnqueueTime, ref CancellationToken cancellationToken)
         where TTarget : IServiceBusSender, IDuckType