Backpressure for reading records

Previously driver would read records from network even if client did
not consume them. This happened because event loop threads were always
in auto-read mode despite the amount of currently buffered records.
This could lead to memory problems when handling large results that
are slowly consumed.

This commit adds simple network-level backpressure for reading records.
Auto-read mode is turned off when amount of buffered records reaches
high watermark. It is turned on again when bunch of records are
consumed and buffer size reaches low watermark. Also added unit tests
for PULL_ALL and RUN handlers.

Author: lutovich

driver/src/main/java/org/neo4j/driver/internal/async/NettyConnection.java

@@ -52,7 +52,6 @@ public class NettyConnection implements Connection
     private final Clock clock;
 
     private final AtomicBoolean open = new AtomicBoolean( true );
-    private final AtomicBoolean autoReadEnabled = new AtomicBoolean( true );
 
     public NettyConnection( Channel channel, ChannelPool channelPool, Clock clock )
     {
@@ -73,8 +72,7 @@ public boolean isOpen()
     @Override
     public void enableAutoRead()
     {
-        assertOpen();
-        if ( autoReadEnabled.compareAndSet( false, true ) )
+        if ( isOpen() )
         {
             setAutoRead( true );
         }
@@ -83,8 +81,7 @@ public void enableAutoRead()
     @Override
     public void disableAutoRead()
     {
-        assertOpen();
-        if ( autoReadEnabled.compareAndSet( true, false ) )
+        if ( isOpen() )
         {
             setAutoRead( false );
         }
@@ -111,6 +108,9 @@ public CompletionStage<Void> release()
     {
         if ( open.compareAndSet( true, false ) )
         {
+            // auto-read could've been disabled, re-enable it to automatically receive response for RESET
+            setAutoRead( true );
+
             Promise<Void> releasePromise = channel.eventLoop().newPromise();
             reset( new ResetResponseHandler( channel, channelPool, messageDispatcher, clock, releasePromise ) );
             return asCompletionStage( releasePromise );
@@ -186,7 +186,7 @@ private void setAutoRead( boolean value )
 
     private void assertOpen()
     {
-        if ( !open.get() )
+        if ( !isOpen() )
         {
             throw new IllegalStateException( "Connection has been released to the pool and can't be reused" );
         }

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

@@ -18,7 +18,7 @@
  */
 package org.neo4j.driver.internal.handlers;
 
-import java.util.LinkedList;
+import java.util.ArrayDeque;
 import java.util.Map;
 import java.util.Queue;
 import java.util.concurrent.CompletableFuture;
@@ -38,16 +38,16 @@
 import static java.util.concurrent.CompletableFuture.completedFuture;
 import static org.neo4j.driver.internal.util.Futures.failedFuture;
 
-// todo: unit tests
 public abstract class PullAllResponseHandler implements ResponseHandler
 {
-    private static final boolean TOUCH_AUTO_READ = false;
+    static final int RECORD_BUFFER_LOW_WATERMARK = Integer.getInteger( "recordBufferLowWatermark", 300 );
+    static final int RECORD_BUFFER_HIGH_WATERMARK = Integer.getInteger( "recordBufferHighWatermark", 1000 );
 
     private final Statement statement;
     private final RunResponseHandler runResponseHandler;
     protected final Connection connection;
 
-    private final Queue<Record> records = new LinkedList<>();
+    private final Queue<Record> records = new ArrayDeque<>();
 
     private boolean finished;
     private Throwable failure;
@@ -199,25 +199,31 @@ else if ( finished )
     private void queueRecord( Record record )
     {
         records.add( record );
-        if ( TOUCH_AUTO_READ )
+
+        boolean shouldBufferAllRecords = summaryFuture != null || failureFuture != null;
+        // when summary or failure is requested we have to buffer all remaining records and then return summary/failure
+        // do not disable auto-read in this case, otherwise records will not be consumed and trailing
+        // SUCCESS or FAILURE message will not arrive as well, so callers will get stuck waiting for summary/failure
+        if ( !shouldBufferAllRecords && records.size() > RECORD_BUFFER_HIGH_WATERMARK )
         {
-            if ( records.size() > 10_000 )
-            {
-                connection.disableAutoRead();
-            }
+            // 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();
         }
     }
 
     private Record dequeueRecord()
     {
         Record record = records.poll();
-        if ( TOUCH_AUTO_READ )
+
+        if ( records.size() < RECORD_BUFFER_LOW_WATERMARK )
         {
-            if ( record != null && records.size() < 100 )
-            {
-                connection.enableAutoRead();
-            }
+            // 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();
         }
+
         return record;
     }
 

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

@@ -30,10 +30,12 @@
 
 public class RunResponseHandler implements ResponseHandler
 {
+    private static final int UNKNOWN = -1;
+
     private final CompletableFuture<Void> runCompletedFuture;
 
     private List<String> statementKeys = emptyList();
-    private long resultAvailableAfter;
+    private long resultAvailableAfter = UNKNOWN;
 
     public RunResponseHandler( CompletableFuture<Void> runCompletedFuture )
     {
@@ -107,6 +109,6 @@ private static long extractResultAvailableAfter( Map<String,Value> metadata )
         {
             return resultAvailableAfterValue.asLong();
         }
-        return -1;
+        return UNKNOWN;
     }
 }

driver/src/test/java/org/neo4j/driver/internal/async/NettyConnectionTest.java

@@ -121,25 +121,15 @@ public void shouldWriteForceReleaseInEventLoopThread() throws Exception
     }
 
     @Test
-    public void shouldNotEnableAutoReadWhenReleased()
+    public void shouldEnableAutoReadWhenReleased()
     {
         EmbeddedChannel channel = new EmbeddedChannel();
         channel.config().setAutoRead( false );
 
         NettyConnection connection = newConnection( channel );
 
         connection.release();
-
-        try
-        {
-            connection.enableAutoRead();
-            fail( "Exception expected" );
-        }
-        catch ( IllegalStateException e )
-        {
-            assertConnectionReleasedError( e );
-        }
-        assertFalse( channel.config().isAutoRead() );
+        assertTrue( channel.config().isAutoRead() );
     }
 
     @Test
@@ -151,16 +141,7 @@ public void shouldNotDisableAutoReadWhenReleased()
         NettyConnection connection = newConnection( channel );
 
         connection.release();
-
-        try
-        {
-            connection.disableAutoRead();
-            fail( "Exception expected" );
-        }
-        catch ( IllegalStateException e )
-        {
-            assertConnectionReleasedError( e );
-        }
+        connection.disableAutoRead(); // does nothing on released connection
         assertTrue( channel.config().isAutoRead() );
     }