Add abstraction to create NIO frame queue

Allows to use other data structure implementations than JDK's
ArrayBlockingQueue and BlockingQueue.

References #410

(cherry picked from commit 8c9577d)

Author: acogoluegnes

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

@@ -0,0 +1,41 @@
+package com.rabbitmq.client.impl.nio;
+
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Bridge between {@link NioQueue} and JDK's {@link BlockingQueue}.
+ *
+ * @see NioQueue
+ * @since 5.5.0
+ */
+public class BlockingQueueNioQueue implements NioQueue {
+
+    private final BlockingQueue<WriteRequest> delegate;
+    private final int writeEnqueuingTimeoutInMs;
+
+    public BlockingQueueNioQueue(BlockingQueue<WriteRequest> delegate, int writeEnqueuingTimeoutInMs) {
+        this.delegate = delegate;
+        this.writeEnqueuingTimeoutInMs = writeEnqueuingTimeoutInMs;
+    }
+
+    @Override
+    public boolean offer(WriteRequest writeRequest) throws InterruptedException {
+        return this.delegate.offer(writeRequest, writeEnqueuingTimeoutInMs, TimeUnit.MILLISECONDS);
+    }
+
+    @Override
+    public int size() {
+        return this.delegate.size();
+    }
+
+    @Override
+    public WriteRequest poll() {
+        return this.delegate.poll();
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return this.delegate.isEmpty();
+    }
+}

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

@@ -20,8 +20,10 @@
 import com.rabbitmq.client.SslEngineConfigurator;
 
 import javax.net.ssl.SSLEngine;
+import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ThreadFactory;
+import java.util.function.Function;
 
 import static com.rabbitmq.client.SslEngineConfigurators.ENABLE_HOSTNAME_VERIFICATION;
 
