Skip to content

Commit

Permalink
Call close with the appropriate flag to commit or rollback on Unmanag…
Browse files Browse the repository at this point in the history 
…edTransaction where possible to avoid double state acquisition (#1065)

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

Calling `close` instead of separate `isOpen` and `commitAsync` requires less lock acquisitions and is safer.

* Update tests
  • Loading branch information
@injectives
injectives authored Nov 10, 2021
1 parent 8929bd4 commit bef5547
Show file tree
Hide file tree
Showing 7 changed files with 77 additions and 95 deletions.
54 changes: 22 additions & 32 deletions driver/src/main/java/org/neo4j/driver/internal/async/InternalAsyncSession.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
{
Expand Down Expand Up @@ -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 );
}
} );
}
}
7 changes: 6 additions & 1 deletion driver/src/main/java/org/neo4j/driver/internal/async/UnmanagedTransaction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
2 changes: 1 addition & 1 deletion driver/src/main/java/org/neo4j/driver/internal/reactive/InternalRxSession.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 );
}

Expand Down
9 changes: 4 additions & 5 deletions driver/src/main/java/org/neo4j/driver/internal/reactive/InternalRxTransaction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
{
Expand Down Expand Up @@ -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 ) );
}
}
16 changes: 11 additions & 5 deletions driver/src/test/java/org/neo4j/driver/internal/async/UnmanagedTransactionTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 );
Expand All @@ -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() );
Expand Down
56 changes: 27 additions & 29 deletions driver/src/test/java/org/neo4j/driver/internal/reactive/InternalRxSessionTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 ) );
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
28 changes: 6 additions & 22 deletions driver/src/test/java/org/neo4j/driver/internal/reactive/InternalRxTransactionTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 );
}
}

0 comments on commit bef5547

Please sign in to comment.