Fail-fast interactions with closed driver

It was previously possible to use closed driver instance, obtain sessions from
it and run queries. This was seen to cause problems in multi-threaded scenarios
when some worker threads try to use driver, which was concurrently closed by
other worker.

This commit makes driver track if it is closed or not and throw exceptions if
on every interaction when closed. Thrown exception is `IllegalStateException`.
Actuall `#close()` method can be called multiple times and is a no-op for
already closed driver.

Author: lutovich

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

@@ -16,19 +16,26 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 package org.neo4j.driver.internal;
 
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
 import org.neo4j.driver.internal.security.SecurityPlan;
+import org.neo4j.driver.v1.AccessMode;
 import org.neo4j.driver.v1.Driver;
 import org.neo4j.driver.v1.Logger;
 import org.neo4j.driver.v1.Logging;
+import org.neo4j.driver.v1.Session;
 
 abstract class BaseDriver implements Driver
 {
+    private final static String DRIVER_LOG_NAME = "Driver";
+
     private final SecurityPlan securityPlan;
     protected final Logger log;
-    private final static String DRIVER_LOG_NAME = "Driver";
+
+    private final ReentrantReadWriteLock closedLock = new ReentrantReadWriteLock();
+    private boolean closed;
 
     BaseDriver( SecurityPlan securityPlan, Logging logging )
     {
@@ -37,8 +44,68 @@ abstract class BaseDriver implements Driver
     }
 
     @Override
-    public boolean isEncrypted()
+    public final boolean isEncrypted()
+    {
+        closedLock.readLock().lock();
+        try
+        {
+            assertOpen();
+            return securityPlan.requiresEncryption();
+        }
+        finally
+        {
+            closedLock.readLock().unlock();
+        }
+    }
+
+    @Override
+    public final Session session()
+    {
+        return session( AccessMode.WRITE );
+    }
+
+    @Override
+    public final Session session( AccessMode mode )
+    {
+        closedLock.readLock().lock();
+        try
+        {
+            assertOpen();
+            return newSessionWithMode( mode );
+        }
+        finally
+        {
+            closedLock.readLock().unlock();
+        }
+    }
+
+    @Override
+    public final void close()
+    {
+        closedLock.writeLock().lock();
+        try
+        {
+            if ( !closed )
+            {
+                closeResources();
+            }
+        }
+        finally
+        {
+            closed = true;
+            closedLock.writeLock().unlock();
+        }
+    }
+
+    protected abstract Session newSessionWithMode( AccessMode mode );
+
+    protected abstract void closeResources();
+
+    private void assertOpen()
     {
-        return securityPlan.requiresEncryption();
+        if ( closed )
+        {
+            throw new IllegalStateException( "This driver instance has already been closed" );
+        }
     }
 }

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

@@ -44,19 +44,13 @@ public DirectDriver(
     }
 
     @Override
-    public Session session()
+    protected Session newSessionWithMode( AccessMode mode )
     {
         return new NetworkSession( connections.acquire( address ) );
     }
 
     @Override
-    public Session session( AccessMode ignore )
-    {
-        return session();
-    }
-
-    @Override
-    public void close()
+    protected void closeResources()
     {
         try
         {

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

@@ -59,16 +59,11 @@ public RoutingDriver(
     }
 
     @Override
-    public Session session()
-    {
-        return session( AccessMode.WRITE );
-    }
-
-    @Override
-    public Session session( final AccessMode mode )
+    protected Session newSessionWithMode( AccessMode mode )
     {
         Connection connection = acquireConnection( mode );
-        return new RoutingNetworkSession( new NetworkSession( connection ), mode, connection.boltServerAddress(), loadBalancer );
+        NetworkSession networkSession = new NetworkSession( connection );
+        return new RoutingNetworkSession( networkSession, mode, connection.boltServerAddress(), loadBalancer );
     }
 
     private Connection acquireConnection( AccessMode role )
@@ -85,7 +80,7 @@ private Connection acquireConnection( AccessMode role )
     }
 
     @Override
-    public void close()
+    protected void closeResources()
     {
         try
         {

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

@@ -20,21 +20,22 @@
 
 import org.junit.Test;
 
-import java.util.concurrent.atomic.AtomicBoolean;
-
 import org.neo4j.driver.internal.spi.Connection;
 import org.neo4j.driver.internal.util.Clock;
+import org.neo4j.driver.internal.util.Supplier;
 import org.neo4j.driver.v1.exceptions.ClientException;
 import org.neo4j.driver.v1.util.Function;
 
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 public class PooledConnectionTest
 {
@@ -151,6 +152,56 @@ public void dispose()
         assertThat( flags[0], equalTo( true ) );
     }
 
+    @Test
+    @SuppressWarnings( "unchecked" )
+    public void shouldDisposeAcquiredConnectionsWhenPoolIsClosed()
+    {
+        PooledConnection connection = mock( PooledConnection.class );
+
+        BlockingPooledConnectionQueue pool = new BlockingPooledConnectionQueue( 5 );
+
+        Supplier<PooledConnection> pooledConnectionFactory = mock( Supplier.class );
+        when( pooledConnectionFactory.get() ).thenReturn( connection );
+
+        PooledConnection acquiredConnection = pool.acquire( pooledConnectionFactory );
+        assertSame( acquiredConnection, connection );
+
+        pool.terminate();
+        verify( connection ).dispose();
+    }
+
+    @Test
+    @SuppressWarnings( "unchecked" )
+    public void shouldDisposeAcquiredAndIdleConnectionsWhenPoolIsClosed()
+    {
+        PooledConnection connection1 = mock( PooledConnection.class );
+        PooledConnection connection2 = mock( PooledConnection.class );
+        PooledConnection connection3 = mock( PooledConnection.class );
+
+        BlockingPooledConnectionQueue pool = new BlockingPooledConnectionQueue( 5 );
+
+        Supplier<PooledConnection> pooledConnectionFactory = mock( Supplier.class );
+        when( pooledConnectionFactory.get() )
+                .thenReturn( connection1 )
+                .thenReturn( connection2 )
+                .thenReturn( connection3 );
+
+        PooledConnection acquiredConnection1 = pool.acquire( pooledConnectionFactory );
+        PooledConnection acquiredConnection2 = pool.acquire( pooledConnectionFactory );
+        PooledConnection acquiredConnection3 = pool.acquire( pooledConnectionFactory );
+        assertSame( acquiredConnection1, connection1 );
+        assertSame( acquiredConnection2, connection2 );
+        assertSame( acquiredConnection3, connection3 );
+
+        pool.offer( acquiredConnection2 );
+
+        pool.terminate();
+
+        verify( connection1 ).dispose();
+        verify( connection2 ).dispose();
+        verify( connection3 ).dispose();
+    }
+
     @Test
     public void shouldDisposeConnectionIfPoolAlreadyClosed() throws Throwable
     {