Handle thread interruption by closing the channel

Blocking API like `Session#run()` is now implemented on top of async
API, i.e. `Session#runAsync()`. Blocking calls simply invoke async
ones and uninterruptibly wait for them to complete. Previously, with
blocking IO, thread interruptions resulted in
`ClosedByInterruptException` which terminated query/transaction in a
harsh way. Right now interrupts will do nothing. Handling them is not
strictly required because `Thread#interrupt()` is not a nice thing
to do. However, interrupts are still used in some environments as a
last resort thing.

This commit makes driver react on thread interrupts by closing the
underlying Netty channel. It's only done in places where channel is
actually available. Closing of channel results in exception and
the waiting thread fails. It will also make database reset state
for this connection, terminate transaction and query.

Such aproach seems very similar to the previous one except it will
result in `ServiceUnavailableException` instead of
`ClosedByInterruptException`.

Author: lutovich

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

@@ -82,7 +82,9 @@ public final Driver newInstance( URI uri, AuthToken authToken, RoutingSettings r
         {
             InternalDriver driver = createDriver( uri, address, connectionPool, config, newRoutingSettings,
                     eventExecutorGroup, securityPlan, retryLogic );
-            Futures.blockingGet( driver.verifyConnectivity() );
+
+            // block to verify connectivity, close connection pool if thread gets interrupted
+            Futures.blockingGet( driver.verifyConnectivity(), connectionPool::close );
             return driver;
         }
         catch ( Throwable driverError )

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

@@ -48,7 +48,6 @@
 
 import static java.util.Collections.emptyMap;
 import static java.util.concurrent.CompletableFuture.completedFuture;
-import static org.neo4j.driver.internal.util.Futures.blockingGet;
 import static org.neo4j.driver.internal.util.Futures.completionErrorCause;
 import static org.neo4j.driver.internal.util.Futures.failedFuture;
 import static org.neo4j.driver.v1.Values.value;
@@ -150,7 +149,8 @@ public void failure()
     @Override
     public void close()
     {
-        blockingGet( closeAsync() );
+        Futures.blockingGet( closeAsync(),
+                () -> terminateConnectionOnThreadInterrupt( "Thread interrupted while closing the transaction" ) );
     }
 
     CompletionStage<Void> closeAsync()
@@ -274,8 +274,9 @@ public CompletionStage<StatementResultCursor> runAsync( String statementTemplate
     @Override
     public StatementResult run( Statement statement )
     {
-        StatementResultCursor cursor = blockingGet( run( statement, false ) );
-        return new InternalStatementResult( cursor );
+        StatementResultCursor cursor = Futures.blockingGet( run( statement, false ),
+                () -> terminateConnectionOnThreadInterrupt( "Thread interrupted while running query in transaction" ) );
+        return new InternalStatementResult( connection, cursor );
     }
 
     @Override
@@ -400,4 +401,9 @@ private BiConsumer<Object,Throwable> transactionClosed( State newState )
             session.setBookmark( bookmark );
         };
     }
+
+    private void terminateConnectionOnThreadInterrupt( String reason )
+    {
+        connection.terminateAndRelease( reason );
+    }
 }

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

@@ -22,14 +22,14 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.neo4j.driver.internal.security.SecurityPlan;
+import org.neo4j.driver.internal.util.Futures;
 import org.neo4j.driver.v1.AccessMode;
 import org.neo4j.driver.v1.Driver;
 import org.neo4j.driver.v1.Logger;
 import org.neo4j.driver.v1.Logging;
 import org.neo4j.driver.v1.Session;
 
 import static java.util.concurrent.CompletableFuture.completedFuture;
