Provide an API to create Batch initiated with supported maximum message size

Author: SeanFeldman

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

@@ -14,7 +14,7 @@ namespace Microsoft.Azure.ServiceBus.Core
     [DebuggerDisplay("{" + nameof(DebuggerDisplay) + ",nq}")]
     public class Batch : IDisposable
     {
-        private readonly long maximumBatchSize;
+        internal readonly ulong maximumBatchSize;
         private AmqpMessage firstMessage;
         private readonly List<Data> datas;
         private AmqpMessage result;
@@ -24,7 +24,7 @@ public class Batch : IDisposable
         /// Construct a new batch with a maximum batch size.
         /// </summary>
         /// <param name="maximumBatchSize">Maximum batch size allowed for batch.</param>
-        public Batch(long maximumBatchSize)
+        public Batch(ulong maximumBatchSize)
         {
             this.maximumBatchSize = maximumBatchSize;
             this.datas = new List<Data>();
@@ -68,7 +68,7 @@ public bool TryAdd(Message message)
         /// </summary>
         public int Length => datas.Count;
 
-        internal long Size => result.SerializedMessageSize;
+        internal ulong Size => (ulong) result.SerializedMessageSize;
 
 
         /// <summary>
@@ -128,6 +128,6 @@ private void ThrowIfDisposed()
             }
         }
 
-        private string DebuggerDisplay => $"Batch: size={Size} message count={datas.Count}";
+        private string DebuggerDisplay => $"Batch: size={Size} message count={datas.Count} maximum size={maximumBatchSize}";
     }
 }

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

@@ -25,11 +25,15 @@ public interface ISenderClient : IClientEntity
         /// </summary>
         Task SendAsync(IList<Message> messageList);
 
-// TODO: extract method into this interface for the next major version
-//        /// <summary>
-//        /// Sends a <see cref="Batch"/> of messages to Service Bus.
-//        /// </summary>
-//        Task SendAsync(Batch batch);
+        // TODO: extract methods into this interface for the next major version
+        //        /// <summary>
+        //        /// Sends a <see cref="Batch"/> of messages to Service Bus.
+        //        /// </summary>
+        //        Task SendAsync(Batch batch);
+        //        /// <summary>
+        //        /// Create a new <see cref="Batch"/> setting maximum size to the maximum message size allowed by the underlying namespace.
+        //        /// </summary>
+        //        Task<Batch> CreateBatch();
 
         /// <summary>
         /// Schedules a message to appear on Service Bus.

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

@@ -43,6 +43,7 @@ public class MessageSender : ClientEntity, IMessageSender
         readonly ServiceBusDiagnosticSource diagnosticSource;
         readonly bool isViaSender;
         readonly string transferDestinationPath;
+        private ulong maxMessageSize = 0;
 
         /// <summary>
         /// Creates a new AMQP MessageSender.
@@ -308,6 +309,40 @@ public async Task SendAsync(Batch batch)
 //            MessagingEventSource.Log.MessageSendStop(this.ClientId);
         }
 
+        /// <summary>
+        /// Create a new <see cref="Batch"/> setting maximum size to the maximum message size allowed by the underlying namespace.
+        /// </summary>
+        public async Task<Batch> CreateBatch()
+        {
+            if (maxMessageSize != 0)
+            {
+                return new Batch(maxMessageSize);
+            }
+
+            var timeoutHelper = new TimeoutHelper(this.OperationTimeout, true);
+            SendingAmqpLink amqpLink = null;
+            try
+            {
+                if (!this.SendLinkManager.TryGetOpenedObject(out amqpLink))
+                {
+                    amqpLink = await this.SendLinkManager.GetOrCreateAsync(timeoutHelper.RemainingTime()).ConfigureAwait(false);
+                }
+
+                if (!amqpLink.Settings.MaxMessageSize.HasValue)
+                {
+                    throw new Exception("Broker didn't provide maximum message size. Batch requires maximum message size to operate.");
+                }
+
+                maxMessageSize = amqpLink.Settings.MaxMessageSize.Value;
+
+                return new Batch(maxMessageSize);
+            }
+            catch (Exception exception)
+            {
+                throw AmqpExceptionHelper.GetClientException(exception, amqpLink?.GetTrackingId(), null, amqpLink?.Session.IsClosing() ?? false);
+            }
+        }
+
         /// <summary>
         /// Schedules a message to appear on Service Bus at a later time.
         /// </summary>

src/Microsoft.Azure.ServiceBus/QueueClient.cs

