Add basicPublish with ByteBuffer

Author: acogoluegnes

pom.xml

@@ -797,6 +797,7 @@
               <include>src/test/java/com/rabbitmq/client/test/ProtocolVersionMismatch.java</include>
               <include>src/test/java/com/rabbitmq/client/test/TestUtils.java</include>
               <include>src/test/java/com/rabbitmq/client/test/RpcTopologyRecordingTest.java</include>
+              <include>src/test/java/com/rabbitmq/client/test/PublishWithByteBufferTest.java</include>
             </includes>
             <googleJavaFormat>
               <version>${google-java-format.version}</version>

src/main/java/com/rabbitmq/client/Channel.java

@@ -19,6 +19,7 @@
 import com.rabbitmq.client.AMQP.*;
 
 import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeoutException;
@@ -299,6 +300,59 @@ void basicPublish(String exchange, String routingKey, boolean mandatory, BasicPr
     void basicPublish(String exchange, String routingKey, boolean mandatory, boolean immediate, BasicProperties props, byte[] body)
             throws IOException;
 
+    /**
+     * Publish a message with a {@link ByteBuffer} body for zero-copy transmission.
+     *
+     * <p>The buffer's content between its current position and its limit will be sent.
+     * The caller must not modify the buffer after this call returns.
+     *
+     * @see #basicPublish(String, String, BasicProperties, byte[])
+     * @param exchange the exchange to publish the message to
+     * @param routingKey the routing key
+     * @param props other properties for the message - routing headers etc
+     * @param body the message body
+     * @throws java.io.IOException if an error is encountered
+     */
+    default void basicPublish(String exchange, String routingKey, BasicProperties props, ByteBuffer body) throws IOException {
+        basicPublish(exchange, routingKey, false, false, props, body);
+    }
+
+    /**
+     * Publish a message with a {@link ByteBuffer} body for zero-copy transmission.
+     *
+     * @see #basicPublish(String, String, boolean, BasicProperties, byte[])
+     * @param exchange the exchange to publish the message to
+     * @param routingKey the routing key
+     * @param mandatory true if the 'mandatory' flag is to be set
+     * @param props other properties for the message - routing headers etc
+     * @param body the message body
+     * @throws java.io.IOException if an error is encountered
+     */
+    default void basicPublish(String exchange, String routingKey, boolean mandatory, BasicProperties props, ByteBuffer body)
+            throws IOException {
+        basicPublish(exchange, routingKey, mandatory, false, props, body);
+    }
+
+    /**
+     * Publish a message with a {@link ByteBuffer} body for zero-copy transmission.
+     *
+     * @see #basicPublish(String, String, boolean, boolean, BasicProperties, byte[])
+     * @param exchange the exchange to publish the message to
+     * @param routingKey the routing key
+     * @param mandatory true if the 'mandatory' flag is to be set
+     * @param immediate true if the 'immediate' flag is to be
+     * set. Note that the RabbitMQ server does not support this flag.
+     * @param props other properties for the message - routing headers etc
+     * @param body the message body
+     * @throws java.io.IOException if an error is encountered
+     */
+    default void basicPublish(String exchange, String routingKey, boolean mandatory, boolean immediate, BasicProperties props, ByteBuffer body)
+            throws IOException {
+        byte[] bytes = new byte[body.remaining()];
+        body.get(bytes);
+        basicPublish(exchange, routingKey, mandatory, immediate, props, bytes);
+    }
+
     /**
      * Actively declare a non-autodelete, non-durable exchange with no extra arguments
      * @see com.rabbitmq.client.AMQP.Exchange.Declare

src/main/java/com/rabbitmq/client/impl/AMQCommand.java

@@ -18,6 +18,7 @@
 import java.io.ByteArrayOutputStream;
 import java.io.DataOutputStream;
 import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 
@@ -86,6 +87,16 @@ public AMQCommand(com.rabbitmq.client.Method method, AMQContentHeader contentHea
         this.assembler = new CommandAssembler((Method) method, contentHeader, body, maxBodyLength);
     }
 
+    /**
+     * Construct a command with a ByteBuffer body for zero-copy transmission.
+     * @param method the wrapped method
+     * @param contentHeader the wrapped content header
+     * @param body the message body as a ByteBuffer
+     */
+    public AMQCommand(com.rabbitmq.client.Method method, AMQContentHeader contentHeader, ByteBuffer body) {
+        this.assembler = new CommandAssembler((Method) method, contentHeader, body);
+    }
+
     /** Public API - {@inheritDoc} */
     @Override
     public Method getMethod() {
@@ -122,29 +133,11 @@ public void transmit(AMQChannel channel) throws IOException {
         try {
             Method m = this.assembler.getMethod();
             if (m.hasContent()) {
-                byte[] body = this.assembler.getContentBody();
-
-                Frame headerFrame = this.assembler.getContentHeader().toFrame(channelNumber, body.length);
-
-                int frameMax = connection.getFrameMax();
-                boolean cappedFrameMax = frameMax > 0;
-                int bodyPayloadMax = cappedFrameMax ? frameMax - EMPTY_FRAME_SIZE : body.length;
-
-                if (cappedFrameMax && headerFrame.size() > frameMax) {
-                    String msg = String.format("Content headers exceeded max frame size: %d > %d", headerFrame.size(), frameMax);
-                    throw new IllegalArgumentException(msg);
-                }
-                connection.writeFrame(m.toFrame(channelNumber));
-                connection.writeFrame(headerFrame);
-
-                for (int offset = 0; offset < body.length; offset += bodyPayloadMax) {
-                    int remaining = body.length - offset;
-
-                    int fragmentLength = (remaining < bodyPayloadMax) ? remaining
-                            : bodyPayloadMax;
-                    Frame frame = Frame.fromBodyFragment(channelNumber, body,
-                            offset, fragmentLength);
-                    connection.writeFrame(frame);
+                ByteBuffer bbBody = this.assembler.getByteBufferBody();
+                if (bbBody != null) {
+                    transmitWithByteBuffer(m, bbBody, channelNumber, connection);
+                } else {
+                    transmitWithByteArray(m, this.assembler.getContentBody(), channelNumber, connection);
                 }
             } else {
                 connection.writeFrame(m.toFrame(channelNumber));
@@ -156,6 +149,52 @@ public void transmit(AMQChannel channel) throws IOException {
         connection.flush();
     }
 
+    private void transmitWithByteArray(Method m, byte[] body, int channelNumber, AMQConnection connection) throws IOException {
+        Frame headerFrame = this.assembler.getContentHeader().toFrame(channelNumber, body.length);
+
+        int frameMax = connection.getFrameMax();
+        boolean cappedFrameMax = frameMax > 0;
+        int bodyPayloadMax = cappedFrameMax ? frameMax - EMPTY_FRAME_SIZE : body.length;
+
+        if (cappedFrameMax && headerFrame.size() > frameMax) {
+            String msg = String.format("Content headers exceeded max frame size: %d > %d", headerFrame.size(), frameMax);
+            throw new IllegalArgumentException(msg);
+        }
+        connection.writeFrame(m.toFrame(channelNumber));
+        connection.writeFrame(headerFrame);
+
+        for (int offset = 0; offset < body.length; offset += bodyPayloadMax) {
+            int remaining = body.length - offset;
+            int fragmentLength = (remaining < bodyPayloadMax) ? remaining : bodyPayloadMax;
+            Frame frame = Frame.fromBodyFragment(channelNumber, body, offset, fragmentLength);
+            connection.writeFrame(frame);
+        }
+    }
+
+    private void transmitWithByteBuffer(Method m, ByteBuffer body, int channelNumber, AMQConnection connection) throws IOException {
+        int bodySize = body.remaining();
+        Frame headerFrame = this.assembler.getContentHeader().toFrame(channelNumber, bodySize);
+
+        int frameMax = connection.getFrameMax();
+        boolean cappedFrameMax = frameMax > 0;
+        int bodyPayloadMax = cappedFrameMax ? frameMax - EMPTY_FRAME_SIZE : bodySize;
+
+        if (cappedFrameMax && headerFrame.size() > frameMax) {
+            String msg = String.format("Content headers exceeded max frame size: %d > %d", headerFrame.size(), frameMax);
+            throw new IllegalArgumentException(msg);
+        }
+        connection.writeFrame(m.toFrame(channelNumber));
+        connection.writeFrame(headerFrame);
+
+        int bodyPosition = body.position();
+        for (int offset = 0; offset < bodySize; offset += bodyPayloadMax) {
+            int remaining = bodySize - offset;
+            int fragmentLength = (remaining < bodyPayloadMax) ? remaining : bodyPayloadMax;
+            Frame frame = Frame.fromBodyFragment(channelNumber, body, bodyPosition + offset, fragmentLength);
+            connection.writeFrame(frame);
+        }
+    }
+
     @Override public String toString() {
         return toString(false);
     }

src/main/java/com/rabbitmq/client/impl/ChannelN.java

@@ -28,6 +28,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.util.*;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CopyOnWriteArrayList;
@@ -745,6 +746,62 @@ public void basicPublish(String exchange, String routingKey,
         metricsCollector.basicPublish(this, deliveryTag);
     }
 
+    /** Public API - {@inheritDoc} */
+    @Override
+    public void basicPublish(String exchange, String routingKey,
+                             BasicProperties props, ByteBuffer body)
+        throws IOException
+    {
+        basicPublish(exchange, routingKey, false, false, props, body);
+    }
+
+    /** Public API - {@inheritDoc} */
+    @Override
+    public void basicPublish(String exchange, String routingKey,
+                             boolean mandatory,
+                             BasicProperties props, ByteBuffer body)
+        throws IOException
+    {
+        basicPublish(exchange, routingKey, mandatory, false, props, body);
+    }
+
+    /** Public API - {@inheritDoc} */
+    @Override
+    public void basicPublish(String exchange, String routingKey,
+                             boolean mandatory, boolean immediate,
+                             BasicProperties props, ByteBuffer body)
+        throws IOException
+    {
+        final long deliveryTag;
+        if (nextPublishSeqNo > 0) {
+            deliveryTag = getNextPublishSeqNo();
+            unconfirmedSet.add(deliveryTag);
+            nextPublishSeqNo++;
+        } else {
+            deliveryTag = 0;
+        }
+        if (props == null) {
+            props = MessageProperties.MINIMAL_BASIC;
+        }
+        AMQP.Basic.Publish publish = new Basic.Publish.Builder()
+                .exchange(exchange)
+                .routingKey(routingKey)
+                .mandatory(mandatory)
+                .immediate(immediate)
+                .build();
+        try {
+            ObservationCollector.PublishCall publishCall = properties -> {
+                AMQCommand command = new AMQCommand(publish, properties, body);
+                transmit(command);
+            };
+            observationCollector.publish(publishCall, publish, props, body, this.connectionInfo());
+        } catch (IOException | AlreadyClosedException e) {
+            metricsCollector.basicPublishFailure(this, e);
+            throw e;
+        }
+        metricsCollector.basicPublish(this, deliveryTag);
+    }
+
     /** Public API - {@inheritDoc} */
     @Override
     public Exchange.DeclareOk exchangeDeclare(String exchange, String type,

src/main/java/com/rabbitmq/client/impl/CommandAssembler.java

@@ -16,6 +16,7 @@
 package com.rabbitmq.client.impl;
 
 import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -50,6 +51,9 @@ private enum CAState {
     /** sum of the lengths of all fragments */
     private int bodyLength;
 
+    /** Zero-copy body buffer, mutually exclusive with bodyN usage on the outbound path */
+    private final ByteBuffer byteBufferBody;
+
     /** No bytes of content body not yet accumulated */
     private long remainingBodyBytes;
 
@@ -61,6 +65,7 @@ public CommandAssembler(Method method, AMQContentHeader contentHeader, byte[] bo
         this.contentHeader = contentHeader;
         this.bodyN = new ArrayList<>(2);
         this.bodyLength = 0;
+        this.byteBufferBody = null;
         this.remainingBodyBytes = 0;
         this.maxBodyLength = maxBodyLength;
         appendBodyFragment(body);
@@ -74,6 +79,24 @@ public CommandAssembler(Method method, AMQContentHeader contentHeader, byte[] bo
         }
     }
 
+    public CommandAssembler(Method method, AMQContentHeader contentHeader, ByteBuffer body) {
+        this.method = method;
+        this.contentHeader = contentHeader;
+        this.bodyN = new ArrayList<>(0);
+        this.bodyLength = body == null ? 0 : body.remaining();
+        this.byteBufferBody = body;
+        this.remainingBodyBytes = 0;
+        this.maxBodyLength = Integer.MAX_VALUE;
+        if (method == null) {
+            this.state = CAState.EXPECTING_METHOD;
+        } else if (contentHeader == null) {
+            this.state = method.hasContent() ? CAState.EXPECTING_CONTENT_HEADER : CAState.COMPLETE;
+        } else {
+            this.remainingBodyBytes = contentHeader.getBodySize() - this.bodyLength;
+            updateContentBodyState();
+        }
+    }
+
     public synchronized Method getMethod() {
         return this.method;
     }
@@ -151,9 +174,19 @@ private byte[] coalesceContentBody() {
     }
 
     public synchronized byte[] getContentBody() {
+        if (byteBufferBody != null) {
+            ByteBuffer dup = byteBufferBody.duplicate();
+            byte[] result = new byte[dup.remaining()];
+            dup.get(result);
+            return result;
+        }
         return coalesceContentBody();
     }
 
+    public synchronized ByteBuffer getByteBufferBody() {
+        return this.byteBufferBody;
+    }
+
     private void appendBodyFragment(byte[] fragment) {
         if (fragment == null || fragment.length == 0) return;
         bodyN.add(fragment);