WIP handling read sessions

Author: pontusmelke

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

@@ -19,13 +19,11 @@
 
 package org.neo4j.driver.internal;
 
-import java.util.Collections;
-import java.util.HashSet;
 import java.util.Set;
-import java.util.concurrent.ThreadLocalRandom;
 
 import org.neo4j.driver.internal.net.BoltServerAddress;
 import org.neo4j.driver.internal.security.SecurityPlan;
+import org.neo4j.driver.internal.spi.ConnectionPool;
 import org.neo4j.driver.v1.Driver;
 import org.neo4j.driver.v1.Logger;
 import org.neo4j.driver.v1.Logging;
@@ -35,11 +33,12 @@ abstract class BaseDriver implements Driver
 {
     private final SecurityPlan securityPlan;
     protected final Logger log;
-    protected final Set<BoltServerAddress> servers = new HashSet<>();
+    protected final ConnectionPool connections;
 
-    BaseDriver( BoltServerAddress address, SecurityPlan securityPlan, Logging logging )
+    BaseDriver( ConnectionPool connections, BoltServerAddress address, SecurityPlan securityPlan, Logging logging )
     {
-        this.servers.add( address );
+        this.connections = connections;
+        this.connections.add( address );
         this.securityPlan = securityPlan;
         this.log = logging.getLog( Session.LOG_NAME );
     }
@@ -52,24 +51,7 @@ public boolean isEncrypted()
 
     Set<BoltServerAddress> servers()
     {
-        return Collections.unmodifiableSet( servers );
-    }
-
-    //This is somewhat silly and has O(n) complexity
-    protected BoltServerAddress randomServer()
-    {
-        ThreadLocalRandom random = ThreadLocalRandom.current();
-        int item = random.nextInt(servers.size());
-        int i = 0;
-        for ( BoltServerAddress server : servers )
-        {
-            if (i == item)
-            {
-                return server;
-            }
-        }
-
-        throw new IllegalStateException( "This cannot happen" );
+        return connections.keys();
     }
 
 }

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

@@ -16,18 +16,15 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 package org.neo4j.driver.internal;
 
-import java.util.LinkedList;
 import java.util.List;
 
 import org.neo4j.driver.internal.net.BoltServerAddress;
 import org.neo4j.driver.internal.net.pooling.PoolSettings;
 import org.neo4j.driver.internal.net.pooling.SocketConnectionPool;
 import org.neo4j.driver.internal.security.SecurityPlan;
 import org.neo4j.driver.internal.spi.Connection;
-import org.neo4j.driver.internal.spi.ConnectionPool;
 import org.neo4j.driver.internal.util.Consumer;
 import org.neo4j.driver.internal.util.Supplier;
 import org.neo4j.driver.v1.Logging;
@@ -36,8 +33,8 @@
 import org.neo4j.driver.v1.SessionMode;
 import org.neo4j.driver.v1.StatementResult;
 import org.neo4j.driver.v1.exceptions.ClientException;
-import org.neo4j.driver.v1.exceptions.ClusterUnavailableException;
 import org.neo4j.driver.v1.exceptions.ConnectionFailureException;
+import org.neo4j.driver.v1.exceptions.ServiceUnavailableException;
 
 import static java.lang.String.format;
 
@@ -47,50 +44,44 @@ public class ClusterDriver extends BaseDriver
     private static final String ACQUIRE_ENDPOINTS = "dbms.cluster.acquireEndpoints";
     private static final int MINIMUM_NUMBER_OF_SERVERS = 3;
 
-    private final ConnectionPool connections;
+    private final Endpoints endpoints = new Endpoints();
 
