Make sure connection is only called once.

Instead of guarding in `PooledConnection` as well as in `NetworkSession` we
make sure we only close connection once in `RoutingDriver`

Author: pontusmelke

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

@@ -241,7 +241,7 @@ public BoltServerAddress apply( Value value )
     //must be called from a synchronized method
     private boolean call( BoltServerAddress address, String procedureName, Consumer<Record> recorder )
     {
-        Connection acquire = null;
+        Connection acquire;
         Session session = null;
         try
         {
@@ -271,11 +271,6 @@ private boolean call( BoltServerAddress address, String procedureName, Consumer<
             {
                 session.close();
             }
-            if ( acquire != null )
-            {
-                acquire.close();
-            }
-
         }
         return true;
     }

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

@@ -19,7 +19,6 @@
 package org.neo4j.driver.internal.net.pooling;
 
 import java.util.Map;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.neo4j.driver.internal.net.BoltServerAddress;
 import org.neo4j.driver.internal.spi.Collector;
@@ -59,7 +58,6 @@ public class PooledConnection implements Connection
     private Runnable onError = null;
     private final Clock clock;
     private long lastUsed;
-    private final AtomicBoolean released = new AtomicBoolean( false );
 
     public PooledConnection( Connection delegate, Consumer<PooledConnection> release, Clock clock )
     {
@@ -69,9 +67,8 @@ public PooledConnection( Connection delegate, Consumer<PooledConnection> release
         this.lastUsed = clock.millis();
     }
 
-    public void setInUse()
+    public void updateTimestamp()
     {
-        released.set(false);
         lastUsed = clock.millis();
     }
 
@@ -200,10 +197,7 @@ public void receiveOne()
      */
     public void close()
     {
-        if ( released.compareAndSet( false, true ))
-        {
-            release.accept( this );
-        }
+        release.accept( this );
         // put the full logic of deciding whether to dispose the connection or to put it back to
         // the pool into the release object
     }

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

@@ -18,6 +18,8 @@
  */
 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;
@@ -39,6 +41,8 @@
 import org.neo4j.driver.v1.Value;
 import org.neo4j.driver.v1.exceptions.ClientException;
 
+import static java.util.Collections.emptyList;
+
 /**
  * The pool is designed to buffer certain amount of free sessions into session pool. When closing a session, we first
  * try to return the session into the session pool, however if we failed to return it back, either because the pool
@@ -115,7 +119,7 @@ public Connection acquire( BoltServerAddress address )
             conn = new PooledConnection( connect( address ), new
                     PooledConnectionReleaseConsumer( connections, stopped, new PooledConnectionValidator( this, poolSettings ) ), clock );
         }
-        conn.setInUse();
+        conn.updateTimestamp();
         return conn;
     }
 
@@ -184,4 +188,19 @@ public void close()
         pools.clear();
     }
 
+    //for testing
+    public List<PooledConnection> connectionsForAddress(BoltServerAddress address)
+    {
+        LinkedBlockingQueue<PooledConnection> pooledConnections =
+                (LinkedBlockingQueue<PooledConnection>) pools.get( address );
+        if (pooledConnections == null)
+        {
+            return emptyList();
+        }
+        else
+        {
+            return new ArrayList<>( pooledConnections );
+        }
+    }
+
 }

driver/src/test/java/org/neo4j/driver/internal/RoutingDriverStubTest.java

@@ -35,6 +35,8 @@
 
 import org.neo4j.driver.internal.logging.ConsoleLogging;
 import org.neo4j.driver.internal.net.BoltServerAddress;
+import org.neo4j.driver.internal.net.pooling.PooledConnection;
+import org.neo4j.driver.internal.net.pooling.SocketConnectionPool;
 import org.neo4j.driver.internal.spi.Connection;
 import org.neo4j.driver.v1.AccessMode;
 import org.neo4j.driver.v1.Config;
@@ -86,6 +88,26 @@ public void shouldDiscoverServers() throws IOException, InterruptedException, St
         assertThat( server.exitStatus(), equalTo( 0 ) );
     }
 
+    @Test
+    public void shouldOnlyPutConnectionInPoolOnce() throws IOException, InterruptedException, StubServer.ForceKilled
+    {
+        // Given
+        StubServer server = StubServer.start( "discover_servers.script", 9001 );
+        URI uri = URI.create( "bolt+routing://127.0.0.1:9001" );
+
+        // When
+        try ( RoutingDriver driver = (RoutingDriver) GraphDatabase.driver( uri, config ) )
+        {
+            // Then
+            SocketConnectionPool pool = (SocketConnectionPool) driver.connectionPool();
+            List<PooledConnection> pooledConnections = pool.connectionsForAddress( address( 9001 ) );
+            assertThat(pooledConnections, hasSize( 1 ));
+        }
+
+        // Finally
+        assertThat( server.exitStatus(), equalTo( 0 ) );
+    }
+
     @Test
     public void shouldDiscoverNewServers() throws IOException, InterruptedException, StubServer.ForceKilled
     {

driver/src/test/java/org/neo4j/driver/internal/RoutingDriverTest.java

@@ -43,7 +43,6 @@
 import org.neo4j.driver.v1.exceptions.NoSuchRecordException;
 import org.neo4j.driver.v1.exceptions.ServiceUnavailableException;
 import org.neo4j.driver.v1.summary.ResultSummary;
-import org.neo4j.driver.v1.util.BiFunction;
 import org.neo4j.driver.v1.util.Function;
 
 import static java.util.Arrays.asList;

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

@@ -120,39 +120,6 @@ public void dispose()
         assertThat( flags[0], equalTo( false ) );
     }
 
-    @Test
-    public void shouldOnlyReturnOnceEventhougCloseIsBeingCalledMultipleTimes() throws Throwable
-    {
-        // Given
-        final BlockingQueue<PooledConnection> pool = new LinkedBlockingQueue<>(2);
-
-        final boolean[] flags = {false};
-
-        Connection conn = mock( Connection.class );
-        PooledConnectionReleaseConsumer releaseConsumer = new PooledConnectionReleaseConsumer( pool,
-                new AtomicBoolean( false ), VALID_CONNECTION );
-
-        PooledConnection pooledConnection = new PooledConnection( conn, releaseConsumer, Clock.SYSTEM )
-        {
-            @Override
-            public void dispose()
-            {
-                flags[0] = true;
-            }
-        };
-
-        // When
-        pooledConnection.close();
-        pooledConnection.close();
-        pooledConnection.close();
-        pooledConnection.close();
-        pooledConnection.close();
-
-        // Then
-        assertThat( pool, hasItem(pooledConnection) );
-        assertThat( pool.size(), equalTo( 1 ) );
-        assertThat( flags[0], equalTo( false ) );
-    }
 
     @Test
     public void shouldDisposeConnectionIfValidConnectionAndIdlePoolIsFull() throws Throwable