Get rid of read retry in NIO for plain connection

A FrameBuilder class has been introduced to be able to read a frame
in several network reads. The frame builder keeps track of
the partial frame state. This avoids pulling from the network
when the channel doesn't come back with any value from a network
read, which is against the NIO philosophy.

References #319

Author: acogoluegnes

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

@@ -0,0 +1,114 @@
+// Copyright (c) 2007-Present Pivotal Software, Inc.  All rights reserved.
+//
+// This software, the RabbitMQ Java client library, is triple-licensed under the
+// Mozilla Public License 1.1 ("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.impl.nio;
+
+import com.rabbitmq.client.AMQP;
+import com.rabbitmq.client.MalformedFrameException;
+import com.rabbitmq.client.impl.Frame;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.ReadableByteChannel;
+
+/**
+ *
+ */
+public class FrameBuilder {
+
+    private static final int PAYLOAD_OFFSET = 1 /* type */ + 2 /* channel */ + 4 /* payload size */;
+
+    private final ReadableByteChannel channel;
+
+    private final ByteBuffer buffer;
+
+    private int frameType;
+    private int frameChannel;
+    private byte [] framePayload;
+
+    private int bytesRead = 0;
+
+    private final int [] frameBuffer = new int[3];
+
+    public FrameBuilder(ReadableByteChannel channel, ByteBuffer buffer) {
+        this.channel = channel;
+        this.buffer = buffer;
+    }
+
+    public Frame readFrame() throws IOException {
+        while(readFromNetworkIfNecessary()) {
+            if (bytesRead == 0) {
+                // type
+                // FIXME check first byte isn't 'A' and thus a header indicating protocol version mismatch
+                frameType = readFromBuffer();
+            } else if (bytesRead == 1) {
+                // channel 1/2
+                frameBuffer[0] = readFromBuffer();
+            } else if (bytesRead == 2) {
+                // channel 2/2
+                frameChannel = (frameBuffer[0] << 8) + (readFromBuffer() << 0);
+            } else if (bytesRead == 3) {
+                // payload size 1/4
+                frameBuffer[0] = readFromBuffer();
+            } else if (bytesRead == 4) {
+                // payload size 2/4
+                frameBuffer[1] = readFromBuffer();
+            } else if (bytesRead == 5) {
+                // payload size 3/4
+                frameBuffer[2] = readFromBuffer();
+            } else if (bytesRead == 6) {
+                // payload size 4/4
+                int framePayloadSize = ((frameBuffer[0] << 24) + (frameBuffer[1] << 16) + (frameBuffer[2] << 8) + (readFromBuffer() << 0));
+                framePayload = new byte[framePayloadSize];
+            } else if (bytesRead >= PAYLOAD_OFFSET && bytesRead < framePayload.length + PAYLOAD_OFFSET) {
+                framePayload[bytesRead - PAYLOAD_OFFSET] = (byte) readFromBuffer();
+            } else if (bytesRead == framePayload.length + PAYLOAD_OFFSET) {
+                int frameEndMarker = readFromBuffer();
+                if (frameEndMarker != AMQP.FRAME_END) {
+                    throw new MalformedFrameException("Bad frame end marker: " + frameEndMarker);
+                }
+                bytesRead = 0;
+                return new Frame(frameType, frameChannel, framePayload);
+            } else {
+                throw new IllegalStateException("Number of read bytes incorrect: " + bytesRead);
+            }
+            bytesRead++;
+        }
+        return null;
+    }
+
+    private int read() throws IOException {
+        return NioHelper.read(channel, buffer);
+    }
+
+    private int readFromBuffer() {
+        return buffer.get() & 0xff;
+    }
+
+    private boolean readFromNetworkIfNecessary() throws IOException {
+        if(!buffer.hasRemaining()) {
+            buffer.clear();
+            int read = read();
+            buffer.flip();
+            if (read > 0) {
+                return true;
+            } else {
+                return false;
+            }
+        } else {
+            return true;
+        }
+    }
+}

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

