Relax connection termination policy in routing driver

Previously routing driver terminated all connections towards a
particular address when one of active connections had a network error.
Connections were also terminated when new routing table did not contain
some address that was present in the previous routing table. Such
behaviour might be problematic because it results in terminated queries.
Network errors might have been temporary but always resulted in
termination of active queries.

This commit makes driver keep connections towards machine that had
failure but remove it from the routing table. This is done to prevent
subsequent queries from using it until rediscovery. After rediscovery
address can either appear again in the routing procedure response or not.
If it is returned then driver will re-add it to routing table and old
pool with all connections will be reused. If address is missing from
routing table and pool has no active connections then driver will close
the pool.

Author: lutovich

driver/src/main/java/org/neo4j/driver/internal/async/NettyConnection.java

@@ -20,9 +20,9 @@
 
 import io.netty.channel.Channel;
 import io.netty.channel.pool.ChannelPool;
-import io.netty.util.concurrent.Promise;
 
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.atomic.AtomicBoolean;
 
@@ -39,16 +39,14 @@
 import org.neo4j.driver.internal.util.ServerVersion;
 import org.neo4j.driver.v1.Value;
 
-import static java.util.concurrent.CompletableFuture.completedFuture;
-import static org.neo4j.driver.internal.util.Futures.asCompletionStage;
-
 public class NettyConnection implements Connection
 {
     private final Channel channel;
     private final InboundMessageDispatcher messageDispatcher;
     private final BoltServerAddress serverAddress;
     private final ServerVersion serverVersion;
     private final ChannelPool channelPool;
+    private final CompletableFuture<Void> releaseFuture;
     private final Clock clock;
 
     private final AtomicBoolean open = new AtomicBoolean( true );
@@ -61,6 +59,7 @@ public NettyConnection( Channel channel, ChannelPool channelPool, Clock clock )
         this.serverAddress = ChannelAttributes.serverAddress( channel );
         this.serverVersion = ChannelAttributes.serverVersion( channel );
         this.channelPool = channelPool;
+        this.releaseFuture = new CompletableFuture<>();
         this.clock = clock;
     }
 
@@ -111,14 +110,9 @@ public CompletionStage<Void> release()
     {
         if ( open.compareAndSet( true, false ) )
         {
-            Promise<Void> releasePromise = channel.eventLoop().newPromise();
-            reset( new ResetResponseHandler( channel, channelPool, messageDispatcher, clock, releasePromise ) );
-            return asCompletionStage( releasePromise );
-        }
-        else
-        {
-            return completedFuture( null );
+            reset( new ResetResponseHandler( channel, channelPool, messageDispatcher, clock, releaseFuture ) );
         }
+        return releaseFuture;
     }
 
     @Override

driver/src/main/java/org/neo4j/driver/internal/async/pool/ConnectionPoolImpl.java

@@ -25,6 +25,7 @@
 import io.netty.util.concurrent.Future;
 
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.ConcurrentHashMap;
@@ -58,10 +59,16 @@ public class ConnectionPoolImpl implements ConnectionPool
 
     public ConnectionPoolImpl( ChannelConnector connector, Bootstrap bootstrap, PoolSettings settings,
             Logging logging, Clock clock )
+    {
+        this( connector, bootstrap, new ActiveChannelTracker( logging ), settings, logging, clock );
+    }
+
+    ConnectionPoolImpl( ChannelConnector connector, Bootstrap bootstrap, ActiveChannelTracker activeChannelTracker,
+            PoolSettings settings, Logging logging, Clock clock )
     {
         this.connector = connector;
         this.bootstrap = bootstrap;
-        this.activeChannelTracker = new ActiveChannelTracker( logging );
+        this.activeChannelTracker = activeChannelTracker;
         this.channelHealthChecker = new NettyChannelHealthChecker( settings, clock, logging );
         this.settings = settings;
         this.clock = clock;
@@ -101,6 +108,30 @@ public void purge( BoltServerAddress address )
         }
     }
 
+    @Override
+    public void retainAll( Set<BoltServerAddress> addressesToRetain )
+    {
+        for ( BoltServerAddress address : pools.keySet() )
+        {
+            if ( !addressesToRetain.contains( address ) )
+            {
+                int activeChannels = activeChannelTracker.activeChannelCount( address );
+                if ( activeChannels == 0 )
+                {
+                    // address is not present in updated routing table and has no active connections
+                    // it's now safe to terminate corresponding connection pool and forget about it
+
+                    ChannelPool pool = pools.remove( address );
+                    if ( pool != null )
+                    {
+                        log.info( "Purging idle connections towards %s", address );
+                        pool.close();
+                    }
+                }
+            }
+        }
+    }
+
     @Override
     public boolean hasAddress( BoltServerAddress address )
     {
@@ -157,7 +188,7 @@ private ChannelPool getOrCreatePool( BoltServerAddress address )
         return pool;
     }
 
