Ensure reactive transaction function gets rolled back on cancellation

This update also ensures that session cancellation does not result in multiple rollback attempts.

Author: injectives

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

@@ -22,6 +22,8 @@
 import java.util.EnumSet;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.CompletionStage;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.BiFunction;
 
 import org.neo4j.driver.Bookmark;
@@ -41,6 +43,7 @@
 
 import static org.neo4j.driver.internal.util.Futures.completedWithNull;
 import static org.neo4j.driver.internal.util.Futures.failedFuture;
+import static org.neo4j.driver.internal.util.LockUtil.executeWithLock;
 
 public class UnmanagedTransaction
 {
@@ -62,66 +65,18 @@ private enum State
         ROLLED_BACK
     }
 
-    /**
-     * This is a holder so that we can have ony the state volatile in the tx without having to synchronize the whole block.
-     */
-    private static final class StateHolder
-    {
-        private static final EnumSet<State> OPEN_STATES = EnumSet.of( State.ACTIVE, State.TERMINATED );
-        private static final StateHolder ACTIVE_HOLDER = new StateHolder( State.ACTIVE, null );
-        private static final StateHolder COMMITTED_HOLDER = new StateHolder( State.COMMITTED, null );
-        private static final StateHolder ROLLED_BACK_HOLDER = new StateHolder( State.ROLLED_BACK, null );
-
-        /**
-         * The actual state.
-         */
-        final State value;
-
-        /**
-         * If this holder contains a state of {@link State#TERMINATED}, this represents the cause if any.
-         */
-        final Throwable causeOfTermination;
-
-        static StateHolder of( State value )
-        {
-            switch ( value )
-            {
-            case ACTIVE:
-                return ACTIVE_HOLDER;
-            case COMMITTED:
-                return COMMITTED_HOLDER;
-            case ROLLED_BACK:
-                return ROLLED_BACK_HOLDER;
-            case TERMINATED:
-            default:
-                throw new IllegalArgumentException( "Cannot provide a default state holder for state " + value );
-            }
-        }
-
-        static StateHolder terminatedWith( Throwable cause )
-        {
-            return new StateHolder( State.TERMINATED, cause );
-        }
-
-        private StateHolder( State value, Throwable causeOfTermination )
-        {
-            this.value = value;
-            this.causeOfTermination = causeOfTermination;
-        }
-
-        boolean isOpen()
-        {
-            return OPEN_STATES.contains( this.value );
-        }
-    }
+    private static final EnumSet<State> OPEN_STATES = EnumSet.of( State.ACTIVE, State.TERMINATED );
 
     private final Connection connection;
     private final BoltProtocol protocol;
     private final BookmarkHolder bookmarkHolder;
     private final ResultCursorsHolder resultCursors;
     private final long fetchSize;
