Merge pull request #466 from lutovich/1.4-fix-flaky-tests

Fix couple newly created flaky CC tests

Author: lutovich

driver/src/test/java/org/neo4j/driver/internal/util/ConnectionTrackingConnector.java

@@ -19,6 +19,7 @@
 package org.neo4j.driver.internal.util;
 
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 
 import org.neo4j.driver.internal.net.BoltServerAddress;
 import org.neo4j.driver.internal.spi.Collector;
@@ -29,11 +30,14 @@
 public class ThrowingConnection implements Connection
 {
     private final Connection realConnection;
-    private RuntimeException nextRunFailure;
+    private final AtomicReference<RuntimeException> nextRunError;
+    private final AtomicReference<RuntimeException> nextResetError;
 
     public ThrowingConnection( Connection realConnection )
     {
         this.realConnection = realConnection;
+        this.nextRunError = new AtomicReference<>();
+        this.nextResetError = new AtomicReference<>();
     }
 
     @Override
@@ -45,12 +49,7 @@ public void init( String clientName, Map<String,Value> authToken )
     @Override
     public void run( String statement, Map<String,Value> parameters, Collector collector )
     {
-        if ( nextRunFailure != null )
-        {
-            RuntimeException error = nextRunFailure;
-            nextRunFailure = null;
-            throw error;
-        }
+        throwErrorIfExists( nextRunError );
         realConnection.run( statement, parameters, collector );
     }
 
@@ -69,6 +68,7 @@ public void pullAll( Collector collector )
     @Override
     public void reset()
     {
+        throwErrorIfExists( nextResetError );
         realConnection.reset();
     }
 
@@ -132,8 +132,22 @@ public BoltServerAddress boltServerAddress()
         return realConnection.boltServerAddress();
     }
 
-    public void setNextRunFailure( RuntimeException nextRunFailure )
+    public void setNextRunError( RuntimeException error )
+    {
+        nextRunError.set( error );
+    }
+
+    public void setNextResetError( RuntimeException error )
     {
-        this.nextRunFailure = nextRunFailure;
+        nextResetError.set( error );
+    }
+
+    private static void throwErrorIfExists( AtomicReference<RuntimeException> errorReference )
+    {
+        RuntimeException error = errorReference.getAndSet( null );
+        if ( error != null )
+        {
+            throw error;
+        }
     }
 }

driver/src/test/java/org/neo4j/driver/internal/util/ConnectionTrackingDriverFactory.java

@@ -30,8 +30,25 @@
 
 public class ThrowingConnectionDriverFactory extends DriverFactory
 {
+    private final Clock clock;
     private final List<ThrowingConnection> connections = new CopyOnWriteArrayList<>();
 
+    public ThrowingConnectionDriverFactory()
+    {
+        this( Clock.SYSTEM );
+    }
+
+    public ThrowingConnectionDriverFactory( Clock clock )
+    {
+        this.clock = clock;
+    }
+
+    @Override
+    protected Clock createClock()
+    {
+        return clock;
+    }
+
     @Override
     protected Connector createConnector( ConnectionSettings connectionSettings, SecurityPlan securityPlan,
             Logging logging )
@@ -44,4 +61,11 @@ public List<ThrowingConnection> getConnections()
     {
         return new ArrayList<>( connections );
     }
+
+    public List<ThrowingConnection> pollConnections()
+    {
+        List<ThrowingConnection> result = new ArrayList<>( connections );
+        connections.clear();
+        return result;
+    }
 }

driver/src/test/java/org/neo4j/driver/internal/util/ThrowingConnection.java

@@ -19,6 +19,7 @@
 package org.neo4j.driver.v1.integration;
 
 import org.junit.After;
+import org.junit.AfterClass;
 import org.junit.Rule;
 import org.junit.Test;
 
@@ -35,8 +36,6 @@
 
 import org.neo4j.driver.internal.cluster.RoutingSettings;
 import org.neo4j.driver.internal.retry.RetrySettings;
-import org.neo4j.driver.internal.util.Clock;
-import org.neo4j.driver.internal.util.ConnectionTrackingDriverFactory;
 import org.neo4j.driver.internal.util.FakeClock;
 import org.neo4j.driver.internal.util.ThrowingConnection;
 import org.neo4j.driver.internal.util.ThrowingConnectionDriverFactory;
@@ -70,6 +69,7 @@
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.startsWith;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThat;
@@ -97,6 +97,12 @@ public void tearDown()
         }
     }
 
+    @AfterClass
+    public static void stopSharedCluster()
+    {
+        ClusterRule.stopSharedCluster();
+    }
+
     @Test
     public void shouldExecuteReadAndWritesWhenDriverSuppliedWithAddressOfLeader() throws Exception
     {
@@ -236,7 +242,7 @@ public Void apply( Session session )
     }
 
     @Test
