Deadlock Issues caused by Netty Blocking ChannelPool#close.

The new netty version provides both blocking pool close `FixedChannelPool#close` and async pool close `FixedChannelPool#closeAsync`.
The blocking close causes deadlock issues in our code.

One deadlock issue happens when `ConnectionPoolImpl#retainAll` trying to close a netty pool inside an IO thread that is used by one of the channels owned by the closing pool.

The deadlock happens because
1) The channel cannot be closed until the IO thread is fully free.
2) The pool cannot be closed until the channel is closed.
3) However the IO thread cannot be free until the pool is closed.
`pool --wait-on--> channel --wait-on--> IO thread --wait-on--> pool`

The fix is to change 3) to let the IO thread to call async pool close to not block the close operation.

Another deadlock issue happens in `Driver#close`.

In `Driver#close`, we need to chain the event loop group shutdown after all pool close.
As pool close tasks are executed on event loop group, if we shutdown event loop group first,
then the pool close will be stuck because of no thread to pick up any shutdown tasks.

Author: Zhen Li

driver/src/main/java/org/neo4j/driver/internal/async/pool/ConnectionPoolImpl.java

@@ -24,8 +24,8 @@
 import io.netty.channel.pool.ChannelPool;
 import io.netty.util.concurrent.Future;
 
-import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.ConcurrentHashMap;
@@ -48,6 +48,8 @@
 import org.neo4j.driver.internal.util.Futures;
 
 import static java.lang.String.format;
+import static org.neo4j.driver.internal.util.Futures.combineErrors;
+import static org.neo4j.driver.internal.util.Futures.completedWithNullIfNonError;
 
 public class ConnectionPoolImpl implements ConnectionPool
 {
@@ -62,6 +64,7 @@ public class ConnectionPoolImpl implements ConnectionPool
 
     private final ConcurrentMap<BoltServerAddress,ExtendedChannelPool> pools = new ConcurrentHashMap<>();
     private final AtomicBoolean closed = new AtomicBoolean();
+    private final CompletableFuture<Void> closeFuture = new CompletableFuture<>();
     private final ConnectionFactory connectionFactory;
 
     public ConnectionPoolImpl( ChannelConnector connector, Bootstrap bootstrap, PoolSettings settings, MetricsListener metricsListener, Logging logging,
@@ -131,8 +134,14 @@ public void retainAll( Set<BoltServerAddress> addressesToRetain )
                     if ( pool != null )
                     {
                         log.info( "Closing connection pool towards %s, it has no active connections " +
-                                  "and is not in the routing table", address );
-                        closePool( pool );
+                                  "and is not in the routing table registry.", address );
+                        // Close in the background
+                        closePool( pool ).whenComplete( ( ignored, error ) -> {
+                            if ( error != null )
+                            {
+                                log.warn( format( "An error occurred while closing connection pool towards %s.", address ), error );
+                            }
+                        } );
                     }
                 }
             }
@@ -156,37 +165,40 @@ public CompletionStage<Void> close()
     {
         if ( closed.compareAndSet( false, true ) )
         {
-            try
-            {
-                nettyChannelTracker.prepareToCloseChannels();
-                for ( Map.Entry<BoltServerAddress,ExtendedChannelPool> entry : pools.entrySet() )
-                {
-                    BoltServerAddress address = entry.getKey();
-                    ExtendedChannelPool pool = entry.getValue();
-                    log.info( "Closing connection pool towards %s", address );
-                    closePool( pool );
-                }
+            nettyChannelTracker.prepareToCloseChannels();
 
-                pools.clear();
-            }
-            finally
-            {
+            CompletableFuture<Void> allPoolClosedFuture = CompletableFuture.allOf(
+                    pools.entrySet().stream().map( entry -> {
+                        BoltServerAddress address = entry.getKey();
+                        ExtendedChannelPool pool = entry.getValue();
+                        log.info( "Closing connection pool towards %s", address );
+                        // Wait for all pools to be closed.
+                        return closePool( pool ).toCompletableFuture();
+                    } ).toArray( CompletableFuture[]::new ) );
 
-                if (ownsEventLoopGroup) {
+            // We can only shutdown event loop group when all netty pools are fully closed,
+            // otherwise the netty pools might missing threads (from event loop group) to execute clean ups.
+            allPoolClosedFuture.whenComplete( ( ignored, pollCloseError ) -> {
+                pools.clear();
+                if ( !ownsEventLoopGroup )
+                {
+                    completedWithNullIfNonError( closeFuture, pollCloseError );
+                }
+                else
+                {
                     // This is an attempt to speed up the shut down procedure of the driver
                     // Feel free return this back to shutdownGracefully() method with default values
                     // if this proves troublesome!!!
-                    eventLoopGroup().shutdownGracefully(200, 15_000, TimeUnit.MILLISECONDS);
+                    eventLoopGroup().shutdownGracefully( 200, 15_000, TimeUnit.MILLISECONDS );
+
+                    Futures.asCompletionStage( eventLoopGroup().terminationFuture() ).whenComplete( ( ignore, eventLoopGroupTerminationError ) -> {
+                        CompletionException combinedErrors = combineErrors( pollCloseError, eventLoopGroupTerminationError );
+                        completedWithNullIfNonError( closeFuture, combinedErrors );
+                    } );
                 }
-            }
-        }
-        if (!ownsEventLoopGroup)
-        {
-            return Futures.completedWithNull();
+            } );
         }
-
-        return Futures.asCompletionStage( eventLoopGroup().terminationFuture() )
-                .thenApply( ignore -> null );
+        return closeFuture;
     }
 
     @Override
@@ -200,11 +212,11 @@ private ExtendedChannelPool getOrCreatePool( BoltServerAddress address )
         return pools.computeIfAbsent( address, this::newPool );
     }
 