@@ -359,6 +359,17 @@ public Task SendAsync(Batch batch)
             return this.innerSender.SendAsync(batch);
         }
 
+        /// <summary>
+        /// Create a new <see cref="Batch"/> setting maximum size to the maximum message size allowed by the underlying namespace.
+        /// </summary>
+        public Task<Batch> CreateBatch()
+        {
+            this.ThrowIfClosed();
+
+            return this.innerSender.CreateBatch();
+        }
+
+
         /// <summary>
         /// Completes a <see cref="Message"/> using its lock token. This will delete the message from the queue.
         /// </summary>

src/Microsoft.Azure.ServiceBus/TopicClient.cs

@@ -194,6 +194,16 @@ public Task SendAsync(Batch batch)
             return this.innerSender.SendAsync(batch);
         }
 
+        /// <summary>
+        /// Create a new <see cref="Batch"/> setting maximum size to the maximum message size allowed by the underlying namespace.
+        /// </summary>
+        public Task<Batch> CreateBatch()
+        {
+            this.ThrowIfClosed();
+
+            return this.innerSender.CreateBatch();
+        }
+
         /// <summary>
         /// Schedules a message to appear on Service Bus at a later time.
         /// </summary>

test/Microsoft.Azure.ServiceBus.UnitTests/API/ApiApprovals.ApproveAzureServiceBus.approved.txt

@@ -227,6 +227,7 @@ namespace Microsoft.Azure.ServiceBus
         public System.Threading.Tasks.Task AbandonAsync(string lockToken, System.Collections.Generic.IDictionary<string, object> propertiesToModify = null) { }
         public System.Threading.Tasks.Task CancelScheduledMessageAsync(long sequenceNumber) { }
         public System.Threading.Tasks.Task CompleteAsync(string lockToken) { }
+        public System.Threading.Tasks.Task<Microsoft.Azure.ServiceBus.Core.Batch> CreateBatch() { }
         public System.Threading.Tasks.Task DeadLetterAsync(string lockToken, System.Collections.Generic.IDictionary<string, object> propertiesToModify = null) { }
         public System.Threading.Tasks.Task DeadLetterAsync(string lockToken, string deadLetterReason, string deadLetterErrorDescription = null) { }
         protected override System.Threading.Tasks.Task OnClosingAsync() { }
@@ -449,6 +450,7 @@ namespace Microsoft.Azure.ServiceBus
         public override Microsoft.Azure.ServiceBus.ServiceBusConnection ServiceBusConnection { get; }
         public string TopicName { get; }
         public System.Threading.Tasks.Task CancelScheduledMessageAsync(long sequenceNumber) { }
+        public System.Threading.Tasks.Task<Microsoft.Azure.ServiceBus.Core.Batch> CreateBatch() { }
         protected override System.Threading.Tasks.Task OnClosingAsync() { }
         public override void RegisterPlugin(Microsoft.Azure.ServiceBus.Core.ServiceBusPlugin serviceBusPlugin) { }
         public System.Threading.Tasks.Task<long> ScheduleMessageAsync(Microsoft.Azure.ServiceBus.Message message, System.DateTimeOffset scheduleEnqueueTimeUtc) { }
@@ -482,7 +484,7 @@ namespace Microsoft.Azure.ServiceBus.Core
     [System.Diagnostics.DebuggerDisplayAttribute("{DebuggerDisplay,nq}")]
     public class Batch : System.IDisposable
     {
-        public Batch(long maximumBatchSize) { }
+        public Batch(ulong maximumBatchSize) { }
         public int Length { get; }
         public void Dispose() { }
         public Microsoft.Azure.Amqp.AmqpMessage ToAmqpMessage() { }
@@ -584,6 +586,7 @@ namespace Microsoft.Azure.ServiceBus.Core
         public override Microsoft.Azure.ServiceBus.ServiceBusConnection ServiceBusConnection { get; }
         public string TransferDestinationPath { get; }
         public System.Threading.Tasks.Task CancelScheduledMessageAsync(long sequenceNumber) { }
+        public System.Threading.Tasks.Task<Microsoft.Azure.ServiceBus.Core.Batch> CreateBatch() { }
         protected override System.Threading.Tasks.Task OnClosingAsync() { }
         public override void RegisterPlugin(Microsoft.Azure.ServiceBus.Core.ServiceBusPlugin serviceBusPlugin) { }
         public System.Threading.Tasks.Task<long> ScheduleMessageAsync(Microsoft.Azure.ServiceBus.Message message, System.DateTimeOffset scheduleEnqueueTimeUtc) { }

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

@@ -67,7 +67,7 @@ public void Should_show_reflect_property_in_batch_size()
 
                 batch.TryAdd(message);
 
-                Assert.Equal(24, batch.Size);
+                Assert.Equal((ulong)24, batch.Size);
             }
 
             using (var batch = new Batch(100))
