Improved return to pool to be thread-safe

Zhen · Zhen · commit 4314207381fd · 2016-06-17T20:05:06.000+02:00
Fixed the bug where the connection is not return to pool or closed when it is invalid
Fixed the bug where the connection is disposed but being put back to pool
Fixed a race condition where the connection is put back to pool after the pool is cleared
Added more tests to make sure that all state changes of a pooled connection related to pooling and disposing are thread-safe
diff --git a/driver/src/main/java/org/neo4j/driver/internal/InternalSession.java b/driver/src/main/java/org/neo4j/driver/internal/InternalSession.java
@@ -19,6 +19,7 @@
 package org.neo4j.driver.internal;
 
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.neo4j.driver.internal.spi.Connection;
 import org.neo4j.driver.internal.spi.Logger;
@@ -52,7 +53,7 @@ public void run()
     };
 
     private InternalTransaction currentTransaction;
-    private boolean isOpen = true;
+    private AtomicBoolean isOpen = new AtomicBoolean( true );
 
     public InternalSession( Connection connection, Logger logger )
     {
@@ -100,19 +101,19 @@ public StatementResult run( Statement statement )
     @Override
     public boolean isOpen()
     {
-        return isOpen;
+        return isOpen.get();
     }
 
     @Override
     public void close()
     {
-        if( !isOpen )
+        // Use atomic operation to protect from closing the connection twice (putting back to the pool twice).
+        if( !isOpen.compareAndSet( true, false ) )
         {
             throw new ClientException( "This session has already been closed." );
         }
         else
         {
-            isOpen = false;
             if ( currentTransaction != null )
             {
                 try
@@ -171,7 +172,7 @@ private void ensureConnectionIsValidBeforeOpeningTransaction()
     @Override
     protected void finalize() throws Throwable
     {
-        if( isOpen )
+        if( isOpen.compareAndSet( true, false ) )
         {
             logger.error( "Neo4j Session object leaked, please ensure that your application calls the `close` " +
                           "method on Sessions before disposing of the objects.", null );
diff --git a/driver/src/main/java/org/neo4j/driver/internal/connector/ConcurrencyGuardingConnection.java b/driver/src/main/java/org/neo4j/driver/internal/connector/ConcurrencyGuardingConnection.java
@@ -157,15 +157,16 @@ public void receiveOne()
 
     @Override
     public void close()
-    {try
     {
-        markAsInUse();
-        delegate.close();
-    }
-    finally
-    {
-        markAsAvailable();
-    }
+        try
+        {
+            markAsInUse();
+            delegate.close();
+        }
+        finally
+        {
+            markAsAvailable();
+        }
     }
 
     @Override
diff --git a/driver/src/main/java/org/neo4j/driver/internal/pool/InternalConnectionPool.java b/driver/src/main/java/org/neo4j/driver/internal/pool/InternalConnectionPool.java
@@ -42,16 +42,16 @@
 import static java.lang.String.format;
 
 /**
- * A basic connection pool that optimizes for threads being long-lived, acquiring/releasing many connections.
- * It uses a global queue as a fallback pool, but tries to avoid coordination by storing connections in a ThreadLocal.
+ * 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
+ * is full or the pool is being cleaned on driver.close, then we directly close the connection attached with the
+ * session.
  *
- * Safety is achieved by tracking thread locals getting garbage collected, returning connections to the global pool
- * when this happens.
+ * The session is NOT meat 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.
  *
- * If threads are long-lived, this pool will achieve linearly scalable performance with overhead equivalent to a
- * hash-map lookup per acquire.
- *
- * If threads are short-lived, this pool is not ideal.
+ * 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.
  */
 public class InternalConnectionPool implements ConnectionPool
 {
diff --git a/driver/src/main/java/org/neo4j/driver/internal/pool/PooledConnection.java b/driver/src/main/java/org/neo4j/driver/internal/pool/PooledConnection.java
@@ -26,7 +26,26 @@
 import org.neo4j.driver.internal.util.Consumer;
 import org.neo4j.driver.v1.Value;
 import org.neo4j.driver.v1.exceptions.Neo4jException;
-
+/**
+ * The state of a pooledConnection from a pool point of view could be one of the following:
+ * Created,
+ * Available,
+ * Claimed,
+ * Closed,
+ * Disposed.
+ *
+ * The state machine looks like:
+ *
+ *                      session.finalize
+ *                       session.close     failed return to pool
+ * Created -------> Claimed  ----------> Closed ---------> Disposed
+ *                    ^                    |                    ^
+ *      pool.acquire  |                    |returned to pool    |
+ *                    |                    |                    |
+ *                    ---- Available <-----                     |
+ *                              |           pool.close          |
+ *                              ---------------------------------
+ */
 public class PooledConnection implements Connection
 {
     /** The real connection who will do all the real jobs */
@@ -157,23 +176,15 @@ public void receiveOne()
     }
 
     @Override
+    /**
+     * Make sure only close the connection once on each session to avoid releasing the connection twice, a.k.a.
+     * adding back the connection twice into the pool.
+     */
     public void close()
     {
-        // In case this session has an open result or transaction or something,
-        // make sure it's reset to a nice state before we reuse it.
-        try
-        {
-            reset( StreamCollector.NO_OP );
-            sync();
-        }
-        catch (Exception ex)
-        {
-            dispose();
-        }
-        finally
-        {
-            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
     }
 
     @Override
@@ -233,6 +244,7 @@ private boolean isClientOrTransientError( RuntimeException e )
 
     public long idleTime()
     {
-        return clock.millis() - lastUsed;
+        long idleTime = clock.millis() - lastUsed;
+        return idleTime;
     }
 }
diff --git a/driver/src/main/java/org/neo4j/driver/internal/pool/PooledConnectionReleaseConsumer.java b/driver/src/main/java/org/neo4j/driver/internal/pool/PooledConnectionReleaseConsumer.java
@@ -64,15 +64,54 @@ else if ( validConnection( pooledConnection ) )
                 // Otherwise, we close the connection directly here.
                 pooledConnection.dispose();
             }
+            else if ( driverStopped.get() )
+            {
+                // If our adding the pooledConnection to the queue was racing with the closing of the driver,
+                // then the loop where the driver is closing all available connections might not observe our newly
+                // added connection. Thus, we must attempt to remove a connection and dispose it. It doesn't matter
+                // 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();
+                }
+            }
+        }
+        else
+        {
+            pooledConnection.dispose();
         }
     }
 
     boolean validConnection( PooledConnection pooledConnection )
     {
-        return !pooledConnection.hasUnrecoverableErrors() &&
+        return reset(pooledConnection) &&
+               !pooledConnection.hasUnrecoverableErrors() &&
                (pooledConnection.idleTime() <= minIdleBeforeConnectionTest || ping( pooledConnection ));
     }
 
+    /**
+     *  In case this session has an open result or transaction or something,
+     *  make sure it's reset to a nice state before we reuse it.
+     * @param conn the PooledConnection
+     * @return true if the connection is reset successfully without any error, otherwise false.
+     */
+    private boolean reset( PooledConnection conn )
+    {
+        try
+        {
+            conn.reset( StreamCollector.NO_OP );
+            conn.sync();
+            return true;
+        }
+        catch ( Throwable e )
+        {
+            return false;
+        }
+    }
+
     private boolean ping( PooledConnection conn )
     {
         try
diff --git a/driver/src/test/java/org/neo4j/driver/internal/InternalSessionTest.java b/driver/src/test/java/org/neo4j/driver/internal/InternalSessionTest.java
@@ -24,10 +24,14 @@
 
 import org.neo4j.driver.internal.logging.DevNullLogger;
 import org.neo4j.driver.internal.spi.Connection;
+import org.neo4j.driver.internal.spi.Logger;
 import org.neo4j.driver.v1.Transaction;
 import org.neo4j.driver.v1.exceptions.ClientException;
 
+import static junit.framework.Assert.fail;
 import static junit.framework.TestCase.assertNotNull;
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.junit.Assert.assertThat;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -135,4 +139,30 @@ public void shouldNotAllowMoreTransactionsInSessionWhileConnectionClosed() throw
         // When
         sess.beginTransaction();
     }
+
+    @Test
+    public void shouldGetExceptionIfTryingToCloseSessionMoreThanOnce() throws Throwable
+    {
+        // Given
+        InternalSession sess = new InternalSession( mock(Connection.class), mock(Logger.class) );
+        try
+        {
+            sess.close();
+        }
+        catch( Exception e )
+        {
+            fail("Should not get any problem to close first time");
+        }
+
+        // When
+        try
+        {
+            sess.close();
+            fail( "Should have received an error to close second time" );
+        }
+        catch( Exception e )
+        {
+           assertThat( e.getMessage(), equalTo("This session has already been closed." ));
+        }
+    }
 }
diff --git a/driver/src/test/java/org/neo4j/driver/internal/pool/ConnectionInvalidationTest.java b/driver/src/test/java/org/neo4j/driver/internal/pool/ConnectionInvalidationTest.java
@@ -22,13 +22,16 @@
 import org.mockito.Mockito;
 
 import java.io.IOException;
+import java.util.HashMap;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.neo4j.driver.internal.spi.Connection;
+import org.neo4j.driver.internal.spi.StreamCollector;
 import org.neo4j.driver.internal.util.Clock;
 import org.neo4j.driver.internal.util.Consumers;
 import org.neo4j.driver.v1.Config;
+import org.neo4j.driver.v1.Value;
 import org.neo4j.driver.v1.exceptions.ClientException;
 import org.neo4j.driver.v1.exceptions.Neo4jException;
 import org.neo4j.driver.v1.exceptions.TransientException;
@@ -38,6 +41,10 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.anyMap;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -57,7 +64,8 @@ public class ConnectionInvalidationTest
     public void shouldInvalidateConnectionThatIsOld() throws Throwable
     {
         // Given a connection that's broken
-        Mockito.doThrow( new ClientException( "That didn't work" ) ).when( delegate ).sync();
+        Mockito.doThrow( new ClientException( "That didn't work" ) )
+                .when( delegate ).run( anyString(), anyMap(), any( StreamCollector.class ) );
         Config config = Config.defaultConfig();
         when( clock.millis() ).thenReturn( 0L, config.idleTimeBeforeConnectionTest() + 1L );
         PooledConnection conn = new PooledConnection( delegate, Consumers.<PooledConnection>noOp(), clock );
@@ -76,7 +84,8 @@ public void shouldInvalidateConnectionThatIsOld() throws Throwable
     public void shouldNotInvalidateConnectionThatIsNotOld() throws Throwable
     {
         // Given a connection that's broken
-        Mockito.doThrow( new ClientException( "That didn't work" ) ).when( delegate ).sync();
+        Mockito.doThrow( new ClientException( "That didn't work" ) )
+                .when( delegate ).run( anyString(), anyMap(), any( StreamCollector.class ) );
         Config config = Config.defaultConfig();
         when( clock.millis() ).thenReturn( 0L, config.idleTimeBeforeConnectionTest() - 1L );
         PooledConnection conn = new PooledConnection( delegate, Consumers.<PooledConnection>noOp(), clock );
@@ -90,6 +99,23 @@ public void shouldNotInvalidateConnectionThatIsNotOld() throws Throwable
         verify( queue ).offer( conn );
     }
 
+    @Test
+    public void shouldInvalidConnectionIfFailedToReset() throws Throwable
+    {
+        // Given a connection that's broken
+        Mockito.doThrow( new ClientException( "That didn't work" ) ).when( delegate ).reset( any( StreamCollector.class ) );
+        Config config = Config.defaultConfig();
+        PooledConnection conn = new PooledConnection( delegate, Consumers.<PooledConnection>noOp(), clock );
+
+        // When/Then
+        BlockingQueue<PooledConnection> queue = mock( BlockingQueue.class );
+        PooledConnectionReleaseConsumer consumer =
+                new PooledConnectionReleaseConsumer( queue, new AtomicBoolean( false ), config );
+        consumer.accept( conn );
+
+        verify( queue, never() ).add( conn );
+    }
+
     @Test
     public void shouldInvalidateOnUnrecoverableProblems() throws Throwable
     {
@@ -115,12 +141,13 @@ public void shouldInvalidateOnProtocolViolationExceptions() throws Throwable
     @SuppressWarnings( "unchecked" )
     private void assertUnrecoverable( Neo4jException exception )
     {
-        doThrow( exception ).when( delegate ).sync();
+        doThrow( exception ).when( delegate )
+                .run( eq("assert unrecoverable"), anyMap(), any( StreamCollector.class ) );
 
         // When
         try
         {
-            conn.sync();
+            conn.run( "assert unrecoverable", new HashMap<String,Value>( ), StreamCollector.NO_OP );
             fail( "Should've rethrown exception" );
         }
         catch ( Neo4jException e )
@@ -141,12 +168,12 @@ private void assertUnrecoverable( Neo4jException exception )
     @SuppressWarnings( "unchecked" )
     private void assertRecoverable( Neo4jException exception )
     {
-        doThrow( exception ).when( delegate ).sync();
+        doThrow( exception ).when( delegate ).run( eq("assert recoverable"), anyMap(), any( StreamCollector.class ) );
 
         // When
         try
         {
-            conn.sync();
+            conn.run( "assert recoverable", new HashMap<String,Value>( ), StreamCollector.NO_OP );
             fail( "Should've rethrown exception" );
         }
         catch ( Neo4jException e )
diff --git a/driver/src/test/java/org/neo4j/driver/internal/pool/InternalConnectionPoolTest.java b/driver/src/test/java/org/neo4j/driver/internal/pool/InternalConnectionPoolTest.java
@@ -31,6 +31,8 @@
 import org.neo4j.driver.v1.Config;
 
 import static java.util.Collections.singletonList;
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.junit.Assert.assertThat;
 import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
@@ -55,10 +57,11 @@ public void shouldAcquireAndRelease() throws Throwable
         conn.close();
 
         // When
-        pool.acquire( uri );
+        Connection acquired = pool.acquire( uri );
 
         // Then
         verify( connector, times( 1 ) ).connect( uri, config, AuthTokens.none() );
+        assertThat( acquired, equalTo(conn) );
     }
 
     private Connector connector( String scheme )
diff --git a/driver/src/test/java/org/neo4j/driver/internal/pool/PooledConnectionTest.java b/driver/src/test/java/org/neo4j/driver/internal/pool/PooledConnectionTest.java

Original file line number	Diff line number	Diff line change
`@@ -19,6 +19,7 @@`
`19`	`19`	`package org.neo4j.driver.internal;`
`20`	`20`
`21`	`21`	`import java.util.Map;`
	`22`	`+import java.util.concurrent.atomic.AtomicBoolean;`
`22`	`23`
`23`	`24`	`import org.neo4j.driver.internal.spi.Connection;`
`24`	`25`	`import org.neo4j.driver.internal.spi.Logger;`
`@@ -52,7 +53,7 @@ public void run()`
`52`	`53`	`};`
`53`	`54`
`54`	`55`	`private InternalTransaction currentTransaction;`
`55`		`- private boolean isOpen = true;`
	`56`	`+ private AtomicBoolean isOpen = new AtomicBoolean( true );`
`56`	`57`
`57`	`58`	`public InternalSession( Connection connection, Logger logger )`
`58`	`59`	`{`
`@@ -100,19 +101,19 @@ public StatementResult run( Statement statement )`
`100`	`101`	`@Override`
`101`	`102`	`public boolean isOpen()`
`102`	`103`	`{`
`103`		`- return isOpen;`
	`104`	`+ return isOpen.get();`
`104`	`105`	`}`
`105`	`106`
`106`	`107`	`@Override`
`107`	`108`	`public void close()`
`108`	`109`	`{`
`109`		`- if( !isOpen )`
	`110`	`+ // Use atomic operation to protect from closing the connection twice (putting back to the pool twice).`
	`111`	`+ if( !isOpen.compareAndSet( true, false ) )`
`110`	`112`	`{`
`111`	`113`	`throw new ClientException( "This session has already been closed." );`
`112`	`114`	`}`
`113`	`115`	`else`
`114`	`116`	`{`
`115`		`- isOpen = false;`
`116`	`117`	`if ( currentTransaction != null )`
`117`	`118`	`{`
`118`	`119`	`try`
`@@ -171,7 +172,7 @@ private void ensureConnectionIsValidBeforeOpeningTransaction()`
`171`	`172`	`@Override`
`172`	`173`	`protected void finalize() throws Throwable`
`173`	`174`	`{`
`174`		`- if( isOpen )`
	`175`	`+ if( isOpen.compareAndSet( true, false ) )`
`175`	`176`	`{`
`176`	`177`	logger.error( "Neo4j Session object leaked, please ensure that your application calls the `close` " +
`177`	`178`	`"method on Sessions before disposing of the objects.", null );`