-    private void closePool( ExtendedChannelPool pool )
+    private CompletionStage<Void> closePool( ExtendedChannelPool pool )
     {
-        pool.close();
-        // after the connection pool is removed/close, I can remove its metrics.
-        metricsListener.removePoolMetrics( pool.id() );
+        return pool.repeatableCloseAsync().whenComplete( ( ignored, error ) ->
+                // after the connection pool is removed/close, I can remove its metrics.
+                metricsListener.removePoolMetrics( pool.id() ) );
     }
 
     ExtendedChannelPool newPool( BoltServerAddress address )

driver/src/main/java/org/neo4j/driver/internal/async/pool/ExtendedChannelPool.java

@@ -20,9 +20,13 @@
 
 import io.netty.channel.pool.ChannelPool;
 
+import java.util.concurrent.CompletionStage;
+
 public interface ExtendedChannelPool extends ChannelPool
 {
     boolean isClosed();
 
     String id();
+
+    CompletionStage<Void> repeatableCloseAsync();
 }

driver/src/main/java/org/neo4j/driver/internal/async/pool/NettyChannelPool.java

@@ -24,6 +24,8 @@
 import io.netty.channel.pool.ChannelHealthChecker;
 import io.netty.channel.pool.FixedChannelPool;
 
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.neo4j.driver.internal.BoltServerAddress;
@@ -32,6 +34,7 @@
 
 import static java.util.Objects.requireNonNull;
 import static org.neo4j.driver.internal.async.connection.ChannelAttributes.setPoolId;
+import static org.neo4j.driver.internal.util.Futures.asCompletionStage;
 
 public class NettyChannelPool extends FixedChannelPool implements ExtendedChannelPool
 {
@@ -49,6 +52,7 @@ public class NettyChannelPool extends FixedChannelPool implements ExtendedChanne
     private final NettyChannelTracker handler;
     private final AtomicBoolean closed = new AtomicBoolean( false );
     private final String id;
+    private final CompletableFuture<Void> closeFuture = new CompletableFuture<>();
 
     public NettyChannelPool( BoltServerAddress address, ChannelConnector connector, Bootstrap bootstrap, NettyChannelTracker handler,
             ChannelHealthChecker healthCheck, long acquireTimeoutMillis, int maxConnections )
@@ -85,19 +89,22 @@ protected ChannelFuture connectChannel( Bootstrap bootstrap )
     }
 
     @Override
-    public void close()
+    public CompletionStage<Void> repeatableCloseAsync()
     {
         if ( closed.compareAndSet( false, true ) )
         {
-            super.close();
+            asCompletionStage( super.closeAsync(), closeFuture );
         }
+        return closeFuture;
     }
 
+    @Override
     public boolean isClosed()
     {
         return closed.get();
     }
 
+    @Override
     public String id()
     {
         return this.id;

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

@@ -45,9 +45,27 @@ public static <T> CompletableFuture<T> completedWithNull()
         return (CompletableFuture) COMPLETED_WITH_NULL;
     }
 