@@ -77,7 +77,7 @@ public void Should_show_reflect_property_in_batch_size()
 
                 batch.TryAdd(message);
 
-                Assert.Equal(45, batch.Size);
+                Assert.Equal((ulong)45, batch.Size);
             }
         }
     }

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

@@ -463,25 +463,31 @@ public async Task ClientsUseGlobalConnectionCloseFirstClientSecoundClientShouldS
 
         [Fact]
         [DisplayTestMethodName]
-        public async Task Sending_batch_with_properties()
+        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 message = new Message("Hello Neeraj".GetBytes());
-                message.UserProperties["custom"] = "value";
+                var message1 = new Message("Hello Neeraj".GetBytes());
+                var message2 = new Message("from".GetBytes());
+                var message3 = new Message("Sean Feldman".GetBytes());
 
                 var batch = new Batch(100);
-                Assert.True(batch.TryAdd(message), "Couldn't add first message");
+                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, 1);
-                var receivedMessage = receivedMessages.FirstOrDefault();
-                Assert.NotNull(receivedMessage);
-                Assert.Equal("value", receivedMessage.UserProperties["custom"]);
+                var receivedMessages = await TestUtility.ReceiveMessagesAsync(receiver, 2);
+                var bodies = receivedMessages.Select(m => m.Body.GetString());
+                var bodiesArray = bodies as string[] ?? bodies.ToArray();
+                Assert.True(bodiesArray.Contains("Hello Neeraj") && bodiesArray.Contains("from"));
+
+                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. Body='{extraMessage?.Body.GetString()}'");
             }
             finally
             {
@@ -492,37 +498,50 @@ public async Task Sending_batch_with_properties()
 
         [Fact]
         [DisplayTestMethodName]
-        public async Task Sending_batch()
+        public async Task Sending_batch_with_properties()
         {
             var sender = new MessageSender(TestUtility.NamespaceConnectionString, TestConstants.PartitionedQueueName);
             var receiver = new MessageReceiver(TestUtility.NamespaceConnectionString, TestConstants.PartitionedQueueName, receiveMode: ReceiveMode.ReceiveAndDelete);
             try
             {
-                var message1 = new Message("Hello Neeraj".GetBytes());
-                var message2 = new Message("from".GetBytes());
-                var message3 = new Message("Sean Feldman".GetBytes());
+                var message = new Message("Hello Neeraj".GetBytes());
+                message.UserProperties["custom"] = "value";
 
                 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");
+                Assert.True(batch.TryAdd(message), "Couldn't add  message");
                 await sender.SendAsync(batch);
                 batch.Dispose();
                 await sender.CloseAsync();
 
-                var receivedMessages = await TestUtility.ReceiveMessagesAsync(receiver, 2);
-                var bodies = receivedMessages.Select(m => m.Body.GetString());
-                var bodiesArray = bodies as string[] ?? bodies.ToArray();
-                Assert.True(bodiesArray.Contains("Hello Neeraj") && bodiesArray.Contains("from"));
-
-                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. Body='{extraMessage?.Body.GetString()}'");
+                var receivedMessages = await TestUtility.ReceiveMessagesAsync(receiver, 1);
+                var receivedMessage = receivedMessages.FirstOrDefault();
+                Assert.NotNull(receivedMessage);
+                Assert.Equal("value", receivedMessage.UserProperties["custom"]);
             }
             finally
             {
                 await sender.CloseAsync().ConfigureAwait(false);
                 await receiver.CloseAsync().ConfigureAwait(false);
             }
         }
+
+        [Fact]
+        [DisplayTestMethodName]
+        public async Task Batch_should_have_maximum_size_set()
+        {
+            var sender = new MessageSender(TestUtility.NamespaceConnectionString, TestConstants.PartitionedQueueName);
+            try
+            {
+                using (var batch = await sender.CreateBatch())
+                {
+                    Assert.True(batch.maximumBatchSize == 256 * 1024 || batch.maximumBatchSize == 1024 * 1024,
+                        $"Maximum batch size was expected to be 256KB or 1MB, but it wasn't. Reported size: {batch.maximumBatchSize}");
+                }
+            }
+            finally
+            {
+                await sender.CloseAsync().ConfigureAwait(false);
+            }
+        }
     }
 }