Refactored out blocking read and writes

Author: pontusmelke

driver/src/main/java/org/neo4j/driver/internal/connector/socket/SocketClient.java

@@ -35,6 +35,8 @@
 import org.neo4j.driver.v1.exceptions.ClientException;
 
 import static java.nio.ByteOrder.BIG_ENDIAN;
+import static org.neo4j.driver.internal.connector.socket.SocketUtils.blockingRead;
+import static org.neo4j.driver.internal.connector.socket.SocketUtils.blockingWrite;
 
 public class SocketClient
 {
@@ -178,29 +180,12 @@ private SocketProtocol negotiateProtocol() throws IOException
         buf.flip();
 
         //Do a blocking write
-        while(buf.hasRemaining())
-        {
-            if (channel.write( buf ) < 0)
-            {
-                throw new ClientException(
-                        "Connection terminated while proposing protocol. This can happen due to network " +
-                        "instabilities, or due to restarts of the database." );
-            }
-        }
+       blockingWrite(channel, buf);
 
         // Read (blocking) back the servers choice
         buf.clear();
         buf.limit( 4 );
-
-        while(buf.hasRemaining())
-        {
-            if ( channel.read( buf ) < 0 )
-            {
-                throw new ClientException(
-                        "Connection terminated while negotiating protocol. This can happen due to network " +
-                        "instabilities, or due to restarts of the database." );
-            }
-        }
+        blockingRead(channel, buf);
 
         // Choose protocol, or fail
         buf.flip();

driver/src/main/java/org/neo4j/driver/internal/connector/socket/SocketUtils.java

@@ -0,0 +1,47 @@
+package org.neo4j.driver.internal.connector.socket;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.ByteChannel;
+
+import org.neo4j.driver.internal.util.BytePrinter;
+import org.neo4j.driver.v1.exceptions.ClientException;
+
+/**
+ * Utility class for common operations.
+ */
+public final class SocketUtils
+{
+    private SocketUtils()
+    {
+        throw new UnsupportedOperationException( "Do not instantiate" );
+    }
+
+    public static void blockingRead(ByteChannel channel, ByteBuffer buf) throws IOException
+    {
+        while(buf.hasRemaining())
+        {
+            if (channel.read( buf ) < 0)
+            {
+                throw new ClientException( String.format(
+                        "Connection terminated while receiving data. This can happen due to network " +
+                        "instabilities, or due to restarts of the database. Expected %s bytes, received %s.",
+                        buf.limit(), BytePrinter.hex( buf ) ) );
+            }
+        }
+    }
+
+    public static void blockingWrite(ByteChannel channel, ByteBuffer buf) throws IOException
+    {
+        while(buf.hasRemaining())
+        {
+            if (channel.write( buf ) < 0)
+            {
+                throw new ClientException( String.format(
+                        "Connection terminated while sending data. This can happen due to network " +
+                        "instabilities, or due to restarts of the database. Expected %s bytes, wrote %s.",
+                        buf.limit(), BytePrinter.hex( buf ) ) );
+            }
+        }
+    }
+}

driver/src/test/java/org/neo4j/driver/internal/connector/socket/SocketUtilsTest.java

@@ -0,0 +1,134 @@
+package org.neo4j.driver.internal.connector.socket;
+
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.ByteChannel;
+import java.util.ArrayList;
+import java.util.List;
+
+import org.neo4j.driver.v1.exceptions.ClientException;
+
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class SocketUtilsTest
+{
+    @Rule
+    public ExpectedException exception = ExpectedException.none();
+
+    @Test
+    public void shouldReadAllBytes() throws IOException
+    {
+        // Given
+        ByteBuffer buffer = ByteBuffer.allocate( 4 );
+        ByteAtATimeChannel channel = new ByteAtATimeChannel( new byte[]{0, 1, 2, 3} );
+
+        // When
+        SocketUtils.blockingRead(channel, buffer );
+        buffer.flip();
+
+        // Then
+        assertThat(buffer.get(), equalTo((byte) 0));
+        assertThat(buffer.get(), equalTo((byte) 1));
+        assertThat(buffer.get(), equalTo((byte) 2));
+        assertThat(buffer.get(), equalTo((byte) 3));
+    }
+
+    @Test
+    public void shouldFailIfConnectionFailsWhileReading() throws IOException
+    {
+        // Given
+        ByteBuffer buffer = ByteBuffer.allocate( 4 );
+        ByteChannel channel = mock( ByteChannel.class );
+        when(channel.read( buffer )).thenReturn( -1 );
+
+        //Expect
+        exception.expect( ClientException.class );
+
+        // When
+        SocketUtils.blockingRead(channel, buffer );
+    }
+
+    @Test
+    public void shouldWriteAllBytes() throws IOException
+    {
+        // Given
+        ByteBuffer buffer = ByteBuffer.wrap(  new byte[]{0, 1, 2, 3});
+        ByteAtATimeChannel channel = new ByteAtATimeChannel( new byte[0] );
+
+        // When
+        SocketUtils.blockingWrite(channel, buffer );
+
+        // Then
+        assertThat(channel.writtenBytes.get(0), equalTo((byte) 0));
+        assertThat(channel.writtenBytes.get(1), equalTo((byte) 1));
+        assertThat(channel.writtenBytes.get(2), equalTo((byte) 2));
+        assertThat(channel.writtenBytes.get(3), equalTo((byte) 3));
+    }
+
+    @Test
+    public void shouldFailIfConnectionFailsWhileWriting() throws IOException
+    {
+        // Given
+        ByteBuffer buffer = ByteBuffer.allocate( 4 );
+        ByteChannel channel = mock( ByteChannel.class );
+        when(channel.write( buffer )).thenReturn( -1 );
+
+        //Expect
+        exception.expect( ClientException.class );
+
+        // When
+        SocketUtils.blockingWrite(channel, buffer );
+    }
+
+    private static class ByteAtATimeChannel implements ByteChannel
+    {
+
+        private final byte[] bytes;
+        private int index = 0;
+        private List<Byte> writtenBytes = new ArrayList<>(  );
+
+        private ByteAtATimeChannel( byte[] bytes )
+        {
+            this.bytes = bytes;
+        }
+
+        @Override
+        public int read( ByteBuffer dst ) throws IOException
+        {
+            if (index >= bytes.length)
+            {
+                return -1;
+            }
+
+            dst.put( bytes[index++]);
+            return 1;
+        }
+
+        @Override
+        public int write( ByteBuffer src ) throws IOException
+        {
+            writtenBytes.add( src.get() );
+            return 1;
+        }
+
+        @Override
+        public boolean isOpen()
+        {
+            return true;
+        }
+
+        @Override
+        public void close() throws IOException
+        {
+
+        }
+    }
+
+}