Avoid marking objects as both live and disposed

In the `ThreadCachingPool` whenever an object is picked from the
`ThreadLocal` we cannot add it to the `disposed` queue directly since the
object will already be present in the `live` queue (or will be on release). Doing
that means that the object will be present both in the  `live` and `disposed`
queue and also means that on reallocation we can multiple copies of the same
reference and can lead to that the `live` queue blows up in size.

Author: pontusmelke

driver/src/main/java/org/neo4j/driver/internal/pool/ThreadCachingPool.java

@@ -106,12 +106,9 @@ public T acquire( long timeout, TimeUnit unit ) throws InterruptedException
                 allocator.onAcquire( slot.value );
                 return slot.value;
             }
-            else
-            {
-                // We've acquired the slot, but the validation strategy says it's time for it to die. Dispose of it,
-                // and go to the global pool.
-                dispose( slot );
-            }
+
+            //The slot was invalidated however we cannot put it to the
+            //disposed queue yet since it already exists in the live queue
         }
 
         // 2. If that fails, acquire from big pool
@@ -133,17 +130,18 @@ private T acquireFromGlobal( long deadline ) throws InterruptedException
             if ( slot != null )
             {
                 // Yay, got a slot - can we keep it?
-                if ( slot.availableToClaimed() )
+                if ( slot.isValid( validationStrategy ) )
                 {
-                    if ( slot.isValid( validationStrategy ) )
+                    if ( slot.availableToClaimed() )
                     {
                         break;
                     }
-                    else
-                    {
-                        // We've acquired the slot, but the validation strategy says it's time for it to die.
-                        dispose( slot );
-                    }
+                }
+                // We've acquired the slot, but the validation strategy says it's time for it to die.
+                // Either the slot is already claimed or if it is available make it claimed
+                else if ( slot.isClaimedOrAvailableToClaimed() )
+                {
+                    dispose( slot );
                 }
             }
             else
