Tracing batch fix (Azure#19091)

* Instrument messages when adding to batch

* Fix tests

* Add docs

Author: JoshLove-msft

sdk/servicebus/Azure.Messaging.ServiceBus/src/Diagnostics/EntityScopeFactory.cs

@@ -2,7 +2,7 @@
 // Licensed under the MIT License.
 
 using System.Collections.Generic;
-using System.Threading.Tasks;
+using System.Diagnostics;
 using Azure.Core.Pipeline;
 
 namespace Azure.Messaging.ServiceBus.Diagnostics
@@ -71,5 +71,22 @@ public DiagnosticScope CreateScope(
             scope.AddAttribute(DiagnosticProperty.EndpointAttribute, _fullyQualifiedNamespace);
             return scope;
         }
+
+        public void InstrumentMessage(ServiceBusMessage message)
+        {
+            if (!message.ApplicationProperties.ContainsKey(DiagnosticProperty.DiagnosticIdAttribute))
+            {
+                using DiagnosticScope messageScope = CreateScope(
+                    DiagnosticProperty.MessageActivityName,
+                    DiagnosticProperty.ProducerKind);
+                messageScope.Start();
+
+                Activity activity = Activity.Current;
+                if (activity != null)
+                {
+                    message.ApplicationProperties[DiagnosticProperty.DiagnosticIdAttribute] = activity.Id;
+                }
+            }
+        }
     }
 }

sdk/servicebus/Azure.Messaging.ServiceBus/src/Primitives/ServiceBusModelFactory.cs

@@ -3,11 +3,11 @@
 
 using System;
 using System.Collections.Generic;
-using System.ComponentModel;
 using Azure.Core.Amqp;
 using Azure.Messaging.ServiceBus.Amqp;
 using Azure.Messaging.ServiceBus.Administration;
 using Azure.Messaging.ServiceBus.Core;
+using Azure.Messaging.ServiceBus.Diagnostics;
 
 namespace Azure.Messaging.ServiceBus
 {
@@ -237,7 +237,7 @@ public static ServiceBusMessageBatch ServiceBusMessageBatch(long batchSizeBytes,
             batchOptions.MaxSizeInBytes ??= long.MaxValue;
 
             var transportBatch = new ListTransportBatch(batchOptions.MaxSizeInBytes.Value, batchSizeBytes, batchMessageStore, tryAddCallback);
-            return new ServiceBusMessageBatch(transportBatch);
+            return new ServiceBusMessageBatch(transportBatch, new EntityScopeFactory("mock", "mock"));
         }
 
         /// <summary>

sdk/servicebus/Azure.Messaging.ServiceBus/src/Sender/ServiceBusMessageBatch.cs

@@ -5,6 +5,7 @@
 using System.Collections.Generic;
 using Azure.Core;
 using Azure.Messaging.ServiceBus.Core;
+using Azure.Messaging.ServiceBus.Diagnostics;
 
 namespace Azure.Messaging.ServiceBus
 {
@@ -29,32 +30,35 @@ public sealed class ServiceBusMessageBatch : IDisposable
         ///   well as any overhead for the batch itself when sent to the Queue/Topic.
         /// </summary>
         ///
-        public long MaxSizeInBytes => InnerBatch.MaxSizeInBytes;
+        public long MaxSizeInBytes => _innerBatch.MaxSizeInBytes;
 
         /// <summary>
         ///   The size of the batch, in bytes, as it will be sent to the Queue/Topic.
         /// </summary>
         ///
-        public long SizeInBytes => InnerBatch.SizeInBytes;
+        public long SizeInBytes => _innerBatch.SizeInBytes;
 
         /// <summary>
         ///   The count of messages contained in the batch.
         /// </summary>
         ///
-        public int Count => InnerBatch.Count;
+        public int Count => _innerBatch.Count;
 
         /// <summary>
         ///   The transport-specific batch responsible for performing the batch operations
         ///   in a manner compatible with the associated <see cref="TransportSender" />.
         /// </summary>
         ///
-        private TransportMessageBatch InnerBatch { get; }
+        private readonly TransportMessageBatch _innerBatch;
+
+        private readonly EntityScopeFactory _scopeFactory;
 
         /// <summary>
         ///   Initializes a new instance of the <see cref="ServiceBusMessageBatch"/> class.
         /// </summary>
         ///
         /// <param name="transportBatch">The  transport-specific batch responsible for performing the batch operations.</param>
+        /// <param name="entityScope">The entity scope used for instrumentation.</param>
         ///
         /// <remarks>
         ///   As an internal type, this class performs only basic sanity checks against its arguments.  It
@@ -65,15 +69,19 @@ public sealed class ServiceBusMessageBatch : IDisposable
         ///   caller.
         /// </remarks>
         ///
-        internal ServiceBusMessageBatch(TransportMessageBatch transportBatch)
+        internal ServiceBusMessageBatch(TransportMessageBatch transportBatch, EntityScopeFactory entityScope)
         {
             Argument.AssertNotNull(transportBatch, nameof(transportBatch));
-            InnerBatch = transportBatch;
+            _innerBatch = transportBatch;
+            _scopeFactory = entityScope;
         }
 
         /// <summary>
         ///   Attempts to add a message to the batch, ensuring that the size
-        ///   of the batch does not exceed its maximum.
+        ///   of the batch does not exceed its maximum. If the message is modified
+        ///   after being added to the batch, the batch will fail to send if the modification
+        ///   caused the batch to exceed the maximum allowable size. Therefore it is best
+        ///   to not modify a message after adding it to the batch.
         /// </summary>
         ///
         /// <param name="message">The message to attempt to add to the batch.</param>
@@ -91,7 +99,9 @@ public bool TryAddMessage(ServiceBusMessage message)
             lock (_syncGuard)
             {
                 AssertNotLocked();
-                return InnerBatch.TryAddMessage(message);
+
+                _scopeFactory.InstrumentMessage(message);
+                return _innerBatch.TryAddMessage(message);
             }
         }
 
@@ -104,7 +114,7 @@ public void Dispose()
             lock (_syncGuard)
             {
                 AssertNotLocked();
-                InnerBatch.Dispose();
+                _innerBatch.Dispose();
             }
         }
 
