Updated handling of getServers response

`getServers` now only returns a single row containing all the information, together with a
`ttl` that is no longer per server but a global `ttl` for the returned servers.

Author: pontusmelke

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

@@ -1,15 +1,15 @@
 /**
  * Copyright (c) 2002-2016 "Neo Technology,"
  * Network Engine for Objects in Lund AB [http://neotechnology.com]
- *
+ * <p>
  * This file is part of Neo4j.
- *
+ * <p>
  * 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
- *
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
  * 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.
@@ -20,6 +20,7 @@
 
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.List;
 import java.util.Set;
 
 import org.neo4j.driver.internal.net.BoltServerAddress;
@@ -34,10 +35,12 @@
 import org.neo4j.driver.v1.Record;
 import org.neo4j.driver.v1.Session;
 import org.neo4j.driver.v1.StatementResult;
+import org.neo4j.driver.v1.Value;
 import org.neo4j.driver.v1.exceptions.ClientException;
 import org.neo4j.driver.v1.exceptions.ConnectionFailureException;
 import org.neo4j.driver.v1.exceptions.ServiceUnavailableException;
 import org.neo4j.driver.v1.util.BiFunction;
+import org.neo4j.driver.v1.util.Function;
 
 import static java.lang.String.format;
 
@@ -50,7 +53,7 @@ public class ClusterDriver extends BaseDriver
         public int compare( BoltServerAddress o1, BoltServerAddress o2 )
         {
             int compare = o1.host().compareTo( o2.host() );
-            if (compare == 0)
+            if ( compare == 0 )
             {
                 compare = Integer.compare( o1.port(), o2.port() );
             }
@@ -60,16 +63,17 @@ public int compare( BoltServerAddress o1, BoltServerAddress o2 )
     };
     private static final int MIN_SERVERS = 2;
     private final ConnectionPool connections;
-    private final BiFunction<Connection,Logger, Session> sessionProvider;
+    private final BiFunction<Connection,Logger,Session> sessionProvider;
 
-    private final ConcurrentRoundRobinSet<BoltServerAddress> routingServers = new ConcurrentRoundRobinSet<>(COMPARATOR);
-    private final ConcurrentRoundRobinSet<BoltServerAddress> readServers = new ConcurrentRoundRobinSet<>(COMPARATOR);
-    private final ConcurrentRoundRobinSet<BoltServerAddress> writeServers = new ConcurrentRoundRobinSet<>(COMPARATOR);
+    private final ConcurrentRoundRobinSet<BoltServerAddress> routingServers =
+            new ConcurrentRoundRobinSet<>( COMPARATOR );
+    private final ConcurrentRoundRobinSet<BoltServerAddress> readServers = new ConcurrentRoundRobinSet<>( COMPARATOR );
+    private final ConcurrentRoundRobinSet<BoltServerAddress> writeServers = new ConcurrentRoundRobinSet<>( COMPARATOR );
 
     public ClusterDriver( BoltServerAddress seedAddress,
             ConnectionPool connections,
             SecurityPlan securityPlan,
-            BiFunction<Connection,Logger, Session> sessionProvider,
+            BiFunction<Connection,Logger,Session> sessionProvider,
             Logging logging )
     {
         super( securityPlan, logging );
@@ -85,7 +89,7 @@ private void checkServers()
         {
             if ( routingServers.size() < MIN_SERVERS ||
                  readServers.isEmpty() ||
-                 writeServers.isEmpty())
+                 writeServers.isEmpty() )
             {
                 getServers();
             }
@@ -107,18 +111,22 @@ private void getServers()
                     @Override
                     public void accept( Record record )
                     {
-                        BoltServerAddress newAddress = new BoltServerAddress( record.get( "address" ).asString() );
-                        switch ( record.get( "mode" ).asString().toUpperCase() )
+                        long ttl = record.get( "ttl" ).asLong();
+                        List<ServerInfo> servers = servers( record );
+                        for ( ServerInfo server : servers )
                         {
-                        case "READ":
-                            readServers.add( newAddress );
-                            break;
-                        case "WRITE":
-                            writeServers.add( newAddress );
-                            break;
-                        case "ROUTE":
-                            routingServers.add( newAddress );
-                            break;
+                            switch ( server.role() )
+                            {
+                            case "READ":
+                                readServers.addAll( server.addresses() );
+                                break;
+                            case "WRITE":
+                                writeServers.addAll( server.addresses() );
+                                break;
+                            case "ROUTE":
+                                routingServers.addAll( server.addresses() );
+                                break;
+                            }
                         }
                     }
                 } );
@@ -137,7 +145,7 @@ public void accept( Record record )
                 this.close();
                 throw new ServiceUnavailableException(
                         String.format( "Server %s couldn't perform discovery",
-                                address == null ? "`UNKNOWN`" : address.toString()),  ex );
+                                address == null ? "`UNKNOWN`" : address.toString() ), ex );
             }
             else
             {
@@ -146,14 +154,55 @@ public void accept( Record record )
         }
     }
 
+    private static class ServerInfo
+    {
+        private final List<BoltServerAddress> addresses;
+        private final String role;
+
+        public ServerInfo( List<BoltServerAddress> addresses, String role )
+        {
+            this.addresses = addresses;
+            this.role = role;
+        }
+
+        public String role()
+        {
+            return role;
+        }
+
+        List<BoltServerAddress> addresses()
+        {
+            return addresses;
+        }
+    }
+
+    private List<ServerInfo> servers( Record record )
+    {
+        return record.get( "servers" ).asList( new Function<Value,ServerInfo>()
+        {
+            @Override
+            public ServerInfo apply( Value value )
+            {
+                return new ServerInfo( value.get("addresses").asList( new Function<Value,BoltServerAddress>()
+                {
+                    @Override
+                    public BoltServerAddress apply( Value value )
+                    {
+                        return new BoltServerAddress( value.asString() );
+                    }
+                } ), value.get("role").asString() );
+            }
+        } );
+    }
+
     //must be called from a synchronized method
     private boolean call( BoltServerAddress address, String procedureName, Consumer<Record> recorder )
     {
         Connection acquire = null;
         Session session = null;
         try
         {
-            acquire = connections.acquire(address);
+            acquire = connections.acquire( address );
             session = sessionProvider.apply( acquire, log );
 
             StatementResult records = session.run( format( "CALL %s", procedureName ) );
@@ -255,13 +304,13 @@ Set<BoltServerAddress> routingServers()
     //For testing
     Set<BoltServerAddress> readServers()
     {
-        return Collections.unmodifiableSet(readServers);
+        return Collections.unmodifiableSet( readServers );
     }
 
     //For testing
     Set<BoltServerAddress> writeServers()
     {
-        return Collections.unmodifiableSet( writeServers);
+        return Collections.unmodifiableSet( writeServers );
     }
 
     //For testing

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

@@ -1,15 +1,15 @@
 /**
  * Copyright (c) 2002-2016 "Neo Technology,"
  * Network Engine for Objects in Lund AB [http://neotechnology.com]
- *
+ * <p>
  * This file is part of Neo4j.
- *
+ * <p>
  * 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
- *
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
  * 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.
@@ -24,14 +24,13 @@
 import org.junit.rules.ExpectedException;
 
 import java.util.Collections;
-import java.util.Iterator;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import org.neo4j.driver.internal.net.BoltServerAddress;
 import org.neo4j.driver.internal.spi.Connection;
 import org.neo4j.driver.internal.spi.ConnectionPool;
-import org.neo4j.driver.internal.value.IntegerValue;
-import org.neo4j.driver.internal.value.StringValue;
 import org.neo4j.driver.v1.AccessMode;
 import org.neo4j.driver.v1.Logger;
 import org.neo4j.driver.v1.Logging;
@@ -56,6 +55,7 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 import static org.neo4j.driver.internal.security.SecurityPlan.insecure;
+import static org.neo4j.driver.v1.Values.value;
 
 public class ClusterDriverTest
 {
@@ -64,17 +64,17 @@ public class ClusterDriverTest
 
     private static final BoltServerAddress SEED = new BoltServerAddress( "localhost", 7687 );
     private static final String GET_SERVERS = "CALL dbms.cluster.routing.getServers";
-    private static final List<BoltServerAddress> NO_ADDRESSES = Collections.<BoltServerAddress>emptyList();
+    private static final List<String> NO_ADDRESSES = Collections.emptyList();
 
     @Test
     public void shouldDoRoutingOnInitialization()
     {
         // Given
         final Session session = mock( Session.class );
         when( session.run( GET_SERVERS ) ).thenReturn(
-                getServers( singletonList( boltAddress( "localhost", 1111 ) ),
-                        singletonList( boltAddress( "localhost", 2222 ) ),
-                        singletonList( boltAddress( "localhost", 3333 ) ) ) );
+                getServers( singletonList( "localhost:1111" ),
+                        singletonList( "localhost:2222" ),
+                        singletonList( "localhost:3333" ) ) );
 
         // When
         ClusterDriver clusterDriver = forSession( session );
@@ -96,11 +96,11 @@ public void shouldDoReRoutingOnSessionAcquisitionIfNecessary()
         final Session session = mock( Session.class );
         when( session.run( GET_SERVERS ) )
                 .thenReturn(
-                        getServers( singletonList( boltAddress( "localhost", 1111 ) ), NO_ADDRESSES, NO_ADDRESSES ) )
+                        getServers( singletonList( "localhost:1111" ), NO_ADDRESSES, NO_ADDRESSES ) )
                 .thenReturn(
-                        getServers( singletonList( boltAddress( "localhost", 1112 ) ),
-                                singletonList( boltAddress( "localhost", 2222 ) ),
-                                singletonList( boltAddress( "localhost", 3333 ) ) ) );
+                        getServers( singletonList( "localhost:1112" ),
+                                singletonList( "localhost:2222" ),
+                                singletonList( "localhost:3333" ) ) );
 
         ClusterDriver clusterDriver = forSession( session );
 
@@ -129,12 +129,11 @@ public void shouldNotDoReRoutingOnSessionAcquisitionIfNotNecessary()
         final Session session = mock( Session.class );
         when( session.run( GET_SERVERS ) )
                 .thenReturn(
-                        getServers( asList( boltAddress( "localhost", 1111 ), boltAddress( "localhost", 1112 ),
-                                boltAddress( "localhost", 1113 ) ),
-                                singletonList( boltAddress( "localhost", 2222 ) ),
-                                singletonList( boltAddress( "localhost", 3333 ) ) ) )
+                        getServers( asList( "localhost:1111", "localhost:1112", "localhost:1113" ),
+                                singletonList( "localhost:2222" ),
+                                singletonList( "localhost:3333" ) ) )
                 .thenReturn(
-                        getServers( singletonList( boltAddress( "localhost", 5555 ) ), NO_ADDRESSES, NO_ADDRESSES ) );
+                        getServers( singletonList( "localhost:5555" ), NO_ADDRESSES, NO_ADDRESSES ) );
 
         ClusterDriver clusterDriver = forSession( session );
 
@@ -180,59 +179,34 @@ private BoltServerAddress boltAddress( String host, int port )
         return new BoltServerAddress( host, port );
     }
 
-    StatementResult getServers( final List<BoltServerAddress> routers, final List<BoltServerAddress> readers,
-            final List<BoltServerAddress> writers )
+    StatementResult getServers( final List<String> routers, final List<String> readers,
+            final List<String> writers )
     {
-
-
         return new StatementResult()
         {
-            private final int totalSize = routers.size() + readers.size() + writers.size();
-            private final Iterator<BoltServerAddress> routeIterator = routers.iterator();
-            private final Iterator<BoltServerAddress> readIterator = readers.iterator();
-            private final Iterator<BoltServerAddress> writeIterator = writers.iterator();
             private int counter = 0;
 
             @Override
             public List<String> keys()
             {
-                return asList( "address", "mode", "expires" );
+                return asList( "ttl", "servers" );
             }
 
             @Override
             public boolean hasNext()
             {
-                return counter++ < totalSize;
+                return counter++ < 1;
             }
 
             @Override
             public Record next()
             {
-                if ( routeIterator.hasNext() )
-                {
-                    return new InternalRecord( asList( "address", "mode", "expires" ),
-                            new Value[]{new StringValue( routeIterator.next().toString() ),
-                                    new StringValue( "ROUTE" ),
-                                    new IntegerValue( Long.MAX_VALUE )} );
-                }
-                else if ( readIterator.hasNext() )
-                {
-                    return new InternalRecord( asList( "address", "mode", "expires" ),
-                            new Value[]{new StringValue( readIterator.next().toString() ),
-                                    new StringValue( "READ" ),
-                                    new IntegerValue( Long.MAX_VALUE )} );
-                }
-                else if ( writeIterator.hasNext() )
-                {
-                    return new InternalRecord( asList( "address", "mode", "expires" ),
-                            new Value[]{new StringValue( writeIterator.next().toString() ),
-                                    new StringValue( "WRITE" ),
-                                    new IntegerValue( Long.MAX_VALUE )} );
-                }
-                else
-                {
-                    return Collections.<Record>emptyIterator().next();
-                }
+                return new InternalRecord( asList( "ttl", "servers" ),
+                        new Value[]{
+                                value( Long.MAX_VALUE ),
+                                value( asList( serverInfo( "ROUTE", routers ), serverInfo( "WRITE", writers ),
+                                        serverInfo( "READ", readers ) ) )
+                        } );
             }
 
             @Override
@@ -268,11 +242,20 @@ public ResultSummary consume()
             @Override
             public void remove()
             {
-                throw new UnsupportedOperationException(  );
+                throw new UnsupportedOperationException();
             }
         };
     }
 
+    private Map<String,Object> serverInfo( String role, List<String> addresses )
+    {
+        Map<String,Object> map = new HashMap<>();
+        map.put( "role", role );
+        map.put( "addresses", addresses );
+
+        return map;
+    }
+
     private ConnectionPool pool()
     {
         ConnectionPool pool = mock( ConnectionPool.class );