Refined the tx run methods using Flux.usingWhen

Added more IT for tx run with retries.

Author: Zhen Li

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

@@ -20,7 +20,6 @@
 
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
-import reactor.core.publisher.Mono;
 
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
@@ -122,14 +121,8 @@ public <T> Publisher<T> writeTransaction( RxTransactionWork<Publisher<T>> work,
 
     private <T> Publisher<T> runTransaction( AccessMode mode, RxTransactionWork<Publisher<T>> work, TransactionConfig config )
     {
-        Publisher<RxTransaction> publisher = beginTransaction( mode, config );
-        Flux<T> txResult = Mono.from( publisher )
-                .flatMapMany( tx -> Flux.from( work.execute( tx ) )
-                        .onErrorResume( error -> Mono.from( tx.rollback() ).then( Mono.error( error ) ) ) // if failed then we rollback and rethrow the error
-                        .concatWith( tx.commit() ) // if succeeded then we commit.
-                );
-
-        return session.retryLogic().retryRx( txResult );
+        Flux<T> repeatableWork = Flux.usingWhen( beginTransaction( mode, config ), work::execute, RxTransaction::commit, RxTransaction::rollback );
+        return session.retryLogic().retryRx( repeatableWork );
     }
 
     @Override

driver/src/test/java/org/neo4j/driver/integration/reactive/RxSessionIT.java

@@ -20,16 +20,33 @@
 
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import org.neo4j.driver.Session;
+import org.neo4j.driver.StatementResult;
+import org.neo4j.driver.exceptions.ClientException;
+import org.neo4j.driver.exceptions.DatabaseException;
+import org.neo4j.driver.exceptions.ServiceUnavailableException;
+import org.neo4j.driver.exceptions.SessionExpiredException;
+import org.neo4j.driver.exceptions.TransientException;
 import org.neo4j.driver.internal.util.EnabledOnNeo4jWith;
 import org.neo4j.driver.reactive.RxResult;
 import org.neo4j.driver.reactive.RxSession;
-import org.neo4j.driver.exceptions.ClientException;
+import org.neo4j.driver.reactive.RxTransaction;
+import org.neo4j.driver.reactive.RxTransactionWork;
 import org.neo4j.driver.util.DatabaseExtension;
 import org.neo4j.driver.util.ParallelizableIT;
 
+import static java.util.Collections.emptyIterator;
+import static org.hamcrest.CoreMatchers.instanceOf;
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.neo4j.driver.internal.util.Neo4jFeature.BOLT_V4;
 
@@ -74,4 +91,146 @@ void shouldBeAbleToReuseSessionAfterFailure()
             assertEquals( record.get("1").asLong(), 1L );
         } ).expectComplete().verify();
     }
