Configurable server address resolver

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

@@ -24,12 +24,14 @@
 import java.net.URI;
 import java.net.UnknownHostException;
 
+import org.neo4j.driver.v1.net.ServerAddress;
+
 import static java.util.Objects.requireNonNull;
 
 /**
  * Holds a host and port pair that denotes a Bolt server address.
  */
-public class BoltServerAddress
+public class BoltServerAddress implements ServerAddress
 {
     public static final int DEFAULT_PORT = 7687;
     public static final BoltServerAddress LOCAL_DEFAULT = new BoltServerAddress( "localhost", DEFAULT_PORT );
@@ -50,11 +52,18 @@ public BoltServerAddress( URI uri )
 
     public BoltServerAddress( String host, int port )
     {
-        this.host = requireNonNull( host );
-        this.port = port;
+        this.host = requireNonNull( host, "host" );
+        this.port = requireValidPort( port );
         this.stringValue = String.format( "%s:%d", host, port );
     }
 
+    public static BoltServerAddress from( ServerAddress address )
+    {
+        return address instanceof BoltServerAddress
+               ? (BoltServerAddress) address
+               : new BoltServerAddress( address.host(), address.port() );
+    }
+
     @Override
     public boolean equals( Object o )
     {
@@ -114,11 +123,13 @@ public BoltServerAddress resolve() throws UnknownHostException
         }
     }
 
+    @Override
     public String host()
     {
         return host;
     }
 
+    @Override
     public int port()
     {
         return port;
@@ -194,4 +205,13 @@ private static RuntimeException invalidAddressFormat( String address )
     {
         return new IllegalArgumentException( "Invalid address format `" + address + "`" );
     }
+
+    private static int requireValidPort( int port )
+    {
+        if ( port >= 0 && port <= 65_535 )
+        {
+            return port;
+        }
+        throw new IllegalArgumentException( "Illegal port: " + port );
+    }
 }

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

@@ -31,6 +31,7 @@
 import org.neo4j.driver.internal.async.ChannelConnectorImpl;
 import org.neo4j.driver.internal.async.pool.ConnectionPoolImpl;
 import org.neo4j.driver.internal.async.pool.PoolSettings;
+import org.neo4j.driver.internal.cluster.DnsResolver;
 import org.neo4j.driver.internal.cluster.RoutingContext;
 import org.neo4j.driver.internal.cluster.RoutingSettings;
 import org.neo4j.driver.internal.cluster.loadbalancing.LeastConnectedLoadBalancingStrategy;
@@ -58,6 +59,7 @@
 import org.neo4j.driver.v1.Logging;
 import org.neo4j.driver.v1.exceptions.ClientException;
 import org.neo4j.driver.v1.exceptions.ServiceUnavailableException;
+import org.neo4j.driver.v1.net.ServerAddressResolver;
 
 import static java.lang.String.format;
 import static org.neo4j.driver.internal.metrics.InternalAbstractMetrics.DEV_NULL_METRICS;
@@ -204,8 +206,9 @@ protected LoadBalancer createLoadBalancer( BoltServerAddress address, Connection
             EventExecutorGroup eventExecutorGroup, Config config, RoutingSettings routingSettings )
     {
         LoadBalancingStrategy loadBalancingStrategy = createLoadBalancingStrategy( config, connectionPool );
+        ServerAddressResolver resolver = createResolver( config );
         return new LoadBalancer( address, routingSettings, connectionPool, eventExecutorGroup, createClock(),
-                config.logging(), loadBalancingStrategy );
+                config.logging(), loadBalancingStrategy, resolver );
     }
 
     private static LoadBalancingStrategy createLoadBalancingStrategy( Config config,
@@ -222,6 +225,12 @@ private static LoadBalancingStrategy createLoadBalancingStrategy( Config config,
         }
     }
 
