Throw ServiceUnavailableException when socket write fails

Previously `ClientException` was thrown which was incorrect and resulted in
connections not being properly disposed.

Author: lutovich

driver/src/main/java/org/neo4j/driver/internal/RoutingTransaction.java

@@ -47,7 +47,7 @@ public class RoutingTransaction implements Transaction
     private final BoltServerAddress address;
     private final RoutingErrorHandler onError;
 
-    RoutingTransaction( Transaction delegate, AccessMode mode, BoltServerAddress address,
+    public RoutingTransaction( Transaction delegate, AccessMode mode, BoltServerAddress address,
             RoutingErrorHandler onError )
     {
         this.delegate = delegate;

driver/src/main/java/org/neo4j/driver/internal/cluster/LoadBalancer.java

@@ -57,10 +57,10 @@ private LoadBalancer(
             RoutingTable routingTable,
             ClusterComposition.Provider provider ) throws ServiceUnavailableException
     {
-        this( log, connections, routingTable, new Rediscovery( settings, clock, log, provider ) );
+        this( routingTable, connections, new Rediscovery( settings, clock, log, provider ), log );
     }
 
-    LoadBalancer( Logger log, ConnectionPool connections, RoutingTable routingTable, Rediscovery rediscovery )
+    LoadBalancer( RoutingTable routingTable, ConnectionPool connections, Rediscovery rediscovery, Logger log )
             throws ServiceUnavailableException
     {
         this.log = log;

driver/src/main/java/org/neo4j/driver/internal/net/SocketConnection.java

@@ -38,6 +38,7 @@
 import org.neo4j.driver.v1.Value;
 import org.neo4j.driver.v1.exceptions.ClientException;
 import org.neo4j.driver.v1.exceptions.Neo4jException;
+import org.neo4j.driver.v1.exceptions.ServiceUnavailableException;
 import org.neo4j.driver.v1.summary.ServerInfo;
 
 import static java.lang.String.format;
@@ -165,8 +166,7 @@ public synchronized void flush()
         }
         catch ( IOException e )
         {
-            String message = e.getMessage();
-            throw new ClientException( "Unable to send messages to server: " + message, e );
+            throw new ServiceUnavailableException( "Unable to send messages to server: " + e.getMessage(), e );
         }
     }
 

driver/src/test/java/org/neo4j/driver/internal/cluster/LoadBalancerTest.java

@@ -29,13 +29,18 @@
 import java.util.List;
 
 import org.neo4j.driver.internal.EventHandler;
+import org.neo4j.driver.internal.RoutingTransaction;
 import org.neo4j.driver.internal.net.BoltServerAddress;
 import org.neo4j.driver.internal.spi.Connection;
+import org.neo4j.driver.internal.spi.ConnectionPool;
 import org.neo4j.driver.internal.spi.StubConnectionPool;
 import org.neo4j.driver.internal.util.FakeClock;
 import org.neo4j.driver.internal.util.MatcherFactory;
+import org.neo4j.driver.v1.AccessMode;
 import org.neo4j.driver.v1.EventLogger;
+import org.neo4j.driver.v1.Transaction;
 import org.neo4j.driver.v1.exceptions.ServiceUnavailableException;
+import org.neo4j.driver.v1.exceptions.SessionExpiredException;
 import org.neo4j.driver.v1.util.Function;
 
 import static org.hamcrest.Matchers.any;
@@ -46,9 +51,13 @@
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.fail;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.neo4j.driver.internal.cluster.ClusterTopology.Role.READ;
 import static org.neo4j.driver.internal.cluster.ClusterTopology.Role.ROUTE;
 import static org.neo4j.driver.internal.cluster.ClusterTopology.Role.WRITE;
+import static org.neo4j.driver.internal.logging.DevNullLogger.DEV_NULL_LOGGER;
 import static org.neo4j.driver.internal.spi.StubConnectionPool.Event.acquire;
 import static org.neo4j.driver.internal.spi.StubConnectionPool.Event.connectionFailure;
 import static org.neo4j.driver.internal.util.FakeClock.Event.sleep;
@@ -100,7 +109,7 @@ private LoadBalancer seedLoadBalancer( String host, int port ) throws Exception
     private LoadBalancer seedLoadBalancer( String host, int port, RoutingSettings settings ) throws Exception
     {
         ClusterRoutingTable routingTable = new ClusterRoutingTable( clock, new BoltServerAddress( host, port ) );
-        return new LoadBalancer( log, connections, routingTable, new Rediscovery( settings, clock, log, cluster ) );
+        return new LoadBalancer( routingTable, connections, new Rediscovery( settings, clock, log, cluster ), log );
     }
 
     @Test
@@ -607,6 +616,36 @@ public void sleep( long timestamp, long millis )
         assertEquals( new BoltServerAddress( "two", 1337 ), connection.boltServerAddress() );
     }
 