-
-    private volatile StateHolder state = StateHolder.of( State.ACTIVE );
+    private final Lock lock = new ReentrantLock();
+    private State state = State.ACTIVE;
+    private CompletionStage<Void> commitStage;
+    private CompletionStage<Void> rollbackStage;
+    private Throwable causeOfTermination;
 
     public UnmanagedTransaction( Connection connection, BookmarkHolder bookmarkHolder, long fetchSize )
     {
@@ -164,50 +119,61 @@ else if ( beginError instanceof ConnectionReadTimeoutException )
 
     public CompletionStage<Void> closeAsync()
     {
-        if ( isOpen() )
-        {
-            return rollbackAsync();
-        }
-        else
-        {
-            return completedWithNull();
-        }
+        return executeWithLock( lock, () -> isOpen() ? rollbackAsync() : completedWithNull() );
     }
 
     public CompletionStage<Void> commitAsync()
     {
-        if ( state.value == State.COMMITTED )
-        {
-            return failedFuture( new ClientException( "Can't commit, transaction has been committed" ) );
-        }
-        else if ( state.value == State.ROLLED_BACK )
+        return executeWithLock( lock, () ->
         {
-            return failedFuture( new ClientException( "Can't commit, transaction has been rolled back" ) );
-        }
-        else
-        {
-            return resultCursors.retrieveNotConsumedError()
-                                .thenCompose( error -> doCommitAsync( error ).handle( handleCommitOrRollback( error ) ) )
-                                .whenComplete( ( ignore, error ) -> handleTransactionCompletion( true, error ) );
-        }
+            if ( state == State.COMMITTED )
+            {
+                return failedFuture( new ClientException( "Can't commit, transaction has been committed" ) );
+            }
+            else if ( state == State.ROLLED_BACK )
+            {
+                return failedFuture( new ClientException( "Can't commit, transaction has been rolled back" ) );
+            }
+            else if ( commitStage != null )
+            {
+                return commitStage;
+            }
+            else
+            {
+                CompletionStage<Void> stage = resultCursors.retrieveNotConsumedError()
+                                                           .thenCompose( error -> doCommitAsync( error ).handle( handleCommitOrRollback( error ) ) );
+                commitStage = stage.whenComplete( ( ignore, error ) -> releaseConnection( error ) );
+                stage.whenComplete( ( ignored, error ) -> updateStateAfterCommitOrRollback( true, error ) );
+                return stage;
+            }
+        } );
     }
 
     public CompletionStage<Void> rollbackAsync()
     {
-        if ( state.value == State.COMMITTED )
+        return executeWithLock( lock, () ->
         {
-            return failedFuture( new ClientException( "Can't rollback, transaction has been committed" ) );
-        }
-        else if ( state.value == State.ROLLED_BACK )
-        {
-            return failedFuture( new ClientException( "Can't rollback, transaction has been rolled back" ) );
-        }
-        else
-        {
-            return resultCursors.retrieveNotConsumedError()
-                                .thenCompose( error -> doRollbackAsync().handle( handleCommitOrRollback( error ) ) )
-                                .whenComplete( ( ignore, error ) -> handleTransactionCompletion( false, error ) );
-        }
+            if ( state == State.COMMITTED )
+            {
+                return failedFuture( new ClientException( "Can't rollback, transaction has been committed" ) );
+            }
+            else if ( state == State.ROLLED_BACK )
+            {
+                return failedFuture( new ClientException( "Can't rollback, transaction has been rolled back" ) );
+            }
+            else if ( rollbackStage != null )
+            {
+                return rollbackStage;
+            }
+            else
+            {
+                CompletionStage<Void> stage = resultCursors.retrieveNotConsumedError()
+                                                           .thenCompose( error -> doRollbackAsync().handle( handleCommitOrRollback( error ) ) );
+                rollbackStage = stage.whenComplete( ( ignore, error ) -> releaseConnection( error ) );
+                stage.whenComplete( ( ignored, error ) -> updateStateAfterCommitOrRollback( false, error ) );
+                return stage;
+            }
+        } );
     }
 
     public CompletionStage<ResultCursor> runAsync( Query query )
@@ -230,22 +196,27 @@ public CompletionStage<RxResultCursor> runRx(Query query)
 
     public boolean isOpen()
     {
-        return state.isOpen();
+        State currentState = executeWithLock( lock, () -> state );
+        return OPEN_STATES.contains( currentState );
     }
 
     public void markTerminated( Throwable cause )
     {
-        if ( state.value == State.TERMINATED )
+        executeWithLock( lock, () ->
         {
-            if ( state.causeOfTermination != null )
+            if ( state == State.TERMINATED )
             {
-                addSuppressedWhenNotCaptured( state.causeOfTermination, cause );
+                if ( causeOfTermination != null )
+                {
+                    addSuppressedWhenNotCaptured( causeOfTermination, cause );
+                }
             }
-        }
-        else
-        {
-            state = StateHolder.terminatedWith( cause );
-        }
+            else
+            {
+                state = State.TERMINATED;
+                causeOfTermination = cause;
+            }
+        } );
     }
 
     private void addSuppressedWhenNotCaptured( Throwable currentCause, Throwable newCause )
