Add ByteBufferFactory for NIO mode

This allows to create heap-based or direct ByteBuffer instances.

References #410

(cherry picked from commit bb2c47a)
(cherry picked from commit 0252e2c)

Author: acogoluegnes

src/main/java/com/rabbitmq/client/impl/nio/ByteBufferFactory.java

@@ -0,0 +1,57 @@
+package com.rabbitmq.client.impl.nio;
+
+import java.nio.ByteBuffer;
+
+/**
+ * Contract to create {@link ByteBuffer}s.
+ *
+ * @see NioParams
+ * @since 5.5.0
+ */
+public interface ByteBufferFactory {
+
+    /**
+     * Create the {@link ByteBuffer} that contains inbound frames.
+     * This buffer is the network buffer for plain connections.
+     * When using SSL/TLS, this buffer isn't directly connected to
+     * the network, the encrypted read buffer is.
+     *
+     * @param nioContext
+     * @return
+     */
+    ByteBuffer createReadBuffer(NioContext nioContext);
+
+    /**
+     * Create the {@link ByteBuffer} that contains outbound frames.
+     * This buffer is the network buffer for plain connections.
+     * When using SSL/TLS, this buffer isn't directed connected to
+     * the network, the encrypted write buffer is.
+     *
+     * @param nioContext
+     * @return
+     */
+    ByteBuffer createWriteBuffer(NioContext nioContext);
+
+    /**
+     * Create the network read {@link ByteBuffer}.
+     * This buffer contains encrypted frames read from the network.
+     * The {@link javax.net.ssl.SSLEngine} decrypts frame and pass them
+     * over to the read buffer.
+     *
+     * @param nioContext
+     * @return
+     */
+    ByteBuffer createEncryptedReadBuffer(NioContext nioContext);
+
+    /**
+     * Create the network write {@link ByteBuffer}.
+     * This buffer contains encrypted outbound frames. These
+     * frames come from the write buffer that sends them through
+     * the {@link javax.net.ssl.SSLContext} for encryption to
+     * this buffer.
+     *
+     * @param nioContext
+     * @return
+     */
+    ByteBuffer createEncryptedWriteBuffer(NioContext nioContext);
+}

src/main/java/com/rabbitmq/client/impl/nio/DefaultByteBufferFactory.java

@@ -0,0 +1,62 @@
+package com.rabbitmq.client.impl.nio;
+
+import java.nio.ByteBuffer;
+import java.util.function.Function;
+
+/**
+ * Default {@link ByteBufferFactory} that creates heap-based {@link ByteBuffer}s.
+ * This behavior can be changed by passing in a custom {@link Function<Integer, ByteBuffer>}
+ * to the constructor.
+ *
+ * @see NioParams
+ * @see ByteBufferFactory
+ * @since 5.5.0
+ */
+public class DefaultByteBufferFactory implements ByteBufferFactory {
+
+    private final Function<Integer, ByteBuffer> allocator;
+
+    public DefaultByteBufferFactory(Function<Integer, ByteBuffer> allocator) {
+        this.allocator = allocator;
+    }
+
+    public DefaultByteBufferFactory() {
+        this(capacity -> ByteBuffer.allocate(capacity));
+    }
+
+    @Override
+    public ByteBuffer createReadBuffer(NioContext nioContext) {
+        if (nioContext.getSslEngine() == null) {
+            return allocator.apply(nioContext.getNioParams().getReadByteBufferSize());
+        } else {
+            return allocator.apply(nioContext.getSslEngine().getSession().getApplicationBufferSize());
+        }
+    }
+
+    @Override
+    public ByteBuffer createWriteBuffer(NioContext nioContext) {
+        if (nioContext.getSslEngine() == null) {
+            return allocator.apply(nioContext.getNioParams().getWriteByteBufferSize());
+        } else {
+            return allocator.apply(nioContext.getSslEngine().getSession().getApplicationBufferSize());
+        }
+    }
+
+    @Override
+    public ByteBuffer createEncryptedReadBuffer(NioContext nioContext) {
+        return createEncryptedByteBuffer(nioContext);
+    }
+
+    @Override
+    public ByteBuffer createEncryptedWriteBuffer(NioContext nioContext) {
+        return createEncryptedByteBuffer(nioContext);
+    }
+
+    protected ByteBuffer createEncryptedByteBuffer(NioContext nioContext) {
+        if (nioContext.getSslEngine() == null) {
+            throw new IllegalArgumentException("Encrypted byte buffer should be created only in SSL/TLS context");
+        } else {
+            return allocator.apply(nioContext.getSslEngine().getSession().getPacketBufferSize());
+        }
+    }
+}

