Call close with the appropriate flag to commit or rollback on UnmanagedTransaction where possible to avoid double state acquisition

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

@@ -146,7 +146,7 @@ private <T> void executeWork(CompletableFuture<T> resultFuture, UnmanagedTransac
             Throwable error = Futures.completionExceptionCause( completionError );
             if ( error != null )
             {
-                rollbackTxAfterFailedTransactionWork( tx, resultFuture, error );
+                closeTxAfterFailedTransactionWork( tx, resultFuture, error );
             }
             else
             {
@@ -174,43 +174,33 @@ private <T> CompletionStage<T> safeExecuteWork(UnmanagedTransaction tx, AsyncTra
         }
     }
 
-    private <T> void rollbackTxAfterFailedTransactionWork(UnmanagedTransaction tx, CompletableFuture<T> resultFuture, Throwable error )
+    private <T> void closeTxAfterFailedTransactionWork( UnmanagedTransaction tx, CompletableFuture<T> resultFuture, Throwable error )
     {
-        if ( tx.isOpen() )
-        {
-            tx.rollbackAsync().whenComplete( ( ignore, rollbackError ) -> {
-                if ( rollbackError != null )
+        tx.closeAsync().whenComplete(
+                ( ignored, rollbackError ) ->
                 {
-                    error.addSuppressed( rollbackError );
-                }
-                resultFuture.completeExceptionally( error );
-            } );
-        }
-        else
-        {
-            resultFuture.completeExceptionally( error );
-        }
+                    if ( rollbackError != null )
+                    {
+                        error.addSuppressed( rollbackError );
+                    }
+                    resultFuture.completeExceptionally( error );
+                } );
     }
 
     private <T> void closeTxAfterSucceededTransactionWork(UnmanagedTransaction tx, CompletableFuture<T> resultFuture, T result )
     {
-        if ( tx.isOpen() )
-        {
-            tx.commitAsync().whenComplete( ( ignore, completionError ) -> {
-                Throwable commitError = Futures.completionExceptionCause( completionError );
-                if ( commitError != null )
+        tx.closeAsync( true ).whenComplete(
+                ( ignored, completionError ) ->
                 {
-                    resultFuture.completeExceptionally( commitError );
-                }
-                else
-                {
-                    resultFuture.complete( result );
-                }
-            } );
-        }
-        else
-        {
-            resultFuture.complete( result );
-        }
+                    Throwable commitError = Futures.completionExceptionCause( completionError );
+                    if ( commitError != null )
+                    {
+                        resultFuture.completeExceptionally( commitError );
+                    }
+                    else
+                    {
+                        resultFuture.complete( result );
+                    }
+                } );
     }
 }

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

@@ -134,7 +134,12 @@ else if ( beginError instanceof ConnectionReadTimeoutException )
 
     public CompletionStage<Void> closeAsync()
     {
-        return closeAsync( false, true );
+        return closeAsync( false );
+    }
+
+    public CompletionStage<Void> closeAsync( boolean commit )
+    {
+        return closeAsync( commit, true );
     }
 
     public CompletionStage<Void> commitAsync()

driver/src/main/java/org/neo4j/driver/internal/reactive/InternalRxSession.java

@@ -130,7 +130,7 @@ public <T> Publisher<T> writeTransaction( RxTransactionWork<? extends Publisher<
     private <T> Publisher<T> runTransaction( AccessMode mode, RxTransactionWork<? extends Publisher<T>> work, TransactionConfig config )
     {
         Flux<T> repeatableWork = Flux.usingWhen( beginTransaction( mode, config ), work::execute,
-                                                 InternalRxTransaction::commitIfOpen, ( tx, error ) -> tx.close(), InternalRxTransaction::close );
+                                                 tx -> tx.close( true ), ( tx, error ) -> tx.close(), InternalRxTransaction::close );
         return session.retryLogic().retryRx( repeatableWork );
     }
 

driver/src/main/java/org/neo4j/driver/internal/reactive/InternalRxTransaction.java

@@ -30,7 +30,6 @@
 import org.neo4j.driver.reactive.RxTransaction;
 
 import static org.neo4j.driver.internal.reactive.RxUtils.createEmptyPublisher;
-import static org.neo4j.driver.internal.util.Futures.completedWithNull;
 
 public class InternalRxTransaction extends AbstractRxQueryRunner implements RxTransaction
 {
@@ -77,13 +76,13 @@ public <T> Publisher<T> rollback()
         return createEmptyPublisher( tx::rollbackAsync );
     }
 
-    Publisher<Void> commitIfOpen()
+    Publisher<Void> close()
     {
-        return createEmptyPublisher( () -> tx.isOpen() ? tx.commitAsync() : completedWithNull() );
+        return close( false );
     }
 
-    Publisher<Void> close()
+    Publisher<Void> close( boolean commit )
     {
-        return createEmptyPublisher( tx::closeAsync );
+        return createEmptyPublisher( () -> tx.closeAsync( commit ) );
     }
 }

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

