Initial (simple) batching implementation

Author: SeanFeldman

src/Microsoft.Azure.ServiceBus/Amqp/AmqpMessageConverter.cs

@@ -1,6 +1,9 @@
 // Copyright (c) Microsoft. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
+using System.Diagnostics;
+using System.Linq;
+
 namespace Microsoft.Azure.ServiceBus.Amqp
 {
     using System;
@@ -23,9 +26,9 @@ static class AmqpMessageConverter
         const string EnqueueSequenceNumberName = "x-opt-enqueue-sequence-number";
         const string LockedUntilName = "x-opt-locked-until";
         const string PublisherName = "x-opt-publisher";
-        const string PartitionKeyName = "x-opt-partition-key";
         const string PartitionIdName = "x-opt-partition-id";
-        const string ViaPartitionKeyName = "x-opt-via-partition-key";
+        internal const string PartitionKeyName = "x-opt-partition-key";
+        internal const string ViaPartitionKeyName = "x-opt-via-partition-key";
         const string DeadLetterSourceName = "x-opt-deadletter-source";
         const string TimeSpanName = AmqpConstants.Vendor + ":timespan";
         const string UriName = AmqpConstants.Vendor + ":uri";
@@ -641,7 +644,7 @@ static ArraySegment<byte> StreamToBytes(Stream stream)
             return buffer;
         }
 
-        private static Data ToData(AmqpMessage message)
+        internal static Data ToData(AmqpMessage message)
         {
             ArraySegment<byte>[] payload = message.GetPayload();
             var buffer = new BufferListStream(payload);

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

@@ -0,0 +1,121 @@
+using System;
+using System.Collections.Generic;
+using System.Diagnostics;
+using Microsoft.Azure.Amqp;
+using Microsoft.Azure.Amqp.Framing;
+using Microsoft.Azure.ServiceBus.Amqp;
+
+namespace Microsoft.Azure.ServiceBus.Core
+{
+    [DebuggerDisplay("{" + nameof(DebuggerDisplay) + ",nq}")]
+    public class Batch : IDisposable
+    {
+        private readonly long maximumBatchSize;
+        private AmqpMessage firstMessage;
+        private readonly List<Data> datas;
+        private AmqpMessage result;
+        private (string messageId, string sessionId, string partitionKey, string viaPartitionKey) originalMessageData;
+
+        /// <summary>
+        /// Construct a new batch with a maximum batch size.
+        /// </summary>
+        /// <param name="maximumBatchSize">Maximum batch size allowed for batch.</param>
+        public Batch(long maximumBatchSize)
+        {
+            this.maximumBatchSize = maximumBatchSize;
+            this.datas = new List<Data>();
+            this.result = AmqpMessage.Create(datas);
+        }
+
+        /// <summary>
+        /// Add <see cref="Message"/> to the batch if the overall size of the batch with the added message is not exceeding the batch maximum.
+        /// </summary>
+        /// <param name="message"><see cref="Message"/> to add to the batch.</param>
+        /// <returns></returns>
+        public bool TryAdd(Message message)
+        {
+            ThrowIfDisposed();
+
+            var amqpMessage = AmqpMessageConverter.SBMessageToAmqpMessage(message);
+
+            if (firstMessage == null)
+            {
+                originalMessageData = (message.MessageId, message.SessionId, message.PartitionKey, message.ViaPartitionKey);
+                firstMessage = amqpMessage;
+            }
+
+            var data = AmqpMessageConverter.ToData(amqpMessage);
+            datas.Add(data);
+
+            if (Size <= maximumBatchSize)
+            {
+                return true;
+            }
+
+            datas.Remove(data);
+            return false;
+
+        }
+
+        private long Size => result.SerializedMessageSize;
+
+        /// <summary>
+        /// Convert batch to AMQP message
+        /// </summary>
+        /// <returns></returns>
+        public AmqpMessage ToAmqpMessage()
+        {
+            ThrowIfDisposed();
+
+            if (datas.Count == 1)
+            {
+                firstMessage.Batchable = true;
+                return firstMessage;
+            }
+
+            if (originalMessageData.messageId != null)
+            {
+                result.Properties.MessageId = originalMessageData.messageId;
+            }
+
+            if (originalMessageData.sessionId != null)
+            {
+                result.Properties.GroupId = originalMessageData.sessionId;
+            }
+
+            if (originalMessageData.partitionKey != null)
+            {
+                result.MessageAnnotations.Map[AmqpMessageConverter.PartitionKeyName] = originalMessageData.partitionKey;
+            }
+
+            if (originalMessageData.viaPartitionKey != null)
+            {
+                result.MessageAnnotations.Map[AmqpMessageConverter.ViaPartitionKeyName] = originalMessageData.viaPartitionKey;
+            }
+
+            result.MessageFormat = AmqpConstants.AmqpBatchedMessageFormat;
+            result.Batchable = true;
+            return result;
+        }
+
+        public void Dispose()
+        {
+            // TODO: review if there's anything else to do
+            firstMessage?.Dispose();
+            result?.Dispose();
+
+            firstMessage = null;
+            result = null;
+        }
+
+        private void ThrowIfDisposed()
+        {
+            if (result == null)
+            {
+                throw new Exception("Batch is has been disposed and cannot be re-used.");
+            }
+        }
+
+        private string DebuggerDisplay => $"Batch: size={Size} message count={datas.Count}";
+    }
+}

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

@@ -247,7 +247,7 @@ public async Task SendAsync(IList<Message> messageList)
             {
                 var processedMessages = await this.ProcessMessages(messageList).ConfigureAwait(false);
 
-                sendTask = this.RetryPolicy.RunOperation(() => this.OnSendAsync(processedMessages), this.OperationTimeout);
+                sendTask = this.RetryPolicy.RunOperation(() => this.OnSendAsync(() => AmqpMessageConverter.BatchSBMessagesAsAmqpMessage(processedMessages)), this.OperationTimeout);
                 await sendTask.ConfigureAwait(false);
             }
             catch (Exception exception)
@@ -267,6 +267,41 @@ public async Task SendAsync(IList<Message> messageList)
 
             MessagingEventSource.Log.MessageSendStop(this.ClientId);
         }
