[Messaging] Reduce allocation for Send by using RedableBuffer (Azure#34591)

anuchandy · Liudmila Molkova · web-flow · commit cdd3daa6ccb9 · 2023-05-03T16:29:56.000-07:00
* Using RedableBuffer to send

* better fun name, revert collapsed imports

* Using rewind to seek the position when retrying

Co-authored-by: Liudmila Molkova &lt;limolkova@microsoft.com&gt;

* Productize the ReadableBuffer based apporach

* remove var added for dev-time testing

* Moving code comment to git description

* changelog update

* let's not new-up ReadableBuffer in non-batch cases

* (review feedback): use monoError to report buffer-overflow and log when a message in the batch is empty.

---------

Co-authored-by: Liudmila Molkova &lt;limolkova@microsoft.com&gt;
diff --git a/sdk/core/azure-core-amqp/CHANGELOG.md b/sdk/core/azure-core-amqp/CHANGELOG.md
@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- Addressing the overhead of batch send allocating byte array equal to the max message size. ([34426](https://github.com/Azure/azure-sdk-for-java/issues/34426))
+
 ### Other Changes
 
 ## 2.8.4 (2023-04-07)
diff --git a/sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/ReactorSender.java b/sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/ReactorSender.java
@@ -14,6 +14,7 @@
 import com.azure.core.amqp.implementation.handler.SendLinkHandler;
 import com.azure.core.util.AsyncCloseable;
 import com.azure.core.util.CoreUtils;
+import com.azure.core.util.FluxUtil;
 import com.azure.core.util.logging.ClientLogger;
 import org.apache.qpid.proton.Proton;
 import org.apache.qpid.proton.amqp.Binary;
@@ -26,6 +27,8 @@
 import org.apache.qpid.proton.amqp.transaction.Declared;
 import org.apache.qpid.proton.amqp.transport.DeliveryState;
 import org.apache.qpid.proton.amqp.transport.ErrorCondition;
+import org.apache.qpid.proton.codec.CompositeReadableBuffer;
+import org.apache.qpid.proton.codec.ReadableBuffer;
 import org.apache.qpid.proton.engine.Delivery;
 import org.apache.qpid.proton.engine.EndpointState;
 import org.apache.qpid.proton.engine.Sender;
@@ -44,6 +47,7 @@
 import java.nio.BufferOverflowException;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.Arrays;
 import java.util.Comparator;
 import java.util.List;
 import java.util.Locale;
@@ -265,62 +269,103 @@ public Mono<Void> send(List<Message> messageBatch, DeliveryState deliveryState)
 
         return getLinkSize()
             .flatMap(maxMessageSize -> {
-                final Message firstMessage = messageBatch.get(0);
+                int totalEncodedSize = 0;
+                final CompositeReadableBuffer buffer = new CompositeReadableBuffer();
+
+                final byte[] envelopBytes = batchEnvelopBytes(messageBatch.get(0), maxMessageSize);
+                if (envelopBytes.length > 0) {
+                    totalEncodedSize += envelopBytes.length;
+                    if (totalEncodedSize > maxMessageSize) {
+                        return batchBufferOverflowError(maxMessageSize);
+                    }
+                    buffer.append(envelopBytes);
+                }
 
-                // proton-j doesn't support multiple dataSections to be part of AmqpMessage
-                // here's the alternate approach provided by them: https://github.com/apache/qpid-proton/pull/54
-                final Message batchMessage = Proton.message();
-                batchMessage.setMessageAnnotations(firstMessage.getMessageAnnotations());
+                for (final Message message : messageBatch) {
+                    final byte[] sectionBytes = batchBinaryDataSectionBytes(message, maxMessageSize);
+                    if (sectionBytes.length > 0) {
+                        totalEncodedSize += sectionBytes.length;
+                        if (totalEncodedSize > maxMessageSize) {
+                            return batchBufferOverflowError(maxMessageSize);
+                        }
+                        buffer.append(sectionBytes);
+                    } else {
+                        logger.info("Ignoring the empty message org.apache.qpid.proton.message.message@{} in the batch.",
+                            Integer.toHexString(System.identityHashCode(message)));
+                    }
+                }
 
-                // Set partition identifier properties of the first message on batch message
-                if ((firstMessage.getMessageId() instanceof String)
-                    && !CoreUtils.isNullOrEmpty((String) firstMessage.getMessageId())) {
+                return send(buffer, AmqpConstants.AMQP_BATCH_MESSAGE_FORMAT, deliveryState);
+            }).then();
+    }
 
-                    batchMessage.setMessageId(firstMessage.getMessageId());
-                }
+    private byte[] batchEnvelopBytes(Message envelopMessage, int maxMessageSize) {
+        // Proton-j doesn't support multiple dataSections to be part of AmqpMessage.
+        // Here's the alternate approach provided: https://github.com/apache/qpid-proton/pull/54
+        final Message message = Proton.message();
+        message.setMessageAnnotations(envelopMessage.getMessageAnnotations());
 
-                if (!CoreUtils.isNullOrEmpty(firstMessage.getGroupId())) {
-                    batchMessage.setGroupId(firstMessage.getGroupId());
-                }
+        // Set partition identifier properties of the first message on batch message
+        if ((envelopMessage.getMessageId() instanceof String)
+            && !CoreUtils.isNullOrEmpty((String) envelopMessage.getMessageId())) {
 
-                final int maxMessageSizeTemp = maxMessageSize;
-
-                final byte[] bytes = new byte[maxMessageSizeTemp];
-                int encodedSize = batchMessage.encode(bytes, 0, maxMessageSizeTemp);
-                int byteArrayOffset = encodedSize;
-
-                for (final Message amqpMessage : messageBatch) {
-                    final Message messageWrappedByData = Proton.message();
-
-                    int payloadSize = messageSerializer.getSize(amqpMessage);
-                    int allocationSize =
-                        Math.min(payloadSize + MAX_AMQP_HEADER_SIZE_BYTES, maxMessageSizeTemp);
-
-                    byte[] messageBytes = new byte[allocationSize];
-                    int messageSizeBytes = amqpMessage.encode(messageBytes, 0, allocationSize);
-                    messageWrappedByData.setBody(new Data(new Binary(messageBytes, 0, messageSizeBytes)));
-
-                    try {
-                        encodedSize =
-                            messageWrappedByData
-                                .encode(bytes, byteArrayOffset, maxMessageSizeTemp - byteArrayOffset - 1);
-                    } catch (BufferOverflowException exception) {
-                        final String message =
-                            String.format(Locale.US,
-                                "Size of the payload exceeded maximum message size: %s kb",
-                                maxMessageSizeTemp / 1024);
-                        final AmqpException error = new AmqpException(false,
-                            AmqpErrorCondition.LINK_PAYLOAD_SIZE_EXCEEDED, message, exception,
-                            handler.getErrorContext(sender));
-
-                        return Mono.error(error);
-                    }
+            message.setMessageId(envelopMessage.getMessageId());
+        }
 
-                    byteArrayOffset = byteArrayOffset + encodedSize;
-                }
+        if (!CoreUtils.isNullOrEmpty(envelopMessage.getGroupId())) {
+            message.setGroupId(envelopMessage.getGroupId());
+        }
 
-                return send(bytes, byteArrayOffset, AmqpConstants.AMQP_BATCH_MESSAGE_FORMAT, deliveryState);
-            }).then();
+        final int size = messageSerializer.getSize(message);
+        final int allocationSize = Math.min(size + MAX_AMQP_HEADER_SIZE_BYTES, maxMessageSize);
+        final byte[] encodedBytes = new byte[allocationSize];
+        final int encodedSize = message.encode(encodedBytes, 0, allocationSize);
+        // This copyOf copying is just few bytes for envelop.
+        return Arrays.copyOf(encodedBytes, encodedSize);
+    }
+
+    private byte[] batchBinaryDataSectionBytes(Message sectionMessage, int maxMessageSize) {
+        final int size = messageSerializer.getSize(sectionMessage);
+        final int allocationSize = Math.min(size + MAX_AMQP_HEADER_SIZE_BYTES, maxMessageSize);
+        final byte[] encodedBytes = new byte[allocationSize];
+        final int encodedSize = sectionMessage.encode(encodedBytes, 0, allocationSize);
+
+        final Message message = Proton.message();
+        final Data binaryData = new Data(new Binary(encodedBytes, 0, encodedSize));
+        message.setBody(binaryData);
+        final int binaryRawSize = binaryData.getValue().getLength();
+        // Precompute the "amqp:data:binary" encoded size -
+        final int binaryEncodedSize = binaryEncodedSize(binaryRawSize);
+        // ^ this pre-computation avoids allocating byte[] 'arr1' of estimated encoded size (to pass to message.encode(arr1,))
+        // and a second allocation of byte[] 'arr2' with exact encoded size (returned from message.encode(arr1,)) then
+        // copying encoded size bytes from 'arr1' to 'arr2'. Skipping extra allocations and CPU cycles for copying.
+        final byte[] binaryEncodedBytes = new byte[binaryEncodedSize];
+        message.encode(binaryEncodedBytes, 0, binaryEncodedSize);
+        return binaryEncodedBytes;
+    }
+
+    private Mono<Void> batchBufferOverflowError(int maxMessageSize) {
+        return FluxUtil.monoError(logger, new AmqpException(false, AmqpErrorCondition.LINK_PAYLOAD_SIZE_EXCEEDED,
+            String.format(Locale.US, "Size of the payload exceeded maximum message size: %s kb", maxMessageSize / 1024),
+            new BufferOverflowException(), handler.getErrorContext(sender)));
+    }
+
+    /**
+     * Compute the encoded size when encoding a binary data of given size per Amqp 1.0 spec "amqp:data:binary" format.
+     *
+     * @param binaryRawSize the length of the binary data.
+     * @return the encoded size.
+     */
+    private int binaryEncodedSize(int binaryRawSize) {
+        if (binaryRawSize <= 255) {
+            // [0x00,0x53,0x75,0xa0,{byte(Data.Binary.Length)},{Data.Binary.bytes}]
+            // The AMQP 1.0 spec format ^ for amqp:data:binary when the raw bytes length is <= 255.
+            return 5 + binaryRawSize;
+        } else {
+            // [0x00,0x53,0x75,0xb0,{int(Data.Binary.Length)},{Data.Binary.bytes}]
+            // The AMQP 1.0 spec format ^ for amqp:data:binary when the raw bytes length is > 255.
+            return  8 + binaryRawSize;
+        }
     }
 
     @Override
