Improve disposal of broken connections

Connection pool can discard broken connections during acquisition
(when connection liveness check timeout is configured) and when
connections are returned to the pool. In both cases connections
should be disposed and removed from the set of active connections.
This is especially important with least connected load balancing
strategy which examines amount of active connections for each address.

This commit makes sure broken connections are disposed through the
connections queue to make sure active set is always updated.

Author: lutovich

driver/src/main/java/org/neo4j/driver/internal/net/pooling/BlockingPooledConnectionQueue.java

@@ -107,6 +107,12 @@ public PooledConnection acquire( Supplier<PooledConnection> supplier )
         return connection;
     }
 
+    void disposeBroken( PooledConnection connection )
+    {
+        acquiredConnections.remove( connection );
+        disposeSafely( connection );
+    }
+
     public boolean isEmpty()
     {
         return queue.isEmpty();
@@ -163,7 +169,7 @@ private void disposeSafely( PooledConnection connection )
         }
         catch ( Throwable disposeError )
         {
-            logger.warn( "Error disposing connection during termination", disposeError );
+            logger.warn( "Error disposing connection", disposeError );
         }
     }
 

driver/src/main/java/org/neo4j/driver/internal/net/pooling/PooledConnectionReleaseConsumer.java

@@ -47,7 +47,7 @@ public void accept( PooledConnection pooledConnection )
         }
         else
         {
-            pooledConnection.dispose();
+            connections.disposeBroken( pooledConnection );
         }
     }
 }

driver/src/main/java/org/neo4j/driver/internal/net/pooling/SocketConnectionPool.java