-    private NettyChannelPool newPool( BoltServerAddress address )
+    ChannelPool newPool( BoltServerAddress address )
     {
         return new NettyChannelPool( address, connector, bootstrap, activeChannelTracker, channelHealthChecker,
                 settings.connectionAcquisitionTimeout(), settings.maxConnectionPoolSize() );

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

@@ -39,54 +39,9 @@ public int size()
         return addresses.length;
     }
 
-    public synchronized void update( Set<BoltServerAddress> addresses, Set<BoltServerAddress> removed )
+    public synchronized void update( Set<BoltServerAddress> addresses )
     {
-        BoltServerAddress[] prev = this.addresses;
-        if ( addresses.isEmpty() )
-        {
-            this.addresses = NONE;
-            return;
-        }
-        if ( prev.length == 0 )
-        {
-            this.addresses = addresses.toArray( NONE );
-            return;
-        }
-        BoltServerAddress[] copy = null;
-        if ( addresses.size() != prev.length )
-        {
-            copy = new BoltServerAddress[addresses.size()];
-        }
-        int j = 0;
-        for ( int i = 0; i < prev.length; i++ )
-        {
-            if ( addresses.remove( prev[i] ) )
-            {
-                if ( copy != null )
-                {
-                    copy[j++] = prev[i];
-                }
-            }
-            else
-            {
-                removed.add( prev[i] );
-                if ( copy == null )
-                {
-                    copy = new BoltServerAddress[prev.length];
-                    System.arraycopy( prev, 0, copy, 0, i );
-                    j = i;
-                }
-            }
-        }
-        if ( copy == null )
-        {
-            return;
-        }
-        for ( BoltServerAddress address : addresses )
-        {
-            copy[j++] = address;
-        }
-        this.addresses = copy;
+        this.addresses = addresses.toArray( NONE );
     }
 
     public synchronized void remove( BoltServerAddress address )

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

@@ -19,6 +19,7 @@
 
 package org.neo4j.driver.internal.cluster;
 
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.Set;
@@ -43,7 +44,7 @@ public class ClusterRoutingTable implements RoutingTable
     public ClusterRoutingTable( Clock clock, BoltServerAddress... routingAddresses )
     {
         this( clock );
-        routers.update( new LinkedHashSet<>( asList( routingAddresses ) ), new HashSet<BoltServerAddress>() );
+        routers.update( new LinkedHashSet<>( asList( routingAddresses ) ) );
     }
 
     private ClusterRoutingTable( Clock clock )
@@ -66,14 +67,12 @@ public boolean isStaleFor( AccessMode mode )
     }
 
     @Override
-    public synchronized Set<BoltServerAddress> update( ClusterComposition cluster )
+    public synchronized void update( ClusterComposition cluster )
     {
         expirationTimeout = cluster.expirationTimestamp();
-        Set<BoltServerAddress> removed = new HashSet<>();
-        readers.update( cluster.readers(), removed );
-        writers.update( cluster.writers(), removed );
-        routers.update( cluster.routers(), removed );
-        return removed;
+        readers.update( cluster.readers() );
+        writers.update( cluster.writers() );
+        routers.update( cluster.routers() );
     }
 
     @Override
@@ -102,6 +101,16 @@ public AddressSet routers()
         return routers;
     }
 
+    @Override
+    public Set<BoltServerAddress> servers()
+    {
+        Set<BoltServerAddress> servers = new HashSet<>();
+        Collections.addAll( servers, readers.toArray() );
+        Collections.addAll( servers, writers.toArray() );
+        Collections.addAll( servers, routers.toArray() );
+        return servers;
+    }
+
     @Override
     public void removeWriter( BoltServerAddress toRemove )
     {

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

@@ -53,18 +53,9 @@ public CompletionStage<RoutingProcedureResponse> run( CompletionStage<Connection
         return connectionStage.thenCompose( connection ->
         {
             Statement procedure = procedureStatement( connection.serverVersion() );
-            return runProcedure( connection, procedure ).handle( ( records, error ) ->
-            {
-                Throwable cause = Futures.completionErrorCause( error );
-                if ( cause != null )
-                {
-                    return handleError( procedure, cause );
-                }
-                else
-                {
-                    return new RoutingProcedureResponse( procedure, records );
-                }
-            } );
+            return runProcedure( connection, procedure )
+                    .thenCompose( records -> releaseConnection( connection, records ) )
+                    .handle( ( records, error ) -> processProcedureResponse( procedure, records, error ) );
         } );
     }
 
@@ -87,6 +78,30 @@ private Statement procedureStatement( ServerVersion serverVersion )
         }
     }
 