src/main/java/com/rabbitmq/client/impl/nio/NioContext.java

@@ -0,0 +1,40 @@
+package com.rabbitmq.client.impl.nio;
+
+import javax.net.ssl.SSLEngine;
+
+/**
+ * Context when creating resources for a NIO-based connection.
+ *
+ * @see ByteBufferFactory
+ * @since 5.5.0
+ */
+public class NioContext {
+
+    private final NioParams nioParams;
+
+    private final SSLEngine sslEngine;
+
+    NioContext(NioParams nioParams, SSLEngine sslEngine) {
+        this.nioParams = nioParams;
+        this.sslEngine = sslEngine;
+    }
+
+    /**
+     * NIO params.
+     *
+     * @return
+     */
+    public NioParams getNioParams() {
+        return nioParams;
+    }
+
+    /**
+     * {@link SSLEngine} for SSL/TLS connection.
+     * Null for plain connection.
+     *
+     * @return
+     */
+    public SSLEngine getSslEngine() {
+        return sslEngine;
+    }
+}

src/main/java/com/rabbitmq/client/impl/nio/NioLoopContext.java

@@ -48,8 +48,9 @@ public NioLoopContext(SocketChannelFrameHandlerFactory socketChannelFrameHandler
         this.socketChannelFrameHandlerFactory = socketChannelFrameHandlerFactory;
         this.executorService = nioParams.getNioExecutor();
         this.threadFactory = nioParams.getThreadFactory();
-        this.readBuffer = ByteBuffer.allocate(nioParams.getReadByteBufferSize());
-        this.writeBuffer = ByteBuffer.allocate(nioParams.getWriteByteBufferSize());
+        NioContext nioContext = new NioContext(nioParams, null);
+        this.readBuffer = nioParams.getByteBufferFactory().createReadBuffer(nioContext);
+        this.writeBuffer = nioParams.getByteBufferFactory().createWriteBuffer(nioContext);
     }
 
     void initStateIfNecessary() throws IOException {

src/main/java/com/rabbitmq/client/impl/nio/NioParams.java

@@ -91,6 +91,14 @@ public void configure(SSLEngine sslEngine) throws IOException {
      */
     private ExecutorService connectionShutdownExecutor;
 
+    /**
+     * The factory to create {@link java.nio.ByteBuffer}s.
+     * The default is to create heap-based {@link java.nio.ByteBuffer}s.
+     *
+     * @since 5.5.0
+     */
+    private ByteBufferFactory byteBufferFactory = new DefaultByteBufferFactory();
+
     public NioParams() {
     }
 
@@ -105,6 +113,7 @@ public NioParams(NioParams nioParams) {
         setSocketChannelConfigurator(nioParams.getSocketChannelConfigurator());
         setSslEngineConfigurator(nioParams.getSslEngineConfigurator());
         setConnectionShutdownExecutor(nioParams.getConnectionShutdownExecutor());
+        setByteBufferFactory(nioParams.getByteBufferFactory());
     }
 
     /**
@@ -364,4 +373,24 @@ public NioParams setConnectionShutdownExecutor(ExecutorService connectionShutdow
         this.connectionShutdownExecutor = connectionShutdownExecutor;
         return this;
     }
+
+    /**
+     * Set the factory to create {@link java.nio.ByteBuffer}s.
+     * <p>
+     * The default implementation creates heap-based {@link java.nio.ByteBuffer}s.
+     *
+     * @param byteBufferFactory the factory to use
+     * @return this {@link NioParams} instance
+     * @see ByteBufferFactory
+     * @see DefaultByteBufferFactory
+     * @since 5.5.0
+     */
+    public NioParams setByteBufferFactory(ByteBufferFactory byteBufferFactory) {
+        this.byteBufferFactory = byteBufferFactory;
+        return this;
+    }
+
+    public ByteBufferFactory getByteBufferFactory() {
+        return byteBufferFactory;
+    }
 }

src/main/java/com/rabbitmq/client/impl/nio/SocketChannelFrameHandlerState.java

@@ -98,10 +98,11 @@ public SocketChannelFrameHandlerState(SocketChannel channel, NioLoopContext nioL
 
         } else {
             this.ssl = true;
-            this.plainOut = ByteBuffer.allocate(sslEngine.getSession().getApplicationBufferSize());
-            this.cipherOut = ByteBuffer.allocate(sslEngine.getSession().getPacketBufferSize());
-            this.plainIn = ByteBuffer.allocate(sslEngine.getSession().getApplicationBufferSize());
-            this.cipherIn = ByteBuffer.allocate(sslEngine.getSession().getPacketBufferSize());
+            NioContext nioContext = new NioContext(nioParams, sslEngine);
+            this.plainOut = nioParams.getByteBufferFactory().createWriteBuffer(nioContext);
+            this.cipherOut = nioParams.getByteBufferFactory().createEncryptedWriteBuffer(nioContext);
+            this.plainIn = nioParams.getByteBufferFactory().createReadBuffer(nioContext);
+            this.cipherIn = nioParams.getByteBufferFactory().createEncryptedReadBuffer(nioContext);
 
             this.outputStream = new DataOutputStream(
                 new SslEngineByteBufferOutputStream(sslEngine, plainOut, cipherOut, channel)

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

@@ -1,14 +1,20 @@
 package com.rabbitmq.client.test;
 
 import com.rabbitmq.client.*;
+import com.rabbitmq.client.impl.nio.DefaultByteBufferFactory;
 import com.rabbitmq.client.impl.nio.NioParams;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 
 import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.List;
 import java.util.concurrent.*;
 
+import static org.hamcrest.Matchers.hasSize;
+import static org.hamcrest.Matchers.isOneOf;
+import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 
 /**
@@ -128,6 +134,38 @@ public void nioLoopCleaning() throws Exception {
         }
     }
 
+    @Test public void byteBufferFactory() throws Exception {
+        ConnectionFactory cf = new ConnectionFactory();
+        cf.useNio();
+        int baseCapacity = 32768;
+        NioParams nioParams = new NioParams();
+        nioParams.setReadByteBufferSize(baseCapacity / 2);
+        nioParams.setWriteByteBufferSize(baseCapacity / 4);
+        List<ByteBuffer> byteBuffers = new CopyOnWriteArrayList<>();
+        cf.setNioParams(nioParams.setByteBufferFactory(new DefaultByteBufferFactory(capacity -> {
+            ByteBuffer bb = ByteBuffer.allocate(capacity);
+            byteBuffers.add(bb);
+            return bb;
+        })));
+
+        try (Connection c = cf.newConnection()) {
+            sendAndVerifyMessage(c, 100);
+        }
+
+        assertThat(byteBuffers, hasSize(2));
+        assertThat(byteBuffers.get(0).capacity(), isOneOf(nioParams.getReadByteBufferSize(), nioParams.getWriteByteBufferSize()));
+        assertThat(byteBuffers.get(1).capacity(), isOneOf(nioParams.getReadByteBufferSize(), nioParams.getWriteByteBufferSize()));
+    }
+
+    @Test public void directByteBuffers() throws Exception {
+        ConnectionFactory cf = new ConnectionFactory();
+        cf.useNio();
+        cf.setNioParams(new NioParams().setByteBufferFactory(new DefaultByteBufferFactory(capacity -> ByteBuffer.allocateDirect(capacity))));
+        try (Connection c = cf.newConnection()) {
+            sendAndVerifyMessage(c, 100);
+        }
+    }
+
     private void sendAndVerifyMessage(Connection connection, int size) throws Exception {
         CountDownLatch latch = new CountDownLatch(1);
         boolean messageReceived = basicGetBasicConsume(connection, QUEUE, latch, size);