Merge pull request #295 from lutovich/1.1-flaky-tests

Fix couple flaky tests

Author: nigelsmall

driver/src/main/java/org/neo4j/driver/internal/net/SocketClient.java

@@ -130,7 +130,7 @@ public void start()
         {
             throw new ServiceUnavailableException( format(
                     "Unable to connect to %s, ensure the database is running and that there is a " +
-                    "working network connection to it.", address ) );
+                    "working network connection to it.", address ), e );
         }
         catch ( IOException e )
         {

driver/src/main/java/org/neo4j/driver/internal/net/pooling/SocketConnectionPool.java

@@ -84,11 +84,7 @@ public PooledConnection get()
         };
         PooledConnection conn = connections.acquire( supplier );
 
-        if ( closed.get() )
-        {
-            connections.terminate();
-            throw poolClosedException();
-        }
+        assertNotClosed( address, connections );
 
         conn.updateTimestamp();
         return conn;
@@ -140,16 +136,21 @@ public void close()
         }
     }
 
-    private void assertNotClosed()
+    private void assertNotClosed( BoltServerAddress address, BlockingPooledConnectionQueue connections )
     {
         if ( closed.get() )
         {
-            throw poolClosedException();
+            connections.terminate();
+            pools.remove( address );
+            assertNotClosed();
         }
     }
 
-    private static RuntimeException poolClosedException()
+    private void assertNotClosed()
     {
-        return new IllegalStateException( "Pool closed" );
+        if ( closed.get() )
+        {
+            throw new IllegalStateException( "Pool closed" );
+        }
     }
 }

driver/src/test/java/org/neo4j/driver/internal/net/pooling/SocketConnectionPoolTest.java

