Add callback parameter to ByteBuffer basicPublish

To be notified when the ByteBuffer can be released.

Author: acogoluegnes

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

@@ -314,7 +314,7 @@ void basicPublish(String exchange, String routingKey, boolean mandatory, boolean
      * @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);
+        basicPublish(exchange, routingKey, false, false, props, body, null);
     }
 
     /**
@@ -330,7 +330,7 @@ default void basicPublish(String exchange, String routingKey, BasicProperties pr
      */
     default void basicPublish(String exchange, String routingKey, boolean mandatory, BasicProperties props, ByteBuffer body)
             throws IOException {
-        basicPublish(exchange, routingKey, mandatory, false, props, body);
+        basicPublish(exchange, routingKey, mandatory, false, props, body, null);
     }
 
     /**
@@ -348,9 +348,70 @@ default void basicPublish(String exchange, String routingKey, boolean mandatory,
      */
     default void basicPublish(String exchange, String routingKey, boolean mandatory, boolean immediate, BasicProperties props, ByteBuffer body)
             throws IOException {
+        basicPublish(exchange, routingKey, mandatory, immediate, props, body, null);
+    }
+
+    /**
+     * Publish a message with a {@link ByteBuffer} body and a {@link WriteListener} for
+     * write completion notification.
+     *
+     * <p>The listener is called when the network layer has finished writing the body to the
+     * socket. This is useful for off-heap buffers that must be released after the write completes.
+     *
+     * <p><strong>Threading:</strong> the listener may be called on the Netty event loop thread.
+     * It must not perform blocking operations.
+     *
+     * @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
+     * @param listener called when the write completes, may be {@code null}
+     * @throws java.io.IOException if an error is encountered
+     */
+    default void basicPublish(String exchange, String routingKey, BasicProperties props, ByteBuffer body, WriteListener listener)
+            throws IOException {
+        basicPublish(exchange, routingKey, false, false, props, body, listener);
+    }
+
+    /**
+     * Publish a message with a {@link ByteBuffer} body and a {@link WriteListener} for
+     * write completion notification.
+     *
+     * @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
+     * @param listener called when the write completes, may be {@code null}
+     * @throws java.io.IOException if an error is encountered
+     */
+    default void basicPublish(String exchange, String routingKey, boolean mandatory, BasicProperties props, ByteBuffer body, WriteListener listener)
+            throws IOException {
+        basicPublish(exchange, routingKey, mandatory, false, props, body, listener);
+    }
+
+    /**
+     * Publish a message with a {@link ByteBuffer} body and a {@link WriteListener} for
+     * write completion notification.
+     *
+     * @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
+     * @param listener called when the write completes, may be {@code null}
+     * @throws java.io.IOException if an error is encountered
+     */
+    default void basicPublish(String exchange, String routingKey, boolean mandatory, boolean immediate, BasicProperties props, ByteBuffer body, WriteListener listener)
+            throws IOException {
         byte[] bytes = new byte[body.remaining()];
         body.get(bytes);
         basicPublish(exchange, routingKey, mandatory, immediate, props, bytes);
+        if (listener != null) {
+            listener.done(true, null);
+        }
     }
 
     /**

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

@@ -0,0 +1,39 @@
+// Copyright (c) 2007-2026 Broadcom. All Rights Reserved. The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
+//
+// This software, the RabbitMQ Java client library, is triple-licensed under the
+// Mozilla Public License 2.0 ("MPL"), the GNU General Public License version 2
+// ("GPL") and the Apache License version 2 ("ASL"). For the MPL, please see
+// LICENSE-MPL-RabbitMQ. For the GPL, please see LICENSE-GPL2.  For the ASL,
+// please see LICENSE-APACHE2.
+//
+// This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
+// either express or implied. See the LICENSE file for specific language governing
+// rights and limitations of this software.
+//
+// If you have any questions regarding licensing, please contact us at
+// info@rabbitmq.com.
+
+package com.rabbitmq.client;
+
+/**
+ * Listener notified when the network layer has finished writing a message body.
+ *
+ * <p>This is primarily useful when publishing with an off-heap {@link java.nio.ByteBuffer}:
+ * the callback tells the caller when the buffer is no longer in use and can be safely released.
+ *
+ * <p><strong>Threading:</strong> the callback may be invoked on the Netty event loop thread.
+ * Implementations must not perform blocking operations.
+ *
+ * @since 5.25.0
+ */
+@FunctionalInterface
+public interface WriteListener {
+
+  /**
+   * Called when the write operation completes.
+   *
+   * @param success {@code true} if the data was written to the socket successfully
+   * @param cause the exception if the write failed, {@code null} on success
+   */
+  void done(boolean success, Throwable cause);
+}

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

