[Service Bus] Add message batching

sdk/servicebus/Azure.Messaging.ServiceBus/src/Amqp/AmqpMessageBatch.cs

@@ -0,0 +1,162 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Collections.Generic;
+using System.Globalization;
+using System.Linq;
+using Azure.Core;
+using Azure.Messaging.ServiceBus.Core;
+using Microsoft.Azure.Amqp;
+
+namespace Azure.Messaging.ServiceBus.Amqp
+{
+    /// <summary>
+    ///   A set of messages with known size constraints, based on messages to be sent
+    ///   using an AMQP-based transport.
+    /// </summary>
+    ///
+    internal class AmqpMessageBatch : TransportMessageBatch
+    {
+        /// <summary>The amount of bytes to reserve as overhead for a small message.</summary>
+        private const byte OverheadBytesSmallMessage = 5;
+
+        /// <summary>The amount of bytes to reserve as overhead for a large message.</summary>
+        private const byte OverheadBytesLargeMessage = 8;
+
+        /// <summary>The maximum number of bytes that a message may be to be considered small.</summary>
+        private const byte MaximumBytesSmallMessage = 255;
+
+        /// <summary>A flag that indicates whether or not the instance has been disposed.</summary>
+        private bool _disposed = false;
+
+        /// <summary>The size of the batch, in bytes, as it will be sent via the AMQP transport.</summary>
+        private long _sizeBytes = 0;
+
+        /// <summary>
+        ///   The maximum size allowed for the batch, in bytes.  This includes the messages in the batch as
+        ///   well as any overhead for the batch itself when sent to the Queue/Topic.
+        /// </summary>
+        ///
+        public override long MaximumSizeInBytes { get; }
+
+        /// <summary>
+        ///   The size of the batch, in bytes, as it will be sent to the Queue/Topic
+        ///   service.
+        /// </summary>
+        ///
+        public override long SizeInBytes => _sizeBytes;
+
+        /// <summary>
+        ///   The count of messages contained in the batch.
+        /// </summary>
+        ///
+        public override int Count => BatchMessages.Count;
+
+        /// <summary>
+        ///   The set of options to apply to the batch.
+        /// </summary>
+        ///
+        private CreateBatchOptions Options { get; }
+
+        /// <summary>
+        ///   The set of messages that have been added to the batch.
+        /// </summary>
+        ///
+        private List<ServiceBusMessage> BatchMessages { get; } = new List<ServiceBusMessage>();
+
+        /// <summary>
+        ///   Initializes a new instance of the <see cref="AmqpMessageBatch"/> class.
+        /// </summary>
+        ///
+        /// <param name="options">The set of options to apply to the batch.</param>
+        ///
+        public AmqpMessageBatch(CreateBatchOptions options)
+        {
+            Argument.AssertNotNull(options, nameof(options));
+            Argument.AssertNotNull(options.MaximumSizeInBytes, nameof(options.MaximumSizeInBytes));
+
+            Options = options;
+            MaximumSizeInBytes = options.MaximumSizeInBytes.Value;
+
+            // Initialize the size by reserving space for the batch envelope.
+
+            using AmqpMessage envelope = AmqpMessageConverter.BatchSBMessagesAsAmqpMessage(Enumerable.Empty<ServiceBusMessage>());
+            _sizeBytes = envelope.SerializedMessageSize;
+        }
+
+        /// <summary>
+        ///   Attempts to add a message to the batch, ensuring that the size
+        ///   of the batch does not exceed its maximum.
+        /// </summary>
+        ///
+        /// <param name="message">The message to attempt to add to the batch.</param>
+        ///
+        /// <returns><c>true</c> if the message was added; otherwise, <c>false</c>.</returns>
+        ///
+        public override bool TryAdd(ServiceBusMessage message)
+        {
+            Argument.AssertNotNull(message, nameof(message));
+            Argument.AssertNotDisposed(_disposed, nameof(ServiceBusMessageBatch));
+
+            AmqpMessage amqpMessage = AmqpMessageConverter.SBMessageToAmqpMessage(message);
+
+            try
+            {
+                // Calculate the size for the message, based on the AMQP message size and accounting for a
+                // bit of reserved overhead size.
+
+                var size = _sizeBytes
+                    + amqpMessage.SerializedMessageSize
+                    + (amqpMessage.SerializedMessageSize <= MaximumBytesSmallMessage
+                        ? OverheadBytesSmallMessage
+                        : OverheadBytesLargeMessage);
+
+                if (size > MaximumSizeInBytes)
+                {
+                    return false;
+                }
+
+                _sizeBytes = size;
+                BatchMessages.Add(message);
+
+                return true;
+            }
+            finally
+            {
+                amqpMessage?.Dispose();
+            }
+        }
+
+        /// <summary>
+        ///   Represents the batch as an enumerable set of transport-specific
+        ///   representations of a message.
+        /// </summary>
+        ///
+        /// <typeparam name="T">The transport-specific message representation being requested.</typeparam>
+        ///
+        /// <returns>The set of messages as an enumerable of the requested type.</returns>
+        ///
+        public override IEnumerable<T> AsEnumerable<T>()
+        {
+            if (typeof(T) != typeof(ServiceBusMessage))
+            {
+                throw new FormatException(string.Format(CultureInfo.CurrentCulture, Resources1.UnsupportedTransportEventType, typeof(T).Name));
+            }
+
+            return (IEnumerable<T>)BatchMessages;
+        }
+
+        /// <summary>
+        ///   Performs the task needed to clean up resources used by the <see cref="AmqpMessageBatch" />.
+        /// </summary>
+        ///
+        public override void Dispose()
+        {
+            _disposed = true;
+
+            BatchMessages.Clear();
+            _sizeBytes = 0;
+        }
+    }
+}

