Update UnmanagedTransaction handling of pending query pulls

Author: injectives

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

@@ -18,6 +18,7 @@
  */
 package org.neo4j.driver.internal.async;
 
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -82,7 +83,10 @@ public CompletionStage<ResultCursor> runAsync( Query query, TransactionConfig co
                 buildResultCursorFactory( query, config ).thenCompose( ResultCursorFactory::asyncResult );
 
         resultCursorStage = newResultCursorStage.exceptionally( error -> null );
-        return newResultCursorStage.thenApply( cursor -> cursor ); // convert the return type
+        return newResultCursorStage.thenCompose(
+                cursor -> cursor.runError()
+                                .map( Futures::<ResultCursor>failedFuture )
+                                .orElseGet( () -> CompletableFuture.completedFuture( cursor ) ) );
     }
 
     public CompletionStage<RxResultCursor> runRx(Query query, TransactionConfig config )

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

@@ -20,6 +20,7 @@
 
 import java.util.Arrays;
 import java.util.EnumSet;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.CompletionStage;
 import java.util.function.BiFunction;
@@ -211,7 +212,10 @@ public CompletionStage<ResultCursor> runAsync( Query query )
         CompletionStage<AsyncResultCursor> cursorStage =
                 protocol.runInUnmanagedTransaction( connection, query, this, fetchSize ).asyncResult();
         resultCursors.add( cursorStage );
-        return cursorStage.thenApply( cursor -> cursor );
+        return cursorStage.thenCompose(
+                cursor -> cursor.runError()
+                                .map( Futures::<ResultCursor>failedFuture )
+                                .orElseGet( () -> CompletableFuture.completedFuture( cursor ) ) );
     }
 
     public CompletionStage<RxResultCursor> runRx(Query query)

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

@@ -18,9 +18,12 @@
  */
 package org.neo4j.driver.internal.cursor;
 
-import org.neo4j.driver.internal.FailableCursor;
+import java.util.Optional;
+
 import org.neo4j.driver.async.ResultCursor;
+import org.neo4j.driver.internal.FailableCursor;
 
 public interface AsyncResultCursor extends ResultCursor, FailableCursor
 {
+    Optional<Throwable> runError();
 }

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

@@ -19,6 +19,7 @@
 package org.neo4j.driver.internal.cursor;
 
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.function.Consumer;
@@ -33,11 +34,13 @@
 
 public class AsyncResultCursorImpl implements AsyncResultCursor
 {
+    private final Throwable runError;
     private final RunResponseHandler runHandler;
     private final PullAllResponseHandler pullAllHandler;
 
-    public AsyncResultCursorImpl(RunResponseHandler runHandler, PullAllResponseHandler pullAllHandler )
+    public AsyncResultCursorImpl( Throwable runError, RunResponseHandler runHandler, PullAllResponseHandler pullAllHandler )
     {
+        this.runError = runError;
         this.runHandler = runHandler;
         this.pullAllHandler = pullAllHandler;
     }
@@ -113,13 +116,38 @@ public <T> CompletionStage<List<T>> listAsync( Function<Record,T> mapFunction )
     @Override
     public CompletionStage<Throwable> discardAllFailureAsync()
     {
-        return consumeAsync().handle( ( summary, error ) -> error );
+        return consumeAsync().handle(
+                ( summary, consumeError ) ->
+                {
+                    if ( runError != null )
+                    {
+                        if ( consumeError != null && runError != consumeError )
+                        {
+                            runError.addSuppressed( consumeError );
+                        }
+                        return runError;
+                    }
+                    return consumeError;
+                } );
     }
 
     @Override
     public CompletionStage<Throwable> pullAllFailureAsync()
     {
-        return pullAllHandler.pullAllFailureAsync();
+        return pullAllHandler.pullAllFailureAsync().handle(
+                ( ignored, pullAllError ) ->
+                {
+                    if ( runError != null )
+                    {
+                        if ( pullAllError != null && runError != pullAllError )
+                        {
+                            runError.addSuppressed( pullAllError );
+                        }
+                        return runError;
+                    }
+                    return pullAllError;
+                }
+        );
     }
 
     private void internalForEachAsync( Consumer<Record> action, CompletableFuture<Void> resultFuture )
@@ -154,4 +182,10 @@ else if ( record != null )
             }
         } );
     }
+
+    @Override
+    public Optional<Throwable> runError()
+    {
+        return Optional.ofNullable( runError );
+    }
 }

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