+    @Test
+    public void shouldForgetAddressAndItsConnectionsOnServiceUnavailable()
+    {
+        Transaction tx = mock( Transaction.class );
+        RoutingTable routingTable = mock( RoutingTable.class );
+        ConnectionPool connectionPool = mock( ConnectionPool.class );
+        Rediscovery rediscovery = mock( Rediscovery.class );
+        LoadBalancer loadBalancer = new LoadBalancer( routingTable, connectionPool, rediscovery, DEV_NULL_LOGGER );
+        BoltServerAddress address = new BoltServerAddress( "host", 42 );
+
+        RoutingTransaction routingTx = new RoutingTransaction( tx, AccessMode.WRITE, address, loadBalancer );
+
+        ServiceUnavailableException txCloseError = new ServiceUnavailableException( "Oh!" );
+        doThrow( txCloseError ).when( tx ).close();
+
+        try
+        {
+            routingTx.close();
+            fail( "Exception expected" );
+        }
+        catch ( Exception e )
+        {
+            assertThat( e, instanceOf( SessionExpiredException.class ) );
+            assertThat( e.getCause(), instanceOf( ServiceUnavailableException.class ) );
+        }
+
+        verify( routingTable ).forget( address );
+        verify( connectionPool ).purge( address );
+    }
+
     private void coreClusterOn( int ttlSeconds, String leader, int port, String... others )
     {
         for ( int i = 0; i <= others.length; i++ )

driver/src/test/java/org/neo4j/driver/internal/net/SocketConnectionTest.java

@@ -22,19 +22,22 @@
 import org.mockito.invocation.InvocationOnMock;
 import org.mockito.stubbing.Answer;
 
+import java.io.IOException;
 import java.net.URI;
 import java.util.ArrayList;
 import java.util.Iterator;
+import java.util.Queue;
 
 import org.neo4j.driver.internal.messaging.Message;
 import org.neo4j.driver.internal.messaging.SuccessMessage;
 import org.neo4j.driver.internal.summary.InternalServerInfo;
-import org.neo4j.driver.v1.Logger;
 import org.neo4j.driver.v1.Values;
+import org.neo4j.driver.v1.exceptions.ServiceUnavailableException;
 import org.neo4j.driver.v1.summary.ServerInfo;
 
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.CoreMatchers.instanceOf;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.fail;
 import static org.mockito.Matchers.any;
@@ -45,16 +48,20 @@
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static org.neo4j.driver.internal.logging.DevNullLogger.DEV_NULL_LOGGER;
+import static org.neo4j.driver.internal.net.BoltServerAddress.LOCAL_DEFAULT;
 import static org.neo4j.driver.v1.Values.parameters;
 
 public class SocketConnectionTest
 {
+    private static final InternalServerInfo SERVER_INFO = new InternalServerInfo( LOCAL_DEFAULT, "test" );
+
     @Test
     public void shouldReceiveServerInfoAfterInit() throws Throwable
     {
         // Given
         SocketClient socket = mock( SocketClient.class );
-        SocketConnection conn = new SocketConnection( socket, mock( InternalServerInfo.class ), mock( Logger.class ) );
+        SocketConnection conn = new SocketConnection( socket, SERVER_INFO, DEV_NULL_LOGGER );
 
         when( socket.address() ).thenReturn( BoltServerAddress.from( URI.create( "http://neo4j.com:9000" ) ) );
 
@@ -98,7 +105,7 @@ public void shouldCloseConnectionIfFailedToCreate() throws Throwable
         // Then
         try
         {
-            SocketConnection conn = new SocketConnection( socket, mock( InternalServerInfo.class ), mock( Logger.class ) );
+            new SocketConnection( socket, SERVER_INFO, DEV_NULL_LOGGER );
             fail( "should have failed with the provided exception" );
         }
         catch( Throwable e )
@@ -108,4 +115,26 @@ public void shouldCloseConnectionIfFailedToCreate() throws Throwable
         }
         verify( socket, times( 1 ) ).stop();
     }
+
+    @Test
+    @SuppressWarnings( "unchecked" )
+    public void flushThrowsWhenSocketIsBroken() throws Exception
+    {
+        SocketClient socket = mock( SocketClient.class );
+        IOException sendError = new IOException( "Unable to send" );
+        doThrow( sendError ).when( socket ).send( any( Queue.class ) );
+
+        SocketConnection connection = new SocketConnection( socket, SERVER_INFO, DEV_NULL_LOGGER );
+
+        try
+        {
+            connection.flush();
+            fail( "Exception expected" );
+        }
+        catch ( Exception e )
+        {
+            assertThat( e, instanceOf( ServiceUnavailableException.class ) );
+            assertSame( sendError, e.getCause() );
+        }
+    }
 }

driver/src/test/java/org/neo4j/driver/v1/integration/CausalClusteringIT.java

@@ -312,8 +312,8 @@ public void shouldHandleGracefulLeaderSwitch() throws Exception
                     parameters( "name", "Webber", "title", "Mr" ) );
             tx1.success();
 
-            closeAndExpectException( tx1, ClientException.class );
-            closeAndExpectException( session1, ClientException.class );
+            closeAndExpectException( tx1, SessionExpiredException.class );
+            session1.close();
 
             String bookmark = inExpirableSession( driver, createSession(), new Function<Session,String>()
             {