sdk/servicebus/Azure.Messaging.ServiceBus/src/Amqp/AmqpMessageConverter.cs

@@ -6,9 +6,10 @@ namespace Azure.Messaging.ServiceBus.Amqp
     using System;
     using System.Collections;
     using System.Collections.Generic;
-    using System.Data;
     using System.IO;
+    using System.Linq;
     using System.Runtime.Serialization;
+    using Azure.Core;
     using Microsoft.Azure.Amqp;
     using Microsoft.Azure.Amqp.Encoding;
     using Microsoft.Azure.Amqp.Framing;
@@ -32,70 +33,122 @@ internal static class AmqpMessageConverter
         private const string DateTimeOffsetName = AmqpConstants.Vendor + ":datetime-offset";
         private const int GuidSize = 16;
 
-        public static AmqpMessage BatchSBMessagesAsAmqpMessage(IEnumerable<SBMessage> sbMessages)
+        /// <summary>The size, in bytes, to use as a buffer for stream operations.</summary>
+        private const int StreamBufferSizeInBytes = 512;
+
+        public static AmqpMessage BatchSBMessagesAsAmqpMessage(IEnumerable<SBMessage> source)
         {
-            if (sbMessages == null)
-            {
-                throw Fx.Exception.ArgumentNull(nameof(sbMessages));
-            }
+            Argument.AssertNotNull(source, nameof(source));
+            return BuildAmqpBatchFromMessage(source);
+        }
 
-            AmqpMessage amqpMessage;
+        /// <summary>
+        ///   Builds a batch <see cref="AmqpMessage" /> from a set of <see cref="SBMessage" />
+        ///   optionally propagating the custom properties.
+        /// </summary>
+        ///
+        /// <param name="source">The set of messages to use as the body of the batch message.</param>
+        ///
+        /// <returns>The batch <see cref="AmqpMessage" /> containing the source messages.</returns>
+        ///
+        private static AmqpMessage BuildAmqpBatchFromMessage(IEnumerable<SBMessage> source)
+        {
             AmqpMessage firstAmqpMessage = null;
             SBMessage firstMessage = null;
-            List<Data> dataList = null;
-            var messageCount = 0;
-            foreach (var sbMessage in sbMessages)
-            {
-                messageCount++;
 
-                amqpMessage = AmqpMessageConverter.SBMessageToAmqpMessage(sbMessage);
-                if (firstAmqpMessage == null)
+            return BuildAmqpBatchFromMessages(
+                source.Select(sbMessage =>
                 {
-                    firstAmqpMessage = amqpMessage;
-                    firstMessage = sbMessage;
-                    continue;
-                }
+                    if (firstAmqpMessage == null)
+                    {
+                        firstAmqpMessage = SBMessageToAmqpMessage(sbMessage);
+                        firstMessage = sbMessage;
+                        return firstAmqpMessage;
+                    }
+                    else
+                    {
+                        return SBMessageToAmqpMessage(sbMessage);
+                    }
+                }), firstMessage);
+        }
 