-    public ClusterDriver( BoltServerAddress seedAddress, ConnectionSettings connectionSettings, SecurityPlan securityPlan,
-                          PoolSettings poolSettings, Logging logging )
+    public ClusterDriver( BoltServerAddress seedAddress, ConnectionSettings connectionSettings,
+            SecurityPlan securityPlan,
+            PoolSettings poolSettings, Logging logging )
     {
-        super( seedAddress, securityPlan, logging );
-        this.connections = new SocketConnectionPool( connectionSettings, securityPlan, poolSettings, logging );
+        super( new SocketConnectionPool( connectionSettings, securityPlan, poolSettings, logging ),seedAddress, securityPlan, logging );
         discover();
     }
 
-    void discover()
+    synchronized void discover()
     {
-        final List<BoltServerAddress> newServers = new LinkedList<>(  );
         try
         {
             boolean success = false;
-            while ( !servers.isEmpty() && !success )
+            while ( !connections.isEmpty() && !success )
             {
                 success = call( DISCOVER_MEMBERS, new Consumer<Record>()
                 {
                     @Override
                     public void accept( Record record )
                     {
-                        newServers.add( new BoltServerAddress( record.get( "address" ).asString() ) );
+                        connections.add(new BoltServerAddress( record.get( "address" ).asString() ));
                     }
                 } );
-
-            }
-            if ( success )
-            {
-                this.servers.clear();
-                this.servers.addAll( newServers );
-                log.debug( "~~ [MEMBERS] -> %s", newServers );
             }
-            else
+            if ( !success )
             {
-                throw new ClusterUnavailableException( "Run out of servers" );
+                throw new ServiceUnavailableException( "Run out of servers" );
             }
         }
         catch ( ClientException ex )
         {
             if ( ex.code().equals( "Neo.ClientError.Procedure.ProcedureNotFound" ) )
             {
-                throw new ClientException( "Discovery failed: could not find procedure %s", DISCOVER_MEMBERS );
+                //no procedure there, not much to do, stick with what we've got
+                //this may happen because server is running in standalone mode
+                log.warn( "Could not find procedure %s", DISCOVER_MEMBERS );
             }
             else
             {
@@ -99,13 +90,15 @@ public void accept( Record record )
         }
     }
 
+    //must be called from a synchronized method
     private boolean call( String procedureName, Consumer<Record> recorder )
     {
+        Connection acquire = null;
+        Session session = null;
+        try {
+            acquire = connections.acquire();
+            session = new NetworkSession( acquire, log );
 
-        BoltServerAddress address = randomServer();
-        Connection acquire =  connections.acquire( address );
-        try ( Session session = new NetworkSession( acquire, log ) )
-        {
             StatementResult records = session.run( format( "CALL %s", procedureName ) );
             while ( records.hasNext() )
             {
@@ -114,65 +107,157 @@ private boolean call( String procedureName, Consumer<Record> recorder )
         }
         catch ( ConnectionFailureException e )
         {
-            forget(address );
+            if (acquire != null)
+            {
+                forget( acquire.address() );
+            }
             return false;
         }
+        finally
+        {
+            if (acquire != null)
+            {
+                acquire.close();
+            }
+            if (session != null)
+            {
+                session.close();
+            }
+        }
         return true;
     }
 
-    private void forget(BoltServerAddress address)
+    //must be called from a synchronized method
+    private void callWithRetry(String procedureName, Consumer<Record> recorder )
+    {
+        while ( !connections.isEmpty() )
+        {
+            Connection acquire = null;
+            Session session = null;
+            try {
+                acquire = connections.acquire();
+                session = new NetworkSession( acquire, log );
+                List<Record> list = session.run( format( "CALL %s", procedureName ) ).list();
+                for ( Record record : list )
+                {
+                    recorder.accept( record );
+                }
+                //we found results give up
+                return;
+            }
+            catch ( ConnectionFailureException e )
+            {
+                if (acquire != null)
+                {
+                    forget( acquire.address() );
+                }
+            }
+            finally
+            {
+                if (acquire != null)
+                {
+                 acquire.close();
+                }
+                if (session != null)
+                {
+                    session.close();
+                }
+            }
+        }
+
+        throw new ServiceUnavailableException( "Failed to communicate with any of the cluster members" );
+    }
+
+    private synchronized void forget( BoltServerAddress address )
     {
-        servers.remove( address );
-        connections.purge(address);
+        address.markInvalid();
+        connections.purge( address );
     }
 
-    //TODO this could return a WRITE session but that may lead to users using the LEADER too much
-    //a `ClientException` may be what we want
     @Override
     public Session session()
     {
-        throw new UnsupportedOperationException();
+        return session( SessionMode.WRITE );
     }
 
     @Override
     public Session session( final SessionMode mode )
     {
-        return new ClusteredSession( new Supplier<Connection>()
+        switch ( mode )
         {
-            @Override
-            public Connection get()
+        case READ:
+            return new ReadNetworkSession( new Supplier<Connection>()
             {
-                return acquireConnection( mode );
-            }
-        }, log );
+                @Override
+                public Connection get()
+                {
+                    return acquireConnection( mode );
+                }
+            }, new Consumer<Connection>()
+            {
+                @Override
+                public void accept( Connection connection )
+                {
+                    forget( connection.address() );
+                }
+            }, log );
+        case WRITE:
+            throw new UnsupportedOperationException();
+        default:
+            throw new UnsupportedOperationException();
+        }
     }
 
-    private Connection acquireConnection( SessionMode mode )
+    private synchronized Connection acquireConnection( SessionMode mode )
     {
         //if we are short on servers, find new ones
-        if ( servers.size() < MINIMUM_NUMBER_OF_SERVERS )
+        if ( connections.size() < MINIMUM_NUMBER_OF_SERVERS )
         {
             discover();
         }
 
-        final BoltServerAddress[] addresses = new BoltServerAddress[2];
-        call( ACQUIRE_ENDPOINTS, new Consumer<Record>()
+        endpoints.clear();
+        try
+        {
+            callWithRetry( ACQUIRE_ENDPOINTS, new Consumer<Record>()
+            {
+                @Override
+                public void accept( Record record )
+                {
+                    String serverMode = record.get( "mode" ).asString();
+                    if ( serverMode.equals( "READ" ) )
+                    {
+                        endpoints.readServer = new BoltServerAddress( record.get( "address" ).asString() );
+                    }
+                    else if ( serverMode.equals( "WRITE" ) )
+                    {
+                        endpoints.writeServer = new BoltServerAddress( record.get( "address" ).asString() );
+                    }
+                }
+            } );
+        }
+        catch (ClientException e)
         {
-            @Override
-            public void accept( Record record )
+            if ( e.code().equals( "Neo.ClientError.Procedure.ProcedureNotFound" ) )
             {
-                addresses[0] = new BoltServerAddress( record.get( "READ" ).asString() );
-                addresses[1] = new BoltServerAddress( record.get( "WRITE" ).asString() );
+                log.warn( "Could not find procedure %s", ACQUIRE_ENDPOINTS );
+                return connections.acquire();
             }
-        } );
+            throw e;
+        }
+
+        if ( !endpoints.valid() )
+        {
+            throw new ServiceUnavailableException("Could not establish any endpoints for the call");
+        }
 
 
         switch ( mode )
         {
         case READ:
-            return connections.acquire( addresses[0] );
+            return connections.acquire( endpoints.readServer );
         case WRITE:
-            return connections.acquire( addresses[0] );
+            return connections.acquire( endpoints.writeServer );
         default:
             throw new ClientException( mode + " is not supported for creating new sessions" );
         }
@@ -191,4 +276,21 @@ public void close()
         }
     }
 
+    private static class Endpoints
+    {
+        BoltServerAddress readServer;
+        BoltServerAddress writeServer;
+
+        public boolean valid()
+        {
+            return readServer != null && writeServer != null;
+        }
+
+        public void clear()
+        {
+            readServer = null;
+            writeServer = null;
+        }
+    }
+
 }