+    private CompletionStage<List<Record>> releaseConnection( Connection connection, List<Record> records )
+    {
+        // It is not strictly required to release connection after routing procedure invocation because it'll
+        // be released by the PULL_ALL response handler after result is fully fetched. Such release will happen
+        // in background. However, releasing it early as part of whole chain makes it easier to reason about
+        // rediscovery in stub server tests. Some of them assume connections to instances not present in new
+        // routing table will be closed immediately.
+        return connection.release().thenApply( ignore -> records );
+    }
+
+    private RoutingProcedureResponse processProcedureResponse( Statement procedure, List<Record> records,
+            Throwable error )
+    {
+        Throwable cause = Futures.completionErrorCause( error );
+        if ( cause != null )
+        {
+            return handleError( procedure, cause );
+        }
+        else
+        {
+            return new RoutingProcedureResponse( procedure, records );
+        }
+    }
+
     private RoutingProcedureResponse handleError( Statement procedure, Throwable error )
     {
         if ( error instanceof ClientException )

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

@@ -27,7 +27,7 @@ public interface RoutingTable
 {
     boolean isStaleFor( AccessMode mode );
 
-    Set<BoltServerAddress> update( ClusterComposition cluster );
+    void update( ClusterComposition cluster );
 
     void forget( BoltServerAddress address );
 
@@ -37,5 +37,7 @@ public interface RoutingTable
 
     AddressSet routers();
 
+    Set<BoltServerAddress> servers();
+
     void removeWriter( BoltServerAddress toRemove );
 }

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

@@ -20,7 +20,6 @@
 
 import io.netty.util.concurrent.EventExecutorGroup;
 
-import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 
@@ -125,10 +124,8 @@ public CompletionStage<Void> close()
 
     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 );
-        // drop all current connections to the address
-        connectionPool.purge( address );
     }
 
     private synchronized CompletionStage<RoutingTable> freshRoutingTable( AccessMode mode )
@@ -171,18 +168,21 @@ else if ( routingTable.isStaleFor( mode ) )
 
     private synchronized void freshClusterCompositionFetched( ClusterComposition composition )
     {
-        Set<BoltServerAddress> removed = routingTable.update( composition );
-
-        for ( BoltServerAddress address : removed )
+        try
         {
-            connectionPool.purge( address );
-        }
+            routingTable.update( composition );
+            connectionPool.retainAll( routingTable.servers() );
 
-        log.info( "Refreshed routing information. %s", routingTable );
+            log.info( "Refreshed routing information. %s", routingTable );
 
-        CompletableFuture<RoutingTable> routingTableFuture = refreshRoutingTableFuture;
-        refreshRoutingTableFuture = null;
-        routingTableFuture.complete( routingTable );
+            CompletableFuture<RoutingTable> routingTableFuture = refreshRoutingTableFuture;
+            refreshRoutingTableFuture = null;
+            routingTableFuture.complete( routingTable );
+        }
+        catch ( Throwable error )
+        {
+            clusterCompositionLookupFailed( error );
+        }
     }
 
     private synchronized void clusterCompositionLookupFailed( Throwable error )

driver/src/main/java/org/neo4j/driver/internal/handlers/ResetResponseHandler.java

@@ -20,9 +20,10 @@
 
 import io.netty.channel.Channel;
 import io.netty.channel.pool.ChannelPool;
-import io.netty.util.concurrent.Promise;
+import io.netty.util.concurrent.Future;
 
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 
 import org.neo4j.driver.internal.async.inbound.InboundMessageDispatcher;
 import org.neo4j.driver.internal.spi.ResponseHandler;
@@ -37,16 +38,16 @@ public class ResetResponseHandler implements ResponseHandler
     private final ChannelPool pool;
     private final InboundMessageDispatcher messageDispatcher;
     private final Clock clock;
-    private final Promise<Void> releasePromise;
+    private final CompletableFuture<Void> releaseFuture;
 
     public ResetResponseHandler( Channel channel, ChannelPool pool, InboundMessageDispatcher messageDispatcher,
-            Clock clock, Promise<Void> releasePromise )
+            Clock clock, CompletableFuture<Void> releaseFuture )
     {
         this.channel = channel;
         this.pool = pool;
         this.messageDispatcher = messageDispatcher;
         this.clock = clock;
-        this.releasePromise = releasePromise;
+        this.releaseFuture = releaseFuture;
     }
 
     @Override
@@ -72,13 +73,7 @@ private void releaseChannel()
         messageDispatcher.unMuteAckFailure();
         setLastUsedTimestamp( channel, clock.millis() );
 
-        if ( releasePromise == null )
-        {
-            pool.release( channel );
-        }
-        else
-        {
-            pool.release( channel, releasePromise );
-        }
+        Future<Void> released = pool.release( channel );
+        released.addListener( ignore -> releaseFuture.complete( null ) );
     }
 }