Ensure a single handler manages channel's auto-read

`PullAllResponseHandler` is responsible for receiving a stream of
records. Stream might be very large and handler tries to not fetch
if fully into memory. It applies a network-level backpressure by
disabling/enabling auto-read property of the underlying Netty channel.
This property tells channel to not read from the network. Auto-read is
disabled when number of queued records goes beyond the threshold.

Management of auto-read turned out to be problematic for nested
queries within a single transaction. Nested queries result in
multiple `PullAllResponseHandler`s being added to the queue of
handlers. They will try to manage auto-read concurrently and can
sometimes disable it completely. Callers would then be blocked
and unable to proceed. This is not a problem for `Session#run()`
because every such call ensures there was a logical SYNC and the
previous query has completed and its result is buffered.

This commit fixes a problem by making only a single installed handler
manage the channel's auto-read property. Making sure there only exists
a single such handler is the responsibility of
`InboundMessageDispatcher`. Auto-read is enabled when new auto-read
managing handler is installed. Previous such handler is disabled.

Author: lutovich

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

@@ -28,6 +28,7 @@
 import org.neo4j.driver.internal.handlers.AckFailureResponseHandler;
 import org.neo4j.driver.internal.logging.ChannelActivityLogger;
 import org.neo4j.driver.internal.messaging.MessageHandler;
+import org.neo4j.driver.internal.spi.AutoReadManagingResponseHandler;
 import org.neo4j.driver.internal.spi.ResponseHandler;
 import org.neo4j.driver.internal.util.ErrorUtil;
 import org.neo4j.driver.v1.Logger;
@@ -48,6 +49,8 @@ public class InboundMessageDispatcher implements MessageHandler
     private boolean fatalErrorOccurred;
     private boolean ackFailureMuted;
 
+    private AutoReadManagingResponseHandler autoReadManagingHandler;
+
     public InboundMessageDispatcher( Channel channel, Logging logging )
     {
         this.channel = requireNonNull( channel );
@@ -63,6 +66,7 @@ public void enqueue( ResponseHandler handler )
         else
         {
             handlers.add( handler );
+            updateAutoReadManagingHandlerIfNeeded( handler );
         }
     }
 
@@ -249,4 +253,20 @@ private void ackFailureIfNeeded()
             channel.writeAndFlush( ACK_FAILURE, channel.voidPromise() );
         }
     }
+
+    private void updateAutoReadManagingHandlerIfNeeded( ResponseHandler handler )
+    {
+        if ( handler instanceof AutoReadManagingResponseHandler )
+        {
+            if ( autoReadManagingHandler != null )
+            {
+                // 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
+                channel.config().setAutoRead( true );
+            }
+            autoReadManagingHandler = (AutoReadManagingResponseHandler) handler;
+        }
+    }
 }

driver/src/main/java/org/neo4j/driver/internal/handlers/PullAllResponseHandler.java

@@ -27,8 +27,8 @@
 import java.util.concurrent.CompletionStage;
 
 import org.neo4j.driver.internal.InternalRecord;
+import org.neo4j.driver.internal.spi.AutoReadManagingResponseHandler;
 import org.neo4j.driver.internal.spi.Connection;
-import org.neo4j.driver.internal.spi.ResponseHandler;
 import org.neo4j.driver.internal.util.Futures;
 import org.neo4j.driver.internal.util.Iterables;
 import org.neo4j.driver.internal.util.MetadataUtil;
@@ -44,7 +44,7 @@
 import static org.neo4j.driver.internal.util.Futures.completedWithNull;
 import static org.neo4j.driver.internal.util.Futures.failedFuture;
 
-public abstract class PullAllResponseHandler implements ResponseHandler
+public abstract class PullAllResponseHandler implements AutoReadManagingResponseHandler
 {
     private static final Queue<Record> UNINITIALIZED_RECORDS = Iterables.emptyQueue();
 
@@ -58,6 +58,7 @@ public abstract class PullAllResponseHandler implements ResponseHandler
     // initialized lazily when first record arrives
     private Queue<Record> records = UNINITIALIZED_RECORDS;
 
+    private boolean autoReadManagementEnabled = true;
     private boolean finished;
     private Throwable failure;
     private ResultSummary summary;
@@ -129,6 +130,12 @@ public synchronized void onRecord( Value[] fields )
         }
     }
 
+    @Override
+    public synchronized void disableAutoReadManagement()
+    {
+        autoReadManagementEnabled = false;
+    }
+
     public synchronized CompletionStage<Record> peekAsync()
     {
         Record record = records.peek();
@@ -209,7 +216,7 @@ else if ( finished )
                 // neither SUCCESS nor FAILURE message has arrived, register future to be notified when it arrives
                 // future will be completed with null on SUCCESS and completed with Throwable on FAILURE
                 // enable auto-read, otherwise we might not read SUCCESS/FAILURE if records are not consumed
-                connection.enableAutoRead();
+                enableAutoRead();
                 failureFuture = new CompletableFuture<>();
             }
             return failureFuture;
@@ -234,7 +241,7 @@ private void enqueueRecord( Record record )
             // more than high watermark records are already queued, tell connection to stop auto-reading from network
             // this is needed to deal with slow consumers, we do not want to buffer all records in memory if they are
             // fetched from network faster than consumed
-            connection.disableAutoRead();
+            disableAutoRead();
         }
     }
 
@@ -246,7 +253,7 @@ private Record dequeueRecord()
         {
             // less than low watermark records are now available in the buffer, tell connection to pre-fetch more
             // and populate queue with new records from network
-            connection.enableAutoRead();
+            enableAutoRead();
         }
 
         return record;