+    public static <T> CompletableFuture<T> completedWithNullIfNonError( CompletableFuture<T> future, Throwable error )
+    {
+        if ( error != null )
+        {
+            future.completeExceptionally( error );
+        }
+        else
+        {
+            future.complete( null );
+        }
+        return future;
+    }
+
     public static <T> CompletionStage<T> asCompletionStage( io.netty.util.concurrent.Future<T> future )
     {
         CompletableFuture<T> result = new CompletableFuture<>();
+        return asCompletionStage( future, result );
+    }
+
+    public static <T> CompletionStage<T> asCompletionStage( io.netty.util.concurrent.Future<T> future, CompletableFuture<T> result )
+    {
         if ( future.isCancelled() )
         {
             result.cancel( true );

driver/src/test/java/org/neo4j/driver/internal/async/pool/ConnectionPoolImplIT.java

@@ -24,6 +24,9 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
+import java.util.Collections;
+import java.util.concurrent.CompletionStage;
+
 import org.neo4j.driver.exceptions.ServiceUnavailableException;
 import org.neo4j.driver.internal.BoltServerAddress;
 import org.neo4j.driver.internal.ConnectionSettings;
@@ -86,6 +89,18 @@ void shouldAcquireIdleConnection()
         assertNotNull( connection2 );
     }
 
+    @Test
+    void shouldBeAbleToClosePoolInIOWorkerThread() throws Throwable
+    {
+        // In the IO worker thread of a channel obtained from a pool, we shall be able to close the pool.
+        CompletionStage<Void> future = pool.acquire( neo4j.address() ).thenCompose( Connection::release )
+                // This shall close all pools
+                .whenComplete( ( ignored, error ) -> pool.retainAll( Collections.emptySet() ) );
+
+        // We should be able to come to this line.
+        await( future );
+    }
+
     @Test
     void shouldFailToAcquireConnectionToWrongAddress()
     {
@@ -118,7 +133,7 @@ void shouldFailToAcquireConnectionWhenPoolIsClosed()
     {
         await( pool.acquire( neo4j.address() ) );
         ExtendedChannelPool channelPool = this.pool.getPool( neo4j.address() );
-        channelPool.close();
+        await( channelPool.repeatableCloseAsync() );
         ServiceUnavailableException error =
                 assertThrows( ServiceUnavailableException.class, () -> await( pool.acquire( neo4j.address() ) ) );
         assertThat( error.getMessage(), containsString( "closed while acquiring a connection" ) );

driver/src/test/java/org/neo4j/driver/internal/async/pool/ConnectionPoolImplTest.java

@@ -20,7 +20,6 @@
 
 import io.netty.bootstrap.Bootstrap;
 import io.netty.channel.Channel;
-import io.netty.channel.pool.ChannelPool;
 import io.netty.util.concurrent.ImmediateEventExecutor;
 import org.junit.jupiter.api.Test;
 
@@ -44,6 +43,7 @@
 import static org.neo4j.driver.internal.BoltServerAddress.LOCAL_DEFAULT;
 import static org.neo4j.driver.internal.logging.DevNullLogging.DEV_NULL_LOGGING;
 import static org.neo4j.driver.internal.metrics.InternalAbstractMetrics.DEV_NULL_METRICS;
+import static org.neo4j.driver.internal.util.Futures.completedWithNull;
 
 class ConnectionPoolImplTest
 {
@@ -73,9 +73,9 @@ void shouldRetainSpecifiedAddresses()
         pool.acquire( ADDRESS_3 );
 
         pool.retainAll( new HashSet<>( asList( ADDRESS_1, ADDRESS_2, ADDRESS_3 ) ) );
-        for ( ChannelPool channelPool : pool.channelPoolsByAddress.values() )
+        for ( ExtendedChannelPool channelPool : pool.channelPoolsByAddress.values() )
         {
-            verify( channelPool, never() ).close();
+            verify( channelPool, never() ).repeatableCloseAsync();
         }
     }
 
@@ -94,9 +94,9 @@ void shouldClosePoolsWhenRetaining()
         when( nettyChannelTracker.inUseChannelCount( ADDRESS_3 ) ).thenReturn( 3 );
 
         pool.retainAll( new HashSet<>( asList( ADDRESS_1, ADDRESS_3 ) ) );
-        verify( pool.getPool( ADDRESS_1 ), never() ).close();
-        verify( pool.getPool( ADDRESS_2 ) ).close();
-        verify( pool.getPool( ADDRESS_3 ), never() ).close();
+        verify( pool.getPool( ADDRESS_1 ), never() ).repeatableCloseAsync();
+        verify( pool.getPool( ADDRESS_2 ) ).repeatableCloseAsync();
+        verify( pool.getPool( ADDRESS_3 ), never() ).repeatableCloseAsync();
     }
 
     @Test
@@ -114,9 +114,9 @@ void shouldNotClosePoolsWithActiveConnectionsWhenRetaining()
         when( nettyChannelTracker.inUseChannelCount( ADDRESS_3 ) ).thenReturn( 0 );
 
         pool.retainAll( singleton( ADDRESS_2 ) );
-        verify( pool.getPool( ADDRESS_1 ), never() ).close();
-        verify( pool.getPool( ADDRESS_2 ), never() ).close();
-        verify( pool.getPool( ADDRESS_3 ) ).close();
+        verify( pool.getPool( ADDRESS_1 ), never() ).repeatableCloseAsync();
+        verify( pool.getPool( ADDRESS_2 ), never() ).repeatableCloseAsync();
+        verify( pool.getPool( ADDRESS_3 ) ).repeatableCloseAsync();
     }
 
     private static PoolSettings newSettings()
@@ -147,6 +147,7 @@ ExtendedChannelPool newPool( BoltServerAddress address )
             ExtendedChannelPool channelPool = mock( ExtendedChannelPool.class );
             Channel channel = mock( Channel.class );
             doReturn( ImmediateEventExecutor.INSTANCE.newSucceededFuture( channel ) ).when( channelPool ).acquire();
+            doReturn( completedWithNull() ).when( channelPool ).repeatableCloseAsync();
             channelPoolsByAddress.put( address, channelPool );
             return channelPool;
         }