Close all connections on driver.close()

Connections/Sessions that were in flight, i.e. acquired from the pool, were not
closed when shutting down the driver, connections remained open until GC
happened. The expected behavior is that `driver.close()` closes down all
sessions it has created.

Author: pontusmelke

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

@@ -139,7 +139,7 @@ public synchronized void reset()
     @Override
     public boolean isOpen()
     {
-        return isOpen.get();
+        return isOpen.get() && connection.isOpen();
     }
 
     @Override
@@ -177,10 +177,6 @@ public void close()
         {
             connection.sync();
         }
-        catch ( Throwable t )
-        {
-            throw t;
-        }
         finally
         {
             closeConnection();
@@ -314,7 +310,7 @@ private void ensureConnectionIsOpen()
 
     private void ensureSessionIsOpen()
     {
-        if ( !isOpen() )
+        if ( !isOpen.get() )
         {
             throw new ClientException(
                     "No more interaction with this session is allowed " +

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

@@ -0,0 +1,121 @@
+/**
+ * Copyright (c) 2002-2016 "Neo Technology,"
+ * Network Engine for Objects in Lund AB [http://neotechnology.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.neo4j.driver.internal.net.pooling;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.LinkedBlockingQueue;
+
+import org.neo4j.driver.internal.util.Supplier;
+
+/**
+ * A blocking queue that also keeps track of connections that are acquired in order
+ * to facilitate termination of all connections.
+ */
+public class BlockingPooledConnectionQueue
+{
+    /** The backing queue, keeps track of connections currently in queue */
+    private final BlockingQueue<PooledConnection> queue;
+
+    /** Keeps track of acquired connections */
+    private final Set<PooledConnection> acquiredConnections =
+            Collections.newSetFromMap(new ConcurrentHashMap<PooledConnection, Boolean>());
+
+    public BlockingPooledConnectionQueue( int capacity )
+    {
+        this.queue = new LinkedBlockingQueue<>( capacity );
+    }
+
+    /**
+     * Offer a connections back to the queue
+     *
+     * @param pooledConnection the connection to put back to the queue
+     * @return <code>true</code> if connections was accepted otherwise <code>false</code>
+     */
+    public boolean offer( PooledConnection pooledConnection )
+    {
+        boolean offer = queue.offer( pooledConnection );
+        if ( offer )
+        {
+            acquiredConnections.remove( pooledConnection );
+        }
+        return offer;
+    }
+
+    public List<PooledConnection> toList()
+    {
+        return new ArrayList<>( queue );
+    }
+
+    public boolean isEmpty()
+    {
+        return queue.isEmpty();
+    }
+
+    public int size()
+    {
+        return queue.size();
+    }
+
+    public boolean contains( PooledConnection pooledConnection )
+    {
+        return queue.contains( pooledConnection );
+    }
+
+    /**
+     * Terminates all connections, both those that are currently in the queue as well
+     * as those that have been acquired.
+     */
+    public void terminate()
+    {
+        while ( !queue.isEmpty() )
+        {
+            PooledConnection conn = queue.poll();
+            if ( conn != null )
+            {
+                //close the underlying connection without adding it back to the queue
+                conn.dispose();
+            }
+        }
+        for ( PooledConnection pooledConnection : acquiredConnections )
+        {
+            pooledConnection.dispose();
+        }
+    }
+
+    /**
+     * Acquire connection or create a new one if the queue is empty
+     * @param supplier used to create a new connection if queue is empty
+     * @return a PooledConnection instance
+     */
+    public PooledConnection acquire( Supplier<PooledConnection> supplier )
+    {
+        PooledConnection poll = queue.poll();
+        if ( poll == null )
+        {
+            poll = supplier.get();
+        }
+        acquiredConnections.add( poll );
+        return poll;
+    }
+}

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

@@ -18,7 +18,6 @@
  */
 package org.neo4j.driver.internal.net.pooling;
 
-import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.neo4j.driver.internal.util.Consumer;
@@ -30,11 +29,11 @@
  */
 class PooledConnectionReleaseConsumer implements Consumer<PooledConnection>
 {
-    private final BlockingQueue<PooledConnection> connections;
+    private final BlockingPooledConnectionQueue connections;
     private final AtomicBoolean driverStopped;
     private final Function<PooledConnection, Boolean> validConnection;
 
-    PooledConnectionReleaseConsumer( BlockingQueue<PooledConnection> connections, AtomicBoolean driverStopped,
+    PooledConnectionReleaseConsumer( BlockingPooledConnectionQueue connections, AtomicBoolean driverStopped,
             Function<PooledConnection, Boolean> validConnection)
     {
         this.connections = connections;
@@ -67,11 +66,7 @@ else if ( driverStopped.get() )
                 // which connection we get back, because other threads might be in the same situation as ours. It only
                 // matters that we added *a* connection that might not be observed by the loop, and that we dispose of
                 // *a* connection in response.
-                PooledConnection conn = connections.poll();
-                if ( conn != null )
-                {
-                    conn.dispose();
-                }
+                connections.terminate();
             }
         }
         else

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

@@ -18,12 +18,9 @@
  */
 package org.neo4j.driver.internal.net.pooling;
 
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.neo4j.driver.internal.ConnectionSettings;
@@ -35,6 +32,7 @@
 import org.neo4j.driver.internal.spi.Connection;
 import org.neo4j.driver.internal.spi.ConnectionPool;
 import org.neo4j.driver.internal.util.Clock;
+import org.neo4j.driver.internal.util.Supplier;
 import org.neo4j.driver.v1.AuthToken;
 import org.neo4j.driver.v1.AuthTokens;
 import org.neo4j.driver.v1.Logging;
@@ -48,10 +46,10 @@
  * try to return the session into the session pool, however if we failed to return it back, either because the pool
  * is full or the pool is being cleaned on driver.close, then we directly close the connection attached with the
  * session.
- *
+ * <p>
  * The session is NOT meant to be thread safe, each thread should have an independent session and close it (return to
  * pool) when the work with the session has been done.
- *
+ * <p>
  * The driver is thread safe. Each thread could try to get a session from the pool and then return it to the pool
  * at the same time.
  */
@@ -60,7 +58,8 @@ public class SocketConnectionPool implements ConnectionPool
     /**
      * Pools, organized by server address.
      */
-    private final ConcurrentHashMap<BoltServerAddress,BlockingQueue<PooledConnection>> pools = new ConcurrentHashMap<>();
+    private final ConcurrentHashMap<BoltServerAddress,BlockingPooledConnectionQueue> pools =
+            new ConcurrentHashMap<>();
 
     private final Clock clock = Clock.SYSTEM;
 
@@ -73,7 +72,7 @@ public class SocketConnectionPool implements ConnectionPool
     private final AtomicBoolean stopped = new AtomicBoolean( false );
 
     public SocketConnectionPool( ConnectionSettings connectionSettings, SecurityPlan securityPlan,
-                                 PoolSettings poolSettings, Logging logging )
+            PoolSettings poolSettings, Logging logging )
     {
         this.connectionSettings = connectionSettings;
         this.securityPlan = securityPlan;
@@ -94,41 +93,48 @@ private Connection connect( BoltServerAddress address ) throws ClientException
 
     private static Map<String,Value> tokenAsMap( AuthToken token )
     {
-        if( token instanceof InternalAuthToken )
+        if ( token instanceof InternalAuthToken )
         {
             return ((InternalAuthToken) token).toMap();
         }
         else
         {
-            throw new ClientException( "Unknown authentication token, `" + token + "`. Please use one of the supported " +
+            throw new ClientException(
+                    "Unknown authentication token, `" + token + "`. Please use one of the supported " +
                     "tokens from `" + AuthTokens.class.getSimpleName() + "`." );
         }
     }
 
     @Override
-    public Connection acquire( BoltServerAddress address )
+    public Connection acquire( final BoltServerAddress address )
     {
         if ( stopped.get() )
         {
             throw new IllegalStateException( "Pool has been closed, cannot acquire new values." );
         }
-        BlockingQueue<PooledConnection> connections = pool( address );
-        PooledConnection conn = connections.poll();
-        if ( conn == null )
+        final BlockingPooledConnectionQueue connections = pool( address );
+        Supplier<PooledConnection> supplier = new Supplier<PooledConnection>()
         {
-            conn = new PooledConnection( connect( address ), new
-                    PooledConnectionReleaseConsumer( connections, stopped, new PooledConnectionValidator( this, poolSettings ) ), clock );
-        }
+            @Override
+            public PooledConnection get()
+            {
+                return new PooledConnection( connect( address ), new
+                        PooledConnectionReleaseConsumer( connections, stopped,
+                        new PooledConnectionValidator( SocketConnectionPool.this, poolSettings ) ), clock );
+
+            }
+        };
+        PooledConnection conn = connections.acquire( supplier );
         conn.updateTimestamp();
         return conn;
     }
 
-    private BlockingQueue<PooledConnection> pool( BoltServerAddress address )
+    private BlockingPooledConnectionQueue pool( BoltServerAddress address )
     {
-        BlockingQueue<PooledConnection> pool = pools.get( address );
+        BlockingPooledConnectionQueue pool = pools.get( address );
         if ( pool == null )
         {
-            pool = new LinkedBlockingQueue<>(poolSettings.maxIdleConnectionPoolSize());
+            pool = new BlockingPooledConnectionQueue( poolSettings.maxIdleConnectionPoolSize() );
 
             if ( pools.putIfAbsent( address, pool ) != null )
             {
@@ -142,19 +148,13 @@ private BlockingQueue<PooledConnection> pool( BoltServerAddress address )
     @Override
     public void purge( BoltServerAddress address )
     {
-        BlockingQueue<PooledConnection> connections = pools.remove( address );
+        BlockingPooledConnectionQueue connections = pools.remove( address );
         if ( connections == null )
         {
             return;
         }
-        while (!connections.isEmpty())
-        {
-            PooledConnection connection = connections.poll();
-            if ( connection != null)
-            {
-                connection.dispose();
-            }
-        }
+
+        connections.terminate();
     }
 
     @Override
@@ -166,41 +166,34 @@ public boolean hasAddress( BoltServerAddress address )
     @Override
     public void close()
     {
-        if( !stopped.compareAndSet( false, true ) )
+        if ( !stopped.compareAndSet( false, true ) )
         {
             // already closed or some other thread already started close
             return;
         }
 
-        for ( BlockingQueue<PooledConnection> pool : pools.values() )
+        for ( BlockingPooledConnectionQueue pool : pools.values() )
         {
-            while ( !pool.isEmpty() )
-            {
-                PooledConnection conn = pool.poll();
-                if ( conn != null )
-                {
-                    //close the underlying connection without adding it back to the queue
-                    conn.dispose();
-                }
-            }
+            pool.terminate();
         }
 
         pools.clear();
     }
 
+
     //for testing
-    public List<PooledConnection> connectionsForAddress(BoltServerAddress address)
+    public List<PooledConnection> connectionsForAddress( BoltServerAddress address )
     {
-        LinkedBlockingQueue<PooledConnection> pooledConnections =
-                (LinkedBlockingQueue<PooledConnection>) pools.get( address );
-        if (pooledConnections == null)
+        BlockingPooledConnectionQueue pooledConnections = pools.get( address );
+        if ( pooledConnections == null )
         {
             return emptyList();
         }
         else
         {
-            return new ArrayList<>( pooledConnections );
+            return pooledConnections.toList();
         }
     }
 
+
 }