Remember initial routing address

Rediscovery procedure will try to contact initial router when all other
routers have failed or current routing table does not have any routers.

Author: lutovich

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

@@ -143,7 +143,7 @@ protected InternalDriver createDriver( Config config, SecurityPlan securityPlan,
     protected LoadBalancer createLoadBalancer( BoltServerAddress address, ConnectionPool connectionPool, Config config,
             RoutingSettings routingSettings )
     {
-        return new LoadBalancer( routingSettings, connectionPool, createClock(), config.logging(), address );
+        return new LoadBalancer( address, routingSettings, connectionPool, createClock(), config.logging() );
     }
 
     /**

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

@@ -41,26 +41,25 @@ public class LoadBalancer implements ConnectionProvider, RoutingErrorHandler, Au
     private final Rediscovery rediscovery;
     private final Logger log;
 
-    public LoadBalancer( RoutingSettings settings, ConnectionPool connections, Clock clock, Logging logging,
-            BoltServerAddress... routingAddresses )
+    public LoadBalancer( BoltServerAddress initialRouter, RoutingSettings settings, ConnectionPool connections,
+            Clock clock, Logging logging )
     {
-        this( settings, new ClusterRoutingTable( clock, routingAddresses ), connections, clock,
+        this( initialRouter, settings, connections, new ClusterRoutingTable( clock, initialRouter ), clock,
                 logging.getLog( LOAD_BALANCER_LOG_NAME ) );
     }
 
-    private LoadBalancer( RoutingSettings settings, RoutingTable routingTable, ConnectionPool connections,
-            Clock clock, Logger log )
+    private LoadBalancer( BoltServerAddress initialRouter, RoutingSettings settings, ConnectionPool connections,
+            RoutingTable routingTable, Clock clock, Logger log )
     {
-        this( routingTable, connections, new Rediscovery( settings, clock, log,
-                new GetServersProcedureClusterCompositionProvider( clock, log ) ), log );
+        this( connections, routingTable, createRediscovery( initialRouter, settings, clock, log ), log );
     }
 
-    LoadBalancer( RoutingTable routingTable, ConnectionPool connections, Rediscovery rediscovery, Logger log )
+    LoadBalancer( ConnectionPool connections, RoutingTable routingTable, Rediscovery rediscovery, Logger log )
     {
-        this.log = log;
         this.connections = connections;
         this.routingTable = routingTable;
         this.rediscovery = rediscovery;
+        this.log = log;
 
         // initialize the routing table
         ensureRouting();
@@ -153,4 +152,11 @@ private RoundRobinAddressSet addressSetFor( AccessMode mode )
             throw new IllegalArgumentException( "Mode '" + mode + "' is not supported" );
         }
     }
+
+    private static Rediscovery createRediscovery( BoltServerAddress initialRouter, RoutingSettings settings,
+            Clock clock, Logger log )
+    {
+        ClusterCompositionProvider clusterComposition = new GetServersProcedureClusterCompositionProvider( clock, log );
+        return new Rediscovery( initialRouter, settings, clock, log, clusterComposition );
+    }
 }

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

@@ -32,13 +32,16 @@ public class Rediscovery
 {
     private static final String NO_ROUTERS_AVAILABLE = "Could not perform discovery. No routing servers available.";
 
+    private final BoltServerAddress initialRouter;
     private final RoutingSettings settings;
     private final Clock clock;
     private final Logger logger;
     private final ClusterCompositionProvider provider;
 
-    public Rediscovery( RoutingSettings settings, Clock clock, Logger logger, ClusterCompositionProvider provider )
+    public Rediscovery( BoltServerAddress initialRouter, RoutingSettings settings, Clock clock, Logger logger,
+            ClusterCompositionProvider provider )
     {
+        this.initialRouter = initialRouter;
         this.settings = settings;
         this.clock = clock;
         this.logger = logger;
@@ -49,7 +52,6 @@ public Rediscovery( RoutingSettings settings, Clock clock, Logger logger, Cluste
     // cluster composition, which would be used to update the routing table and connection pool
     public ClusterComposition lookupClusterComposition( ConnectionPool connections, RoutingTable routingTable )
     {
-        assertHasRouters( routingTable );
         int failures = 0;
 
         for ( long start = clock.millis(), delay = 0; ; delay = Math.max( settings.retryTimeoutDelay, delay * 2 ) )
@@ -74,13 +76,19 @@ public ClusterComposition lookupClusterComposition( ConnectionPool connections,
     private ClusterComposition lookupClusterCompositionOnKnownRouters( ConnectionPool connections,
             RoutingTable routingTable )
     {
+        boolean triedInitialRouter = false;
         int size = routingTable.routerSize();
         for ( int i = 0; i < size; i++ )
         {
             BoltServerAddress address = routingTable.nextRouter();
             if ( address == null )
             {
-                throw new ServiceUnavailableException( NO_ROUTERS_AVAILABLE );
+                break;
+            }
+
+            if ( address.equals( initialRouter ) )
+            {
+                triedInitialRouter = true;
             }
 
             ClusterComposition composition = lookupClusterCompositionOnRouter( address, connections, routingTable );
@@ -89,7 +97,12 @@ private ClusterComposition lookupClusterCompositionOnKnownRouters( ConnectionPoo
                 return composition;
             }
         }
-        return null;
+
+        if ( triedInitialRouter )
+        {
+            return null;
+        }
+        return lookupClusterCompositionOnRouter( initialRouter, connections, routingTable );
     }
 
     private ClusterComposition lookupClusterCompositionOnRouter( BoltServerAddress routerAddress,
@@ -138,12 +151,4 @@ private void sleep( long millis )
             }
         }
     }
-
-    private void assertHasRouters( RoutingTable table )
-    {
-        if ( table.routerSize() == 0 )
-        {
-            throw new ServiceUnavailableException( NO_ROUTERS_AVAILABLE );
-        }
-    }
 }

driver/src/test/java/org/neo4j/driver/internal/RoutingDriverBoltKitTest.java

@@ -833,6 +833,30 @@ public void shouldRetryWriteTransactionAndPerformRediscoveryUntilSuccess() throw
         }
     }
 
+    @Test
+    public void shouldUseInitialRouterForRediscoveryWhenAllOtherRoutersAreDead() throws Exception
+    {
+        // initial router does not have itself in the returned set of routers
+        StubServer router = StubServer.start( "acquire_endpoints.script", 9010 );
+
+        try ( Driver driver = GraphDatabase.driver( "bolt+routing://127.0.0.1:9010", config ) )
+        {
+            try ( Session session = driver.session( AccessMode.READ ) )
+            {
+                // restart router on the same port with different script that contains itself as reader
+                assertEquals( 0, router.exitStatus() );
+                router = StubServer.start( "rediscover_using_initial_router.script", 9010 );
+
+                List<Record> records = session.run( "MATCH (n) RETURN n.name AS name" ).list();
+                assertEquals( 2, records.size() );
+                assertEquals( "Bob", records.get( 0 ).get( "name" ).asString() );
+                assertEquals( "Alice", records.get( 1 ).get( "name" ).asString() );
+            }
+        }
+
+        assertEquals( 0, router.exitStatus() );
+    }
+
     private static Driver newDriverWithSleeplessClock( String uriString )
     {
         DriverFactory driverFactory = new DriverFactoryWithClock( new SleeplessClock() );

driver/src/test/java/org/neo4j/driver/internal/RoutingDriverTest.java

@@ -345,7 +345,7 @@ private final Driver driverWithServers( long ttl, Map<String,Object>... serverIn
     private Driver driverWithPool( ConnectionPool pool )
     {
         RoutingSettings settings = new RoutingSettings( 10, 5_000 );
-        ConnectionProvider connectionProvider = new LoadBalancer( settings, pool, clock, logging, SEED );
+        ConnectionProvider connectionProvider = new LoadBalancer( SEED, settings, pool, clock, logging );
         Config config = Config.build().withLogging( logging ).toConfig();
         SessionFactory sessionFactory = new NetworkSessionWithAddressFactory( connectionProvider, config );
         return new InternalDriver( insecure(), sessionFactory, logging );

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

@@ -73,7 +73,7 @@ public void ensureRoutingShouldUpdateRoutingTableAndPurgeConnectionPoolWhenStale
         when( routingTable.update( any( ClusterComposition.class ) ) ).thenReturn( set );
 
         // when
-        LoadBalancer balancer = new LoadBalancer( routingTable, conns, rediscovery, DEV_NULL_LOGGER );
+        LoadBalancer balancer = new LoadBalancer( conns, routingTable, rediscovery, DEV_NULL_LOGGER );
 
         // then
         assertNotNull( balancer );
@@ -88,7 +88,7 @@ public void shouldEnsureRoutingOnInitialization() throws Exception
     {
         // given & when
         final AtomicInteger ensureRoutingCounter = new AtomicInteger( 0 );
-        LoadBalancer balancer = new LoadBalancer( mock( RoutingTable.class ), mock( ConnectionPool.class ),
+        LoadBalancer balancer = new LoadBalancer( mock( ConnectionPool.class ), mock( RoutingTable.class ),
                 mock( Rediscovery.class ), DEV_NULL_LOGGER )
         {
             @Override
@@ -143,7 +143,7 @@ public void shouldForgetAddressAndItsConnectionsOnServiceUnavailableWhileClosing
         RoutingTable routingTable = mock( RoutingTable.class );
         ConnectionPool connectionPool = mock( ConnectionPool.class );
         Rediscovery rediscovery = mock( Rediscovery.class );
-        LoadBalancer loadBalancer = new LoadBalancer( routingTable, connectionPool, rediscovery, DEV_NULL_LOGGER );
+        LoadBalancer loadBalancer = new LoadBalancer( connectionPool, routingTable, rediscovery, DEV_NULL_LOGGER );
         BoltServerAddress address = new BoltServerAddress( "host", 42 );
 
         PooledConnection connection = newConnectionWithFailingSync( address );
@@ -174,7 +174,7 @@ public void shouldForgetAddressAndItsConnectionsOnServiceUnavailableWhileClosing
         PooledConnection connectionWithFailingSync = newConnectionWithFailingSync( address );
         when( connectionPool.acquire( any( BoltServerAddress.class ) ) ).thenReturn( connectionWithFailingSync );
         Rediscovery rediscovery = mock( Rediscovery.class );
-        LoadBalancer loadBalancer = new LoadBalancer( routingTable, connectionPool, rediscovery, DEV_NULL_LOGGER );
+        LoadBalancer loadBalancer = new LoadBalancer( connectionPool, routingTable, rediscovery, DEV_NULL_LOGGER );
 
         Session session = newSession( loadBalancer );
         // begin transaction to make session obtain a connection
@@ -205,7 +205,7 @@ private LoadBalancer setupLoadBalancer( PooledConnection writerConn, PooledConne
         when( routingTable.readers() ).thenReturn( readerAddrs );
         when( routingTable.writers() ).thenReturn( writerAddrs );
 
-        return new LoadBalancer( routingTable, connPool, mock( Rediscovery.class ), DEV_NULL_LOGGER );
+        return new LoadBalancer( connPool, routingTable, mock( Rediscovery.class ), DEV_NULL_LOGGER );
     }
 
     private static Session newSession( LoadBalancer loadBalancer )