@@ -265,7 +265,8 @@ public void closeWithConcurrentAcquisitionsEmptiesThePool() throws InterruptedEx
 
         for ( int i = 0; i < port.intValue(); i++ )
         {
-            assertFalse( pool.hasAddress( new BoltServerAddress( "localhost", i ) ) );
+            boolean hasAddress = pool.hasAddress( new BoltServerAddress( "localhost", i ) );
+            assertFalse( "Pool still has connection queues" + pool, hasAddress );
         }
         for ( Connection connection : createdConnections )
         {

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

@@ -18,6 +18,8 @@
  */
 package org.neo4j.driver.v1.stress;
 
+import org.junit.After;
+import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 
@@ -28,6 +30,7 @@
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 
 import org.neo4j.driver.v1.Config;
 import org.neo4j.driver.v1.Driver;
@@ -36,6 +39,7 @@
 import org.neo4j.driver.v1.util.TestNeo4j;
 
 import static java.util.Arrays.asList;
+import static org.junit.Assert.assertTrue;
 import static org.neo4j.driver.v1.GraphDatabase.driver;
 
 public class SessionPoolingStressIT
@@ -44,65 +48,107 @@ public class SessionPoolingStressIT
     public TestNeo4j neo4j = new TestNeo4j();
 
     private static final int N_THREADS = 50;
-    private final ExecutorService executor = Executors.newFixedThreadPool( N_THREADS );
-    private static final List<String> QUERIES = asList( "RETURN 1295 + 42", "UNWIND range(1,10000) AS x CREATE (n {prop:x}) DELETE n " );
     private static final int MAX_TIME = 10000;
-    private final AtomicBoolean hasFailed = new AtomicBoolean( false );
+
+    private static final List<String> QUERIES = asList(
+            "RETURN 1295 + 42", "UNWIND range(1,10000) AS x CREATE (n {prop:x}) DELETE n " );
+
+    private Driver driver;
+    private ExecutorService executor;
+
+    @Before
+    public void setUp() throws Exception
+    {
+        executor = Executors.newFixedThreadPool( N_THREADS );
+    }
+
+    @After
+    public void tearDown() throws Exception
+    {
+        if ( executor != null )
+        {
+            executor.shutdownNow();
+        }
+
+        if ( driver != null )
+        {
+            driver.close();
+        }
+    }
 
     @Test
-    public void shouldWorkFine() throws InterruptedException
+    public void shouldWorkFine() throws Throwable
     {
-        Driver driver = driver( neo4j.uri(),
-                Config.build()
-                        .withEncryptionLevel( Config.EncryptionLevel.NONE )
-                        .withMaxSessions( N_THREADS ).toConfig() );
-
-        doWork( driver );
-        executor.awaitTermination( MAX_TIME + (int)(MAX_TIME * 0.2), TimeUnit.MILLISECONDS );
-        driver.close();
+        Config config = Config.build()
+                .withEncryptionLevel( Config.EncryptionLevel.NONE )
+                .withMaxSessions( N_THREADS )
+                .toConfig();
+
+        driver = driver( neo4j.uri(), config );
+
+        AtomicBoolean stop = new AtomicBoolean();
+        AtomicReference<Throwable> failureReference = new AtomicReference<>();
+
+        doWork( stop, failureReference );
+
+        Thread.sleep( MAX_TIME );
+
+        stop.set( true );
+        executor.shutdown();
+        assertTrue( executor.awaitTermination( MAX_TIME, TimeUnit.MILLISECONDS ) );
+
+        Throwable failure = failureReference.get();
+        if ( failure != null )
+        {
+            throw new AssertionError( "Some workers have failed", failure );
+        }
     }
 
-    private void doWork( final Driver driver )
+    private void doWork( AtomicBoolean stop, AtomicReference<Throwable> failure )
     {
         for ( int i = 0; i < N_THREADS; i++ )
         {
-            executor.execute( new Worker( driver ) );
+            executor.execute( new Worker( driver, stop, failure ) );
         }
     }
 
     private class Worker implements Runnable
     {
         private final Random random = ThreadLocalRandom.current();
         private final Driver driver;
+        private final AtomicBoolean stop;
+        private final AtomicReference<Throwable> failureReference;
 
-        public Worker( Driver driver )
+        Worker( Driver driver, AtomicBoolean stop, AtomicReference<Throwable> failureReference )
         {
             this.driver = driver;
+            this.stop = stop;
+            this.failureReference = failureReference;
         }
 
         @Override
         public void run()
         {
             try
             {
-                long deadline = System.currentTimeMillis() + MAX_TIME;
-                for (;;)
+                while ( !stop.get() )
                 {
                     for ( String query : QUERIES )
                     {
                         runQuery( query );
                     }
-                    long left = deadline - System.currentTimeMillis();
-                    if ( left <= 0 )
-                    {
-                        break;
-                    }
                 }
             }
-            catch ( Throwable e )
+            catch ( Throwable failure )
             {
-                e.printStackTrace();
-                hasFailed.set( true );
+                if ( !failureReference.compareAndSet( null, failure ) )
+                {
+                    Throwable firstFailure = failureReference.get();
+                    synchronized ( firstFailure )
+                    {
+                        firstFailure.addSuppressed( failure );
+                    }
+                }
             }
         }
 

driver/src/test/java/org/neo4j/driver/v1/stress/SessionPoolingStressIT.java

@@ -20,6 +20,9 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.net.SocketAddress;
+import java.nio.channels.SocketChannel;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -33,6 +36,8 @@
 
 public class StubServer
 {
+    private static final int SOCKET_CONNECT_ATTEMPTS = 20;
+
     public static final Config INSECURE_CONFIG = Config.build()
             .withEncryptionLevel( Config.EncryptionLevel.NONE ).toConfig();
 
@@ -50,7 +55,7 @@ private StubServer( String script, int port ) throws IOException, InterruptedExc
         command.addAll( asList( Integer.toString( port ), script ) );
         ProcessBuilder server = new ProcessBuilder().inheritIO().command( command );
         process = server.start();
-        sleep( 500 );  // might take a moment for the socket to start listening
+        waitForSocket( port );
     }
 
     public static StubServer start( String resource, int port ) throws IOException, InterruptedException
@@ -100,4 +105,22 @@ private static boolean boltKitAvailable()
             return false;
         }
     }
+
+    private static void waitForSocket( int port ) throws InterruptedException
+    {
+        SocketAddress address = new InetSocketAddress( "localhost", port );
+        for ( int i = 0; i < SOCKET_CONNECT_ATTEMPTS; i++ )
+        {
+            try
+            {
+                SocketChannel.open( address );
+                return;
+            }
+            catch ( Exception e )
+            {
+                sleep( 300 );
+            }
+        }
+        throw new AssertionError( "Can't connect to " + address );
+    }
 }