@@ -413,15 +413,21 @@ void shouldReturnFailingStageOnConflictingCompletingAction( boolean protocolComm
     private static Stream<Arguments> closingNotActionTransactionArgs()
     {
         return Stream.of(
-                Arguments.of( true, 1, "commit" ),
-                Arguments.of( false, 1, "rollback" ),
-                Arguments.of( false, 0, "terminate" )
+                Arguments.of( true, 1, "commit", null ),
+                Arguments.of( false, 1, "rollback", null ),
+                Arguments.of( false, 0, "terminate", null ),
+                Arguments.of( true, 1, "commit", true ),
+                Arguments.of( false, 1, "rollback", true ),
+                Arguments.of( true, 1, "commit", false ),
+                Arguments.of( false, 1, "rollback", false ),
+                Arguments.of( false, 0, "terminate", false )
         );
     }
 
     @ParameterizedTest
     @MethodSource( "closingNotActionTransactionArgs" )
-    void shouldReturnCompletedWithNullStageOnClosingNotActiveTransaction( boolean protocolCommit, int expectedProtocolInvocations, String originalAction )
+    void shouldReturnCompletedWithNullStageOnClosingInactiveTransactionExceptCommittingAborted(
+            boolean protocolCommit, int expectedProtocolInvocations, String originalAction, Boolean commitOnClose )
     {
         Connection connection = mock( Connection.class );
         BoltProtocol protocol = mock( BoltProtocol.class );
@@ -431,7 +437,7 @@ void shouldReturnCompletedWithNullStageOnClosingNotActiveTransaction( boolean pr
         UnmanagedTransaction tx = new UnmanagedTransaction( connection, new DefaultBookmarkHolder(), UNLIMITED_FETCH_SIZE );
 
         CompletionStage<Void> originalActionStage = mapTransactionAction( originalAction, tx ).get();
-        CompletionStage<Void> closeStage = tx.closeAsync();
+        CompletionStage<Void> closeStage = commitOnClose != null ? tx.closeAsync( commitOnClose ) : tx.closeAsync();
 
         assertTrue( originalActionStage.toCompletableFuture().isDone() );
         assertFalse( originalActionStage.toCompletableFuture().isCompletedExceptionally() );

driver/src/test/java/org/neo4j/driver/internal/reactive/InternalRxSessionTest.java

@@ -199,9 +199,7 @@ void shouldDelegateRunTx( Function<RxSession,Publisher<String>> runTx ) throws T
         // Given
         NetworkSession session = mock( NetworkSession.class );
         UnmanagedTransaction tx = mock( UnmanagedTransaction.class );
-        when( tx.isOpen() ).thenReturn( true );
-        when( tx.commitAsync() ).thenReturn( completedWithNull() );
-        when( tx.rollbackAsync() ).thenReturn( completedWithNull() );
+        when( tx.closeAsync( true ) ).thenReturn( completedWithNull() );
 
         when( session.beginTransactionAsync( any( AccessMode.class ), any( TransactionConfig.class ) ) ).thenReturn( completedFuture( tx ) );
         when( session.retryLogic() ).thenReturn( new FixedRetryLogic( 1 ) );
@@ -213,7 +211,7 @@ void shouldDelegateRunTx( Function<RxSession,Publisher<String>> runTx ) throws T
 
         // Then
         verify( session ).beginTransactionAsync( any( AccessMode.class ), any( TransactionConfig.class ) );
-        verify( tx ).commitAsync();
+        verify( tx ).closeAsync( true );
     }
 
     @Test
@@ -223,25 +221,24 @@ void shouldRetryOnError() throws Throwable
         int retryCount = 2;
         NetworkSession session = mock( NetworkSession.class );
         UnmanagedTransaction tx = mock( UnmanagedTransaction.class );
-        when( tx.isOpen() ).thenReturn( true );
-        when( tx.commitAsync() ).thenReturn( completedWithNull() );
-        when( tx.rollbackAsync() ).thenReturn( completedWithNull() );
+        when( tx.closeAsync( false ) ).thenReturn( completedWithNull() );
 
         when( session.beginTransactionAsync( any( AccessMode.class ), any( TransactionConfig.class ) ) ).thenReturn( completedFuture( tx ) );
         when( session.retryLogic() ).thenReturn( new FixedRetryLogic( retryCount ) );
         InternalRxSession rxSession = new InternalRxSession( session );
 
         // When
-        Publisher<String> strings = rxSession.readTransaction( t ->
-                Flux.just( "a" ).then( Mono.error( new RuntimeException( "Errored" ) ) ) );
+        Publisher<String> strings = rxSession.readTransaction(
+                t ->
+                        Flux.just( "a" ).then( Mono.error( new RuntimeException( "Errored" ) ) ) );
         StepVerifier.create( Flux.from( strings ) )
-                // we lost the "a"s too as the user only see the last failure
-                .expectError( RuntimeException.class )
-                .verify();
+                    // we lost the "a"s too as the user only see the last failure
+                    .expectError( RuntimeException.class )
+                    .verify();
 
         // Then
         verify( session, times( retryCount + 1 ) ).beginTransactionAsync( any( AccessMode.class ), any( TransactionConfig.class ) );
-        verify( tx, times( retryCount + 1 ) ).closeAsync();
+        verify( tx, times( retryCount + 1 ) ).closeAsync( false );
     }
 
     @Test
@@ -251,33 +248,34 @@ void shouldObtainResultIfRetrySucceed() throws Throwable
         int retryCount = 2;
         NetworkSession session = mock( NetworkSession.class );
         UnmanagedTransaction tx = mock( UnmanagedTransaction.class );
-        when( tx.isOpen() ).thenReturn( true );
-        when( tx.commitAsync() ).thenReturn( completedWithNull() );
-        when( tx.rollbackAsync() ).thenReturn( completedWithNull() );
+        when( tx.closeAsync( false ) ).thenReturn( completedWithNull() );
+        when( tx.closeAsync( true ) ).thenReturn( completedWithNull() );
 
         when( session.beginTransactionAsync( any( AccessMode.class ), any( TransactionConfig.class ) ) ).thenReturn( completedFuture( tx ) );
         when( session.retryLogic() ).thenReturn( new FixedRetryLogic( retryCount ) );
         InternalRxSession rxSession = new InternalRxSession( session );
 
         // When
         AtomicInteger count = new AtomicInteger();
-        Publisher<String> strings = rxSession.readTransaction( t -> {
-            // we fail for the first few retries, and then success on the last run.
-            if ( count.getAndIncrement() == retryCount )
-            {
-                return Flux.just( "a" );
-            }
-            else
-            {
-                return Flux.just( "a" ).then( Mono.error( new RuntimeException( "Errored" ) ) );
-            }
-        } );
+        Publisher<String> strings = rxSession.readTransaction(
+                t ->
+                {
+                    // we fail for the first few retries, and then success on the last run.
+                    if ( count.getAndIncrement() == retryCount )
+                    {
+                        return Flux.just( "a" );
+                    }
+                    else
+                    {
+                        return Flux.just( "a" ).then( Mono.error( new RuntimeException( "Errored" ) ) );
+                    }
+                } );
         StepVerifier.create( Flux.from( strings ) ).expectNext( "a" ).verifyComplete();
 
         // Then
         verify( session, times( retryCount + 1 ) ).beginTransactionAsync( any( AccessMode.class ), any( TransactionConfig.class ) );
-        verify( tx, times( retryCount ) ).closeAsync();
-        verify( tx ).commitAsync();
+        verify( tx, times( retryCount ) ).closeAsync( false );
+        verify( tx ).closeAsync( true );
     }
 
     @Test