@@ -319,4 +326,20 @@ private ResultSummary extractResultSummary( Map<String,Value> metadata )
         long resultAvailableAfter = runResponseHandler.resultAvailableAfter();
         return MetadataUtil.extractSummary( statement, connection, resultAvailableAfter, metadata );
     }
+
+    private void enableAutoRead()
+    {
+        if ( autoReadManagementEnabled )
+        {
+            connection.enableAutoRead();
+        }
+    }
+
+    private void disableAutoRead()
+    {
+        if ( autoReadManagementEnabled )
+        {
+            connection.disableAutoRead();
+        }
+    }
 }

driver/src/main/java/org/neo4j/driver/internal/spi/AutoReadManagingResponseHandler.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright (c) 2002-2018 "Neo4j,"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.neo4j.driver.internal.spi;
+
+import org.neo4j.driver.internal.async.inbound.InboundMessageDispatcher;
+
+/**
+ * A type of {@link ResponseHandler handler} that manages auto-read of the underlying connection using {@link Connection#enableAutoRead()} and
+ * {@link Connection#disableAutoRead()}.
+ * <p>
+ * Implementations can use auto-read management to apply network-level backpressure when receiving a stream of records.
+ * There should only be a single such handler active for a connection at one point in time. Otherwise, handlers can interfere and turn on/off auto-read
+ * racing with each other. {@link InboundMessageDispatcher} is responsible for tracking these handlers and disabling auto-read management to maintain just
+ * a single auto-read managing handler per connection.
+ */
+public interface AutoReadManagingResponseHandler extends ResponseHandler
+{
+    /**
+     * Tell this handler that it should stop changing auto-read setting for the connection.
+     */
+    void disableAutoReadManagement();
+}

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

@@ -19,6 +19,7 @@
 package org.neo4j.driver.internal.async.inbound;
 
 import io.netty.channel.Channel;
+import io.netty.channel.ChannelConfig;
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
 import org.mockito.InOrder;
@@ -27,6 +28,7 @@
 import java.util.HashMap;
 import java.util.Map;
 
+import org.neo4j.driver.internal.spi.AutoReadManagingResponseHandler;
 import org.neo4j.driver.internal.spi.ResponseHandler;
 import org.neo4j.driver.internal.value.IntegerValue;
 import org.neo4j.driver.v1.Value;
@@ -49,6 +51,7 @@
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyZeroInteractions;
+import static org.mockito.Mockito.when;
 import static org.neo4j.driver.internal.logging.DevNullLogging.DEV_NULL_LOGGING;
 import static org.neo4j.driver.internal.messaging.AckFailureMessage.ACK_FAILURE;
 import static org.neo4j.driver.v1.Values.value;
@@ -443,6 +446,25 @@ public void shouldMuteAndUnMuteAckFailure()
         assertFalse( dispatcher.isAckFailureMuted() );
     }
 
+    @Test
+    public void shouldKeepSingleAutoReadManagingHandler()
+    {
+        InboundMessageDispatcher dispatcher = newDispatcher();
+
+        AutoReadManagingResponseHandler handler1 = mock( AutoReadManagingResponseHandler.class );
+        AutoReadManagingResponseHandler handler2 = mock( AutoReadManagingResponseHandler.class );
+        AutoReadManagingResponseHandler handler3 = mock( AutoReadManagingResponseHandler.class );
+
+        dispatcher.enqueue( handler1 );
+        dispatcher.enqueue( handler2 );
+        dispatcher.enqueue( handler3 );
+
+        InOrder inOrder = inOrder( handler1, handler2, handler3 );
+        inOrder.verify( handler1 ).disableAutoReadManagement();
+        inOrder.verify( handler2 ).disableAutoReadManagement();
+        inOrder.verify( handler3, never() ).disableAutoReadManagement();
+    }
+
     private static void verifyFailure( ResponseHandler handler )
     {
         ArgumentCaptor<Neo4jException> captor = ArgumentCaptor.forClass( Neo4jException.class );
@@ -453,7 +475,10 @@ private static void verifyFailure( ResponseHandler handler )
 
     private static InboundMessageDispatcher newDispatcher()
     {
-        return newDispatcher( mock( Channel.class ) );
+        Channel channel = mock( Channel.class );
+        ChannelConfig channelConfig = mock( ChannelConfig.class );
+        when( channel.config() ).thenReturn( channelConfig );
+        return newDispatcher( channel );
     }
 
     private static InboundMessageDispatcher newDispatcher( Channel channel )

driver/src/test/java/org/neo4j/driver/internal/handlers/PullAllResponseHandlerTest.java

@@ -823,6 +823,21 @@ public void shouldEnableAutoReadOnConnectionWhenSummaryRequestedButNotAvailable(
         assertNotNull( summaryFuture.get() );
     }
 
+    @Test
+    public void shouldNotDisableAutoReadWhenAutoReadManagementDisabled()
+    {
+        Connection connection = connectionMock();
+        PullAllResponseHandler handler = newHandler( asList( "key1", "key2" ), connection );
+        handler.disableAutoReadManagement();
+
+        for ( int i = 0; i < PullAllResponseHandler.RECORD_BUFFER_HIGH_WATERMARK + 1; i++ )
+        {
+            handler.onRecord( values( 100, 200 ) );
+        }
+
+        verify( connection, never() ).disableAutoRead();
+    }
+
     @Test
     public void shouldPropagateFailureFromListAsync()
     {