-                if (dataList == null)
-                {
-                    dataList = new List<Data> { ToData(firstAmqpMessage) };
-                }
+        /// <summary>
+        ///   Builds a batch <see cref="AmqpMessage" /> from a set of <see cref="AmqpMessage" />.
+        /// </summary>
+        ///
+        /// <param name="source">The set of messages to use as the body of the batch message.</param>
+        /// <param name="firstMessage"></param>
+        ///
+        /// <returns>The batch <see cref="AmqpMessage" /> containing the source messages.</returns>
+        ///
+        private static AmqpMessage BuildAmqpBatchFromMessages(
+            IEnumerable<AmqpMessage> source,
+            SBMessage firstMessage = null)
+        {
+            AmqpMessage batchEnvelope;
 
-                dataList.Add(ToData(amqpMessage));
-            }
+            var batchMessages = source.ToList();
 
-            if (messageCount == 1 && firstAmqpMessage != null)
+            if (batchMessages.Count == 1)
             {
-                firstAmqpMessage.Batchable = true;
-                return firstAmqpMessage;
+                batchEnvelope = batchMessages[0];
             }
+            else
+            {
+                batchEnvelope = AmqpMessage.Create(batchMessages.Select(message =>
+                {
+                    message.Batchable = true;
+                    using var messageStream = message.ToStream();
+                    return new Data { Value = ReadStreamToArraySegment(messageStream) };
+                }));
 
-            amqpMessage = AmqpMessage.Create(dataList);
-            amqpMessage.MessageFormat = AmqpConstants.AmqpBatchedMessageFormat;
+                batchEnvelope.MessageFormat = AmqpConstants.AmqpBatchedMessageFormat;
+            }
 
-            if (firstMessage.MessageId != null)
+            if (firstMessage?.MessageId != null)
             {
-                amqpMessage.Properties.MessageId = firstMessage.MessageId;
+                batchEnvelope.Properties.MessageId = firstMessage.MessageId;
             }
-            if (firstMessage.SessionId != null)
+            if (firstMessage?.SessionId != null)
             {
-                amqpMessage.Properties.GroupId = firstMessage.SessionId;
+                batchEnvelope.Properties.GroupId = firstMessage.SessionId;
             }
 
-            if (firstMessage.PartitionKey != null)
+            if (firstMessage?.PartitionKey != null)
             {
-                amqpMessage.MessageAnnotations.Map[AmqpMessageConverter.PartitionKeyName] =
+                batchEnvelope.MessageAnnotations.Map[AmqpMessageConverter.PartitionKeyName] =
                     firstMessage.PartitionKey;
             }
 
-            if (firstMessage.ViaPartitionKey != null)
+            if (firstMessage?.ViaPartitionKey != null)
             {
-                amqpMessage.MessageAnnotations.Map[AmqpMessageConverter.ViaPartitionKeyName] =
+                batchEnvelope.MessageAnnotations.Map[AmqpMessageConverter.ViaPartitionKeyName] =
                     firstMessage.ViaPartitionKey;
             }
 
-            amqpMessage.Batchable = true;
-            return amqpMessage;
+            batchEnvelope.Batchable = true;
+            return batchEnvelope;
+        }
+
+        /// <summary>
+        ///   Converts a stream to an <see cref="ArraySegment{T}" /> representation.
+        /// </summary>
+        ///
+        /// <param name="stream">The stream to read and capture in memory.</param>
+        ///
+        /// <returns>The <see cref="ArraySegment{T}" /> containing the stream data.</returns>
+        ///
+        private static ArraySegment<byte> ReadStreamToArraySegment(Stream stream)
+        {
+            if (stream == null)
+            {
+                return new ArraySegment<byte>();
+            }
+
+            using var memStream = new MemoryStream(StreamBufferSizeInBytes);
+            stream.CopyTo(memStream, StreamBufferSizeInBytes);
+
+            return new ArraySegment<byte>(memStream.ToArray());
         }
 
         public static AmqpMessage SBMessageToAmqpMessage(SBMessage sbMessage)