Forget auto-read managing handler when it is dequeued

Make sure `InboundMessageDispatcher` does not hold on to the last
auto-read management handler after it is removed from the queue of
all handlers.

Author: lutovich

driver/src/main/java/org/neo4j/driver/internal/async/inbound/InboundMessageDispatcher.java

@@ -119,7 +119,7 @@ public void handleAckFailureMessage()
     public void handleSuccessMessage( Map<String,Value> meta )
     {
         log.debug( "S: SUCCESS %s", meta );
-        ResponseHandler handler = handlers.remove();
+        ResponseHandler handler = removeHandler();
         handler.onSuccess( meta );
     }
 
@@ -152,7 +152,7 @@ public void handleFailureMessage( String code, String message )
         // try to write ACK_FAILURE before notifying the next response handler
         ackFailureIfNeeded();
 
-        ResponseHandler handler = handlers.remove();
+        ResponseHandler handler = removeHandler();
         handler.onFailure( currentError );
     }
 
@@ -161,7 +161,7 @@ public void handleIgnoredMessage()
     {
         log.debug( "S: IGNORED" );
 
-        ResponseHandler handler = handlers.remove();
+        ResponseHandler handler = removeHandler();
 
         Throwable error;
         if ( currentError != null )
@@ -189,7 +189,7 @@ public void handleFatalError( Throwable error )
 
         while ( !handlers.isEmpty() )
         {
-            ResponseHandler handler = handlers.remove();
+            ResponseHandler handler = removeHandler();
             handler.onFailure( currentError );
         }
     }
@@ -245,6 +245,14 @@ public boolean isAckFailureMuted()
         return ackFailureMuted;
     }
 
+    /**
+     * <b>Visible for testing</b>
+     */
+    AutoReadManagingResponseHandler autoReadManagingHandler()
+    {
+        return autoReadManagingHandler;
+    }
+
     private void ackFailureIfNeeded()
     {
         if ( !ackFailureMuted )
@@ -254,6 +262,18 @@ private void ackFailureIfNeeded()
         }
     }
 
+    private ResponseHandler removeHandler()
+    {
+        ResponseHandler handler = handlers.remove();
+        if ( autoReadManagingHandler == handler )
+        {
+            // handler that is being removed is the auto-read managing handler
+            // make sure this dispatcher does not hold on to a removed handler
+            autoReadManagingHandler = null;
+        }
+        return handler;
+    }
+
     private void updateAutoReadManagingHandlerIfNeeded( ResponseHandler handler )
     {
         if ( handler instanceof AutoReadManagingResponseHandler )
@@ -263,7 +283,7 @@ private void updateAutoReadManagingHandlerIfNeeded( ResponseHandler handler )
                 // there already exists a handler that manages channel's auto-read
                 // make it stop because new managing handler is being added and there should only be a single such handler
                 autoReadManagingHandler.disableAutoReadManagement();
-                // restore the original auto-read value
+                // restore the default value of auto-read
                 channel.config().setAutoRead( true );
             }
             autoReadManagingHandler = (AutoReadManagingResponseHandler) handler;

driver/src/test/java/org/neo4j/driver/internal/async/inbound/InboundMessageDispatcherTest.java

@@ -465,6 +465,44 @@ public void shouldKeepSingleAutoReadManagingHandler()
         inOrder.verify( handler3, never() ).disableAutoReadManagement();
     }
 
+    @Test
+    public void shouldKeepTrackOfAutoReadManagingHandler()
+    {
+        InboundMessageDispatcher dispatcher = newDispatcher();
+
+        AutoReadManagingResponseHandler handler1 = mock( AutoReadManagingResponseHandler.class );
+        AutoReadManagingResponseHandler handler2 = mock( AutoReadManagingResponseHandler.class );
+
+        assertNull( dispatcher.autoReadManagingHandler() );
+
+        dispatcher.enqueue( handler1 );
+        assertEquals( handler1, dispatcher.autoReadManagingHandler() );
+
+        dispatcher.enqueue( handler2 );
+        assertEquals( handler2, dispatcher.autoReadManagingHandler() );
+    }
+
+    @Test
+    public void shouldForgetAutoReadManagingHandlerWhenItIsRemoved()
+    {
+        InboundMessageDispatcher dispatcher = newDispatcher();
+
+        ResponseHandler handler1 = mock( ResponseHandler.class );
+        ResponseHandler handler2 = mock( ResponseHandler.class );
+        AutoReadManagingResponseHandler handler3 = mock( AutoReadManagingResponseHandler.class );
+
+        dispatcher.enqueue( handler1 );
+        dispatcher.enqueue( handler2 );
+        dispatcher.enqueue( handler3 );
+        assertEquals( handler3, dispatcher.autoReadManagingHandler() );
+
+        dispatcher.handleSuccessMessage( emptyMap() );
+        dispatcher.handleSuccessMessage( emptyMap() );
+        dispatcher.handleSuccessMessage( emptyMap() );
+
+        assertNull( dispatcher.autoReadManagingHandler() );
+    }
+
     private static void verifyFailure( ResponseHandler handler )
     {
         ArgumentCaptor<Neo4jException> captor = ArgumentCaptor.forClass( Neo4jException.class );
@@ -475,14 +513,19 @@ private static void verifyFailure( ResponseHandler handler )
 
     private static InboundMessageDispatcher newDispatcher()
     {
-        Channel channel = mock( Channel.class );
-        ChannelConfig channelConfig = mock( ChannelConfig.class );
-        when( channel.config() ).thenReturn( channelConfig );
-        return newDispatcher( channel );
+        return newDispatcher( newChannelMock() );
     }
 
     private static InboundMessageDispatcher newDispatcher( Channel channel )
     {
         return new InboundMessageDispatcher( channel, DEV_NULL_LOGGING );
     }
+
+    private static Channel newChannelMock()
+    {
+        Channel channel = mock( Channel.class );
+        ChannelConfig channelConfig = mock( ChannelConfig.class );
+        when( channel.config() ).thenReturn( channelConfig );
+        return channel;
+    }
 }