@@ -186,15 +184,13 @@ private T acquireFromGlobal( long deadline ) throws InterruptedException
 
     private void dispose( Slot<T> slot )
     {
-        if ( !slot.claimedToDisposed() )
+        if ( slot.claimedToDisposed() )
         {
-            throw new IllegalStateException( "Cannot dispose unclaimed pool object: " + slot );
+            // Done before below, in case dispose call fails. This is safe since objects on the
+            // pool are used for read-only operations
+            disposed.add( slot );
+            allocator.onDispose( slot.value );
         }
-
-        // Done before below, in case dispose call fails. This is safe since objects on the
-        // pool are used for read-only operations
-        disposed.add( slot );
-        allocator.onDispose( slot.value );
     }
 
     /**
@@ -234,33 +230,30 @@ private Consumer<T> createDisposeCallback( final Slot<T> slot )
             public void accept( T t )
             {
                 slot.updateUsageTimestamp();
-                if ( !slot.isValid( validationStrategy ) )
+                if ( !slot.isValid( validationStrategy) )
                 {
-                    // The value has for some reason become invalid, dispose of it
                     dispose( slot );
                     return;
                 }
 
-                if ( !slot.claimedToAvailable() )
-                {
-                    throw new IllegalStateException( "Failed to release pooled object: " + slot );
-                }
-
-                // Make sure the pool isn't being stopped in the middle of all these shenanigans
-                if ( !stopped.get() )
-                {
-                    // All good, as you were.
-                    live.add( slot );
-                }
-                else
+                if ( slot.claimedToAvailable() )
                 {
-                    // Another thread concurrently closing the pool may have started closing before we
-                    // set our slot to "available". In that case, the slot will not be disposed of by the closing thread
-                    // We mitigate this by trying to claim the slot back - if we are able to, we dispose the slot.
-                    // If we can't claim the slot back, that means another thread is dealing with it.
-                    if ( slot.availableToClaimed() )
+                    // Make sure the pool isn't being stopped in the middle of all these shenanigans
+                    if ( !stopped.get() )
                     {
-                        dispose( slot );
+                        // All good, as you were.
+                        live.add( slot );
+                    }
+                    else
+                    {
+                        // Another thread concurrently closing the pool may have started closing before we
+                        // set our slot to "available". In that case, the slot will not be disposed of by the closing thread
+                        // We mitigate this by trying to claim the slot back - if we are able to, we dispose the slot.
+                        // If we can't claim the slot back, that means another thread is dealing with it.
+                        if ( slot.availableToClaimed() )
+                        {
+                            dispose( slot );
+                        }
                     }
                 }
             }
@@ -304,13 +297,6 @@ enum State
     long lastUsed;
     T value;
 
-    public static <T> Slot<T> disposed( int index, Clock clock )
-    {
-        Slot<T> slot = new Slot<>( index, clock );
-        slot.claimedToDisposed();
-        return slot;
-    }
-
     /**
      * @param index the index into the {@link ThreadCachingPool#all all} array, used to re-use that slot when this is
      * disposed
@@ -343,6 +329,16 @@ public boolean claimedToDisposed()
         return state.compareAndSet( State.CLAIMED, State.DISPOSED );
     }
 
+    public boolean isClaimedOrAvailableToClaimed()
+    {
+        return availableToClaimed() || state.get() == State.CLAIMED;
+    }
+
+    public boolean disposed()
+    {
+        return state.get() == State.DISPOSED;
+    }
+
     public void updateUsageTimestamp()
     {
         lastUsed = clock.millis();

driver/src/test/java/org/neo4j/driver/internal/pool/ThreadCachingPoolTest.java

@@ -18,33 +18,45 @@
  */
 package org.neo4j.driver.internal.pool;
 
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.reflect.Field;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.rules.ExpectedException;
-
 import org.neo4j.driver.internal.util.Clock;
 import org.neo4j.driver.internal.util.Consumer;
 import org.neo4j.driver.v1.exceptions.ClientException;
 
 import static junit.framework.TestCase.fail;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.hasSize;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
 
 public class ThreadCachingPoolTest
 {
     private final List<PooledObject> inUse = new LinkedList<>();
     private final List<PooledObject> inPool = new LinkedList<>();
     private final List<PooledObject> disposed = new LinkedList<>();
-
+    private final MethodHandle liveQueueGet = queueGetter( "live" );
+    private final MethodHandle disposedQueueGet = queueGetter( "disposed" );
+    private final ExecutorService executor = Executors.newFixedThreadPool( 10 );
     private static AtomicInteger IDGEN = new AtomicInteger();
 
     @Rule
@@ -127,7 +139,7 @@ public void shouldDisposeOfInvalidItems() throws Throwable
     {
         // Given
         ThreadCachingPool<PooledObject>
-                pool = new ThreadCachingPool<>( 4, trackAllocator, invalidIfIdIs(0), Clock.SYSTEM );
+                pool = new ThreadCachingPool<>( 4, trackAllocator, invalidIfIdIs( 0 ), Clock.SYSTEM );
 
         // And given we've allocated/releasd object with id 0 once (no validation on first allocation)
         // TODO: Is that the right thing to do? I assume the allocator will allocate healthy objects..
@@ -137,8 +149,8 @@ public void shouldDisposeOfInvalidItems() throws Throwable
         pool.acquire( 10, TimeUnit.SECONDS );
 
         // Then object with id 0 should've been disposed of, and we should have one live object with id 1
-        assertThat( inPool,   equalTo( none() ) );
-        assertThat( inUse,    equalTo( items( 1 ) ) );
+        assertThat( inPool, equalTo( none() ) );
+        assertThat( inUse, equalTo( items( 1 ) ) );
         assertThat( disposed, equalTo( items( 0 ) ) );
     }
 
@@ -171,8 +183,8 @@ public void shouldDisposeOfObjectsThatBecomeInvalidWhileInUse() throws Throwable
         val.invalidate().release();
 
         // Then
-        assertThat( inPool,   equalTo( none() ) );
-        assertThat( inUse,    equalTo( none() ) );
+        assertThat( inPool, equalTo( none() ) );
+        assertThat( inUse, equalTo( none() ) );
         assertThat( disposed, equalTo( items( val ) ) );
     }
 
@@ -191,9 +203,9 @@ public void shouldRecoverFromItemCreationFailure() throws Throwable
             try
             {
                 pool.acquire( 10, TimeUnit.SECONDS );
-                fail("Should not succeed at allocating any item here.");
+                fail( "Should not succeed at allocating any item here." );
             }
-            catch( ClientException e )
+            catch ( ClientException e )
             {
                 // Expected
             }
@@ -207,8 +219,8 @@ public void shouldRecoverFromItemCreationFailure() throws Throwable
         {
             pool.acquire( 10, TimeUnit.SECONDS );
         }
-        assertThat( inPool,   equalTo( none() ) );
-        assertThat( inUse,    equalTo( items( 0, 1, 2, 3 ) ) );
+        assertThat( inPool, equalTo( none() ) );
+        assertThat( inUse, equalTo( items( 0, 1, 2, 3 ) ) );
         assertThat( disposed, equalTo( none() ) ); // because allocation fails, onDispose is not called
     }
 
