Decouple load balancing logic from address set

driver/src/main/java/org/neo4j/driver/internal/cluster/RoundRobinAddressSet.java → driver/src/main/java/org/neo4j/driver/internal/cluster/AddressSet.java

@@ -20,39 +20,23 @@
 
 import java.util.Arrays;
 import java.util.Set;
-import java.util.concurrent.atomic.AtomicInteger;
 
 import org.neo4j.driver.internal.net.BoltServerAddress;
 
-public class RoundRobinAddressSet
+public class AddressSet
 {
     private static final BoltServerAddress[] NONE = {};
-    private final AtomicInteger offset = new AtomicInteger();
-    private volatile BoltServerAddress[] addresses = NONE;
 
-    public int size()
-    {
-        return addresses.length;
-    }
+    private volatile BoltServerAddress[] addresses = NONE;
 
-    public BoltServerAddress next()
+    public BoltServerAddress[] toArray()
     {
-        BoltServerAddress[] addresses = this.addresses;
-        if ( addresses.length == 0 )
-        {
-            return null;
-        }
-        return addresses[next( addresses.length )];
+        return addresses;
     }
 
-    int next( int divisor )
+    public int size()
     {
-        int index = offset.getAndIncrement();
-        for ( ; index == Integer.MAX_VALUE; index = offset.getAndIncrement() )
-        {
-            offset.compareAndSet( Integer.MIN_VALUE, index % divisor );
-        }
-        return index % divisor;
+        return addresses.length;
     }
 
     public synchronized void update( Set<BoltServerAddress> addresses, Set<BoltServerAddress> removed )
@@ -132,12 +116,6 @@ public synchronized void remove( BoltServerAddress address )
     @Override
     public String toString()
     {
-        return "RoundRobinAddressSet=" + Arrays.toString( addresses );
-    }
-
-    /** breaking encapsulation in order to perform white-box testing of boundary case */
-    void setOffset( int target )
-    {
-        offset.set( target );
+        return "AddressSet=" + Arrays.toString( addresses );
     }
 }

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

@@ -36,9 +36,9 @@ public class ClusterRoutingTable implements RoutingTable
 
     private final Clock clock;
     private volatile long expirationTimeout;
-    private final RoundRobinAddressSet readers;
-    private final RoundRobinAddressSet writers;
-    private final RoundRobinAddressSet routers;
+    private final AddressSet readers;
+    private final AddressSet writers;
+    private final AddressSet routers;
 
     public ClusterRoutingTable( Clock clock, BoltServerAddress... routingAddresses )
     {
@@ -51,9 +51,9 @@ private ClusterRoutingTable( Clock clock )
         this.clock = clock;
         this.expirationTimeout = clock.millis() - 1;
 
-        this.readers = new RoundRobinAddressSet();
-        this.writers = new RoundRobinAddressSet();
-        this.routers = new RoundRobinAddressSet();
+        this.readers = new AddressSet();
+        this.writers = new AddressSet();
+        this.routers = new AddressSet();
     }
 
     @Override
@@ -85,27 +85,21 @@ public synchronized void forget( BoltServerAddress address )
     }
 
     @Override
-    public RoundRobinAddressSet readers()
+    public AddressSet readers()
     {
         return readers;
     }
 
     @Override
-    public RoundRobinAddressSet writers()
+    public AddressSet writers()
     {
         return writers;
     }
 
     @Override
-    public BoltServerAddress nextRouter()
+    public AddressSet routers()
     {
-        return routers.next();
-    }
-
-    @Override
-    public int routerSize()
-    {
-        return routers.size();
+        return routers;
     }
 
     @Override

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

@@ -39,6 +39,7 @@ public class LoadBalancer implements ConnectionProvider, RoutingErrorHandler, Au
     private final ConnectionPool connections;
     private final RoutingTable routingTable;
     private final Rediscovery rediscovery;