@@ -457,16 +502,24 @@ Mono<Void> isClosed() {
 
     @Override
     public Mono<DeliveryState> send(byte[] bytes, int arrayOffset, int messageFormat, DeliveryState deliveryState) {
-        final Flux<EndpointState> activeEndpointFlux = RetryUtil.withRetry(
-            handler.getEndpointStates().takeUntil(state -> state == EndpointState.ACTIVE), retryOptions,
-            activeTimeoutMessage);
-
-        return activeEndpointFlux.then(Mono.create(sink -> {
+        return onEndpointActive().then(Mono.create(sink -> {
             sendWork(new RetriableWorkItem(bytes, arrayOffset, messageFormat, sink, retryOptions.getTryTimeout(),
                 deliveryState, metricsProvider));
         }));
     }
 
+    Mono<DeliveryState> send(ReadableBuffer buffer, int messageFormat, DeliveryState deliveryState) {
+        return onEndpointActive().then(Mono.create(sink -> {
+            sendWork(new RetriableWorkItem(buffer, messageFormat, sink, retryOptions.getTryTimeout(),
+                deliveryState, metricsProvider));
+        }));
+    }
+
+    private Flux<EndpointState> onEndpointActive() {
+        return RetryUtil.withRetry(handler.getEndpointStates().takeUntil(state -> state == EndpointState.ACTIVE),
+            retryOptions, activeTimeoutMessage);
+    }
+
     /**
      * Add the work item in pending send to be processed on {@link ReactorDispatcher} thread.
      *
@@ -536,10 +589,7 @@ private void processSendWork() {
                 if (workItem.isDeliveryStateProvided()) {
                     delivery.disposition(workItem.getDeliveryState());
                 }
-                sentMsgSize = sender.send(workItem.getMessage(), 0, workItem.getEncodedMessageSize());
-                assert sentMsgSize == workItem.getEncodedMessageSize()
-                    : "Contract of the ProtonJ library for Sender. Send API changed";
-
+                workItem.send(sender);
                 linkAdvance = sender.advance();
             } catch (Exception exception) {
                 sendException = exception;
diff --git a/sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/RetriableWorkItem.java b/sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/RetriableWorkItem.java
@@ -4,6 +4,8 @@
 package com.azure.core.amqp.implementation;
 
 import org.apache.qpid.proton.amqp.transport.DeliveryState;
+import org.apache.qpid.proton.codec.ReadableBuffer;
+import org.apache.qpid.proton.engine.Sender;
 import reactor.core.publisher.MonoSink;
 
 import java.time.Duration;
@@ -17,7 +19,8 @@ class RetriableWorkItem {
     private final AtomicInteger retryAttempts = new AtomicInteger();
     private final MonoSink<DeliveryState> monoSink;
     private final TimeoutTracker timeoutTracker;
-    private final byte[] amqpMessage;
+    private final ReadableBuffer encodedBuffer;
+    private final byte[] encodedBytes;
     private final int messageFormat;
     private final int encodedMessageSize;
     private final DeliveryState deliveryState;
@@ -27,27 +30,30 @@ class RetriableWorkItem {
     private final AmqpMetricsProvider metricsProvider;
     private long tryStartTime = 0;
 
-    RetriableWorkItem(byte[] amqpMessage, int encodedMessageSize, int messageFormat, MonoSink<DeliveryState> monoSink,
-                      Duration timeout, DeliveryState deliveryState, AmqpMetricsProvider metricsProvider) {
-        this(amqpMessage, encodedMessageSize, messageFormat, monoSink, new TimeoutTracker(timeout,
-            false), deliveryState, metricsProvider);
+    RetriableWorkItem(ReadableBuffer buffer, int messageFormat, MonoSink<DeliveryState> monoSink, Duration timeout,
+        DeliveryState deliveryState, AmqpMetricsProvider metricsProvider) {
+        this.encodedBuffer = buffer;
+        this.encodedBytes = null;
+        this.encodedMessageSize = buffer.remaining();
+        this.messageFormat = messageFormat;
+        this.monoSink = monoSink;
+        this.timeoutTracker = new TimeoutTracker(timeout, false);
+        this.deliveryState = deliveryState;
+        this.metricsProvider = metricsProvider;
     }
 
-    private RetriableWorkItem(byte[] amqpMessage, int encodedMessageSize, int messageFormat, MonoSink<DeliveryState>
-        monoSink, TimeoutTracker timeout, DeliveryState deliveryState, AmqpMetricsProvider metricsProvider) {
-        this.amqpMessage = amqpMessage;
+    RetriableWorkItem(byte[] bytes, int encodedMessageSize, int messageFormat, MonoSink<DeliveryState> monoSink, Duration timeout,
+        DeliveryState deliveryState, AmqpMetricsProvider metricsProvider) {
+        this.encodedBytes = bytes;
+        this.encodedBuffer = null;
         this.encodedMessageSize = encodedMessageSize;
         this.messageFormat = messageFormat;
         this.monoSink = monoSink;
-        this.timeoutTracker = timeout;
+        this.timeoutTracker = new TimeoutTracker(timeout, false);
         this.deliveryState = deliveryState;
         this.metricsProvider = metricsProvider;
     }
 
-    byte[] getMessage() {
-        return amqpMessage;
-    }
-
     DeliveryState getDeliveryState() {
         return deliveryState;
     }
@@ -108,6 +114,17 @@ boolean isWaitingForAck() {
         return this.waitingForAck;
     }
 
+    void send(Sender sender) {
+        final int sentMsgSize;
+        if (encodedBytes != null) {
+            sentMsgSize = sender.send(encodedBytes, 0, encodedMessageSize);
+        } else {
+            encodedBuffer.rewind();
+            sentMsgSize = sender.send(encodedBuffer);
+        }
+        assert sentMsgSize == encodedMessageSize : "Contract of the ProtonJ library for Sender. Send API changed";
+    }
+
     private void reportMetrics(DeliveryState deliveryState) {
         if (metricsProvider.isSendDeliveryEnabled()) {
             metricsProvider.recordSend(tryStartTime, deliveryState == null ? null : deliveryState.getType());
diff --git a/sdk/core/azure-core-amqp/src/test/java/com/azure/core/amqp/implementation/ReactorSenderTest.java b/sdk/core/azure-core-amqp/src/test/java/com/azure/core/amqp/implementation/ReactorSenderTest.java
@@ -25,6 +25,7 @@
 import org.apache.qpid.proton.amqp.transaction.TransactionalState;
 import org.apache.qpid.proton.amqp.transport.DeliveryState;
 import org.apache.qpid.proton.amqp.transport.ErrorCondition;
+import org.apache.qpid.proton.codec.ReadableBuffer;
 import org.apache.qpid.proton.engine.Delivery;
 import org.apache.qpid.proton.engine.EndpointState;
 import org.apache.qpid.proton.engine.Sender;
@@ -107,7 +108,7 @@ public class ReactorSenderTest {
     private Scheduler scheduler;
 
     @Captor
-    private ArgumentCaptor<byte[]> messageBytesCaptor;
+    private ArgumentCaptor<ReadableBuffer> messageBufferCaptor;
 
     @Captor
     private  ArgumentCaptor<DeliveryState> deliveryStateArgumentCaptor;
@@ -318,7 +319,7 @@ public void testSendBatch() {
             reactorProvider, tokenManager, messageSerializer, options, scheduler, AmqpMetricsProvider.noop());
         final ReactorSender spyReactorSender = spy(reactorSender);
 
-        doReturn(Mono.empty()).when(spyReactorSender).send(any(byte[].class), anyInt(), anyInt(), isNull());
+        doReturn(Mono.empty()).when(spyReactorSender).send(any(ReadableBuffer.class), anyInt(), isNull());
 
         // Act
         StepVerifier.create(spyReactorSender.send(Arrays.asList(message, message2)))
@@ -330,14 +331,13 @@ public void testSendBatch() {
 
         // Assert
         verify(sender, times(1)).getRemoteMaxMessageSize();
-        verify(spyReactorSender, times(2)).send(messageBytesCaptor.capture(), anyInt(),
-            anyInt(), isNull());
+        verify(spyReactorSender, times(2)).send(messageBufferCaptor.capture(), anyInt(), isNull());
 
-        assertFalse(messageBytesCaptor.getAllValues().isEmpty());
+        assertFalse(messageBufferCaptor.getAllValues().isEmpty());
 
-        messageBytesCaptor.getAllValues().forEach(delivery -> {
+        messageBufferCaptor.getAllValues().forEach(delivery -> {
             final Message actual = Proton.message();
-            actual.decode(delivery, 0, delivery.length);
+            actual.decode(delivery);
 
             final Object actualMessageId = actual.getMessageId();
             final String actualGroupId = actual.getGroupId();