@@ -256,10 +268,88 @@ public void shouldRecovedDisposedItemReallocationFailing() throws Throwable
         assertThat( inPool, equalTo( none() ) );
         assertThat( inUse, equalTo( items( 2, 3 ) ) );
         // only the first two items value onDispose called, since allocation fails after that
-        assertThat( disposed, equalTo( items( 0, 1) ) );
+        assertThat( disposed, equalTo( items( 0, 1 ) ) );
+    }
+
+    @SuppressWarnings( "unchecked" )
+    @Test
+    public void shouldNotHaveReferenceAsBothLiveAndDisposed() throws Throwable
+    {
+        // Given
+        final ThreadCachingPool<PooledObject>
+                pool = new ThreadCachingPool<>( 4, trackAllocator, checkInvalidateFlag, Clock.SYSTEM );
+
+        // This object will be cached in ThreadLocal
+        final PooledObject obj1 = pool.acquire( 10, TimeUnit.SECONDS );
+
+        //This will add another object to the live queue
+        assertTrue( acquireInSeparateThread( pool ) );
+
+        //Now we release the first object, meaning that it will be added
+        //to the live queue (as well as being cached as ThreadLocal in this thread)
+        obj1.release();
+        //Now we invalidate the object
+        obj1.invalidate();
+
+        // When
+        //Now the cached object is invalidated, we should now pick the object
+        //from the live objects created in the background thread
+        PooledObject obj2 = pool.acquire( 10, TimeUnit.SECONDS );
+
+        //THEN
+        assertThat( obj1.id, equalTo( 0 ) );
+        assertThat( obj2.id, equalTo( 1 ) );
+        BlockingQueue<Slot<PooledObject>> liveQueue = (BlockingQueue<Slot<PooledObject>>) liveQueueGet.invoke( pool );
+        BlockingQueue<Slot<PooledObject>> disposedQueue =
+                (BlockingQueue<Slot<PooledObject>>) disposedQueueGet.invoke( pool );
+
+        assertThat( disposedQueue, empty() );
+        assertThat( liveQueue, hasSize( 1 ) );
+        assertThat( liveQueue.poll( 10, TimeUnit.SECONDS ).value.id, equalTo( 0 ) );
+    }
+
+    private boolean acquireInSeparateThread( final ThreadCachingPool<PooledObject> pool ) throws InterruptedException
+    {
+        final AtomicBoolean succeeded = new AtomicBoolean( true );
+        executor.execute( new Runnable()
+        {
+            @Override
+            public void run()
+            {
+                try
+                {
+                    PooledObject obj = pool.acquire( 10, TimeUnit.MINUTES );
+                    obj.release();
+                }
+                catch ( InterruptedException e )
+                {
+                    succeeded.set( false );
+                }
+            }
+        } );
+        executor.awaitTermination( 2, TimeUnit.SECONDS );
+        return succeeded.get();
+    }
+
+
+    //This is terrible hack, but I really want to keep the queues private in
+    //ThreadCachingPool
+    private static MethodHandle queueGetter( String name )
+    {
+        try
+        {
+            MethodHandles.Lookup lookup = MethodHandles.lookup();
+            Field value = ThreadCachingPool.class.getDeclaredField( name );
+            value.setAccessible( true );
+            return lookup.unreflectGetter( value );
+        }
+        catch ( NoSuchFieldException | IllegalAccessException e )
+        {
+            throw new AssertionError( e );
+        }
     }
 
-    private List<PooledObject> items( int ... objects )
+    private List<PooledObject> items( int... objects )
     {
         List<PooledObject> out = new LinkedList<>();
         for ( int id : objects )
@@ -269,9 +359,9 @@ private List<PooledObject> items( int ... objects )
         return out;
     }
 
-    private List<PooledObject> items( PooledObject ... objects )
+    private List<PooledObject> items( PooledObject... objects )
     {
-        return Arrays.asList(objects);
+        return Arrays.asList( objects );
     }
 
     private List<PooledObject> none()
@@ -305,7 +395,7 @@ private class PooledObject
 
         public PooledObject( Consumer<PooledObject> release )
         {
-            this(IDGEN.getAndIncrement(), release);
+            this( IDGEN.getAndIncrement(), release );
         }
 
         public PooledObject( int id, Consumer<PooledObject> release )
@@ -362,7 +452,7 @@ private class TestAllocator implements Allocator<PooledObject>
         @Override
         public PooledObject allocate( Consumer<PooledObject> release )
         {
-            if( creationException != null )
+            if ( creationException != null )
             {
                 throw creationException;
             }