+    private static ServerAddressResolver createResolver( Config config )
+    {
+        ServerAddressResolver configuredResolver = config.resolver();
+        return configuredResolver != null ? configuredResolver : new DnsResolver( config.logging() );
+    }
+
     /**
      * Creates new {@link Clock}.
      * <p>

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

@@ -20,42 +20,40 @@
 
 import java.net.InetAddress;
 import java.net.UnknownHostException;
-import java.util.HashSet;
 import java.util.Set;
+import java.util.stream.Stream;
 
 import org.neo4j.driver.internal.BoltServerAddress;
 import org.neo4j.driver.v1.Logger;
+import org.neo4j.driver.v1.Logging;
+import org.neo4j.driver.v1.net.ServerAddress;
+import org.neo4j.driver.v1.net.ServerAddressResolver;
 
-public class DnsResolver implements HostNameResolver
+import static java.util.Collections.singleton;
+import static java.util.stream.Collectors.toSet;
+
+public class DnsResolver implements ServerAddressResolver
 {
     private final Logger logger;
 
-    public DnsResolver( Logger logger )
+    public DnsResolver( Logging logging )
     {
-        this.logger = logger;
+        this.logger = logging.getLog( DnsResolver.class.getSimpleName() );
     }
 
     @Override
-    public Set<BoltServerAddress> resolve( BoltServerAddress initialRouter )
+    public Set<ServerAddress> resolve( ServerAddress initialRouter )
     {
-        Set<BoltServerAddress> addresses = new HashSet<>();
         try
         {
-            InetAddress[] ipAddresses = InetAddress.getAllByName( initialRouter.host() );
-
-            for ( InetAddress ipAddress : ipAddresses )
-            {
-                addresses.add( new BoltServerAddress( ipAddress.getHostAddress(), initialRouter.port() ) );
-            }
-
-            return addresses;
+            return Stream.of( InetAddress.getAllByName( initialRouter.host() ) )
+                    .map( address -> new BoltServerAddress( address.getHostAddress(), initialRouter.port() ) )
+                    .collect( toSet() );
         }
         catch ( UnknownHostException e )
         {
-            logger.error( "Failed to resolve URI `" + initialRouter + "` to IPs due to error: " + e.getMessage(), e );
-
-            addresses.add( initialRouter );
-            return addresses;
+            logger.error( "Failed to resolve address `" + initialRouter + "` to IPs due to error: " + e.getMessage(), e );
+            return singleton( initialRouter );
         }
     }
 }

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

@@ -20,8 +20,8 @@
 
 import io.netty.util.concurrent.EventExecutorGroup;
 
-import java.util.Collections;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
@@ -35,9 +35,13 @@
 import org.neo4j.driver.v1.Logger;
 import org.neo4j.driver.v1.exceptions.SecurityException;
 import org.neo4j.driver.v1.exceptions.ServiceUnavailableException;
+import org.neo4j.driver.v1.net.ServerAddressResolver;
 
 import static java.lang.String.format;
+import static java.util.Collections.emptySet;
+import static java.util.Collections.singletonList;
 import static java.util.concurrent.CompletableFuture.completedFuture;
+import static java.util.stream.Collectors.toList;
 import static org.neo4j.driver.internal.util.Futures.completedWithNull;
 
 public class Rediscovery
@@ -48,27 +52,26 @@ public class Rediscovery
     private final RoutingSettings settings;
     private final Logger logger;
     private final ClusterCompositionProvider provider;
-    private final HostNameResolver hostNameResolver;
+    private final ServerAddressResolver resolver;
     private final EventExecutorGroup eventExecutorGroup;
 
     private volatile boolean useInitialRouter;
 
     public Rediscovery( BoltServerAddress initialRouter, RoutingSettings settings, ClusterCompositionProvider provider,
-            EventExecutorGroup eventExecutorGroup, HostNameResolver hostNameResolver, Logger logger )
+            EventExecutorGroup eventExecutorGroup, ServerAddressResolver resolver, Logger logger )
     {
-        this( initialRouter, settings, provider, hostNameResolver, eventExecutorGroup, logger, true );
+        this( initialRouter, settings, provider, resolver, eventExecutorGroup, logger, true );
     }
 
     // Test-only constructor
-    public Rediscovery( BoltServerAddress initialRouter, RoutingSettings settings, ClusterCompositionProvider provider,
-            HostNameResolver hostNameResolver, EventExecutorGroup eventExecutorGroup, Logger logger,
-            boolean useInitialRouter )
+    Rediscovery( BoltServerAddress initialRouter, RoutingSettings settings, ClusterCompositionProvider provider,
+            ServerAddressResolver resolver, EventExecutorGroup eventExecutorGroup, Logger logger, boolean useInitialRouter )
     {
         this.initialRouter = initialRouter;
         this.settings = settings;
         this.logger = logger;
         this.provider = provider;
-        this.hostNameResolver = hostNameResolver;
+        this.resolver = resolver;
         this.eventExecutorGroup = eventExecutorGroup;
         this.useInitialRouter = useInitialRouter;
     }
@@ -163,7 +166,7 @@ private CompletionStage<ClusterComposition> lookupOnKnownRoutersThenOnInitialRou
     private CompletionStage<ClusterComposition> lookupOnInitialRouterThenOnKnownRouters( RoutingTable routingTable,
             ConnectionPool connectionPool )
     {
-        Set<BoltServerAddress> seenServers = Collections.emptySet();
+        Set<BoltServerAddress> seenServers = emptySet();
         return lookupOnInitialRouter( routingTable, connectionPool, seenServers ).thenCompose( composition ->
         {
             if ( composition != null )
@@ -201,7 +204,7 @@ private CompletionStage<ClusterComposition> lookupOnKnownRouters( RoutingTable r
     private CompletionStage<ClusterComposition> lookupOnInitialRouter( RoutingTable routingTable,
             ConnectionPool connectionPool, Set<BoltServerAddress> seenServers )
     {
-        Set<BoltServerAddress> addresses = hostNameResolver.resolve( initialRouter );
+        List<BoltServerAddress> addresses = resolve( initialRouter );
         addresses.removeAll( seenServers );
 
         CompletableFuture<ClusterComposition> result = completedWithNull();
@@ -255,4 +258,19 @@ private ClusterComposition handleRoutingProcedureError( Throwable error, Routing
         }
     }
 
+    private List<BoltServerAddress> resolve( BoltServerAddress address )
+    {
+        try
+        {
+            return resolver.resolve( address )
+                    .stream()
+                    .map( BoltServerAddress::from )
+                    .collect( toList() ); // collect to list to preserve the order
+        }
+        catch ( Throwable error )
+        {
+            logger.error( "Resolver function failed to resolve '" + address + "'. The address will be used as is", error );
+            return singletonList( address );
+        }
+    }
 }

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

@@ -30,7 +30,6 @@
 import org.neo4j.driver.internal.cluster.ClusterComposition;
 import org.neo4j.driver.internal.cluster.ClusterCompositionProvider;
 import org.neo4j.driver.internal.cluster.ClusterRoutingTable;
-import org.neo4j.driver.internal.cluster.DnsResolver;
 import org.neo4j.driver.internal.cluster.Rediscovery;
 import org.neo4j.driver.internal.cluster.RoutingProcedureClusterCompositionProvider;
 import org.neo4j.driver.internal.cluster.RoutingSettings;
@@ -45,6 +44,7 @@
 import org.neo4j.driver.v1.Logging;
 import org.neo4j.driver.v1.exceptions.ServiceUnavailableException;
 import org.neo4j.driver.v1.exceptions.SessionExpiredException;
+import org.neo4j.driver.v1.net.ServerAddressResolver;
 
 import static java.util.concurrent.CompletableFuture.completedFuture;
 
@@ -63,15 +63,15 @@ public class LoadBalancer implements ConnectionProvider, RoutingErrorHandler
 
     public LoadBalancer( BoltServerAddress initialRouter, RoutingSettings settings, ConnectionPool connectionPool,
             EventExecutorGroup eventExecutorGroup, Clock clock, Logging logging,
-            LoadBalancingStrategy loadBalancingStrategy )
+            LoadBalancingStrategy loadBalancingStrategy, ServerAddressResolver resolver )
     {
         this( connectionPool, new ClusterRoutingTable( clock, initialRouter ),
-                createRediscovery( initialRouter, settings, eventExecutorGroup, clock, logging ),
+                createRediscovery( initialRouter, settings, eventExecutorGroup, resolver, clock, logging ),
                 loadBalancerLogger( logging ), loadBalancingStrategy, eventExecutorGroup );
     }
 
     // Used only in testing
-    public LoadBalancer( ConnectionPool connectionPool, RoutingTable routingTable, Rediscovery rediscovery,
+    LoadBalancer( ConnectionPool connectionPool, RoutingTable routingTable, Rediscovery rediscovery,
             EventExecutorGroup eventExecutorGroup, Logging logging )
     {
         this( connectionPool, routingTable, rediscovery, loadBalancerLogger( logging ),
@@ -264,13 +264,11 @@ private BoltServerAddress selectAddress( AccessMode mode, AddressSet servers )
     }
 
     private static Rediscovery createRediscovery( BoltServerAddress initialRouter, RoutingSettings settings,
-            EventExecutorGroup eventExecutorGroup, Clock clock, Logging logging )
+            EventExecutorGroup eventExecutorGroup, ServerAddressResolver resolver, Clock clock, Logging logging )
     {
         Logger log = loadBalancerLogger( logging );
-        ClusterCompositionProvider clusterCompositionProvider =
-                new RoutingProcedureClusterCompositionProvider( clock, settings );
-        return new Rediscovery( initialRouter, settings, clusterCompositionProvider, eventExecutorGroup,
-                new DnsResolver( log ), log );
+        ClusterCompositionProvider clusterCompositionProvider = new RoutingProcedureClusterCompositionProvider( clock, settings );
+        return new Rediscovery( initialRouter, settings, clusterCompositionProvider, eventExecutorGroup, resolver, log );
     }
 
     private static Logger loadBalancerLogger( Logging logging )