Fixed issue with connection pool deactivation

Author: lutovich

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

@@ -64,7 +64,8 @@ public boolean isStaleFor( AccessMode mode )
                mode == AccessMode.WRITE && writers.size() == 0;
     }
 
-    private Set<BoltServerAddress> servers()
+    @Override
+    public Set<BoltServerAddress> servers()
     {
         Set<BoltServerAddress> servers = new HashSet<>();
         servers.addAll( readers.servers() );

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

@@ -110,16 +110,13 @@ private PooledConnection acquireConnection( AccessMode mode, RoundRobinAddressSe
 
     private synchronized void forget( BoltServerAddress address )
     {
-        // First remove from the load balancer, to prevent concurrent threads from making connections to them.
+        // remove from the routing table, to prevent concurrent threads from making connections to this address
         routingTable.forget( address );
+
         if ( PURGE_ON_ERROR )
         {
             connections.purge( address );
         }
-        else
-        {
-            connections.deactivate( address );
-        }
     }
 
     synchronized void ensureRouting( AccessMode mode )
@@ -153,15 +150,7 @@ private void updateConnectionPool( RoutingTableChange routingTableChange )
         }
         else
         {
-            for ( BoltServerAddress addedAddress : routingTableChange.added() )
-            {
-                connections.activate( addedAddress );
-            }
-            for ( BoltServerAddress removedAddress : routingTableChange.removed() )
-            {
-                connections.deactivate( removedAddress );
-            }
-            connections.compact();
+            connections.retainAll( routingTable.servers() );
         }
     }
 

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

@@ -18,6 +18,8 @@
  */
 package org.neo4j.driver.internal.cluster;
 
+import java.util.Set;
+
 import org.neo4j.driver.internal.net.BoltServerAddress;
 import org.neo4j.driver.v1.AccessMode;
 
@@ -37,5 +39,7 @@ public interface RoutingTable
 
     int routerSize();
 
+    Set<BoltServerAddress> servers();
+
     void removeWriter( BoltServerAddress toRemove );
 }

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

@@ -23,7 +23,7 @@
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.LinkedBlockingQueue;
-import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.neo4j.driver.internal.logging.DelegatingLogger;
 import org.neo4j.driver.internal.net.BoltServerAddress;
@@ -40,15 +40,11 @@ public class BlockingPooledConnectionQueue
 {
     public static final String LOG_NAME = "ConnectionQueue";
 
-    private static final int ACTIVE = 1;
-    private static final int INACTIVE = 2;
-    private static final int TERMINATED = 3;
-
     /** The backing queue, keeps track of connections currently in queue */
     private final BlockingQueue<PooledConnection> queue;
     private final Logger logger;
 
-    private final AtomicInteger state = new AtomicInteger( ACTIVE );
+    private final AtomicBoolean terminated = new AtomicBoolean();
 
     /** Keeps track of acquired connections */
     private final Set<PooledConnection> acquiredConnections =
@@ -75,7 +71,7 @@ public boolean offer( PooledConnection pooledConnection )
         {
             disposeSafely( pooledConnection );
         }
-        if ( state.get() != ACTIVE )
+        if ( terminated.get() )
         {
             terminateIdleConnections();
         }
@@ -96,13 +92,11 @@ public PooledConnection acquire( Supplier<PooledConnection> supplier )
         }
         acquiredConnections.add( connection );
 
-        int poolState = state.get();
-        if ( poolState != ACTIVE )
+        if ( terminated.get() )
         {
             acquiredConnections.remove( connection );
             disposeSafely( connection );
-            throw new IllegalStateException( "Pool is " + (poolState == INACTIVE ? "deactivated" : "terminated") +
-                                             ", new connections can't be acquired" );
+            throw new IllegalStateException( "Pool is terminated, new connections can't be acquired" );
         }
         else
         {
@@ -131,24 +125,6 @@ public boolean contains( PooledConnection pooledConnection )
         return queue.contains( pooledConnection );
     }
 