+
+    @Test
+    void shouldRunAsyncTransactionWithoutRetries()
+    {
+        RxSession session = neo4j.driver().rxSession();
+        InvocationTrackingWork work = new InvocationTrackingWork( "CREATE (:Apa) RETURN 42" );
+        Publisher<Integer> publisher = session.writeTransaction( work );
+
+        StepVerifier.create( publisher ).expectNext( 42 ).verifyComplete();
+
+        assertEquals( 1, work.invocationCount() );
+        assertEquals( 1, countNodesByLabel( "Apa" ) );
+    }
+
+    @Test
+    void shouldRunAsyncTransactionWithRetriesOnAsyncFailures()
+    {
+        RxSession session = neo4j.driver().rxSession();
+        InvocationTrackingWork work = new InvocationTrackingWork( "CREATE (:Node) RETURN 24" ).withAsyncFailures(
+                new ServiceUnavailableException( "Oh!" ),
+                new SessionExpiredException( "Ah!" ),
+                new TransientException( "Code", "Message" ) );
+
+        Publisher<Integer> publisher = session.writeTransaction( work );
+        StepVerifier.create( publisher ).expectNext( 24 ).verifyComplete();
+
+        assertEquals( 4, work.invocationCount() );
+        assertEquals( 1, countNodesByLabel( "Node" ) );
+    }
+
+    @Test
+    void shouldRunAsyncTransactionWithRetriesOnSyncFailures()
+    {
+        RxSession session = neo4j.driver().rxSession();
+        InvocationTrackingWork work = new InvocationTrackingWork( "CREATE (:Test) RETURN 12" ).withSyncFailures(
+                new TransientException( "Oh!", "Deadlock!" ),
+                new ServiceUnavailableException( "Oh! Network Failure" ) );
+
+        Publisher<Integer> publisher = session.writeTransaction( work );
+        StepVerifier.create( publisher ).expectNext( 12 ).verifyComplete();
+
+        assertEquals( 3, work.invocationCount() );
+        assertEquals( 1, countNodesByLabel( "Test" ) );
+    }
+
+    @Test
+    void shouldRunAsyncTransactionThatCanNotBeRetried()
+    {
+        RxSession session = neo4j.driver().rxSession();
+        InvocationTrackingWork work = new InvocationTrackingWork( "UNWIND [10, 5, 0] AS x CREATE (:Hi) RETURN 10/x" );
+        Publisher<Integer> publisher = session.writeTransaction( work );
+
+        StepVerifier.create( publisher )
+                .expectNext( 1 ).expectNext( 2 )
+                .expectErrorSatisfies( error -> assertThat( error, instanceOf( ClientException.class ) ) )
+                .verify();
+
+        assertEquals( 1, work.invocationCount() );
+        assertEquals( 0, countNodesByLabel( "Hi" ) );
+    }
+
+    @Test
+    void shouldRunAsyncTransactionThatCanNotBeRetriedAfterATransientFailure()
+    {
+        RxSession session = neo4j.driver().rxSession();
+        // first throw TransientException directly from work, retry can happen afterwards
+        // then return a future failed with DatabaseException, retry can't happen afterwards
+        InvocationTrackingWork work = new InvocationTrackingWork( "CREATE (:Person) RETURN 1" )
+                .withSyncFailures( new TransientException( "Oh!", "Deadlock!" ) )
+                .withAsyncFailures( new DatabaseException( "Oh!", "OutOfMemory!" ) );
+        Publisher<Integer> publisher = session.writeTransaction( work );
+
+        StepVerifier.create( publisher )
+                .expectErrorSatisfies( e -> {
+                    assertThat( e, instanceOf( DatabaseException.class ) );
+                    assertEquals( 1, e.getSuppressed().length );
+                    assertThat( e.getSuppressed()[0], instanceOf( TransientException.class ) );
+                } )
+                .verify();
+
+        assertEquals( 2, work.invocationCount() );
+        assertEquals( 0, countNodesByLabel( "Person" ) );
+    }
+
+    private long countNodesByLabel( String label )
+    {
+        try ( Session session = neo4j.driver().session() )
+        {
+            StatementResult result = session.run( "MATCH (n:" + label + ") RETURN count(n)" );
+            return result.single().get( 0 ).asLong();
+        }
+    }
+
+    private static class InvocationTrackingWork implements RxTransactionWork<Publisher<Integer>>
+    {
+        final String query;
+        final AtomicInteger invocationCount;
+
+        Iterator<RuntimeException> asyncFailures = emptyIterator();
+        Iterator<RuntimeException> syncFailures = emptyIterator();
+
+        InvocationTrackingWork( String query )
+        {
+            this.query = query;
+            this.invocationCount = new AtomicInteger();
+        }
+
+        InvocationTrackingWork withAsyncFailures( RuntimeException... failures )
+        {
+            asyncFailures = Arrays.asList( failures ).iterator();
+            return this;
+        }
+
+        InvocationTrackingWork withSyncFailures( RuntimeException... failures )
+        {
+            syncFailures = Arrays.asList( failures ).iterator();
+            return this;
+        }
+
+        int invocationCount()
+        {
+            return invocationCount.get();
+        }
+
+        @Override
+        public Publisher<Integer> execute( RxTransaction tx )
+        {
+            invocationCount.incrementAndGet();
+
+            if ( syncFailures.hasNext() )
+            {
+                throw syncFailures.next();
+            }
+
+            if ( asyncFailures.hasNext() )
+            {
+                return Mono.error( asyncFailures.next() );
+            }
+
+            return Flux.from( tx.run( query ).records() ).map( r -> r.get( 0 ).asInt() );
+        }
+    }
 }