@@ -24,6 +24,7 @@
 
 import com.rabbitmq.client.AMQP;
 import com.rabbitmq.client.Command;
+import com.rabbitmq.client.WriteListener;
 
 /**
  * AMQP 0-9-1-specific implementation of {@link Command} which accumulates
@@ -47,10 +48,12 @@ public class AMQCommand implements Command {
     /** The assembler for this command - synchronised on - contains all the state */
     private final CommandAssembler assembler;
     private final Lock assemblerLock = new ReentrantLock();
+    private final WriteListener writeListener;
 
     /** Construct a command for inbound frame assembly with a max body length. */
     AMQCommand(int maxBodyLength) {
         this.assembler = new CommandAssembler(null, null, maxBodyLength);
+        this.writeListener = null;
     }
 
     /** Construct a command ready to fill in by reading frames. */
@@ -64,6 +67,7 @@ public AMQCommand() {
      */
     public AMQCommand(com.rabbitmq.client.Method method) {
         this.assembler = new CommandAssembler((Method) method, null, Integer.MAX_VALUE);
+        this.writeListener = null;
     }
 
     /**
@@ -73,7 +77,20 @@ public AMQCommand(com.rabbitmq.client.Method method) {
      * @param body the message body as a ByteBuffer
      */
     public AMQCommand(com.rabbitmq.client.Method method, AMQContentHeader contentHeader, ByteBuffer body) {
+        this(method, contentHeader, body, null);
+    }
+
+    /**
+     * Construct a command with a ByteBuffer body and a write completion listener.
+     * @param method the wrapped method
+     * @param contentHeader the wrapped content header
+     * @param body the message body as a ByteBuffer
+     * @param writeListener called when the network layer finishes writing, may be {@code null}
+     */
+    public AMQCommand(com.rabbitmq.client.Method method, AMQContentHeader contentHeader, ByteBuffer body,
+                      WriteListener writeListener) {
         this.assembler = new CommandAssembler((Method) method, contentHeader, body);
+        this.writeListener = writeListener;
     }
 
     /** Public API - {@inheritDoc} */
@@ -143,7 +160,7 @@ public void transmit(AMQChannel channel) throws IOException {
             assemblerLock.unlock();
         }
 
-        connection.flush();
+        connection.flush(this.writeListener);
     }
 
     @Override public String toString() {

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

@@ -655,8 +655,12 @@ void writeFrame(Frame f) throws IOException {
      * Public API - flush the output buffers
      */
     public void flush() throws IOException {
+        flush(null);
+    }
+
+    void flush(WriteListener listener) throws IOException {
         try {
-            _frameHandler.flush();
+            _frameHandler.flush(listener);
         } catch (IOException ioe) {
             this.errorOnWriteListener.handle(this, ioe);
         }

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

@@ -717,7 +717,7 @@ public void basicPublish(String exchange, String routingKey,
         throws IOException
     {
         basicPublish(exchange, routingKey, mandatory, immediate, props,
-                body == null ? null : ByteBuffer.wrap(body));
+                body == null ? null : ByteBuffer.wrap(body), null);
     }
 
     /** Public API - {@inheritDoc} */
@@ -726,7 +726,7 @@ public void basicPublish(String exchange, String routingKey,
                              BasicProperties props, ByteBuffer body)
         throws IOException
     {
-        basicPublish(exchange, routingKey, false, false, props, body);
+        basicPublish(exchange, routingKey, false, false, props, body, null);
     }
 
     /** Public API - {@inheritDoc} */
@@ -736,7 +736,7 @@ public void basicPublish(String exchange, String routingKey,
                              BasicProperties props, ByteBuffer body)
         throws IOException
     {
-        basicPublish(exchange, routingKey, mandatory, false, props, body);
+        basicPublish(exchange, routingKey, mandatory, false, props, body, null);
     }
 
     /** Public API - {@inheritDoc} */
@@ -745,6 +745,38 @@ public void basicPublish(String exchange, String routingKey,
                              boolean mandatory, boolean immediate,
                              BasicProperties props, ByteBuffer body)
         throws IOException
+    {
+        basicPublish(exchange, routingKey, mandatory, immediate, props, body, null);
+    }
+
+    /** Public API - {@inheritDoc} */
+    @Override
+    public void basicPublish(String exchange, String routingKey,
+                             BasicProperties props, ByteBuffer body,
+                             WriteListener listener)
+        throws IOException
+    {
+        basicPublish(exchange, routingKey, false, false, props, body, listener);
+    }
+
+    /** Public API - {@inheritDoc} */
+    @Override
+    public void basicPublish(String exchange, String routingKey,
+                             boolean mandatory,
+                             BasicProperties props, ByteBuffer body,
+                             WriteListener listener)
+        throws IOException
+    {
+        basicPublish(exchange, routingKey, mandatory, false, props, body, listener);
+    }
+
+    /** Public API - {@inheritDoc} */
+    @Override
+    public void basicPublish(String exchange, String routingKey,
+                             boolean mandatory, boolean immediate,
+                             BasicProperties props, ByteBuffer body,
+                             WriteListener listener)
+        throws IOException
     {
         final long deliveryTag;
         if (nextPublishSeqNo > 0) {
@@ -765,7 +797,7 @@ public void basicPublish(String exchange, String routingKey,
                 .build();
         try {
             ObservationCollector.PublishCall publishCall = properties -> {
-                AMQCommand command = new AMQCommand(publish, properties, body);
+                AMQCommand command = new AMQCommand(publish, properties, body, listener);
                 transmit(command);
             };
             observationCollector.publish(publishCall, publish, props, body, this.connectionInfo());

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

@@ -15,6 +15,7 @@
 
 package com.rabbitmq.client.impl;
 
+import com.rabbitmq.client.WriteListener;
 import java.io.IOException;
 import java.net.SocketException;
 import java.net.SocketTimeoutException;
@@ -79,6 +80,19 @@ default void finishConnectionNegotiation() {
      */
     void flush() throws IOException;
 
+    /**
+     * Flush the underlying data connection and notify the listener when the write completes.
+     * The default implementation flushes synchronously and calls the listener immediately.
+     * @param listener called when the flush completes, may be {@code null}
+     * @throws IOException if there is a problem accessing the connection
+     */
+    default void flush(WriteListener listener) throws IOException {
+        flush();
+        if (listener != null) {
+            listener.done(true, null);
+        }
+    }
+
     /** Close the underlying data connection (complaint not permitted). */
     void close();
 

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

@@ -25,6 +25,7 @@
 import com.rabbitmq.client.MalformedFrameException;
 import com.rabbitmq.client.ShutdownSignalException;
 import com.rabbitmq.client.SocketConfigurator;
+import com.rabbitmq.client.WriteListener;
 import io.netty.bootstrap.Bootstrap;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
@@ -34,6 +35,7 @@
 import io.netty.channel.ChannelInboundHandlerAdapter;
 import io.netty.channel.ChannelInitializer;
 import io.netty.channel.ChannelOption;
+import io.netty.channel.ChannelPromise;
 import io.netty.channel.EventLoopGroup;
 import io.netty.channel.socket.SocketChannel;
 import io.netty.channel.socket.nio.NioSocketChannel;
@@ -408,6 +410,17 @@ public void flush() {
       this.channel.flush();
     }
 
+    @Override
+    public void flush(WriteListener listener) {
+      if (listener == null) {
+        this.channel.flush();
+      } else {
+        ChannelPromise promise = this.channel.newPromise();
+        promise.addListener(future -> listener.done(future.isSuccess(), future.cause()));
+        this.channel.writeAndFlush(Unpooled.EMPTY_BUFFER, promise);
+      }
+    }
+
     @Override
     public void close() {
       if (this.closed.compareAndSet(false, true)) {

src/main/java/com/rabbitmq/client/impl/recovery/AutorecoveringChannel.java

@@ -237,6 +237,21 @@ public void basicPublish(String exchange, String routingKey, boolean mandatory,
         delegate.basicPublish(exchange, routingKey, mandatory, immediate, props, body);
     }
 
+    @Override
+    public void basicPublish(String exchange, String routingKey, AMQP.BasicProperties props, ByteBuffer body, WriteListener listener) throws IOException {
+        delegate.basicPublish(exchange, routingKey, props, body, listener);
+    }
+
+    @Override
+    public void basicPublish(String exchange, String routingKey, boolean mandatory, AMQP.BasicProperties props, ByteBuffer body, WriteListener listener) throws IOException {
+        delegate.basicPublish(exchange, routingKey, mandatory, props, body, listener);
+    }
+
+    @Override
+    public void basicPublish(String exchange, String routingKey, boolean mandatory, boolean immediate, AMQP.BasicProperties props, ByteBuffer body, WriteListener listener) throws IOException {
+        delegate.basicPublish(exchange, routingKey, mandatory, immediate, props, body, listener);
+    }
+
     @Override
     public AMQP.Exchange.DeclareOk exchangeDeclare(String exchange, String type) throws IOException {
         return exchangeDeclare(exchange, type, false, false, null);

src/test/java/com/rabbitmq/client/test/ClientTestSuite.java

@@ -74,7 +74,9 @@
     BlockedConnectionTest.class,
     NettyTest.class,
     IoDeadlockOnConnectionClosing.class,
-    ProtocolVersionMismatch.class
+    ProtocolVersionMismatch.class,
+    ByteBufferPublishTest.class,
+    PublishWithByteBufferTest.class
 })
 public class ClientTestSuite {