Better handle sync failures in async tx functions

Given work might fail by throwing exception or by returning a failed
future. This commit makes retries gracefully handle exceptions.
Previously retrying code would only log an exception and hang.

Author: lutovich

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

@@ -453,7 +453,7 @@ else if ( future.isSuccess() )
     private <T> void executeWork( final InternalPromise<T> resultPromise, final ExplicitTransaction tx,
             TransactionWork<Response<T>> work )
     {
-        Response<T> workResponse = work.execute( tx );
+        Response<T> workResponse = safeExecuteWork( tx, work );
         workResponse.addListener( new ResponseListener<T>()
         {
             @Override
@@ -471,6 +471,22 @@ public void operationCompleted( T result, Throwable error )
         } );
     }
 
+    private <T> Response<T> safeExecuteWork( ExplicitTransaction tx, TransactionWork<Response<T>> work )
+    {
+        // given work might fail in both async and sync way
+        // async failure will result in a failed future being returned
+        // sync failure will result in an exception being thrown
+        try
+        {
+            return work.execute( tx );
+        }
+        catch ( Throwable workError )
+        {
+            // work threw an exception, wrap it in a future and proceed
+            return new InternalPromise<T>( eventExecutorGroup ).setFailure( workError );
+        }
+    }
+
     private <T> void rollbackTxAfterFailedTransactionWork( ExplicitTransaction tx,
             final InternalPromise<T> resultPromise, final Throwable error )
     {

driver/src/test/java/org/neo4j/driver/v1/integration/SessionAsyncIT.java

@@ -28,7 +28,6 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.concurrent.Future;
@@ -46,6 +45,7 @@
 import org.neo4j.driver.v1.TransactionWork;
 import org.neo4j.driver.v1.Value;
 import org.neo4j.driver.v1.exceptions.ClientException;
+import org.neo4j.driver.v1.exceptions.DatabaseException;
 import org.neo4j.driver.v1.exceptions.ServiceUnavailableException;
 import org.neo4j.driver.v1.exceptions.SessionExpiredException;
 import org.neo4j.driver.v1.exceptions.TransientException;
@@ -54,6 +54,7 @@
 import org.neo4j.driver.v1.types.Node;
 import org.neo4j.driver.v1.util.TestNeo4j;
 
+import static java.util.Collections.emptyIterator;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.instanceOf;
@@ -371,11 +372,13 @@ public void shouldRunAsyncTransactionWithoutRetries()
     }
 
     @Test
-    public void shouldRunAsyncTransactionWithRetries()
+    public void shouldRunAsyncTransactionWithRetriesOnAsyncFailures()
     {
-        List<Throwable> failures = Arrays.<Throwable>asList( new ServiceUnavailableException( "Oh!" ),
-                new SessionExpiredException( "Ah!" ), new TransientException( "Code", "Message" ) );
-        InvocationTrackingWork work = new InvocationTrackingWork( "CREATE (:Node) RETURN 24", failures );
+        InvocationTrackingWork work = new InvocationTrackingWork( "CREATE (:Node) RETURN 24" ).withAsyncFailures(
+                new ServiceUnavailableException( "Oh!" ),
+                new SessionExpiredException( "Ah!" ),
+                new TransientException( "Code", "Message" ) );
+
         Response<Record> txResponse = session.writeTransactionAsync( work );
 
         Record record = await( txResponse );
@@ -386,6 +389,23 @@ public void shouldRunAsyncTransactionWithRetries()
         assertEquals( 1, countNodesByLabel( "Node" ) );
     }
 
+    @Test
+    public void shouldRunAsyncTransactionWithRetriesOnSyncFailures()
+    {
+        InvocationTrackingWork work = new InvocationTrackingWork( "CREATE (:Test) RETURN 12" ).withSyncFailures(
+                new TransientException( "Oh!", "Deadlock!" ),
+                new ServiceUnavailableException( "Oh! Network Failure" ) );
+
+        Response<Record> txResponse = session.writeTransactionAsync( work );
+
+        Record record = await( txResponse );
+        assertNotNull( record );
+        assertEquals( 12L, record.get( 0 ).asLong() );
+
+        assertEquals( 3, work.invocationCount() );
+        assertEquals( 1, countNodesByLabel( "Test" ) );
+    }
+
     @Test
     public void shouldRunAsyncTransactionThatCanNotBeRetried()
     {
@@ -406,6 +426,32 @@ public void shouldRunAsyncTransactionThatCanNotBeRetried()
         assertEquals( 0, countNodesByLabel( "Hi" ) );
     }
 