driver/src/test/java/org/neo4j/driver/internal/reactive/InternalRxTransactionTest.java

@@ -48,7 +48,6 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 import static org.neo4j.driver.Values.parameters;
@@ -140,43 +139,28 @@ void shouldMarkTxIfFailedToRun( Function<RxTransaction, RxResult> runReturnOne )
     }
 
     @Test
-    void shouldCommitWhenOpen()
+    void shouldDelegateConditionalClose()
     {
         UnmanagedTransaction tx = mock( UnmanagedTransaction.class );
-        when( tx.isOpen() ).thenReturn( true );
-        when( tx.commitAsync() ).thenReturn( Futures.completedWithNull() );
-
-        InternalRxTransaction rxTx = new InternalRxTransaction( tx );
-        Publisher<Void> publisher = rxTx.commitIfOpen();
-        StepVerifier.create( publisher ).verifyComplete();
-
-        verify( tx ).commitAsync();
-    }
-
-    @Test
-    void shouldNotCommitWhenNotOpen()
-    {
-        UnmanagedTransaction tx = mock( UnmanagedTransaction.class );
-        when( tx.isOpen() ).thenReturn( false );
-        when( tx.commitAsync() ).thenReturn( Futures.completedWithNull() );
+        when( tx.closeAsync( true ) ).thenReturn( Futures.completedWithNull() );
 
         InternalRxTransaction rxTx = new InternalRxTransaction( tx );
-        Publisher<Void> publisher = rxTx.commitIfOpen();
+        Publisher<Void> publisher = rxTx.close( true );
         StepVerifier.create( publisher ).verifyComplete();
 
-        verify( tx, never() ).commitAsync();
+        verify( tx ).closeAsync( true );
     }
 
     @Test
     void shouldDelegateClose()
     {
         UnmanagedTransaction tx = mock( UnmanagedTransaction.class );
-        when( tx.closeAsync() ).thenReturn( Futures.completedWithNull() );
+        when( tx.closeAsync( false ) ).thenReturn( Futures.completedWithNull() );
 
         InternalRxTransaction rxTx = new InternalRxTransaction( tx );
         Publisher<Void> publisher = rxTx.close();
         StepVerifier.create( publisher ).verifyComplete();
 
-        verify( tx ).closeAsync();
+        verify( tx ).closeAsync( false );
     }
 }