Fixes from code review

Author: pontusmelke

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

@@ -25,6 +25,7 @@
 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.util.Supplier;
 
@@ -37,6 +38,8 @@ public class BlockingPooledConnectionQueue
     /** The backing queue, keeps track of connections currently in queue */
     private final BlockingQueue<PooledConnection> queue;
 
+    private final AtomicBoolean isTerminating = new AtomicBoolean( false );
+
     /** Keeps track of acquired connections */
     private final Set<PooledConnection> acquiredConnections =
             Collections.newSetFromMap(new ConcurrentHashMap<PooledConnection, Boolean>());
@@ -54,12 +57,30 @@ public BlockingPooledConnectionQueue( int capacity )
      */
     public boolean offer( PooledConnection pooledConnection )
     {
-        boolean offer = queue.offer( pooledConnection );
-        if ( offer )
+        acquiredConnections.remove( pooledConnection );
+        //Don't offer to the queue if we are terminating.
+        return !isTerminating.get() && queue.offer( pooledConnection );
+
+    }
+
+    /**
+     * 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 )
         {
-            acquiredConnections.remove( pooledConnection );
+            poll = supplier.get();
         }
-        return offer;
+        //We are closing down, don't put anything back
+        if ( isTerminating.get() )
+        {
+            acquiredConnections.add( poll );
+        }
+        return poll;
     }
 
     public List<PooledConnection> toList()
@@ -88,34 +109,21 @@ public boolean contains( PooledConnection pooledConnection )
      */
     public void terminate()
     {
-        while ( !queue.isEmpty() )
+        if (isTerminating.compareAndSet( false, true ))
         {
-            PooledConnection conn = queue.poll();
-            if ( conn != null )
+            while ( !queue.isEmpty() )
             {
-                //close the underlying connection without adding it back to the queue
-                conn.dispose();
+                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();
             }
         }
-        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;
     }
 }