Add diagnostics

Author: SeanFeldman

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

@@ -38,6 +38,7 @@ public Batch(ulong maximumBatchSize, Func<Message, Task<Message>> pluginsCallbac
             this.pluginsCallback = pluginsCallback;
             this.datas = new List<Data>();
             this.result = AmqpMessage.Create(datas);
+            OriginalMessageList = new List<Message>();
         }
 
         /// <summary>
@@ -66,6 +67,7 @@ public async Task<bool> TryAdd(Message message)
 
             if (Size <= maximumBatchSize)
             {
+                OriginalMessageList.Add(message);
                 return true;
             }
 
@@ -127,8 +129,8 @@ public void Dispose()
             firstMessage?.Dispose();
             result?.Dispose();
 
-            firstMessage = null;
-            result = null;
+            datas.Clear();
+            OriginalMessageList.Clear();
         }
 
         private void ThrowIfDisposed()
@@ -139,6 +141,8 @@ private void ThrowIfDisposed()
             }
         }
 
+        internal List<Message> OriginalMessageList { get; }
+
         private string DebuggerDisplay => $"Batch: size={Size} message count={datas.Count} maximum size={maximumBatchSize}";
     }
 }

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

@@ -280,33 +280,31 @@ public async Task SendAsync(Batch batch)
 
             MessagingEventSource.Log.MessageSendStart(this.ClientId, batch.Length);
 
-//            var isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
-//            var activity = isDiagnosticSourceEnabled ? this.diagnosticSource.SendStart(messageList) : null;
+            var isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+            var activity = isDiagnosticSourceEnabled ? this.diagnosticSource.SendStart(batch.OriginalMessageList) : null;
             Task sendTask = null;
 
             try
             {
-                //var processedMessages = await this.ProcessMessages(messageList).ConfigureAwait(false);
-
                 sendTask = this.RetryPolicy.RunOperation(() => this.OnSendAsync(batch.ToAmqpMessage), this.OperationTimeout);
                 await sendTask.ConfigureAwait(false);
             }
             catch (Exception exception)
             {
-//                if (isDiagnosticSourceEnabled)
-//                {
-//                    this.diagnosticSource.ReportException(exception);
-//                }
+                if (isDiagnosticSourceEnabled)
+                {
+                    this.diagnosticSource.ReportException(exception);
+                }
 
                 MessagingEventSource.Log.MessageSendException(this.ClientId, exception);
                 throw;
             }
             finally
             {
-//                this.diagnosticSource.SendStop(activity, messageList, sendTask?.Status);
+                this.diagnosticSource.SendStop(activity, batch.OriginalMessageList, sendTask?.Status);
             }
 
-//            MessagingEventSource.Log.MessageSendStop(this.ClientId);
+            MessagingEventSource.Log.MessageSendStop(this.ClientId);
         }
 
         /// <summary>

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

@@ -309,6 +309,40 @@ async Task BatchSendReceiveFireEvents()
             Assert.True(this.events.IsEmpty);
         }
 
+        [Fact]
+        [DisplayTestMethodName]
+        async Task BatchSendWithBatchReceiveFireEvents()
+        {
+            this.queueClient = new QueueClient(TestUtility.NamespaceConnectionString, TestConstants.NonPartitionedQueueName,
+                ReceiveMode.ReceiveAndDelete);
+
+            this.listener.Enable( (name, queuName, arg) => name.Contains("Send") || name.Contains("Receive") );
+            await TestUtility.SendMessagesUsingBatchAsync(this.queueClient.InnerSender, 5);
+            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, 5);
+
+            Assert.True(this.events.TryDequeue(out var sendStop1));
+            AssertSendStop(sendStop1.eventName, sendStop1.payload, sendStop1.activity, sendStop1.activity, 5);
+
+            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.True(this.events.IsEmpty);
+        }
+
         [Fact]
         [DisplayTestMethodName]
         async Task PeekFireEvents()

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

@@ -66,6 +66,25 @@ internal static async Task SendMessagesAsync(IMessageSender messageSender, int m
             await messageSender.SendAsync(messagesToSend);
             Log($"Sent {messageCount} messages");
         }
+        internal static async Task SendMessagesUsingBatchAsync(IMessageSender messageSender, int messageCount)
+        {
+            if (messageCount == 0)
+            {
+                await Task.FromResult(false);
+            }
+
+            var sender = (MessageSender)messageSender;
+            var batch = await sender.CreateBatch();
+            for (var i = 0; i < messageCount; i++)
+            {
+                var message = new Message(Encoding.UTF8.GetBytes("test" + i));
+                message.Label = "test" + i;
+                await batch.TryAdd(message);
+            }
+
+            await sender.SendAsync(batch);
+            Log($"Sent {messageCount} messages");
+        }
 
         internal static async Task<IList<Message>> ReceiveMessagesAsync(IMessageReceiver messageReceiver, int messageCount)
         {
@@ -98,7 +117,7 @@ internal static async Task<IList<Message>> ReceiveDeferredMessagesAsync(IMessage
                 var msg = await messageReceiver.ReceiveDeferredMessageAsync(sequenceNumber);
                 if (msg != null)
                 {
-                    messagesToReturn.Add(msg); 
+                    messagesToReturn.Add(msg);
                 }
             }