Back pressure support for async and blocking API with BOLT V4.

Author: Zhen Li

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

@@ -34,19 +34,19 @@
 
 public class AsyncStatementResultCursor implements InternalStatementResultCursor
 {
-    private final RunResponseHandler runResponseHandler;
+    private final RunResponseHandler runHandler;
     private final PullAllResponseHandler pullAllHandler;
 
-    public AsyncStatementResultCursor( RunResponseHandler runResponseHandler, PullAllResponseHandler pullAllHandler )
+    public AsyncStatementResultCursor( RunResponseHandler runHandler, PullAllResponseHandler pullAllHandler )
     {
-        this.runResponseHandler = runResponseHandler;
+        this.runHandler = runHandler;
         this.pullAllHandler = pullAllHandler;
     }
 
     @Override
     public List<String> keys()
     {
-        return runResponseHandler.statementKeys();
+        return runHandler.statementKeys();
     }
 
     @Override
@@ -94,7 +94,7 @@ public CompletionStage<Record> singleAsync()
     @Override
     public CompletionStage<ResultSummary> consumeAsync()
     {
-        return pullAllHandler.consumeAsync();
+        return pullAllHandler.summaryAsync();
     }
 
     @Override
@@ -120,7 +120,7 @@ public <T> CompletionStage<List<T>> listAsync( Function<Record,T> mapFunction )
     @Override
     public CompletionStage<Throwable> failureAsync()
     {
-        return pullAllHandler.failureAsync();
+        return pullAllHandler.summaryAsync().handle( ( summary, error ) -> error );
     }
 
     private void internalForEachAsync( Consumer<Record> action, CompletableFuture<Void> resultFuture )

driver/src/main/java/org/neo4j/driver/internal/cursor/InternalStatementResultCursorFactory.java

@@ -23,9 +23,8 @@
 import org.neo4j.driver.internal.async.AsyncStatementResultCursor;
 import org.neo4j.driver.internal.handlers.PullAllResponseHandler;
 import org.neo4j.driver.internal.handlers.RunResponseHandler;
-import org.neo4j.driver.internal.handlers.pulln.BasicPullResponseHandler;
+import org.neo4j.driver.internal.handlers.pulln.PullResponseHandler;
 import org.neo4j.driver.internal.messaging.Message;
-import org.neo4j.driver.internal.messaging.request.PullMessage;
 import org.neo4j.driver.internal.spi.Connection;
 
 import static java.util.Objects.requireNonNull;
@@ -36,12 +35,12 @@ public class InternalStatementResultCursorFactory implements StatementResultCurs
     private final RunResponseHandler runHandler;
     private final Connection connection;
 
-    private final BasicPullResponseHandler pullHandler;
+    private final PullResponseHandler pullHandler;
     private final PullAllResponseHandler pullAllHandler;
     private final boolean waitForRunResponse;
     private final Message runMessage;
 
-    public InternalStatementResultCursorFactory( Connection connection, Message runMessage, RunResponseHandler runHandler, BasicPullResponseHandler pullHandler,
+    public InternalStatementResultCursorFactory( Connection connection, Message runMessage, RunResponseHandler runHandler, PullResponseHandler pullHandler,
             PullAllResponseHandler pullAllHandler, boolean waitForRunResponse )
     {
         requireNonNull( connection );
@@ -62,7 +61,8 @@ public InternalStatementResultCursorFactory( Connection connection, Message runM
     public CompletionStage<InternalStatementResultCursor> asyncResult()
     {
         // only write and flush messages when async result is wanted.
-        connection.writeAndFlush( runMessage, runHandler, PullMessage.PULL_ALL, pullAllHandler );
+        connection.writeAndFlush( runMessage, runHandler );
+        pullAllHandler.prePull();
 
         if ( waitForRunResponse )
         {

driver/src/main/java/org/neo4j/driver/internal/cursor/RxStatementResultCursor.java

@@ -28,26 +28,26 @@
 
 import org.neo4j.driver.internal.FailableCursor;
 import org.neo4j.driver.internal.handlers.RunResponseHandler;
-import org.neo4j.driver.internal.handlers.pulln.BasicPullResponseHandler;
+import org.neo4j.driver.internal.handlers.pulln.PullResponseHandler;
 import org.neo4j.driver.Record;
 import org.neo4j.driver.summary.ResultSummary;
 
-import static org.neo4j.driver.internal.handlers.pulln.AbstractBasicPullResponseHandler.DISCARD_RECORD_CONSUMER;
+import static org.neo4j.driver.internal.handlers.pulln.BasicPullResponseHandler.DISCARD_RECORD_CONSUMER;
 
 public class RxStatementResultCursor implements Subscription, FailableCursor
 {
     private final RunResponseHandler runHandler;
-    private final BasicPullResponseHandler pullHandler;
+    private final PullResponseHandler pullHandler;
     private final Throwable runResponseError;
     private final CompletableFuture<ResultSummary> summaryFuture = new CompletableFuture<>();
-    private boolean isRecordHandlerInstalled = false;
+    private BiConsumer<Record,Throwable> recordConsumer;
 
-    public RxStatementResultCursor( RunResponseHandler runHandler, BasicPullResponseHandler pullHandler )
+    public RxStatementResultCursor( RunResponseHandler runHandler, PullResponseHandler pullHandler )
     {
         this( null, runHandler, pullHandler );
     }
 
-    public RxStatementResultCursor( Throwable runError, RunResponseHandler runHandler, BasicPullResponseHandler pullHandler )
+    public RxStatementResultCursor( Throwable runError, RunResponseHandler runHandler, PullResponseHandler pullHandler )
     {
         Objects.requireNonNull( runHandler );
         Objects.requireNonNull( pullHandler );
@@ -66,15 +66,20 @@ public List<String> keys()
 
     public void installRecordConsumer( BiConsumer<Record,Throwable> recordConsumer )
     {
-        if ( isRecordHandlerInstalled )
+        if ( isRecordConsumerInstalled() )
         {
             return;
         }
-        isRecordHandlerInstalled = true;
-        pullHandler.installRecordConsumer( recordConsumer );
+        this.recordConsumer = recordConsumer;
+        pullHandler.installRecordConsumer( this.recordConsumer );
         assertRunCompletedSuccessfully();
     }
 
+    private boolean isRecordConsumerInstalled()
+    {
+        return this.recordConsumer != null;
+    }
+
     public void request( long n )
     {
         pullHandler.request( n );
@@ -120,8 +125,10 @@ private void assertRunCompletedSuccessfully()
     private void installSummaryConsumer()
     {
         pullHandler.installSummaryConsumer( ( summary, error ) -> {
-            if ( error != null )
+            if ( error != null && recordConsumer == DISCARD_RECORD_CONSUMER )
             {
+                // We will only report the error to summary if there is no user record consumer installed
+                // When a user record consumer is installed, the error will be reported to record consumer instead.
                 summaryFuture.completeExceptionally( error );
             }
             else if ( summary != null )