+    private final LoadBalancingStrategy loadBalancingStrategy;
     private final Logger log;
 
     public LoadBalancer( BoltServerAddress initialRouter, RoutingSettings settings, ConnectionPool connections,
@@ -59,6 +60,7 @@ private LoadBalancer( BoltServerAddress initialRouter, RoutingSettings settings,
         this.connections = connections;
         this.routingTable = routingTable;
         this.rediscovery = rediscovery;
+        this.loadBalancingStrategy = new RoundRobinLoadBalancingStrategy();
         this.log = log;
 
         refreshRoutingTable();
@@ -67,7 +69,7 @@ private LoadBalancer( BoltServerAddress initialRouter, RoutingSettings settings,
     @Override
     public PooledConnection acquireConnection( AccessMode mode )
     {
-        RoundRobinAddressSet addressSet = addressSetFor( mode );
+        AddressSet addressSet = addressSetFor( mode );
         PooledConnection connection = acquireConnection( mode, addressSet );
         return new RoutingPooledConnection( connection, this, mode );
     }
@@ -90,10 +92,10 @@ public void close() throws Exception
         connections.close();
     }
 
-    private PooledConnection acquireConnection( AccessMode mode, RoundRobinAddressSet servers )
+    private PooledConnection acquireConnection( AccessMode mode, AddressSet servers )
     {
         ensureRouting( mode );
-        for ( BoltServerAddress address; (address = servers.next()) != null; )
+        for ( BoltServerAddress address; (address = selectAddress( mode, servers )) != null; )
         {
             try
             {
@@ -141,7 +143,7 @@ synchronized void refreshRoutingTable()
         log.info( "Refreshed routing information. %s", routingTable );
     }
 
-    private RoundRobinAddressSet addressSetFor( AccessMode mode )
+    private AddressSet addressSetFor( AccessMode mode )
     {
         switch ( mode )
         {
@@ -150,7 +152,22 @@ private RoundRobinAddressSet addressSetFor( AccessMode mode )
         case WRITE:
             return routingTable.writers();
         default:
-            throw new IllegalArgumentException( "Mode '" + mode + "' is not supported" );
+            throw unknownMode( mode );
+        }
+    }
+
+    private BoltServerAddress selectAddress( AccessMode mode, AddressSet servers )
+    {
+        BoltServerAddress[] addresses = servers.toArray();
+
+        switch ( mode )
+        {
+        case READ:
+            return loadBalancingStrategy.selectReader( addresses );
+        case WRITE:
+            return loadBalancingStrategy.selectWriter( addresses );
+        default:
+            throw unknownMode( mode );
         }
     }
 
@@ -161,4 +178,9 @@ private static Rediscovery createRediscovery( BoltServerAddress initialRouter, R
                 new RoutingProcedureClusterCompositionProvider( clock, log, settings );
         return new Rediscovery( initialRouter, settings, clock, log, clusterComposition, new DnsResolver( log ) );
     }
+
+    private static RuntimeException unknownMode( AccessMode mode )
+    {
+        return new IllegalArgumentException( "Mode '" + mode + "' is not supported" );
+    }
 }

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

@@ -0,0 +1,28 @@
+/*
+ * 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 org.neo4j.driver.internal.net.BoltServerAddress;
+
+public interface LoadBalancingStrategy
+{
+    BoltServerAddress selectReader( BoltServerAddress[] knownReaders );
+
+    BoltServerAddress selectWriter( BoltServerAddress[] knownWriters );
+}

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

@@ -136,15 +136,10 @@ private ClusterComposition lookupOnInitialRouterThenOnKnownRouters( RoutingTable
     private ClusterComposition lookupOnKnownRouters( RoutingTable routingTable, ConnectionPool connections,
             Set<BoltServerAddress> seenServers )
     {
-        int size = routingTable.routerSize();
-        for ( int i = 0; i < size; i++ )
-        {
-            BoltServerAddress address = routingTable.nextRouter();
-            if ( address == null )
-            {
-                break;
-            }
+        BoltServerAddress[] addresses = routingTable.routers().toArray();
 
+        for ( BoltServerAddress address : addresses )
+        {
             ClusterComposition composition = lookupOnRouter( address, routingTable, connections );
             if ( composition != null )
             {

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

@@ -0,0 +1,53 @@
+/*
+ * 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 java.util.concurrent.atomic.AtomicInteger;
+
+public class RoundRobinArrayIndex
+{
+    private final AtomicInteger offset;
+
+    RoundRobinArrayIndex()
+    {
+        this( 0 );
+    }
+
+    // only for testing
+    RoundRobinArrayIndex( int initialOffset )
+    {
+        this.offset = new AtomicInteger( initialOffset );
+    }
+
+    public int next( int arrayLength )
+    {
+        if ( arrayLength == 0 )
+        {
+            return -1;
+        }
+
+        int nextOffset;
+        while ( (nextOffset = offset.getAndIncrement()) < 0 )
+        {
+            // overflow, try resetting back to zero
+            offset.compareAndSet( nextOffset + 1, 0 );
+        }
+        return nextOffset % arrayLength;
+    }
+}

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

@@ -0,0 +1,50 @@
+/*
+ * 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 org.neo4j.driver.internal.net.BoltServerAddress;
+
+public class RoundRobinLoadBalancingStrategy implements LoadBalancingStrategy
+{
+    private final RoundRobinArrayIndex readersIndex = new RoundRobinArrayIndex();
+    private final RoundRobinArrayIndex writersIndex = new RoundRobinArrayIndex();
+
+    @Override
+    public BoltServerAddress selectReader( BoltServerAddress[] knownReaders )
+    {
+        return select( knownReaders, readersIndex );
+    }
+
+    @Override
+    public BoltServerAddress selectWriter( BoltServerAddress[] knownWriters )
+    {
+        return select( knownWriters, writersIndex );
+    }
+
+    private BoltServerAddress select( BoltServerAddress[] addresses, RoundRobinArrayIndex roundRobinIndex )
+    {
+        int length = addresses.length;
+        if ( length == 0 )
+        {
+            return null;
+        }
+        int index = roundRobinIndex.next( length );
+        return addresses[index];
+    }
+}

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

@@ -31,13 +31,11 @@ public interface RoutingTable
 
     void forget( BoltServerAddress address );
 
-    RoundRobinAddressSet readers();
+    AddressSet readers();
 
-    RoundRobinAddressSet writers();
+    AddressSet writers();
 
-    BoltServerAddress nextRouter();
-
-    int routerSize();
+    AddressSet routers();
 
     void removeWriter( BoltServerAddress toRemove );
 }

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

@@ -808,11 +808,11 @@ public void shouldRetryReadTransactionAndPerformRediscoveryUntilSuccess() throws
     @Test
     public void shouldRetryWriteTransactionAndPerformRediscoveryUntilSuccess() throws Exception
     {
-        StubServer router1 = StubServer.start( "acquire_endpoints.script", 9010 );
-        StubServer brokenWriter1 = StubServer.start( "dead_write_server.script", 9007 );
+        StubServer router1 = StubServer.start( "discover_servers.script", 9010 );
+        StubServer brokenWriter1 = StubServer.start( "dead_write_server.script", 9001 );
+        StubServer router2 = StubServer.start( "acquire_endpoints.script", 9002 );
         StubServer brokenWriter2 = StubServer.start( "dead_write_server.script", 9008 );
-        StubServer router2 = StubServer.start( "discover_servers.script", 9002 );
-        StubServer writer = StubServer.start( "write_server.script", 9001 );
+        StubServer writer = StubServer.start( "write_server.script", 9007 );
 
         try ( Driver driver = newDriverWithSleeplessClock( "bolt+routing://127.0.0.1:9010" );
               Session session = driver.session() )
@@ -827,9 +827,9 @@ public void shouldRetryWriteTransactionAndPerformRediscoveryUntilSuccess() throw
         {
             assertEquals( 0, router1.exitStatus() );
             assertEquals( 0, brokenWriter1.exitStatus() );
-            assertEquals( 0, brokenWriter2.exitStatus() );
             assertEquals( 0, router2.exitStatus() );
             assertEquals( 0, writer.exitStatus() );
+            assertEquals( 0, brokenWriter2.exitStatus() );
         }
     }