@@ -149,22 +149,29 @@ public void run() {
                                 state.prepareForReadSequence();
 
                                 while (state.continueReading()) {
-                                    Frame frame = Frame.readFrom(inputStream);
+                                    final Frame frame;
+                                    if (state.frameBuilder == null) {
+                                        frame = Frame.readFrom(inputStream);
+                                    } else {
+                                        frame = state.frameBuilder.readFrame();
+                                    }
 
-                                    try {
-                                        boolean noProblem = state.getConnection().handleReadFrame(frame);
-                                        if (noProblem && (!state.getConnection().isRunning() || state.getConnection().hasBrokerInitiatedShutdown())) {
-                                            // looks like the frame was Close-Ok or Close
-                                            dispatchShutdownToConnection(state);
+                                    if (frame != null) {
+                                        try {
+                                            boolean noProblem = state.getConnection().handleReadFrame(frame);
+                                            if (noProblem && (!state.getConnection().isRunning() || state.getConnection().hasBrokerInitiatedShutdown())) {
+                                                // looks like the frame was Close-Ok or Close
+                                                dispatchShutdownToConnection(state);
+                                                key.cancel();
+                                                break;
+                                            }
+                                        } catch (Throwable ex) {
+                                            // problem during frame processing, tell connection, and
+                                            // we can stop for this channel
+                                            handleIoError(state, ex);
                                             key.cancel();
                                             break;
                                         }
-                                    } catch (Throwable ex) {
-                                        // problem during frame processing, tell connection, and
-                                        // we can stop for this channel
-                                        handleIoError(state, ex);
-                                        key.cancel();
-                                        break;
                                     }
                                 }
 

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

@@ -77,6 +77,8 @@ public class SocketChannelFrameHandlerState {
 
     final DataInputStream inputStream;
 
+    final FrameBuilder frameBuilder;
+
     public SocketChannelFrameHandlerState(SocketChannel channel, NioLoopContext nioLoopsState, NioParams nioParams, SSLEngine sslEngine) {
         this.channel = channel;
         this.readSelectorState = nioLoopsState.readSelectorState;
@@ -98,6 +100,8 @@ public SocketChannelFrameHandlerState(SocketChannel channel, NioLoopContext nioL
                 new ByteBufferInputStream(channel, plainIn)
             );
 
+            this.frameBuilder = new FrameBuilder(channel, plainIn);
+
         } else {
             this.ssl = true;
             this.plainOut = ByteBuffer.allocate(sslEngine.getSession().getApplicationBufferSize());
@@ -111,6 +115,8 @@ public SocketChannelFrameHandlerState(SocketChannel channel, NioLoopContext nioL
             this.inputStream = new DataInputStream(
                 new SslEngineByteBufferInputStream(sslEngine, plainIn, cipherIn, channel)
             );
+
+            this.frameBuilder = null;
         }
 
     }

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

@@ -16,6 +16,7 @@
 
 package com.rabbitmq.client.test;
 
+import com.rabbitmq.client.impl.nio.FrameBuilder;
 import com.rabbitmq.utility.IntAllocatorTests;
 import org.junit.runner.RunWith;
 import org.junit.runners.Suite;
@@ -49,7 +50,8 @@
     ConnectionFactoryTest.class,
     RecoveryAwareAMQConnectionFactoryTest.class,
     RpcTest.class,
-    RecoveryDelayHandlerTest.class
+    RecoveryDelayHandlerTest.class,
+    FrameBuilder.class
 })
 public class ClientTests {
 

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

@@ -0,0 +1,109 @@
+// Copyright (c) 2007-Present Pivotal Software, Inc.  All rights reserved.
+//
+// This software, the RabbitMQ Java client library, is triple-licensed under the
+// Mozilla Public License 1.1 ("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.test;
+
+import com.rabbitmq.client.AMQP;
+import com.rabbitmq.client.impl.Frame;
+import com.rabbitmq.client.impl.nio.FrameBuilder;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.ReadableByteChannel;
+
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.notNullValue;
+import static org.hamcrest.Matchers.nullValue;
+import static org.junit.Assert.assertThat;
+
+/**
+ *
+ */
+@RunWith(MockitoJUnitRunner.class)
+public class FrameBuilderTest {
+
+    @Mock
+    ReadableByteChannel channel;
+
+    ByteBuffer buffer;
+
+    FrameBuilder builder;
+
+    @Test
+    public void buildFrameInOneGo() throws IOException {
+        buffer = ByteBuffer.wrap(new byte[]{1, 0, 0, 0, 0, 0, 3, 1, 2, 3, end()});
+        builder = new FrameBuilder(channel, buffer);
+        Frame frame = builder.readFrame();
+        assertThat(frame, notNullValue());
+        assertThat(frame.type, is(1));
+        assertThat(frame.channel, is(0));
+        assertThat(frame.getPayload().length, is(3));
+    }
+
+    @Test
+    public void buildFramesInOneGo() throws IOException {
+        byte[] frameContent = new byte[]{1, 0, 0, 0, 0, 0, 3, 1, 2, 3, end()};
+        int nbFrames = 13;
+        byte[] frames = new byte[frameContent.length * nbFrames];
+        for(int i = 0; i < nbFrames; i++) {
+            for (int j = 0; j < frameContent.length; j++) {
+                frames[i * frameContent.length + j] = frameContent[j];
+            }
+        }
+        buffer = ByteBuffer.wrap(frames);
+        builder = new FrameBuilder(channel, buffer);
+        int frameCount = 0;
+        Frame frame;
+        while ((frame = builder.readFrame()) != null) {
+            assertThat(frame, notNullValue());
+            assertThat(frame.type, is(1));
+            assertThat(frame.channel, is(0));
+            assertThat(frame.getPayload().length, is(3));
+            frameCount++;
+        }
+        assertThat(frameCount, is(nbFrames));
+    }
+
+    @Test
+    public void buildFrameInSeveralCalls() throws IOException {
+        buffer = ByteBuffer.wrap(new byte[]{1, 0, 0, 0, 0, 0, 3, 1, 2});
+        builder = new FrameBuilder(channel, buffer);
+        Frame frame = builder.readFrame();
+        assertThat(frame, nullValue());
+
+        buffer.clear();
+        buffer.put(b(3)).put(end());
+        buffer.flip();
+
+        frame = builder.readFrame();
+        assertThat(frame, notNullValue());
+        assertThat(frame.type, is(1));
+        assertThat(frame.channel, is(0));
+        assertThat(frame.getPayload().length, is(3));
+    }
+
+    byte b(int b) {
+        return (byte) b;
+    }
+
+    byte end() {
+        return (byte) AMQP.FRAME_END;
+    }
+
+}