@@ -142,10 +142,16 @@ private PooledConnection acquireConnection( BoltServerAddress address,
     {
         ConnectionSupplier connectionSupplier = new ConnectionSupplier( connectionQueue, address );
 
-        PooledConnection connection;
+        PooledConnection connection = null;
         boolean connectionCreated;
         do
         {
+            // dispose previous connection that can't be acquired
+            if ( connection != null )
+            {
+                connectionQueue.disposeBroken( connection );
+            }
+
             connection = connectionQueue.acquire( connectionSupplier );
             connectionCreated = connectionSupplier.connectionCreated();
         }

driver/src/test/java/org/neo4j/driver/internal/net/pooling/BlockingPooledConnectionQueueTest.java

@@ -285,6 +285,21 @@ public void shouldReportActiveConnections()
         assertEquals( 0, queue.activeConnections() );
     }
 
+    @Test
+    @SuppressWarnings( "unchecked" )
+    public void shouldDisposeBrokenConnections()
+    {
+        BlockingPooledConnectionQueue queue = newConnectionQueue( 5 );
+
+        queue.offer( mock( PooledConnection.class ) );
+        PooledConnection connection = queue.acquire( mock( Supplier.class ) );
+        assertEquals( 1, queue.activeConnections() );
+
+        queue.disposeBroken( connection );
+        assertEquals( 0, queue.activeConnections() );
+        verify( connection ).dispose();
+    }
+
     private static BlockingPooledConnectionQueue newConnectionQueue( int capacity )
     {
         return newConnectionQueue( capacity, mock( Logging.class, RETURNS_MOCKS ) );

driver/src/test/java/org/neo4j/driver/internal/net/pooling/PooledConnectionReleaseConsumerTest.java

@@ -0,0 +1,100 @@
+/*
+ * Copyright (c) 2002-2017 "Neo Technology,"
+ * Network Engine for Objects in Lund AB [http://neotechnology.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.neo4j.driver.internal.net.pooling;
+
+import org.junit.Test;
+
+import org.neo4j.driver.internal.spi.ConnectionPool;
+import org.neo4j.driver.internal.spi.PooledConnection;
+import org.neo4j.driver.internal.util.Supplier;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.neo4j.driver.internal.logging.DevNullLogging.DEV_NULL_LOGGING;
+import static org.neo4j.driver.internal.net.BoltServerAddress.LOCAL_DEFAULT;
+
+public class PooledConnectionReleaseConsumerTest
+{
+    @Test
+    public void shouldOfferReusableConnectionsBackToTheConnectionsQueue()
+    {
+        BlockingPooledConnectionQueue queue = newConnectionQueue();
+        PooledConnection connection = acquireConnection( queue );
+
+        PooledConnectionValidator validator = newConnectionValidator( true );
+        PooledConnectionReleaseConsumer releaseConsumer = new PooledConnectionReleaseConsumer( queue, validator );
+
+        releaseConsumer.accept( connection );
+
+        // connection should now be idle
+        assertEquals( 0, queue.activeConnections() );
+        assertEquals( 1, queue.size() );
+
+        verify( connection ).reset();
+        verify( connection ).sync();
+    }
+
+    @Test
+    public void shouldAskConnectionsQueueToDisposeNotReusableConnections()
+    {
+        BlockingPooledConnectionQueue queue = newConnectionQueue();
+        PooledConnection connection = acquireConnection( queue );
+
+        PooledConnectionValidator validator = newConnectionValidator( false );
+        PooledConnectionReleaseConsumer releaseConsumer = new PooledConnectionReleaseConsumer( queue, validator );
+
+        releaseConsumer.accept( connection );
+
+        // connection should've been disposed
+        assertEquals( 0, queue.activeConnections() );
+        assertEquals( 0, queue.size() );
+
+        verify( connection ).dispose();
+    }
+
+    private static BlockingPooledConnectionQueue newConnectionQueue()
+    {
+        return new BlockingPooledConnectionQueue( LOCAL_DEFAULT, 5, DEV_NULL_LOGGING );
+    }
+
+    @SuppressWarnings( "unchecked" )
+    private static PooledConnection acquireConnection( BlockingPooledConnectionQueue queue )
+    {
+        queue.offer( newConnectionMock() );
+        PooledConnection connection = queue.acquire( mock( Supplier.class ) );
+        assertEquals( 1, queue.activeConnections() );
+        return connection;
+    }
+
+    private static PooledConnectionValidator newConnectionValidator( boolean allowsConnections )
+    {
+        ConnectionPool pool = mock( ConnectionPool.class );
+        when( pool.hasAddress( LOCAL_DEFAULT ) ).thenReturn( allowsConnections );
+        return new PooledConnectionValidator( pool );
+    }
+
+    private static PooledConnection newConnectionMock()
+    {
+        PooledConnection connection = mock( PooledConnection.class );
+        when( connection.boltServerAddress() ).thenReturn( LOCAL_DEFAULT );
+        return connection;
+    }
+}

driver/src/test/java/org/neo4j/driver/internal/net/pooling/SocketConnectionPoolTest.java

@@ -24,7 +24,9 @@
 import org.mockito.stubbing.Answer;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentHashMap;
@@ -42,6 +44,7 @@
 import org.neo4j.driver.internal.util.Clock;
 import org.neo4j.driver.internal.util.FakeClock;
 import org.neo4j.driver.v1.Logging;
+import org.neo4j.driver.v1.Value;
 
 import static java.util.Collections.newSetFromMap;
 import static org.hamcrest.Matchers.instanceOf;
@@ -57,6 +60,7 @@
 import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.RETURNS_MOCKS;
 import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -65,6 +69,7 @@
 import static org.mockito.Mockito.when;
 import static org.neo4j.driver.internal.net.BoltServerAddress.DEFAULT_PORT;
 import static org.neo4j.driver.internal.net.BoltServerAddress.LOCAL_DEFAULT;
+import static org.neo4j.driver.v1.Values.value;
 
 public class SocketConnectionPoolTest
 {
@@ -536,6 +541,52 @@ public void reportActiveConnectionsWhenHasIdleConnections()
         assertEquals( 0, pool.activeConnections( ADDRESS_1 ) );
     }
 
+    @Test
+    public void shouldForgetBrokenIdleConnection()
+    {
+        Connection connection1 = newConnectionMock( ADDRESS_1 );
+        Connection connection2 = newConnectionMock( ADDRESS_1 );
+
+        doNothing().doThrow( new RuntimeException() ).when( connection1 ).reset();
+
+        int idleTimeBeforeConnectionTest = 42;
+        FakeClock clock = new FakeClock();
+        Connector connector = newMockConnector( connection1, connection2 );
+        SocketConnectionPool pool = newPool( connector, clock, idleTimeBeforeConnectionTest );
+
+        // acquire and release one connection
+        pool.acquire( ADDRESS_1 ).close();
+        // make this connection seem idle for too long
+        clock.progress( idleTimeBeforeConnectionTest + 42 );
+
+        PooledConnection acquiredConnection = pool.acquire( ADDRESS_1 );
+
+        Map<String,Value> auth = Collections.singletonMap( "Key", value( "Value" ) );
+        acquiredConnection.init( "DummyClient", auth );
+        verify( connection1, never() ).init( "DummyClient", auth );
+        verify( connection2 ).init( "DummyClient", auth );
+
+        assertEquals( 1, pool.activeConnections( ADDRESS_1 ) );
+        acquiredConnection.close();
+        assertEquals( 0, pool.activeConnections( ADDRESS_1 ) );
+    }
+
+    @Test
+    public void shouldForgetIdleConnection()
+    {
+        Connection connection = newConnectionMock( ADDRESS_1 );
+        doThrow( new RuntimeException() ).when( connection ).reset();
+
+        SocketConnectionPool pool = newPool( newMockConnector( connection ), new FakeClock(), 42 );
+        PooledConnection pooledConnection = pool.acquire( ADDRESS_1 );
+
+        // release the connection, it should fail to reset and be disposed
+        pooledConnection.close();
+
+        assertEquals( 0, pool.activeConnections( ADDRESS_1 ) );
+        verify( connection ).close();
+    }
+
     private static Answer<Connection> createConnectionAnswer( final Set<Connection> createdConnections )
     {
         return new Answer<Connection>()