@@ -32,6 +34,12 @@
  */
 public class NioParams {
 
+    static Function<? super NioContext, ? extends NioQueue> DEFAULT_WRITE_QUEUE_FACTORY =
+        ctx -> new BlockingQueueNioQueue(
+            new ArrayBlockingQueue<>(ctx.getNioParams().getWriteQueueCapacity(), true),
+            ctx.getNioParams().getWriteEnqueuingTimeoutInMs()
+        );
+
     /**
      * size of the byte buffer used for inbound data
      */
@@ -93,6 +101,14 @@ public class NioParams {
      */
     private ByteBufferFactory byteBufferFactory = new DefaultByteBufferFactory();
 
+    /**
+     * Factory to create a {@link NioQueue}.
+     *
+     * @since 5.5.0
+     */
+    private Function<? super NioContext, ? extends NioQueue> writeQueueFactory =
+        DEFAULT_WRITE_QUEUE_FACTORY;
+
     public NioParams() {
     }
 
@@ -108,6 +124,7 @@ public NioParams(NioParams nioParams) {
         setSslEngineConfigurator(nioParams.getSslEngineConfigurator());
         setConnectionShutdownExecutor(nioParams.getConnectionShutdownExecutor());
         setByteBufferFactory(nioParams.getByteBufferFactory());
+        setWriteQueueFactory(nioParams.getWriteQueueFactory());
     }
 
     /**
@@ -384,4 +401,24 @@ public NioParams setByteBufferFactory(ByteBufferFactory byteBufferFactory) {
     public ByteBufferFactory getByteBufferFactory() {
         return byteBufferFactory;
     }
+
+    /**
+     * Set the factory to create {@link NioQueue}s.
+     * <p>
+     * The default uses a {@link ArrayBlockingQueue}.
+     *
+     * @param writeQueueFactory the factory to use
+     * @return this {@link NioParams} instance
+     * @see NioQueue
+     * @since 5.5.0
+     */
+    public NioParams setWriteQueueFactory(
+        Function<? super NioContext, ? extends NioQueue> writeQueueFactory) {
+        this.writeQueueFactory = writeQueueFactory;
+        return this;
+    }
+
+    public Function<? super NioContext, ? extends NioQueue> getWriteQueueFactory() {
+        return writeQueueFactory;
+    }
 }

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

@@ -0,0 +1,45 @@
+package com.rabbitmq.client.impl.nio;
+
+/**
+ * Contract to exchange frame between application threads and NIO thread.
+ * <p>
+ * This is a simplified subset of {@link java.util.concurrent.BlockingQueue}.
+ * This interface is considered a SPI and is likely to move between
+ * minor and patch releases.
+ *
+ * @see NioParams
+ * @since 5.5.0
+ */
+public interface NioQueue {
+
+    /**
+     * Enqueue a frame, block if the queue is full.
+     *
+     * @param writeRequest
+     * @return
+     * @throws InterruptedException
+     */
+    boolean offer(WriteRequest writeRequest) throws InterruptedException;
+
+    /**
+     * Get the current size of the queue.
+     *
+     * @return
+     */
+    int size();
+
+    /**
+     * Retrieves and removes the head of this queue,
+     * or returns {@code null} if this queue is empty.
+     *
+     * @return the head of this queue, or {@code null} if this queue is empty
+     */
+    WriteRequest poll();
+
+    /**
+     * Returns {@code true} if the queue contains no element.
+     *
+     * @return {@code true} if the queue contains no element
+     */
+    boolean isEmpty();
+}

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

@@ -26,10 +26,6 @@
 import java.nio.ByteBuffer;
 import java.nio.channels.SelectionKey;
 import java.nio.channels.SocketChannel;
-import java.util.Queue;
-import java.util.concurrent.ArrayBlockingQueue;
-import java.util.concurrent.BlockingQueue;
-import java.util.concurrent.TimeUnit;
 
 /**
  *
@@ -43,7 +39,7 @@ public class SocketChannelFrameHandlerState {
 
     private final SocketChannel channel;
 
-    private final BlockingQueue<WriteRequest> writeQueue;
+    private final NioQueue writeQueue;
 
     private volatile AMQConnection connection;
 
@@ -54,8 +50,6 @@ public class SocketChannelFrameHandlerState {
 
     private final SelectorHolder readSelectorState;
 
-    private final int writeEnqueuingTimeoutInMs;
-
     final boolean ssl;
 
     final SSLEngine sslEngine;
@@ -80,8 +74,13 @@ public SocketChannelFrameHandlerState(SocketChannel channel, NioLoopContext nioL
         this.channel = channel;
         this.readSelectorState = nioLoopsState.readSelectorState;
         this.writeSelectorState = nioLoopsState.writeSelectorState;
-        this.writeQueue = new ArrayBlockingQueue<WriteRequest>(nioParams.getWriteQueueCapacity(), true);
-        this.writeEnqueuingTimeoutInMs = nioParams.getWriteEnqueuingTimeoutInMs();
+
+        NioContext nioContext = new NioContext(nioParams, sslEngine);
+
+        this.writeQueue = nioParams.getWriteQueueFactory() == null ?
+            NioParams.DEFAULT_WRITE_QUEUE_FACTORY.apply(nioContext) :
+            nioParams.getWriteQueueFactory().apply(nioContext);
+
         this.sslEngine = sslEngine;
         if(this.sslEngine == null) {
             this.ssl = false;
@@ -98,7 +97,6 @@ public SocketChannelFrameHandlerState(SocketChannel channel, NioLoopContext nioL
 
         } else {
             this.ssl = true;
-            NioContext nioContext = new NioContext(nioParams, sslEngine);
             this.plainOut = nioParams.getByteBufferFactory().createWriteBuffer(nioContext);
             this.cipherOut = nioParams.getByteBufferFactory().createEncryptedWriteBuffer(nioContext);
             this.plainIn = nioParams.getByteBufferFactory().createReadBuffer(nioContext);
@@ -116,7 +114,7 @@ public SocketChannel getChannel() {
         return channel;
     }
 
-    public Queue<WriteRequest> getWriteQueue() {
+    public NioQueue getWriteQueue() {
         return writeQueue;
     }
 
@@ -130,7 +128,7 @@ public void write(Frame frame) throws IOException {
 
     private void sendWriteRequest(WriteRequest writeRequest) throws IOException {
         try {
-            boolean offered = this.writeQueue.offer(writeRequest, writeEnqueuingTimeoutInMs, TimeUnit.MILLISECONDS);
+            boolean offered = this.writeQueue.offer(writeRequest);
             if(offered) {
                 this.writeSelectorState.registerFrameHandlerState(this, SelectionKey.OP_WRITE);
                 this.readSelectorState.selector.wakeup();

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

@@ -1,6 +1,7 @@
 package com.rabbitmq.client.test;
 
 import com.rabbitmq.client.*;
+import com.rabbitmq.client.impl.nio.BlockingQueueNioQueue;
 import com.rabbitmq.client.impl.nio.DefaultByteBufferFactory;
 import com.rabbitmq.client.impl.nio.NioParams;
 import org.junit.After;
@@ -11,9 +12,11 @@
 import java.nio.ByteBuffer;
 import java.util.List;
 import java.util.concurrent.*;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.isOneOf;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 
@@ -166,6 +169,23 @@ public void nioLoopCleaning() throws Exception {
         }
     }
 
+    @Test public void customWriteQueue() throws Exception {
+        ConnectionFactory cf = new ConnectionFactory();
+        cf.useNio();
+        AtomicInteger count = new AtomicInteger(0);
+        cf.setNioParams(new NioParams().setWriteQueueFactory(ctx -> {
+            count.incrementAndGet();
+            return new BlockingQueueNioQueue(
+                new LinkedBlockingQueue<>(ctx.getNioParams().getWriteQueueCapacity()),
+                ctx.getNioParams().getWriteEnqueuingTimeoutInMs()
+            );
+        }));
+        try (Connection c = cf.newConnection()) {
+            sendAndVerifyMessage(c, 100);
+        }
+        assertEquals(1, count.get());
+    }
+
     private void sendAndVerifyMessage(Connection connection, int size) throws Exception {
         CountDownLatch latch = new CountDownLatch(1);
         boolean messageReceived = basicGetBasicConsume(connection, QUEUE, latch, size);