Don't add everything back to live queue

Previously whenever an object was picked off the thread local cache it was
added back to the live queue even thought it was never removed from there.
This lead to indefinite growth of the live queue.

Author: pontusmelke

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

@@ -94,18 +94,22 @@ public ThreadCachingPool( int targetSize, Allocator<T> allocator, ValidationStra
 
     public T acquire( long timeout, TimeUnit unit ) throws InterruptedException
     {
+        assert live.size() <= maxSize;
         long deadline = clock.millis() + unit.toMillis( timeout );
 
         // 1. Try and value an object from our local slot
         Slot<T> slot = local.get();
 
-        if ( slot != null && slot.availableToClaimed() )
+        if ( slot != null && slot.availableToThreadLocalClaimed() )
         {
             if ( slot.isValid( validationStrategy ) )
             {
                 allocator.onAcquire( slot.value );
                 return slot.value;
             }
+            else {
+                dispose( slot );
+            }
 
             //The slot was invalidated however we cannot put it to the
             //disposed queue yet since it already exists in the live queue
@@ -177,14 +181,32 @@ else if ( slot.isClaimedOrAvailableToClaimed() )
 
         // Keep this slot cached with our thread, so that we can grab this value quickly next time,
         // assuming threads generally availableToClaimed one instance at a time
-        local.set( slot );
+        updateThreadLocal( slot );
         allocator.onAcquire( slot.value );
         return slot.value;
     }
 
+    private void updateThreadLocal(Slot<T> slot)
+    {
+        Slot<T> localSlot = local.get();
+        if ( localSlot != null )
+        {
+            //The old slot is no longer in the tread local
+            localSlot.threadLocalClaimedToClaimed();
+        }
+        else
+        {
+            //There was nothing stored in thread local
+            //no we must also add this slot to the live queue
+            live.add( slot );
+        }
+        slot.claimByThreadLocal();
+        local.set( slot );
+    }
+
     private void dispose( Slot<T> slot )
     {
-        if ( slot.claimedToDisposed() )
+        if ( slot.claimedToDisposed() || slot.threadLocalClaimedToDisposed() )
         {
             // Done before below, in case dispose call fails. This is safe since objects on the
             // pool are used for read-only operations
@@ -213,7 +235,7 @@ private Slot<T> allocate( int slotIndex )
             // Return it :)
             return slot;
         }
-        catch( Neo4jException e )
+        catch ( Neo4jException e )
         {
             // Failed to allocate slot, return it to the list of disposed slots, rethrow exception.
             slot.claimedToDisposed();
@@ -230,7 +252,7 @@ private Consumer<T> createDisposeCallback( final Slot<T> slot )
             public void accept( T t )
             {
                 slot.updateUsageTimestamp();
-                if ( !slot.isValid( validationStrategy) )
+                if ( !slot.isValid( validationStrategy ) )
                 {
                     dispose( slot );
                     return;
@@ -256,6 +278,17 @@ public void accept( T t )
                         }
                     }
                 }
+
+                // If we are claimed by thread local we are already in the live queue
+                if ( slot.threadLocalClaimedToAvailable() && stopped.get() )
+                {
+                    // As above, try to claim the slot back and dispose
+                    if ( slot.availableToClaimed() )
+                    {
+                        dispose( slot );
+                    }
+                }
+
             }
         };
     }
@@ -286,6 +319,7 @@ class Slot<T>
     enum State
     {
         AVAILABLE,
+        THREAD_LOCAL_CLAIMED,
         CLAIMED,
         DISPOSED
     }
@@ -318,6 +352,11 @@ public boolean availableToClaimed()
         return state.compareAndSet( State.AVAILABLE, State.CLAIMED );
     }
 
+    public boolean availableToThreadLocalClaimed()
+    {
+        return state.compareAndSet( State.AVAILABLE, State.THREAD_LOCAL_CLAIMED );
+    }
+
     public boolean claimedToAvailable()
     {
         updateUsageTimestamp();
@@ -329,6 +368,26 @@ public boolean claimedToDisposed()
         return state.compareAndSet( State.CLAIMED, State.DISPOSED );
     }
 
+    public boolean threadLocalClaimedToDisposed()
+    {
+        return state.compareAndSet( State.THREAD_LOCAL_CLAIMED, State.DISPOSED );
+    }
+
+    public boolean threadLocalClaimedToClaimed()
+    {
+        return state.compareAndSet( State.THREAD_LOCAL_CLAIMED, State.CLAIMED );
+    }
+
+    public boolean threadLocalClaimedToAvailable()
+    {
+        return state.compareAndSet( State.THREAD_LOCAL_CLAIMED, State.AVAILABLE );
+    }
+
+    public void claimByThreadLocal()
+    {
+        state.set( State.THREAD_LOCAL_CLAIMED );
+    }
+
     public boolean isClaimedOrAvailableToClaimed()
     {
         return availableToClaimed() || state.get() == State.CLAIMED;

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

@@ -303,9 +303,47 @@ public void shouldNotHaveReferenceAsBothLiveAndDisposed() throws Throwable
         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 ) );
+        assertThat( disposedQueue, hasSize( 1 ) );
+        assertThat( disposedQueue.poll( 10, TimeUnit.SECONDS ).value.id, equalTo( 0 ) );
+        assertThat( liveQueue, empty() );
+    }
+
+    @SuppressWarnings( "unchecked" )
+    @Test
+    public void shouldNotAddToLiveQueueTwice() throws Throwable
+    {
+        // Given
+        ThreadCachingPool<PooledObject>
+                pool = new ThreadCachingPool<>( 4, trackAllocator, checkInvalidateFlag, Clock.SYSTEM );
+
+        // When
+        PooledObject o1 = pool.acquire( 10, TimeUnit.SECONDS );
+        o1.release();
+        PooledObject o2 = pool.acquire( 10, TimeUnit.SECONDS );
+        o2.release();
+
+        // Then
+        BlockingQueue<Slot<PooledObject>> liveQueue = (BlockingQueue<Slot<PooledObject>>) liveQueueGet.invoke( pool );
+        assertThat(liveQueue, hasSize( 1 ));
+    }
+
+    @SuppressWarnings( "unchecked" )
+    @Test
+    public void shouldNotAddToLiveQueueTwice2() throws Throwable
+    {
+        // Given
+        ThreadCachingPool<PooledObject>
+                pool = new ThreadCachingPool<>( 4, trackAllocator, checkInvalidateFlag, Clock.SYSTEM );
+        PooledObject o1 = pool.acquire( 10, TimeUnit.SECONDS );
+        PooledObject o2 = pool.acquire( 10, TimeUnit.SECONDS );
+
+        // When
+        o1.release();
+        o2.release();
+
+        // Then
+        BlockingQueue<Slot<PooledObject>> liveQueue = (BlockingQueue<Slot<PooledObject>>) liveQueueGet.invoke( pool );
+        assertThat(liveQueue, hasSize( 1 ));
     }
 
     private boolean acquireInSeparateThread( final ThreadCachingPool<PooledObject> pool ) throws InterruptedException
@@ -318,7 +356,7 @@ public void run()
             {
                 try
                 {
-                    PooledObject obj = pool.acquire( 10, TimeUnit.MINUTES );
+                    PooledObject obj = pool.acquire( 10, TimeUnit.SECONDS );
                     obj.release();
                 }
                 catch ( InterruptedException e )