@@ -64,7 +64,7 @@ public CompletionStage<AsyncResultCursor> asyncResult()
         connection.write( runMessage, runHandler ); // queues the run message, will be flushed with pull message together
         pullAllHandler.prePopulateRecords();
 
-        return runFuture.thenApply( ignore -> new DisposableAsyncResultCursor( new AsyncResultCursorImpl( runHandler, pullAllHandler ) ) );
+        return runFuture.handle( ( ignored, error ) -> new DisposableAsyncResultCursor( new AsyncResultCursorImpl( error, runHandler, pullAllHandler ) ) );
     }
 
     public CompletionStage<RxResultCursor> rxResult()

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

@@ -19,6 +19,7 @@
 package org.neo4j.driver.internal.cursor;
 
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.function.Consumer;
@@ -118,4 +119,10 @@ boolean isDisposed()
     {
         return this.isDisposed;
     }
+
+    @Override
+    public Optional<Throwable> runError()
+    {
+        return this.delegate.runError();
+    }
 }

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

@@ -66,7 +66,7 @@ public CompletionStage<AsyncResultCursor> asyncResult()
         // only write and flush messages when async result is wanted.
         connection.write( runMessage, runHandler ); // queues the run message, will be flushed with pull message together
         pullAllHandler.prePopulateRecords();
-        return runFuture.thenApply( ignore -> new DisposableAsyncResultCursor( new AsyncResultCursorImpl( runHandler, pullAllHandler ) ) );
+        return runFuture.handle( ( ignored, error ) -> new DisposableAsyncResultCursor( new AsyncResultCursorImpl( error, runHandler, pullAllHandler ) ) );
     }
 
     @Override

driver/src/test/java/org/neo4j/driver/integration/RoutingDriverBoltKitIT.java

@@ -31,13 +31,15 @@
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.logging.Level;
 
 import org.neo4j.driver.AccessMode;
 import org.neo4j.driver.AuthToken;
 import org.neo4j.driver.AuthTokens;
 import org.neo4j.driver.Config;
 import org.neo4j.driver.Driver;
 import org.neo4j.driver.GraphDatabase;
+import org.neo4j.driver.Logging;
 import org.neo4j.driver.Record;
 import org.neo4j.driver.Session;
 import org.neo4j.driver.TransactionWork;
@@ -104,7 +106,8 @@ void shouldHandleLeaderSwitchAndRetryWhenWritingInTxFunctionAsync() throws IOExc
         StubServer writeServer = stubController.startStub( "not_able_to_write_server_tx_func_retries.script", 9007 );
         URI uri = URI.create( "neo4j://127.0.0.1:9001" );
 
-        Driver driver = GraphDatabase.driver( uri, Config.builder().withMaxTransactionRetryTime( 1, TimeUnit.MILLISECONDS ).build() );
+        Driver driver = GraphDatabase
+                .driver( uri, Config.builder().withLogging( Logging.console( Level.FINE ) ).withMaxTransactionRetryTime( 1, TimeUnit.MILLISECONDS ).build() );
         AsyncSession session = driver.asyncSession( builder().withDatabase( "mydatabase" ).build() );
         List<String> names = Futures.blockingGet( session.writeTransactionAsync(
                 tx -> tx.runAsync( "RETURN 1" )

driver/src/test/java/org/neo4j/driver/internal/InternalResultTest.java

@@ -373,7 +373,7 @@ private Result createResult(int numberOfRecords )
         }
         pullAllHandler.onSuccess( emptyMap() );
 
-        AsyncResultCursor cursor = new AsyncResultCursorImpl( runHandler, pullAllHandler );
+        AsyncResultCursor cursor = new AsyncResultCursorImpl( null, runHandler, pullAllHandler );
         return new InternalResult( connection, new DisposableAsyncResultCursor( cursor ) );
     }
 

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

@@ -404,12 +404,12 @@ void shouldPropagateFailureInConsumeAsync()
 
     private static AsyncResultCursorImpl newCursor(PullAllResponseHandler pullAllHandler )
     {
-        return new AsyncResultCursorImpl( newRunResponseHandler(), pullAllHandler );
+        return new AsyncResultCursorImpl( null, newRunResponseHandler(), pullAllHandler );
     }
 
     private static AsyncResultCursorImpl newCursor(RunResponseHandler runHandler, PullAllResponseHandler pullAllHandler )
     {
-        return new AsyncResultCursorImpl( runHandler, pullAllHandler );
+        return new AsyncResultCursorImpl( null, runHandler, pullAllHandler );
     }
 
     private static RunResponseHandler newRunResponseHandler()