Loop to acquire valid connection in LoadBalancer

Previously load balancer selected a single address and tried to acquire
connection for it. Connection failure would then surface to the user.
This was done even though routing table could possibly contain other
valid machines that might respond.

This commit makes load balancer chain connection acquisition futures to
probe all known machines before failure.

Author: lutovich

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

@@ -51,7 +51,6 @@
 import org.neo4j.driver.v1.exceptions.SessionExpiredException;
 
 import static java.util.concurrent.CompletableFuture.completedFuture;
-import static org.neo4j.driver.internal.async.Futures.failedFuture;
 
 public class LoadBalancer implements ConnectionProvider, RoutingErrorHandler, AutoCloseable
 {
@@ -62,6 +61,7 @@ public class LoadBalancer implements ConnectionProvider, RoutingErrorHandler, Au
     private final RoutingTable routingTable;
     private final Rediscovery rediscovery;
     private final LoadBalancingStrategy loadBalancingStrategy;
+    private final EventExecutorGroup eventExecutorGroup;
     private final Logger log;
 
     private CompletableFuture<RoutingTable> refreshRoutingTableFuture;
@@ -72,26 +72,28 @@ public LoadBalancer( BoltServerAddress initialRouter, RoutingSettings settings,
     {
         this( connections, asyncConnectionPool, new ClusterRoutingTable( clock, initialRouter ),
                 createRediscovery( initialRouter, settings, eventExecutorGroup, clock, logging ),
-                loadBalancerLogger( logging ), loadBalancingStrategy );
+                loadBalancerLogger( logging ), loadBalancingStrategy, eventExecutorGroup );
     }
 
     // Used only in testing
     public LoadBalancer( ConnectionPool connections, AsyncConnectionPool asyncConnectionPool,
-            RoutingTable routingTable, Rediscovery rediscovery, Logging logging )
+            RoutingTable routingTable, Rediscovery rediscovery, EventExecutorGroup eventExecutorGroup, Logging logging )
     {
         this( connections, asyncConnectionPool, routingTable, rediscovery, loadBalancerLogger( logging ),
-                new LeastConnectedLoadBalancingStrategy( connections, asyncConnectionPool, logging ) );
+                new LeastConnectedLoadBalancingStrategy( connections, asyncConnectionPool, logging ),
+                eventExecutorGroup );
     }
 
     private LoadBalancer( ConnectionPool connections, AsyncConnectionPool asyncConnectionPool,
             RoutingTable routingTable, Rediscovery rediscovery, Logger log,
-            LoadBalancingStrategy loadBalancingStrategy )
+            LoadBalancingStrategy loadBalancingStrategy, EventExecutorGroup eventExecutorGroup )
     {
         this.connections = connections;
         this.asyncConnectionPool = asyncConnectionPool;
         this.routingTable = routingTable;
         this.rediscovery = rediscovery;
         this.loadBalancingStrategy = loadBalancingStrategy;
+        this.eventExecutorGroup = eventExecutorGroup;
         this.log = log;
 
         if ( connections != null )
@@ -111,22 +113,51 @@ public PooledConnection acquireConnection( AccessMode mode )
     @Override
     public CompletionStage<AsyncConnection> acquireAsyncConnection( AccessMode mode )
     {
-        return freshRoutingTable( mode ).thenCompose( routingTable ->
-        {
-            AddressSet addressSet = addressSet( mode, routingTable );
-            BoltServerAddress address = selectAddressAsync( mode, addressSet );
+        return freshRoutingTable( mode )
+                .thenCompose( routingTable -> acquire( mode, routingTable ) )
+                .thenApply( connection -> new RoutingAsyncConnection( connection, mode, this ) );
+    }
+
+    private CompletionStage<AsyncConnection> acquire( AccessMode mode, RoutingTable routingTable )
+    {
+        AddressSet addresses = addressSet( mode, routingTable );
+        CompletableFuture<AsyncConnection> result = new CompletableFuture<>();
+        acquire( mode, addresses, result );
+        return result;
+    }
 
-            // todo: loop like in sync version until we get a successful connection
+    private void acquire( AccessMode mode, AddressSet addresses, CompletableFuture<AsyncConnection> result )
+    {
+        BoltServerAddress address = selectAddressAsync( mode, addresses );
+
+        if ( address == null )
+        {
+            result.completeExceptionally( new SessionExpiredException(
+                    "Failed to obtain connection towards " + mode + " server. " +
+                    "Known routing table is: " + routingTable ) );
+            return;
+        }
 
-            if ( address == null )
+        asyncConnectionPool.acquire( address ).whenComplete( ( connection, error ) ->
+        {
+            if ( error != null )
             {
-                return failedFuture( new SessionExpiredException(
-                        "Failed to obtain connection towards " + mode + " server. " +
-                        "Known routing table is: " + routingTable ) );
+                if ( error instanceof ServiceUnavailableException )
+                {
+                    log.error( "Failed to obtain a connection towards address " + address, error );
+                    forget( address );
+                    eventExecutorGroup.next().execute( () -> acquire( mode, addresses, result ) );
+                }
+                else
+                {
+                    result.completeExceptionally( error );
+                }
             }
-
-            return asyncConnectionPool.acquire( address );
-        } ).thenApply( connection -> new RoutingAsyncConnection( connection, mode, this ) );
+            else
+            {
+                result.complete( connection );
+            }
+        } );
     }
 
     @Override
@@ -172,7 +203,10 @@ private synchronized void forget( BoltServerAddress address )
         // First remove from the load balancer, to prevent concurrent threads from making connections to them.
         routingTable.forget( address );
         // drop all current connections to the address
-        connections.purge( address );
+        if ( connections != null )
+        {
+            connections.purge( address );
+        }
         asyncConnectionPool.purge( address );
     }
 

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

@@ -1,3 +1,21 @@
+/*
+ * 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.cluster;
 
 import io.netty.util.concurrent.GlobalEventExecutor;

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

@@ -18,6 +18,7 @@
  */
 package org.neo4j.driver.internal.cluster;
 
+import io.netty.util.concurrent.GlobalEventExecutor;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
@@ -381,7 +382,8 @@ private static LoadBalancer newLoadBalancer( ClusterComposition clusterCompositi
         Rediscovery rediscovery = mock( Rediscovery.class );
         when( rediscovery.lookupClusterComposition( routingTable, connectionPool ) ).thenReturn( clusterComposition );
         AsyncConnectionPool asyncConnectionPool = mock( AsyncConnectionPool.class );
-        return new LoadBalancer( connectionPool, asyncConnectionPool, routingTable, rediscovery, DEV_NULL_LOGGING );
+        return new LoadBalancer( connectionPool, asyncConnectionPool, routingTable, rediscovery,
+                GlobalEventExecutor.INSTANCE, DEV_NULL_LOGGING );
     }
 
     private interface ConnectionMethod