Merge pull request #411 from lutovich/1.5-async-routing-improvements

Async routing improvements

Author: zhenlineo

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

@@ -19,7 +19,6 @@
 package org.neo4j.driver.internal;
 
 import io.netty.bootstrap.Bootstrap;
-import io.netty.channel.EventLoopGroup;
 import io.netty.util.concurrent.EventExecutorGroup;
 
 import java.io.IOException;
@@ -76,17 +75,17 @@ public final Driver newInstance( URI uri, AuthToken authToken, RoutingSettings r
         SecurityPlan securityPlan = createSecurityPlan( address, config );
         ConnectionPool connectionPool = createConnectionPool( authToken, securityPlan, config );
 
-        Bootstrap bootstrap = BootstrapFactory.newBootstrap();
-        EventLoopGroup eventLoopGroup = bootstrap.config().group();
-        RetryLogic retryLogic = createRetryLogic( retrySettings, eventLoopGroup, config.logging() );
+        Bootstrap bootstrap = createBootstrap();
+        EventExecutorGroup eventExecutorGroup = bootstrap.config().group();
+        RetryLogic retryLogic = createRetryLogic( retrySettings, eventExecutorGroup, config.logging() );
 
         AsyncConnectionPool asyncConnectionPool = createAsyncConnectionPool( authToken, securityPlan, bootstrap,
                 config );
 
         try
         {
-            return createDriver( uri, address, connectionPool, config, newRoutingSettings, securityPlan, retryLogic,
-                    asyncConnectionPool );
+            return createDriver( uri, address, connectionPool, asyncConnectionPool, config, newRoutingSettings,
+                    eventExecutorGroup, securityPlan, retryLogic );
         }
         catch ( Throwable driverError )
         {
@@ -121,8 +120,8 @@ private AsyncConnectionPool createAsyncConnectionPool( AuthToken authToken, Secu
     }
 
     private Driver createDriver( URI uri, BoltServerAddress address, ConnectionPool connectionPool,
-            Config config, RoutingSettings routingSettings, SecurityPlan securityPlan, RetryLogic retryLogic,
-            AsyncConnectionPool asyncConnectionPool )
+            AsyncConnectionPool asyncConnectionPool, Config config, RoutingSettings routingSettings,
+            EventExecutorGroup eventExecutorGroup, SecurityPlan securityPlan, RetryLogic retryLogic )
     {
         String scheme = uri.getScheme().toLowerCase();
         switch ( scheme )
@@ -131,7 +130,8 @@ private Driver createDriver( URI uri, BoltServerAddress address, ConnectionPool
             assertNoRoutingContext( uri, routingSettings );
             return createDirectDriver( address, connectionPool, config, securityPlan, retryLogic, asyncConnectionPool );
         case BOLT_ROUTING_URI_SCHEME:
-            return createRoutingDriver( address, connectionPool, config, routingSettings, securityPlan, retryLogic );
+            return createRoutingDriver( address, connectionPool, asyncConnectionPool, config, routingSettings,
+                    securityPlan, retryLogic, eventExecutorGroup );
         default:
             throw new ClientException( format( "Unsupported URI scheme: %s", scheme ) );
         }
@@ -158,13 +158,15 @@ protected Driver createDirectDriver( BoltServerAddress address, ConnectionPool c
      * <b>This method is protected only for testing</b>
      */
     protected Driver createRoutingDriver( BoltServerAddress address, ConnectionPool connectionPool,
-            Config config, RoutingSettings routingSettings, SecurityPlan securityPlan, RetryLogic retryLogic )
+            AsyncConnectionPool asyncConnectionPool, Config config, RoutingSettings routingSettings,
+            SecurityPlan securityPlan, RetryLogic retryLogic, EventExecutorGroup eventExecutorGroup )
     {
         if ( !securityPlan.isRoutingCompatible() )
         {
             throw new IllegalArgumentException( "The chosen security plan is not compatible with a routing driver" );
         }
-        ConnectionProvider connectionProvider = createLoadBalancer( address, connectionPool, config, routingSettings );
+        ConnectionProvider connectionProvider = createLoadBalancer( address, connectionPool, asyncConnectionPool,
+                eventExecutorGroup, config, routingSettings );
         SessionFactory sessionFactory = createSessionFactory( connectionProvider, retryLogic, config );
         return createDriver( config, securityPlan, sessionFactory );
     }
@@ -184,21 +186,25 @@ protected InternalDriver createDriver( Config config, SecurityPlan securityPlan,
      * <p>
      * <b>This method is protected only for testing</b>
      */
-    protected LoadBalancer createLoadBalancer( BoltServerAddress address, ConnectionPool connectionPool, Config config,
-            RoutingSettings routingSettings )
+    protected LoadBalancer createLoadBalancer( BoltServerAddress address, ConnectionPool connectionPool,
+            AsyncConnectionPool asyncConnectionPool, EventExecutorGroup eventExecutorGroup,
+            Config config, RoutingSettings routingSettings )
     {
-        return new LoadBalancer( address, routingSettings, connectionPool, createClock(), config.logging(),
-                createLoadBalancingStrategy( config, connectionPool ) );
+        LoadBalancingStrategy loadBalancingStrategy =
+                createLoadBalancingStrategy( config, connectionPool, asyncConnectionPool );
+        return new LoadBalancer( address, routingSettings, connectionPool, asyncConnectionPool, eventExecutorGroup,
+                createClock(), config.logging(), loadBalancingStrategy );
     }
 
-    private static LoadBalancingStrategy createLoadBalancingStrategy( Config config, ConnectionPool connectionPool )
+    private static LoadBalancingStrategy createLoadBalancingStrategy( Config config, ConnectionPool connectionPool,
+            AsyncConnectionPool asyncConnectionPool )
     {
         switch ( config.loadBalancingStrategy() )
         {
         case ROUND_ROBIN:
             return new RoundRobinLoadBalancingStrategy( config.logging() );
         case LEAST_CONNECTED:
-            return new LeastConnectedLoadBalancingStrategy( connectionPool, config.logging() );
+            return new LeastConnectedLoadBalancingStrategy( connectionPool, asyncConnectionPool, config.logging() );
         default:
             throw new IllegalArgumentException( "Unknown load balancing strategy: " + config.loadBalancingStrategy() );
         }
@@ -253,7 +259,7 @@ protected SessionFactory createSessionFactory( ConnectionProvider connectionProv
     }
 
     /**
-     * Creates new {@link RetryLogic >}.
+     * Creates new {@link RetryLogic}.
      * <p>
      * <b>This method is protected only for testing</b>
      */
@@ -263,6 +269,16 @@ protected RetryLogic createRetryLogic( RetrySettings settings, EventExecutorGrou
         return new ExponentialBackoffRetryLogic( settings, eventExecutorGroup, createClock(), logging );
     }
 
+    /**
+     * Creates new {@link Bootstrap}.
+     * <p>
+     * <b>This method is protected only for testing</b>
+     */
+    protected Bootstrap createBootstrap()
+    {
+        return BootstrapFactory.newBootstrap();
+    }
+
     private static SecurityPlan createSecurityPlan( BoltServerAddress address, Config config )
     {
         try

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

@@ -21,9 +21,10 @@
 import java.util.Map;
 import java.util.concurrent.CompletionStage;
 
+import org.neo4j.driver.internal.net.BoltServerAddress;
 import org.neo4j.driver.internal.spi.ResponseHandler;
+import org.neo4j.driver.internal.util.ServerVersion;
 import org.neo4j.driver.v1.Value;
-import org.neo4j.driver.v1.summary.ServerInfo;
 
 public interface AsyncConnection
 {
@@ -43,5 +44,7 @@ void runAndFlush( String statement, Map<String,Value> parameters, ResponseHandle
 
     CompletionStage<Void> forceRelease();
 
-    ServerInfo serverInfo();
+    BoltServerAddress serverAddress();
+
+    ServerVersion serverVersion();
 }

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

@@ -23,31 +23,42 @@
 
 import org.neo4j.driver.internal.async.inbound.InboundMessageDispatcher;
 import org.neo4j.driver.internal.net.BoltServerAddress;
+import org.neo4j.driver.internal.util.ServerVersion;
 
 import static io.netty.util.AttributeKey.newInstance;
 
 public final class ChannelAttributes
 {
-    private static final AttributeKey<BoltServerAddress> ADDRESS = newInstance( "address" );
+    private static final AttributeKey<BoltServerAddress> ADDRESS = newInstance( "serverAddress" );
+    private static final AttributeKey<ServerVersion> SERVER_VERSION = newInstance( "serverVersion" );
     private static final AttributeKey<Long> CREATION_TIMESTAMP = newInstance( "creationTimestamp" );
     private static final AttributeKey<Long> LAST_USED_TIMESTAMP = newInstance( "lastUsedTimestamp" );
     private static final AttributeKey<InboundMessageDispatcher> MESSAGE_DISPATCHER = newInstance( "messageDispatcher" );
-    private static final AttributeKey<String> SERVER_VERSION = newInstance( "serverVersion" );
 
     private ChannelAttributes()
     {
     }
 
-    public static BoltServerAddress address( Channel channel )
+    public static BoltServerAddress serverAddress( Channel channel )
     {
         return get( channel, ADDRESS );
     }
 
-    public static void setAddress( Channel channel, BoltServerAddress address )
+    public static void setServerAddress( Channel channel, BoltServerAddress address )
     {
         setOnce( channel, ADDRESS, address );
     }
 
+    public static ServerVersion serverVersion( Channel channel )
+    {
+        return get( channel, SERVER_VERSION );
+    }
+
+    public static void setServerVersion( Channel channel, ServerVersion version )
+    {
+        setOnce( channel, SERVER_VERSION, version );
+    }
+
     public static long creationTimestamp( Channel channel )
     {
         return get( channel, CREATION_TIMESTAMP );
@@ -78,16 +89,6 @@ public static void setMessageDispatcher( Channel channel, InboundMessageDispatch
         setOnce( channel, MESSAGE_DISPATCHER, messageDispatcher );
     }
 
-    public static String serverVersion( Channel channel )
-    {
-        return get( channel, SERVER_VERSION );
-    }
-
-    public static void setServerVersion( Channel channel, String serverVersion )
-    {
-        setOnce( channel, SERVER_VERSION, serverVersion );
-    }
-
     private static <T> T get( Channel channel, AttributeKey<T> key )
     {
         return channel.attr( key ).get();