-import static org.neo4j.driver.internal.util.Futures.blockingGet;
 
 public class InternalDriver implements Driver
 {
@@ -105,7 +105,7 @@ private Session newSession( AccessMode mode, Bookmark bookmark )
     @Override
     public void close()
     {
-        blockingGet( closeAsync() );
+        Futures.blockingGet( closeAsync() );
     }
 
     @Override

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

@@ -19,7 +19,10 @@
 package org.neo4j.driver.internal;
 
 import java.util.List;
+import java.util.concurrent.CompletionStage;
 
+import org.neo4j.driver.internal.spi.Connection;
+import org.neo4j.driver.internal.util.Futures;
 import org.neo4j.driver.v1.Record;
 import org.neo4j.driver.v1.StatementResult;
 import org.neo4j.driver.v1.StatementResultCursor;
@@ -28,15 +31,15 @@
 import org.neo4j.driver.v1.summary.ResultSummary;
 import org.neo4j.driver.v1.util.Function;
 
-import static org.neo4j.driver.internal.util.Futures.blockingGet;
-
 public class InternalStatementResult implements StatementResult
 {
+    private final Connection connection;
     private final StatementResultCursor cursor;
     private List<String> keys;
 
-    public InternalStatementResult( StatementResultCursor cursor )
+    public InternalStatementResult( Connection connection, StatementResultCursor cursor )
     {
+        this.connection = connection;
         this.cursor = cursor;
     }
 
@@ -114,4 +117,14 @@ public void remove()
     {
         throw new ClientException( "Removing records from a result is not supported." );
     }
+
+    private <T> T blockingGet( CompletionStage<T> stage )
+    {
+        return Futures.blockingGet( stage, this::terminateConnectionOnThreadInterrupt );
+    }
+
+    private void terminateConnectionOnThreadInterrupt()
+    {
+        connection.terminateAndRelease( "Thread interrupted while waiting for result to arrive" );
+    }
 }

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

@@ -47,7 +47,6 @@
 import org.neo4j.driver.v1.types.TypeSystem;
 
 import static java.util.concurrent.CompletableFuture.completedFuture;
-import static org.neo4j.driver.internal.util.Futures.blockingGet;
 import static org.neo4j.driver.internal.util.Futures.failedFuture;
 import static org.neo4j.driver.v1.Values.value;
 
@@ -132,8 +131,12 @@ public CompletionStage<StatementResultCursor> runAsync( String statementText, Va
     @Override
     public StatementResult run( Statement statement )
     {
-        StatementResultCursor cursor = blockingGet( runAsync( statement, false ) );
-        return new InternalStatementResult( cursor );
+        StatementResultCursor cursor = Futures.blockingGet( runAsync( statement, false ),
+                () -> terminateConnectionOnThreadInterrupt( "Thread interrupted while running query in session" ) );
+
+        // query executed, it is safe to obtain a connection in a blocking way
+        Connection connection = Futures.getNow( connectionStage );
+        return new InternalStatementResult( connection, cursor );
     }
 
     @Override
@@ -152,7 +155,8 @@ public boolean isOpen()
     @Override
     public void close()
     {
-        blockingGet( closeAsync() );
+        Futures.blockingGet( closeAsync(),
+                () -> terminateConnectionOnThreadInterrupt( "Thread interrupted while closing the session" ) );
     }
 
     @Override
@@ -189,7 +193,7 @@ public CompletionStage<Void> closeAsync()
     @Override
     public Transaction beginTransaction()
     {
-        return blockingGet( beginTransactionAsync( mode ) );
+        return beginTransaction( mode );
     }
 
     @Deprecated
@@ -248,7 +252,8 @@ public String lastBookmark()
     @Override
     public void reset()
     {
-        blockingGet( resetAsync() );
+        Futures.blockingGet( resetAsync(),
+                () -> terminateConnectionOnThreadInterrupt( "Thread interrupted while resetting the session" ) );
     }
 
     private CompletionStage<Void> resetAsync()
@@ -288,7 +293,7 @@ private <T> T transaction( AccessMode mode, TransactionWork<T> work )
         // event loop thread will bock and wait for itself to read some data
         return retryLogic.retry( () ->
         {
-            try ( Transaction tx = blockingGet( beginTransactionAsync( mode ) ) )
+            try ( Transaction tx = beginTransaction( mode ) )
             {
                 try
                 {
@@ -433,6 +438,12 @@ private CompletionStage<InternalStatementResultCursor> runAsync( Statement state
         return newResultCursorStage;
     }
 
+    private Transaction beginTransaction( AccessMode mode )
+    {
+        return Futures.blockingGet( beginTransactionAsync( mode ),
+                () -> terminateConnectionOnThreadInterrupt( "Thread interrupted while starting a transaction" ) );
+    }
+
     private CompletionStage<ExplicitTransaction> beginTransactionAsync( AccessMode mode )
     {
         ensureSessionIsOpen();
@@ -526,6 +537,25 @@ private CompletionStage<Void> releaseConnection()
         } );
     }
 
+    private void terminateConnectionOnThreadInterrupt( String reason )
+    {
+        // try to get current connection in a blocking fashion
+        Connection connection = null;
+        try
+        {
+            connection = Futures.getNow( connectionStage );
+        }
+        catch ( Throwable ignore )
+        {
+            // ignore errors because handing interruptions is best effort
+        }
+
+        if ( connection != null )
+        {
+            connection.terminateAndRelease( reason );
+        }
+    }
+
     private CompletionStage<Void> ensureNoOpenTxBeforeRunningQuery()
     {
         return ensureNoOpenTx( "Statements cannot be run directly on a session with an open transaction; " +

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

@@ -34,6 +34,7 @@ public final class ChannelAttributes
     private static final AttributeKey<Long> CREATION_TIMESTAMP = newInstance( "creationTimestamp" );
     private static final AttributeKey<Long> LAST_USED_TIMESTAMP = newInstance( "lastUsedTimestamp" );
     private static final AttributeKey<InboundMessageDispatcher> MESSAGE_DISPATCHER = newInstance( "messageDispatcher" );
+    private static final AttributeKey<String> TERMINATION_REASON = newInstance( "terminationReason" );
 
     private ChannelAttributes()
     {
@@ -89,6 +90,16 @@ public static void setMessageDispatcher( Channel channel, InboundMessageDispatch
         setOnce( channel, MESSAGE_DISPATCHER, messageDispatcher );
     }
 
+    public static String terminationReason( Channel channel )
+    {
+        return get( channel, TERMINATION_REASON );
+    }
+
+    public static void setTerminationReason( Channel channel, String reason )
+    {
+        setOnce( channel, TERMINATION_REASON, reason );
+    }
+
     private static <T> T get( Channel channel, AttributeKey<T> key )
     {
         return channel.attr( key ).get();

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

@@ -39,6 +39,8 @@
 import org.neo4j.driver.internal.util.ServerVersion;
 import org.neo4j.driver.v1.Value;
 
+import static org.neo4j.driver.internal.async.ChannelAttributes.setTerminationReason;
+
 public class NettyConnection implements Connection
 {
     private final Channel channel;
@@ -115,6 +117,18 @@ public CompletionStage<Void> release()
         return releaseFuture;
     }
 
+    @Override
+    public void terminateAndRelease( String reason )
+    {
+        if ( open.compareAndSet( true, false ) )
+        {
+            setTerminationReason( channel, reason );
+            channel.close();
+            channelPool.release( channel );
+            releaseFuture.complete( null );
+        }
+    }
+
     @Override
     public BoltServerAddress serverAddress()
     {

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

@@ -83,6 +83,12 @@ public CompletionStage<Void> release()
         return delegate.release();
     }
 
+    @Override
+    public void terminateAndRelease( String reason )
+    {
+        delegate.terminateAndRelease( reason );
+    }
+
     @Override
     public BoltServerAddress serverAddress()
     {

driver/src/main/java/org/neo4j/driver/internal/async/inbound/ChannelErrorHandler.java

@@ -32,6 +32,7 @@
 
 import static java.util.Objects.requireNonNull;
 import static org.neo4j.driver.internal.async.ChannelAttributes.messageDispatcher;
+import static org.neo4j.driver.internal.async.ChannelAttributes.terminationReason;
 
 public class ChannelErrorHandler extends ChannelInboundHandlerAdapter
 {
@@ -69,8 +70,10 @@ public void channelInactive( ChannelHandlerContext ctx )
         if ( !failed )
         {
             // channel became inactive not because of a fatal exception that came from exceptionCaught
-            // it is most likely inactive because actual network connection broke
-            ServiceUnavailableException error = ErrorUtil.newConnectionTerminatedError();
+            // it is most likely inactive because actual network connection broke or was explicitly closed by the driver
+
+            String terminationReason = terminationReason( ctx.channel() );
+            ServiceUnavailableException error = ErrorUtil.newConnectionTerminatedError( terminationReason );
             fail( ctx, error );
         }
     }

driver/src/main/java/org/neo4j/driver/internal/spi/Connection.java

@@ -41,6 +41,8 @@ void runAndFlush( String statement, Map<String,Value> parameters, ResponseHandle
 
     CompletionStage<Void> release();
 
+    void terminateAndRelease( String reason );
+
     BoltServerAddress serverAddress();
 
     ServerVersion serverVersion();

driver/src/main/java/org/neo4j/driver/internal/util/ErrorUtil.java

@@ -31,6 +31,15 @@ private ErrorUtil()
     {
     }
 
+    public static ServiceUnavailableException newConnectionTerminatedError( String reason )
+    {
+        if ( reason == null )
+        {
+            return newConnectionTerminatedError();
+        }
+        return new ServiceUnavailableException( "Connection to the database terminated. " + reason );
+    }
+
     public static ServiceUnavailableException newConnectionTerminatedError()
     {
         return new ServiceUnavailableException( "Connection to the database terminated. " +

driver/src/main/java/org/neo4j/driver/internal/util/Futures.java

@@ -78,6 +78,11 @@ public static <T> CompletableFuture<T> failedFuture( Throwable error )
     }
 
     public static <V> V blockingGet( CompletionStage<V> stage )
+    {
+        return blockingGet( stage, Futures::noOpInterruptHandler );
+    }
+
+    public static <V> V blockingGet( CompletionStage<V> stage, Runnable interruptHandler )
     {
         EventLoopGroupFactory.assertNotInEventLoopThread();
 
@@ -93,7 +98,14 @@ public static <V> V blockingGet( CompletionStage<V> stage )
                 }
                 catch ( InterruptedException e )
                 {
+                    // this thread was interrupted while waiting
+                    // computation denoted by the future might still be running
+
                     interrupted = true;
+
+                    // run the interrupt handler and ignore if it throws
+                    // need to wait for IO thread to actually finish, can't simply re-rethrow
+                    safeRun( interruptHandler );
                 }
                 catch ( ExecutionException e )
                 {
@@ -110,6 +122,11 @@ public static <V> V blockingGet( CompletionStage<V> stage )
         }
     }
 
+    public static <T> T getNow( CompletionStage<T> stage )
+    {
+        return stage.toCompletableFuture().getNow( null );
+    }
+
     // todo: test all call sites
     public static Throwable completionErrorCause( Throwable error )
     {
@@ -119,4 +136,19 @@ public static Throwable completionErrorCause( Throwable error )
         }
         return error;
     }
+
+    private static void noOpInterruptHandler()
+    {
+    }
+
+    private static void safeRun( Runnable runnable )
+    {
+        try
+        {
+            runnable.run();
+        }
+        catch ( Throwable ignore )
+        {
+        }
+    }
 }