-    public void activate()
-    {
-        state.compareAndSet( INACTIVE, ACTIVE );
-    }
-
-    public void deactivate()
-    {
-        if ( state.compareAndSet( ACTIVE, INACTIVE ) )
-        {
-            terminateIdleConnections();
-        }
-    }
-
-    public boolean isActive()
-    {
-        return state.get() == ACTIVE;
-    }
-
     /**
      * Terminates all connections, both those that are currently in the queue as well
      * as those that have been acquired.
@@ -157,7 +133,7 @@ public boolean isActive()
      */
     public void terminate()
     {
-        if ( state.getAndSet( TERMINATED ) != TERMINATED )
+        if ( terminated.compareAndSet( false, true ) )
         {
             terminateIdleConnections();
             terminateAcquiredConnections();

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

@@ -18,7 +18,6 @@
  */
 package org.neo4j.driver.internal.net.pooling;
 
-import java.util.Collections;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
@@ -33,6 +32,7 @@
 import org.neo4j.driver.internal.spi.PooledConnection;
 import org.neo4j.driver.internal.util.Clock;
 import org.neo4j.driver.internal.util.Supplier;
+import org.neo4j.driver.v1.Logger;
 import org.neo4j.driver.v1.Logging;
 
 /**
@@ -61,6 +61,7 @@ public class SocketConnectionPool implements ConnectionPool
     private final ConnectionValidator<PooledConnection> connectionValidator;
     private final Clock clock;
     private final Logging logging;
+    private final Logger log;
 
     public SocketConnectionPool( PoolSettings poolSettings, Connector connector, Clock clock, Logging logging )
     {
@@ -69,6 +70,7 @@ public SocketConnectionPool( PoolSettings poolSettings, Connector connector, Clo
         this.connectionValidator = new PooledConnectionValidator( this );
         this.clock = clock;
         this.logging = logging;
+        this.log = logging.getLog( getClass().getSimpleName() );
     }
 
     @Override
@@ -93,46 +95,33 @@ public void purge( BoltServerAddress address )
     }
 
     @Override
-    public void activate( BoltServerAddress address )
-    {
-        BlockingPooledConnectionQueue connectionQueue = pools.get( address );
-        if ( connectionQueue != null )
-        {
-            connectionQueue.activate();
-        }
-    }
-
-    @Override
-    public void deactivate( BoltServerAddress address )
-    {
-        BlockingPooledConnectionQueue connections = pools.get( address );
-        if ( connections != null )
-        {
-            connections.deactivate();
-        }
-    }
-
-    @Override
-    public void compact()
+    public void retainAll( Set<BoltServerAddress> addressesToRetain )
     {
         for ( Map.Entry<BoltServerAddress,BlockingPooledConnectionQueue> entry : pools.entrySet() )
         {
             BoltServerAddress address = entry.getKey();
-            BlockingPooledConnectionQueue queue = entry.getValue();
+            BlockingPooledConnectionQueue pool = entry.getValue();
 
-            if ( !queue.isActive() && queue.activeConnections() == 0 )
+            if ( !addressesToRetain.contains( address ) && pool.activeConnections() == 0 )
             {
-                // queue has been in deactivated state and has no open connections by now
-                pools.remove( address );
+                // address is not present in the updated routing table and has no active connections
+                // it's now safe to terminate corresponding connection pool and forget about it
+
+                BlockingPooledConnectionQueue removedPool = pools.remove( address );
+                if ( removedPool != null )
+                {
+                    log.info( "Closing connection pool towards %s, it has no active connections " +
+                              "and is not in the routing table", address );
+                    removedPool.terminate();
+                }
             }
         }
     }
 
     @Override
     public boolean hasAddress( BoltServerAddress address )
     {
-        BlockingPooledConnectionQueue connectionQueue = pools.get( address );
-        return connectionQueue != null && connectionQueue.isActive();
+        return pools.containsKey( address );
     }
 
     @Override
@@ -152,17 +141,7 @@ public void close()
     public int activeConnections( BoltServerAddress address )
     {
         BlockingPooledConnectionQueue connectionQueue = pools.get( address );
-        if ( connectionQueue == null || !connectionQueue.isActive() )
-        {
-            return 0;
-        }
-        return connectionQueue.activeConnections();
-    }
-
-    // test-only accessor
-    Set<BoltServerAddress> addresses()
-    {
-        return Collections.unmodifiableSet( pools.keySet() );
+        return connectionQueue == null ? 0 : connectionQueue.activeConnections();
     }
 
     private BlockingPooledConnectionQueue pool( BoltServerAddress address )

driver/src/main/java/org/neo4j/driver/internal/security/TLSSocketChannel.java

@@ -372,6 +372,8 @@ private HandshakeStatus wrap( ByteBuffer buffer ) throws IOException, ClientExce
                 cipherOut.compact();
             }
             break;
+        case CLOSED:
+            throw new IOException( "TLS socket channel is closed" );
         default:
             throw new ClientException( "Got unexpected status " + status );
         }

driver/src/main/java/org/neo4j/driver/internal/spi/ConnectionPool.java

@@ -18,6 +18,8 @@
  */
 package org.neo4j.driver.internal.spi;
 
+import java.util.Set;
+
 import org.neo4j.driver.internal.net.BoltServerAddress;
 
 public interface ConnectionPool extends AutoCloseable
@@ -36,11 +38,7 @@ public interface ConnectionPool extends AutoCloseable
      */
     void purge( BoltServerAddress address );
 
-    void activate( BoltServerAddress address );
-
-    void deactivate( BoltServerAddress address );
-
-    void compact();
+    void retainAll( Set<BoltServerAddress> addressesToRetain );
 
     boolean hasAddress( BoltServerAddress address );
 }

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

@@ -18,7 +18,7 @@
  */
 package org.neo4j.driver.internal.cluster;
 
-import java.util.ArrayList;
+import java.util.Collections;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Set;
@@ -41,7 +41,7 @@ private ClusterCompositionUtil() {}
     public static final BoltServerAddress E = new BoltServerAddress( "5555:55" );
     public static final BoltServerAddress F = new BoltServerAddress( "6666:66" );
 
-    public static final List<BoltServerAddress> EMPTY = new ArrayList<>();
+    public static final List<BoltServerAddress> EMPTY = Collections.emptyList();
 
     public static final ClusterComposition VALID_CLUSTER_COMPOSITION =
             createClusterComposition( asList( A, B ), asList( C ), asList( D, E ) );

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

@@ -20,7 +20,9 @@
 
 import org.junit.Test;
 
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 
 import org.neo4j.driver.internal.net.BoltServerAddress;
 import org.neo4j.driver.internal.util.FakeClock;
@@ -231,4 +233,25 @@ public void shouldNotRemoveServerIfPreWriterNowReader()
         assertEquals( 2, change.removed().size() );
         assertThat( change.removed(), containsInAnyOrder( A, C ) );
     }
+
+    @Test
+    public void shouldReturnNoServersWhenEmpty()
+    {
+        ClusterRoutingTable routingTable = new ClusterRoutingTable( new FakeClock() );
+
+        Set<BoltServerAddress> servers = routingTable.servers();
+
+        assertEquals( 0, servers.size() );
+    }
+
+    @Test
+    public void shouldReturnAllServers()
+    {
+        ClusterRoutingTable routingTable = new ClusterRoutingTable( new FakeClock() );
+        routingTable.update( createClusterComposition( asList( A, B, C ), asList( B, C, D ), asList( C, D, E, F ) ) );
+
+        Set<BoltServerAddress> servers = routingTable.servers();
+
+        assertEquals( new HashSet<>( asList( A, B, C, D, E, F ) ), servers );
+    }
 }