Make sessioin.reset thread-safe

Author: Zhen

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

@@ -54,6 +54,7 @@ public void run()
 
     private InternalTransaction currentTransaction;
     private AtomicBoolean isOpen = new AtomicBoolean( true );
+    private AtomicBoolean markedToClose = new AtomicBoolean( false );
 
     public InternalSession( Connection connection, Logger logger )
     {
@@ -103,13 +104,16 @@ public void reset()
         ensureNoUnrecoverableError();
         ensureConnectionIsOpen();
 
-        connection.resetAsync();
+        if( markedToClose.compareAndSet( false, true ) )
+        {
+            connection.resetAsync();
+        }
     }
 
     @Override
     public boolean isOpen()
     {
-        return isOpen.get();
+        return isOpen.get() && !markedToClose.get();
     }
 
     @Override
@@ -245,4 +249,16 @@ private void ensureConnectionIsOpen()
                                        "Please close this session and retry your statement in another new session." );
         }
     }
+
+    private void ensureSessionIsOpen()
+    {
+        if( !isOpen() )
+        {
+            throw new ClientException(
+                    "No more interaction with this session is allowed " +
+                    "as the current session is already closed or marked as closed. " +
+                    "You get this error either because you have a bad reference to a session that has already be closed " +
+                    "or you are trying to reuse a session that you have called `reset` on it." );
+        }
+    }
 }

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

@@ -113,7 +113,7 @@ public void sync()
     }
 
     @Override
-    public void flush()
+    public synchronized void flush()
     {
         try
         {
@@ -180,7 +180,7 @@ else if ( e instanceof SocketTimeoutException )
         }
     }
 
-    private void queueMessage( Message msg, StreamCollector collector )
+    private synchronized void queueMessage( Message msg, StreamCollector collector )
     {
         pendingMessages.add( msg );
         responseHandler.appendResultCollector( collector );
@@ -213,18 +213,8 @@ public boolean hasUnrecoverableErrors()
     @Override
     public void resetAsync()
     {
-        if( interrupted.compareAndSet( false, true ) )
-        {
-            queueMessage( RESET, new StreamCollector.ResetStreamCollector( new Runnable()
-            {
-                @Override
-                public void run()
-                {
-                    interrupted.set( false );
-                }
-            } ) );
-            flush();
-        }
+        queueMessage( RESET, StreamCollector.RESET );
+        flush();
     }
 
     @Override

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

@@ -18,9 +18,9 @@
  */
 package org.neo4j.driver.internal.net;
 
-import java.util.LinkedList;
 import java.util.Map;
 import java.util.Queue;
+import java.util.concurrent.ConcurrentLinkedQueue;
 
 import org.neo4j.driver.internal.messaging.MessageHandler;
 import org.neo4j.driver.internal.spi.StreamCollector;
@@ -39,7 +39,7 @@
 
 public class SocketResponseHandler implements MessageHandler
 {
-    private final Queue<StreamCollector> collectors = new LinkedList<>();
+    private final Queue<StreamCollector> collectors = new ConcurrentLinkedQueue<>();
 
     /** If a failure occurs, the error gets stored here */
     private Neo4jException error;

driver/src/test/java/org/neo4j/driver/v1/integration/SessionIT.java

@@ -21,8 +21,12 @@
 import org.junit.Rule;
 import org.junit.Test;
 
+import java.util.logging.Level;
+
+import org.neo4j.driver.internal.logging.ConsoleLogging;
 import org.neo4j.driver.v1.AuthToken;
 import org.neo4j.driver.v1.AuthTokens;
+import org.neo4j.driver.v1.Config;
 import org.neo4j.driver.v1.Driver;
 import org.neo4j.driver.v1.GraphDatabase;
 import org.neo4j.driver.v1.Session;
@@ -99,13 +103,15 @@ public void shouldKillLongRunningStatement() throws Throwable
     {
         neo4j.ensureProcedures( "longRunningStatement.jar" );
         // Given
-        Driver driver = GraphDatabase.driver( neo4j.uri() );
+        Driver driver = GraphDatabase.driver( neo4j.uri(), Config.build().withLogging( new ConsoleLogging( Level.ALL
+        ) ).toConfig() );
 
         int executionTimeout = 10; // 10s
         final int killTimeout = 1; // 1s
         long startTime = -1, endTime;
 
-        try( final Session session = driver.session() )
+        final Session session = driver.session();
+        try
         {
             StatementResult result =
                     session.run( "CALL test.driver.longRunningStatement({seconds})",
@@ -123,13 +129,18 @@ public void shouldKillLongRunningStatement() throws Throwable
         {
             endTime = System.currentTimeMillis();
             assertTrue( startTime > 0 );
-            assertTrue( endTime - startTime > killTimeout * 1000 ); // get killed by session.kill
+            assertTrue( endTime - startTime > killTimeout * 1000 ); // get reset by session.reset
             assertTrue( endTime - startTime < executionTimeout * 1000 / 2 ); // finished before execution finished
+            assertFalse( session.isOpen() );
         }
         catch ( Exception e )
         {
             fail( "Should be a Neo4jException" );
         }
+        finally
+        {
+            session.close();
+        }
     }
 
     @Test
@@ -168,7 +179,7 @@ public void shouldKillLongStreamingResult() throws Throwable
             assertThat( recordCount, greaterThan(1) );
 
             assertTrue( startTime > 0 );
-            assertTrue( endTime - startTime > killTimeout * 1000 ); // get killed by session.kill
+            assertTrue( endTime - startTime > killTimeout * 1000 ); // get reset by session.reset
             assertTrue( endTime - startTime < executionTimeout * 1000 / 2 ); // finished before execution finished
         }
     }
@@ -190,7 +201,7 @@ public void run()
                 }
                 finally
                 {
-                    session.reset(); // kill the session after timeout
+                    session.reset(); // reset the session after timeout
                 }
             }
         } ).start();