@@ -267,39 +238,40 @@ public Connection connection()
 
     private void ensureCanRunQueries()
     {
-        if ( state.value == State.COMMITTED )
-        {
-            throw new ClientException( "Cannot run more queries in this transaction, it has been committed" );
-        }
-        else if ( state.value == State.ROLLED_BACK )
-        {
-            throw new ClientException( "Cannot run more queries in this transaction, it has been rolled back" );
-        }
-        else if ( state.value == State.TERMINATED )
+        executeWithLock( lock, () ->
         {
-            throw new ClientException( "Cannot run more queries in this transaction, " +
-                    "it has either experienced an fatal error or was explicitly terminated", state.causeOfTermination );
-        }
+            if ( state == State.COMMITTED )
+            {
+                throw new ClientException( "Cannot run more queries in this transaction, it has been committed" );
+            }
+            else if ( state == State.ROLLED_BACK )
+            {
+                throw new ClientException( "Cannot run more queries in this transaction, it has been rolled back" );
+            }
+            else if ( state == State.TERMINATED )
+            {
+                throw new ClientException( "Cannot run more queries in this transaction, " +
+                                           "it has either experienced an fatal error or was explicitly terminated", causeOfTermination );
+            }
+        } );
     }
 
     private CompletionStage<Void> doCommitAsync( Throwable cursorFailure )
     {
-        if ( state.value == State.TERMINATED )
-        {
-            return failedFuture( new ClientException( "Transaction can't be committed. " +
-                                                      "It has been rolled back either because of an error or explicit termination",
-                                                      cursorFailure != state.causeOfTermination ? state.causeOfTermination : null ) );
-        }
-        return protocol.commitTransaction( connection ).thenAccept( bookmarkHolder::setBookmark );
+        ClientException exception = executeWithLock(
+                lock, () -> state == State.TERMINATED
+                            ? new ClientException( "Transaction can't be committed. " +
+                                                   "It has been rolled back either because of an error or explicit termination",
+                                                   cursorFailure != causeOfTermination ? causeOfTermination : null )
+                            : null
+        );
+        return exception != null ? failedFuture( exception ) : protocol.commitTransaction( connection ).thenAccept( bookmarkHolder::setBookmark );
     }
 
     private CompletionStage<Void> doRollbackAsync()
     {
-        if ( state.value == State.TERMINATED )
-        {
-            return completedWithNull();
-        }
-        return protocol.rollbackTransaction( connection );
+        State currentState = executeWithLock( lock, () -> state );
+        return currentState == State.TERMINATED ? completedWithNull() : protocol.rollbackTransaction( connection );
     }
 
     private static BiFunction<Void,Throwable,Void> handleCommitOrRollback( Throwable cursorFailure )
@@ -315,17 +287,8 @@ private static BiFunction<Void,Throwable,Void> handleCommitOrRollback( Throwable
         };
     }
 
-    private void handleTransactionCompletion( boolean commitOnSuccess, Throwable throwable )
+    private void releaseConnection( Throwable throwable )
     {
-        if ( commitOnSuccess && throwable == null )
-        {
-            state = StateHolder.of( State.COMMITTED );
-        }
-        else
-        {
-            state = StateHolder.of( State.ROLLED_BACK );
-        }
-
         if ( throwable instanceof AuthorizationExpiredException )
         {
             connection.terminateAndRelease( AuthorizationExpiredException.DESCRIPTION );
@@ -339,4 +302,27 @@ else if ( throwable instanceof ConnectionReadTimeoutException )
             connection.release(); // release in background
         }
     }
+
+    private void updateStateAfterCommitOrRollback( boolean commitAttempt, Throwable throwable )
+    {
+        executeWithLock( lock, () ->
+        {
+            if ( commitAttempt && throwable == null )
+            {
+                state = State.COMMITTED;
+            }
+            else
+            {
+                state = State.ROLLED_BACK;
+            }
+            if ( commitAttempt )
+            {
+                commitStage = null;
+            }
+            else
+            {
+                rollbackStage = null;
+            }
+        } );
+    }
 }

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(), null );
+                                                 InternalRxTransaction::commitIfOpen, ( tx, error ) -> tx.close(), InternalRxTransaction::close );
         return session.retryLogic().retryRx( repeatableWork );
     }