@@ -118,7 +128,7 @@ internal void Clear()
             lock (_syncGuard)
             {
                 AssertNotLocked();
-                InnerBatch.Clear();
+                _innerBatch.Clear();
             }
         }
 
@@ -130,7 +140,7 @@ internal void Clear()
         ///
         /// <returns>The set of messages as an enumerable of the requested type.</returns>
         ///
-        internal IEnumerable<T> AsEnumerable<T>() => InnerBatch.AsEnumerable<T>();
+        internal IEnumerable<T> AsEnumerable<T>() => _innerBatch.AsEnumerable<T>();
 
         /// <summary>
         ///   Locks the batch to prevent new messages from being added while a service

sdk/servicebus/Azure.Messaging.ServiceBus/src/Sender/ServiceBusSender.cs

@@ -246,7 +246,10 @@ private async Task ApplyPlugins(IReadOnlyList<ServiceBusMessage> messages)
 
         private DiagnosticScope CreateDiagnosticScope(IEnumerable<ServiceBusMessage> messages, string activityName)
         {
-            InstrumentMessages(messages);
+            foreach (ServiceBusMessage message in messages)
+            {
+                _scopeFactory.InstrumentMessage(message);
+            }
 
             // create a new scope for the specified operation
             DiagnosticScope scope = _scopeFactory.CreateScope(
@@ -257,32 +260,6 @@ private DiagnosticScope CreateDiagnosticScope(IEnumerable<ServiceBusMessage> mes
             return scope;
         }
 
-        /// <summary>
-        ///   Performs the actions needed to instrument a set of messages.
-        /// </summary>
-        ///
-        /// <param name="messages">The messages to instrument.</param>
-        ///
-        private void InstrumentMessages(IEnumerable<ServiceBusMessage> messages)
-        {
-            foreach (ServiceBusMessage message in messages)
-            {
-                if (!message.ApplicationProperties.ContainsKey(DiagnosticProperty.DiagnosticIdAttribute))
-                {
-                    using DiagnosticScope messageScope = _scopeFactory.CreateScope(
-                        DiagnosticProperty.MessageActivityName,
-                        DiagnosticProperty.ProducerKind);
-                    messageScope.Start();
-
-                    Activity activity = Activity.Current;
-                    if (activity != null)
-                    {
-                        message.ApplicationProperties[DiagnosticProperty.DiagnosticIdAttribute] = activity.Id;
-                    }
-                }
-            }
-        }
-
         /// <summary>
         ///   Creates a size-constraint batch to which <see cref="ServiceBusMessage" /> may be added using
         ///   a <see cref="ServiceBusMessageBatch.TryAddMessage"/>. If a message would exceed the maximum
@@ -329,7 +306,7 @@ public virtual async ValueTask<ServiceBusMessageBatch> CreateMessageBatchAsync(
             try
             {
                 TransportMessageBatch transportBatch = await _innerSender.CreateMessageBatchAsync(options, cancellationToken).ConfigureAwait(false);
-                batch = new ServiceBusMessageBatch(transportBatch);
+                batch = new ServiceBusMessageBatch(transportBatch, _scopeFactory);
             }
             catch (Exception ex)
             {

sdk/servicebus/Azure.Messaging.ServiceBus/tests/Sender/SenderLiveTests.cs

@@ -6,8 +6,6 @@
 using System.Linq;
 using System.Net;
 using System.Threading.Tasks;
-using Azure.Messaging.ServiceBus;
-using Azure.Messaging.ServiceBus.Tests;
 using NUnit.Framework;
 
 namespace Azure.Messaging.ServiceBus.Tests.Sender
@@ -131,48 +129,6 @@ public async Task CanSendAnEmptyBodyMessageBatch()
             }
         }
 
-        [Test]
-        public async Task CanSendLargeMessageBatch()
-        {
-            await using (var scope = await ServiceBusScope.CreateWithQueue(enablePartitioning: true, enableSession: true))
-            {
-                await using var client = new ServiceBusClient(TestEnvironment.ServiceBusConnectionString);
-                ServiceBusSender sender = client.CreateSender(scope.QueueName);
-                using ServiceBusMessageBatch batch = await sender.CreateMessageBatchAsync();
-
-                await AddAndSendMessages();
-
-                batch.Clear();
-                Assert.AreEqual(0, batch.Count);
-                Assert.AreEqual(0, batch.SizeInBytes);
-
-                await AddAndSendMessages();
-
-                async Task AddAndSendMessages()
-                {
-                    // service limits to 4500 messages but we have not added this to our client validation yet
-                    while (batch.Count < 4500 && batch.TryAddMessage(
-                        new ServiceBusMessage(new byte[50])
-                        {
-                            MessageId = "new message ID that takes up some space",
-                            SessionId = "sessionId",
-                            PartitionKey = "sessionId",
-                            ApplicationProperties = { { "key", "value" } }
-                        }))
-                    {
-                    }
-
-                    if (batch.Count < 4500)
-                    {
-                        // the difference in size from the max allowable size should be less than the size of 1 message
-                        Assert.IsTrue(batch.MaxSizeInBytes - batch.SizeInBytes < 180);
-                    }
-                    Assert.Greater(batch.Count, 0);
-                    await sender.SendMessagesAsync(batch);
-                }
-            }
-        }
-
         [Test]
         public async Task CannotSendLargerThanMaximumSize()
         {

sdk/servicebus/Azure.Messaging.ServiceBus/tests/Sender/SenderTests.cs

@@ -8,6 +8,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Messaging.ServiceBus.Core;
+using Azure.Messaging.ServiceBus.Diagnostics;
 using Azure.Messaging.ServiceBus.Plugins;
 using Moq;
 using NUnit.Framework;
@@ -145,7 +146,8 @@ public async Task SendBatchManagesLockingTheBatch()
 
             var completionSource = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
             var mockTransportBatch = new Mock<TransportMessageBatch>();
-            var batch = new ServiceBusMessageBatch(mockTransportBatch.Object);
+            var mockScope = new EntityScopeFactory("mock", "mock");
+            var batch = new ServiceBusMessageBatch(mockTransportBatch.Object, mockScope);
             var mockTransportSender = new Mock<TransportSender>();
             var mockConnection = new Mock<ServiceBusConnection>();
 

sdk/servicebus/Azure.Messaging.ServiceBus/tests/Sender/ServiceBusMessageBatchLiveTests.cs

@@ -0,0 +1,71 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System.Threading.Tasks;
+using Azure.Core.Tests;
+using Azure.Messaging.ServiceBus.Diagnostics;
+using NUnit.Framework;
+
+namespace Azure.Messaging.ServiceBus.Tests.Sender
+{
+    [NonParallelizable]
+    public class ServiceBusMessageBatchLiveTests : ServiceBusLiveTestBase
+    {
+        [Test]
+        [TestCase(true)]
+        [TestCase(false)]
+        public async Task CanSendLargeMessageBatch(bool enableTracing)
+        {
+            await using (var scope = await ServiceBusScope.CreateWithQueue(enablePartitioning: true, enableSession: true))
+            {
+                TestDiagnosticListener listener = null;
+                if (enableTracing)
+                {
+                    listener = new TestDiagnosticListener(EntityScopeFactory.DiagnosticNamespace);
+                }
+                try
+                {
+                    await using var client = new ServiceBusClient(TestEnvironment.ServiceBusConnectionString);
+                    ServiceBusSender sender = client.CreateSender(scope.QueueName);
+                    using ServiceBusMessageBatch batch = await sender.CreateMessageBatchAsync();
+
+                    await AddAndSendMessages();
+
+                    batch.Clear();
+                    Assert.AreEqual(0, batch.Count);
+                    Assert.AreEqual(0, batch.SizeInBytes);
+
+                    await AddAndSendMessages();
+
+                    async Task AddAndSendMessages()
+                    {
+                        // service limits to 4500 messages but we have not added this to our client validation yet
+                        while (batch.Count < 4500 && batch.TryAddMessage(
+                            new ServiceBusMessage(new byte[50])
+                            {
+                                MessageId = "new message ID that takes up some space",
+                                SessionId = "sessionId",
+                                PartitionKey = "sessionId",
+                                ApplicationProperties = { { "key", "value" } }
+                            }))
+                        {
+                        }
+
+                        if (batch.Count < 4500)
+                        {
+                            var diff = batch.MaxSizeInBytes - batch.SizeInBytes;
+                            // the difference in size from the max allowable size should be less than the size of 1 message
+                            Assert.IsTrue(diff < 220, diff.ToString());
+                        }
+                        Assert.Greater(batch.Count, 0);
+                        await sender.SendMessagesAsync(batch);
+                    }
+                }
+                finally
+                {
+                    listener?.Dispose();
+                }
+            }
+        }
+    }
+}