+        public async Task SendAsync(Batch batch)
+        {
+            this.ThrowIfClosed();
+
+            //var count = MessageSender.ValidateMessages(messageList);
+//            MessagingEventSource.Log.MessageSendStart(this.ClientId, count);
+
+//            var isDiagnosticSourceEnabled = ServiceBusDiagnosticSource.IsEnabled();
+//            var activity = isDiagnosticSourceEnabled ? this.diagnosticSource.SendStart(messageList) : 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);
+//                }
+
+                MessagingEventSource.Log.MessageSendException(this.ClientId, exception);
+                throw;
+            }
+            finally
+            {
+//                this.diagnosticSource.SendStop(activity, messageList, sendTask?.Status);
+            }
+
+//            MessagingEventSource.Log.MessageSendStop(this.ClientId);
+        }
 
         /// <summary>
         /// Schedules a message to appear on Service Bus at a later time.
@@ -521,10 +556,10 @@ async Task<IList<Message>> ProcessMessages(IList<Message> messageList)
             return processedMessageList;
         }
 
-        async Task OnSendAsync(IList<Message> messageList)
+        async Task OnSendAsync(Func<AmqpMessage> amqpMessageProvider)
         {
             var timeoutHelper = new TimeoutHelper(this.OperationTimeout, true);
-            using (var amqpMessage = AmqpMessageConverter.BatchSBMessagesAsAmqpMessage(messageList))
+            using (var amqpMessage = amqpMessageProvider())
             {
                 SendingAmqpLink amqpLink = null;
                 try

test/Microsoft.Azure.ServiceBus.UnitTests/Primitives/BatchTests.cs

@@ -0,0 +1,33 @@
+// 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.Primitives
+{
+    using System;
+    using System.Text;
+    using Microsoft.Azure.ServiceBus.Core;
+    using Xunit;
+
+    public class BatchTests
+    {
+        [Fact]
+        public void Should_return_false_when_is_about_to_exceed_max_batch_size()
+        {
+            using (var batch = new Batch(1))
+            {
+                var wasAdded = batch.TryAdd(new Message(Encoding.UTF8.GetBytes("hello")));
+                Assert.False(wasAdded, "Message should not have been added, but it was.");
+            }
+        }
+
+        [Fact]
+        public void Should_throw_if_batch_disposed()
+        {
+            using (var batch = new Batch(1))
+            {
+                batch.Dispose();
+                Assert.Throws<Exception>(() => batch.TryAdd(new Message()));
+            }
+        }
+    }
+}

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

@@ -13,9 +13,9 @@ namespace Microsoft.Azure.ServiceBus.UnitTests
 
     public class SenderReceiverTests : SenderReceiverClientTestBase
     {
-        private static TimeSpan TwoSeconds = TimeSpan.FromSeconds(2);
+        private static readonly TimeSpan TwoSeconds = TimeSpan.FromSeconds(2);
 
-        public static IEnumerable<object[]> TestPermutations => new object[][]
+        public static IEnumerable<object[]> TestPermutations => new[]
         {
             new object[] {TestConstants.NonPartitionedQueueName},
             new object[] {TestConstants.PartitionedQueueName}
@@ -376,7 +376,7 @@ public async Task ClientThrowsObjectDisposedExceptionWhenUserCloseConnectionAndW
 
                 var recivedMessage = await receiver.ReceiveAsync().ConfigureAwait(false);
                 Assert.True(Encoding.UTF8.GetString(recivedMessage.Body) == Encoding.UTF8.GetString(messageBody));
-            
+
                 var connection = sender.ServiceBusConnection;
                 Assert.Throws<ObjectDisposedException>(() => new MessageSender(connection, TestConstants.PartitionedQueueName));
             }
@@ -413,7 +413,7 @@ public async Task SendMesageCloseConnectionCreateAnotherConnectionSendAgainMessa
 
                 messageBody = Encoding.UTF8.GetBytes("Message 2");
                 message = new Message(messageBody);
-                await sender.SendAsync(message); 
+                await sender.SendAsync(message);
 
                 recivedMessage = await receiver.ReceiveAsync().ConfigureAwait(false);
                 Assert.True(Encoding.UTF8.GetString(recivedMessage.Body) == Encoding.UTF8.GetString(messageBody));
@@ -459,5 +459,38 @@ public async Task ClientsUseGlobalConnectionCloseFirstClientSecoundClientShouldS
                 await receiver.CloseAsync().ConfigureAwait(false);
             }
         }
+
+        [Fact]
+        [DisplayTestMethodName]
+        public async Task Sending_batch()
+        {
+            var sender = new MessageSender(TestUtility.NamespaceConnectionString, TestConstants.PartitionedQueueName);
+            var receiver = new MessageReceiver(TestUtility.NamespaceConnectionString, TestConstants.PartitionedQueueName, receiveMode: ReceiveMode.ReceiveAndDelete);
+            try
+            {
+                var message1 = new Message(Encoding.UTF8.GetBytes("Hello Neeraj"));
+                var message2 = new Message(Encoding.UTF8.GetBytes("from"));
+                var message3 = new Message(Encoding.UTF8.GetBytes("Sean Feldman"));
+
+                var batch = new Batch(100);
+                Assert.True(batch.TryAdd(message1), "Couldn't add first message");
+                Assert.True(batch.TryAdd(message2), "Couldn't add second message");
+                Assert.False(batch.TryAdd(message3), "Shouldn't be able to add third message");
+                await sender.SendAsync(batch);
+                //batch.dispose()
+                await sender.CloseAsync();
+
+                var receivedMessages = await TestUtility.ReceiveMessagesAsync(receiver, 2);
+                Assert.Equal("Hello Neeraj", Encoding.UTF8.GetString(receivedMessages[0].Body));
+                Assert.Equal("from", Encoding.UTF8.GetString(receivedMessages[1].Body));
+                var extraMessage = await TestUtility.PeekMessageAsync(receiver);
+                Assert.True(extraMessage == null, "Should not have any messages other than the two, but an extra message is found");
+            }
+            finally
+            {
+                await sender.CloseAsync().ConfigureAwait(false);
+                await receiver.CloseAsync().ConfigureAwait(false);
+            }
+        }
     }
 }