Merge pull request #283 from lutovich/1.1-closed-driver

Better handling of closed driver

Author: Zhen Li

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
     {