+    @Test
+    public void shouldRunAsyncTransactionThatCanNotBeRetriedAfterATransientFailure()
+    {
+        // 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!" ) );
+        Response<Record> txResponse = session.writeTransactionAsync( work );
+
+        try
+        {
+            await( txResponse );
+            fail( "Exception expected" );
+        }
+        catch ( Exception e )
+        {
+            assertThat( e, instanceOf( DatabaseException.class ) );
+            assertEquals( 1, e.getSuppressed().length );
+            assertThat( e.getSuppressed()[0], instanceOf( TransientException.class ) );
+        }
+
+        assertEquals( 2, work.invocationCount() );
+        assertEquals( 0, countNodesByLabel( "Person" ) );
+    }
+
     private Future<List<Future<Boolean>>> runNestedQueries( StatementResultCursor inputCursor )
     {
         Promise<List<Future<Boolean>>> resultPromise = GlobalEventExecutor.INSTANCE.newPromise();
@@ -524,19 +570,27 @@ void killDb()
     private static class InvocationTrackingWork implements TransactionWork<Response<Record>>
     {
         final String query;
-        final Iterator<Throwable> failures;
         final AtomicInteger invocationCount;
 
+        Iterator<RuntimeException> asyncFailures = emptyIterator();
+        Iterator<RuntimeException> syncFailures = emptyIterator();
+
         InvocationTrackingWork( String query )
         {
-            this( query, Collections.<Throwable>emptyList() );
+            this.query = query;
+            this.invocationCount = new AtomicInteger();
+        }
+
+        InvocationTrackingWork withAsyncFailures( RuntimeException... failures )
+        {
+            asyncFailures = Arrays.asList( failures ).iterator();
+            return this;
         }
 
-        InvocationTrackingWork( String query, List<Throwable> failures )
+        InvocationTrackingWork withSyncFailures( RuntimeException... failures )
         {
-            this.query = query;
-            this.failures = failures.iterator();
-            this.invocationCount = new AtomicInteger();
+            syncFailures = Arrays.asList( failures ).iterator();
+            return this;
         }
 
         int invocationCount()
@@ -549,6 +603,11 @@ public Response<Record> execute( Transaction tx )
         {
             invocationCount.incrementAndGet();
 
+            if ( syncFailures.hasNext() )
+            {
+                throw syncFailures.next();
+            }
+
             final InternalPromise<Record> resultPromise = new InternalPromise<>( GlobalEventExecutor.INSTANCE );
 
             tx.runAsync( query ).addListener( new ResponseListener<StatementResultCursor>()
@@ -597,9 +656,9 @@ private void processFetchResult( Boolean recordAvailable, Throwable error,
                 return;
             }
 
-            if ( failures.hasNext() )
+            if ( asyncFailures.hasNext() )
             {
-                resultPromise.setFailure( failures.next() );
+                resultPromise.setFailure( asyncFailures.next() );
             }
             else
             {

driver/src/test/java/org/neo4j/driver/v1/util/TestUtil.java

@@ -68,13 +68,7 @@ public static <T, F extends Future<T>> T await( F future )
         }
         catch ( ExecutionException e )
         {
-            Throwable cause = e.getCause();
-            StackTraceElement[] originalStackTrace = cause.getStackTrace();
-            RuntimeException exceptionWithOriginalStackTrace = new RuntimeException();
-            cause.setStackTrace( exceptionWithOriginalStackTrace.getStackTrace() );
-            exceptionWithOriginalStackTrace.setStackTrace( originalStackTrace );
-            cause.addSuppressed( exceptionWithOriginalStackTrace );
-            throwException( cause );
+            throwException( e.getCause() );
             return null;
         }
         catch ( TimeoutException e )