LoadBalancer implements ConnectionProvider { @Override public CompletionStage<Connection> acquireConnection( ConnectionContext context ) { return routingTables.ensureRoutingTable( context ) .thenCompose( handler -> acquire( context.mode(), handler.routingTable() ) .thenApply( connection -> new RoutingConnection( connection, context.databaseName(), context.mode(), handler ) ) ); } LoadBalancer( BoltServerAddress initialRouter, RoutingSettings settings, ConnectionPool connectionPool,
            EventExecutorGroup eventExecutorGroup, Clock clock, Logging logging,
            LoadBalancingStrategy loadBalancingStrategy, ServerAddressResolver resolver ); private  LoadBalancer( ConnectionPool connectionPool, Rediscovery rediscovery, RoutingSettings settings, LoadBalancingStrategy loadBalancingStrategy,
            EventExecutorGroup eventExecutorGroup, Clock clock, Logger log );   LoadBalancer( ConnectionPool connectionPool, RoutingTableRegistry routingTables, Rediscovery rediscovery, LoadBalancingStrategy loadBalancingStrategy,
            EventExecutorGroup eventExecutorGroup, Logger log ); @Override CompletionStage<Connection> acquireConnection( ConnectionContext context ); @Override CompletionStage<Void> verifyConnectivity(); @Override CompletionStage<Void> close(); @Override CompletionStage<Boolean> supportsMultiDb();  }
@Test void shouldThrowWhenRediscoveryReturnsNoSuitableServers() { ConnectionPool connectionPool = newConnectionPoolMock(); RoutingTable routingTable = mock( RoutingTable.class ); when( routingTable.readers() ).thenReturn( new AddressSet() ); when( routingTable.writers() ).thenReturn( new AddressSet() ); LoadBalancer loadBalancer = newLoadBalancer( connectionPool, routingTable ); SessionExpiredException error1 = assertThrows( SessionExpiredException.class, () -> await( loadBalancer.acquireConnection( contextWithMode( READ ) ) ) ); assertThat( error1.getMessage(), startsWith( "Failed to obtain connection towards READ server" ) ); SessionExpiredException error2 = assertThrows( SessionExpiredException.class, () -> await( loadBalancer.acquireConnection( contextWithMode( WRITE ) ) ) ); assertThat( error2.getMessage(), startsWith( "Failed to obtain connection towards WRITE server" ) ); }
@Test void shouldSelectLeastConnectedAddress() { ConnectionPool connectionPool = newConnectionPoolMock(); when( connectionPool.inUseConnections( A ) ).thenReturn( 0 ); when( connectionPool.inUseConnections( B ) ).thenReturn( 20 ); when( connectionPool.inUseConnections( C ) ).thenReturn( 0 ); RoutingTable routingTable = mock( RoutingTable.class ); AddressSet readerAddresses = mock( AddressSet.class ); when( readerAddresses.toArray() ).thenReturn( new BoltServerAddress[]{A, B, C} ); when( routingTable.readers() ).thenReturn( readerAddresses ); LoadBalancer loadBalancer = newLoadBalancer( connectionPool, routingTable ); Set<BoltServerAddress> seenAddresses = new HashSet<>(); for ( int i = 0; i < 10; i++ ) { Connection connection = await( loadBalancer.acquireConnection( newBoltV4ConnectionContext() ) ); seenAddresses.add( connection.serverAddress() ); } assertEquals( 2, seenAddresses.size() ); assertTrue( seenAddresses.containsAll( asList( A, C ) ) ); }
@Test void shouldRoundRobinWhenNoActiveConnections() { ConnectionPool connectionPool = newConnectionPoolMock(); RoutingTable routingTable = mock( RoutingTable.class ); AddressSet readerAddresses = mock( AddressSet.class ); when( readerAddresses.toArray() ).thenReturn( new BoltServerAddress[]{A, B, C} ); when( routingTable.readers() ).thenReturn( readerAddresses ); LoadBalancer loadBalancer = newLoadBalancer( connectionPool, routingTable ); Set<BoltServerAddress> seenAddresses = new HashSet<>(); for ( int i = 0; i < 10; i++ ) { Connection connection = await( loadBalancer.acquireConnection( newBoltV4ConnectionContext() ) ); seenAddresses.add( connection.serverAddress() ); } assertEquals( 3, seenAddresses.size() ); assertTrue( seenAddresses.containsAll( asList( A, B, C ) ) ); }
@Test void shouldTryMultipleServersAfterRediscovery() { Set<BoltServerAddress> unavailableAddresses = asOrderedSet( A ); ConnectionPool connectionPool = newConnectionPoolMockWithFailures( unavailableAddresses ); RoutingTable routingTable = new ClusterRoutingTable( defaultDatabase(), new FakeClock() ); routingTable.update( new ClusterComposition( -1, new LinkedHashSet<>( Arrays.asList( A, B ) ), emptySet(), emptySet() ) ); LoadBalancer loadBalancer = newLoadBalancer( connectionPool, routingTable ); Connection connection = await( loadBalancer.acquireConnection( newBoltV4ConnectionContext() ) ); assertNotNull( connection ); assertEquals( B, connection.serverAddress() ); assertArrayEquals( new BoltServerAddress[]{B}, routingTable.readers().toArray() ); }