-    public void shouldDropBrokenOldSessions() throws Exception
+    public void shouldDropBrokenOldConnections() throws Exception
     {
         Cluster cluster = clusterRule.getCluster();
 
@@ -249,19 +255,24 @@ public void shouldDropBrokenOldSessions() throws Exception
                 .toConfig();
 
         FakeClock clock = new FakeClock();
-        ConnectionTrackingDriverFactory driverFactory = new ConnectionTrackingDriverFactory( clock );
+        ThrowingConnectionDriverFactory driverFactory = new ThrowingConnectionDriverFactory( clock );
 
         URI routingUri = cluster.leader().getRoutingUri();
         AuthToken auth = clusterRule.getDefaultAuthToken();
         RetrySettings retrySettings = RetrySettings.DEFAULT;
 
         try ( Driver driver = driverFactory.newInstance( routingUri, auth, defaultRoutingSettings(), retrySettings, config ) )
         {
-            // create nodes in different threads using different sessions
+            // create nodes in different threads using different sessions and connections
             createNodesInDifferentThreads( concurrentSessionsCount, driver );
 
-            // now pool contains many sessions, make them all invalid
-            driverFactory.closeConnections();
+            // now pool contains many connections, make them all invalid
+            List<ThrowingConnection> oldConnections = driverFactory.pollConnections();
+            for ( ThrowingConnection oldConnection : oldConnections )
+            {
+                oldConnection.setNextResetError( new ServiceUnavailableException( "Unable to reset" ) );
+            }
+
             // move clock forward more than configured liveness check timeout
             clock.progress( MINUTES.toMillis( livenessCheckTimeoutMinutes + 1 ) );
 
@@ -273,6 +284,12 @@ public void shouldDropBrokenOldSessions() throws Exception
                 assertEquals( 1, records.size() );
                 assertEquals( concurrentSessionsCount, records.get( 0 ).get( 0 ).asInt() );
             }
+
+            // all old connections failed to reset and should be closed
+            for ( ThrowingConnection connection : oldConnections )
+            {
+                assertFalse( connection.isOpen() );
+            }
         }
     }
 
@@ -573,7 +590,7 @@ public void shouldRediscoverWhenConnectionsToAllCoresBreak()
             // make all those connections throw and seem broken
             for ( ThrowingConnection connection : driverFactory.getConnections() )
             {
-                connection.setNextRunFailure( new ServiceUnavailableException( "Disconnected" ) );
+                connection.setNextRunError( new ServiceUnavailableException( "Disconnected" ) );
             }
 
             // observe that connection towards writer is broken
@@ -620,7 +637,7 @@ public void shouldKeepOperatingWhenConnectionsBreak() throws Exception
         String value = "Tony Stark";
         Cluster cluster = clusterRule.getCluster();
 
-        ConnectionTrackingDriverFactory driverFactory = new ConnectionTrackingDriverFactory( Clock.SYSTEM );
+        ThrowingConnectionDriverFactory driverFactory = new ThrowingConnectionDriverFactory();
         AtomicBoolean stop = new AtomicBoolean();
         executor = newExecutor();
 
@@ -632,18 +649,22 @@ public void shouldKeepOperatingWhenConnectionsBreak() throws Exception
             // launch writers and readers that use transaction functions and thus should never fail
             for ( int i = 0; i < 3; i++ )
             {
-                results.add( executor.submit( countNodesCallable( driver, label, property, value, stop ) ) );
+                results.add( executor.submit( readNodesCallable( driver, label, property, value, stop ) ) );
             }
             for ( int i = 0; i < 2; i++ )
             {
                 results.add( executor.submit( createNodesCallable( driver, label, property, value, stop ) ) );
             }
 
-            // terminate connections while reads and writes are in progress
+            // make connections throw while reads and writes are in progress
             long deadline = System.currentTimeMillis() + MINUTES.toMillis( 1 );
             while ( System.currentTimeMillis() < deadline && !stop.get() )
             {
-                driverFactory.closeConnections();
+                List<ThrowingConnection> connections = driverFactory.pollConnections();
+                for ( ThrowingConnection connection : connections )
+                {
+                    connection.setNextRunError( new ServiceUnavailableException( "Unable to execute query" ) );
+                }
                 SECONDS.sleep( 5 ); // sleep a bit to allow readers and writers to progress
             }
             stop.set( true );
@@ -676,7 +697,7 @@ private static void setupLastConnectionToThrow( ThrowingConnectionDriverFactory
     {
         List<ThrowingConnection> connections = factory.getConnections();
         ThrowingConnection lastConnection = connections.get( connections.size() - 1 );
-        lastConnection.setNextRunFailure( error );
+        lastConnection.setNextRunError( error );
     }
 
     private int executeWriteAndReadThroughBolt( ClusterMember member ) throws TimeoutException, InterruptedException
@@ -951,7 +972,7 @@ public Void call() throws Exception
         };
     }
 
-    private static Callable<Void> countNodesCallable( final Driver driver, final String label, final String property, final String value,
+    private static Callable<Void> readNodesCallable( final Driver driver, final String label, final String property, final String value,
             final AtomicBoolean stop )
     {
         return new Callable<Void>()
@@ -963,7 +984,7 @@ public Void call() throws Exception
                 {
                     try ( Session session = driver.session( AccessMode.READ ) )
                     {
-                        countNodes( session, label, property, value );
+                        readNodeIds( session, label, property, value );
                     }
                     catch ( Throwable t )
                     {
@@ -989,6 +1010,26 @@ public Void execute( Transaction tx )
         } );
     }
 
+    private static List<Long> readNodeIds( final Session session, final String label, final String property, final String value )
+    {
+        return session.readTransaction( new TransactionWork<List<Long>>()
+        {
+            @Override
+            public List<Long> execute( Transaction tx )
+            {
+                StatementResult result = tx.run( "MATCH (n:" + label + " {" + property + ": $value}) RETURN n LIMIT 5",
+                        parameters( "value", value ) );
+
+                List<Long> ids = new ArrayList<>();
+                while ( result.hasNext() )
+                {
+                    ids.add( result.next().get( 0 ).asNode().id() );
+                }
+                return ids;
+            }
+        } );
+    }
+
     private static void updateNode( final Session session, final String label, final String property, final String oldValue, final String newValue )
     {
         session.writeTransaction( new TransactionWork<Void>()