Exposing only the methods ConnectionPoolImpl really needs

Author: Zhen Li

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

@@ -19,7 +19,6 @@
 package org.neo4j.driver.internal.async;
 
 import io.netty.channel.Channel;
-import io.netty.channel.pool.ChannelPool;
 
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
@@ -28,6 +27,7 @@
 import org.neo4j.driver.internal.BoltServerAddress;
 import org.neo4j.driver.internal.async.connection.ChannelAttributes;
 import org.neo4j.driver.internal.async.inbound.InboundMessageDispatcher;
+import org.neo4j.driver.internal.async.pool.ExtendedChannelPool;
 import org.neo4j.driver.internal.handlers.ChannelReleasingResetResponseHandler;
 import org.neo4j.driver.internal.handlers.ResetResponseHandler;
 import org.neo4j.driver.internal.messaging.BoltProtocol;
@@ -57,15 +57,15 @@ public class NetworkConnection implements Connection
     private final BoltServerAddress serverAddress;
     private final ServerVersion serverVersion;
     private final BoltProtocol protocol;
-    private final ChannelPool channelPool;
+    private final ExtendedChannelPool channelPool;
     private final CompletableFuture<Void> releaseFuture;
     private final Clock clock;
 
     private final AtomicReference<Status> status = new AtomicReference<>( Status.OPEN );
     private final MetricsListener metricsListener;
     private final ListenerEvent inUseEvent;
 
-    public NetworkConnection( Channel channel, ChannelPool channelPool, Clock clock, MetricsListener metricsListener )
+    public NetworkConnection( Channel channel, ExtendedChannelPool channelPool, Clock clock, MetricsListener metricsListener )
     {
         this.channel = channel;
         this.messageDispatcher = ChannelAttributes.messageDispatcher( channel );

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

@@ -21,8 +21,6 @@
 import io.netty.bootstrap.Bootstrap;
 import io.netty.channel.Channel;
 import io.netty.channel.EventLoopGroup;
-import io.netty.channel.pool.ChannelPool;
-import io.netty.util.concurrent.Future;
 
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
@@ -49,7 +47,7 @@
 
 import static java.lang.String.format;
 import static org.neo4j.driver.internal.util.Futures.combineErrors;
-import static org.neo4j.driver.internal.util.Futures.completedWithNullIfNonError;
+import static org.neo4j.driver.internal.util.Futures.completeWithNullIfNoError;
 
 public class ConnectionPoolImpl implements ConnectionPool
 {
@@ -98,9 +96,9 @@ public CompletionStage<Connection> acquire( BoltServerAddress address )
 
         ListenerEvent acquireEvent = metricsListener.createListenerEvent();
         metricsListener.beforeAcquiringOrCreating( pool.id(), acquireEvent );
-        Future<Channel> connectionFuture = pool.acquire();
+        CompletionStage<Channel> connectionFuture = pool.acquire();
 
-        return Futures.asCompletionStage( connectionFuture ).handle( ( channel, error ) ->
+        return connectionFuture.handle( ( channel, error ) ->
         {
             try
             {
@@ -135,13 +133,7 @@ public void retainAll( Set<BoltServerAddress> addressesToRetain )
                     {
                         log.info( "Closing connection pool towards %s, it has no active connections " +
                                   "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 );
-                            }
-                        } );
+                        closePoolInBackground( address, pool );
                     }
                 }
             }
@@ -166,35 +158,19 @@ public CompletionStage<Void> close()
         if ( closed.compareAndSet( false, true ) )
         {
             nettyChannelTracker.prepareToCloseChannels();
-
-            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 ) );
+            CompletableFuture<Void> allPoolClosedFuture = closeAllPools();
 
             // 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 );
+                    completeWithNullIfNoError( 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 );
-
-                    Futures.asCompletionStage( eventLoopGroup().terminationFuture() ).whenComplete( ( ignore, eventLoopGroupTerminationError ) -> {
-                        CompletionException combinedErrors = combineErrors( pollCloseError, eventLoopGroupTerminationError );
-                        completedWithNullIfNonError( closeFuture, combinedErrors );
-                    } );
+                    shutdownEventLoopGroup( pollCloseError );
                 }
             } );
         }
@@ -207,31 +183,10 @@ public boolean isOpen( BoltServerAddress address )
         return pools.containsKey( address );
     }
 
-    private ExtendedChannelPool getOrCreatePool( BoltServerAddress address )
-    {
-        return pools.computeIfAbsent( address, this::newPool );
-    }
-
-    private CompletionStage<Void> closePool( ExtendedChannelPool pool )
-    {
-        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 )
-    {
-        NettyChannelPool pool =
-                new NettyChannelPool( address, connector, bootstrap, nettyChannelTracker, channelHealthChecker, settings.connectionAcquisitionTimeout(),
-                        settings.maxConnectionPoolSize() );
-        // before the connection pool is added I can add the metrics for the pool.
-        metricsListener.putPoolMetrics( pool.id(), address, this );
-        return pool;
-    }
-
-    private EventLoopGroup eventLoopGroup()
+    @Override
+    public String toString()
     {
-        return bootstrap.config().group();
+        return "ConnectionPoolImpl{" + "pools=" + pools + '}';
     }
 
     private void processAcquisitionError( ExtendedChannelPool pool, BoltServerAddress serverAddress, Throwable error )
@@ -271,26 +226,83 @@ private void assertNotClosed()
         }
     }
 
-    private void assertNotClosed( BoltServerAddress address, Channel channel, ChannelPool pool )
+    private void assertNotClosed( BoltServerAddress address, Channel channel, ExtendedChannelPool pool )
     {
         if ( closed.get() )
         {
             pool.release( channel );
-            pool.close();
+            closePoolInBackground( address, pool );
             pools.remove( address );
             assertNotClosed();
         }
     }
 
-    @Override
-    public String toString()
-    {
-        return "ConnectionPoolImpl{" + "pools=" + pools + '}';
-    }
-
     // for testing only
     ExtendedChannelPool getPool( BoltServerAddress address )
     {
         return pools.get( address );
     }
+
+    ExtendedChannelPool newPool( BoltServerAddress address )
+    {
+        NettyChannelPool pool =
+                new NettyChannelPool( address, connector, bootstrap, nettyChannelTracker, channelHealthChecker, settings.connectionAcquisitionTimeout(),
+                        settings.maxConnectionPoolSize() );
+        // before the connection pool is added I can add the metrics for the pool.
+        metricsListener.putPoolMetrics( pool.id(), address, this );
+        return pool;
+    }
+
+    private ExtendedChannelPool getOrCreatePool( BoltServerAddress address )
+    {
+        return pools.computeIfAbsent( address, this::newPool );
+    }
+
+    private CompletionStage<Void> closePool( ExtendedChannelPool pool )
+    {
+        return pool.close().whenComplete( ( ignored, error ) ->
+                // after the connection pool is removed/close, I can remove its metrics.
+                metricsListener.removePoolMetrics( pool.id() ) );
+    }
+
+    private void closePoolInBackground( BoltServerAddress address, ExtendedChannelPool pool )
+    {
+        // 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 );
+            }
+        } );
+    }
+
+    private EventLoopGroup eventLoopGroup()
+    {
+        return bootstrap.config().group();
+    }
+
+    private void shutdownEventLoopGroup( Throwable pollCloseError )
+    {
+        // This is an attempt to speed up the shut down procedure of the driver
+        // This timeout is needed for `closePoolInBackground` to finish background job, especially for races between `acquire` and `close`.
+        eventLoopGroup().shutdownGracefully( 200, 15_000, TimeUnit.MILLISECONDS );
+
+        Futures.asCompletionStage( eventLoopGroup().terminationFuture() )
+                .whenComplete( ( ignore, eventLoopGroupTerminationError ) -> {
+                    CompletionException combinedErrors = combineErrors( pollCloseError, eventLoopGroupTerminationError );
+                    completeWithNullIfNoError( closeFuture, combinedErrors );
+                } );
+    }
+
+    private CompletableFuture<Void> closeAllPools()
+    {
+        return 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 ) );
+    }
 }

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

@@ -18,15 +18,19 @@
  */
 package org.neo4j.driver.internal.async.pool;
 
-import io.netty.channel.pool.ChannelPool;
+import io.netty.channel.Channel;
 
 import java.util.concurrent.CompletionStage;
 
-public interface ExtendedChannelPool extends ChannelPool
+public interface ExtendedChannelPool
 {
+    CompletionStage<Channel> acquire();
+
+    CompletionStage<Void> release( Channel channel );
+
     boolean isClosed();
 
     String id();
 
-    CompletionStage<Void> repeatableCloseAsync();
+    CompletionStage<Void> close();
 }

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

@@ -36,7 +36,7 @@
 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
+public class NettyChannelPool implements ExtendedChannelPool
 {
     /**
      * Unlimited amount of parties are allowed to request channels from the pool.
@@ -47,57 +47,66 @@ public class NettyChannelPool extends FixedChannelPool implements ExtendedChanne
      */
     private static final boolean RELEASE_HEALTH_CHECK = false;
 
-    private final BoltServerAddress address;
-    private final ChannelConnector connector;
-    private final NettyChannelTracker handler;
+    private final FixedChannelPool delegate;
     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,
+    NettyChannelPool( BoltServerAddress address, ChannelConnector connector, Bootstrap bootstrap, NettyChannelTracker handler,
             ChannelHealthChecker healthCheck, long acquireTimeoutMillis, int maxConnections )
     {
-        super( bootstrap, handler, healthCheck, AcquireTimeoutAction.FAIL, acquireTimeoutMillis, maxConnections,
-                MAX_PENDING_ACQUIRES, RELEASE_HEALTH_CHECK );
-
-        this.address = requireNonNull( address );
-        this.connector = requireNonNull( connector );
-        this.handler = requireNonNull( handler );
+        requireNonNull( address );
+        requireNonNull( connector );
+        requireNonNull( handler );
         this.id = poolId( address );
-    }
-
-    @Override
-    protected ChannelFuture connectChannel( Bootstrap bootstrap )
-    {
-        ListenerEvent creatingEvent = handler.channelCreating( this.id );
-        ChannelFuture channelFuture = connector.connect( address, bootstrap );
-        channelFuture.addListener( future ->
+        this.delegate = new FixedChannelPool( bootstrap, handler, healthCheck, FixedChannelPool.AcquireTimeoutAction.FAIL, acquireTimeoutMillis, maxConnections,
+                MAX_PENDING_ACQUIRES, RELEASE_HEALTH_CHECK )
         {
-            if ( future.isSuccess() )
-            {
-                // notify pool handler about a successful connection
-                Channel channel = channelFuture.channel();
-                setPoolId( channel, this.id );
-                handler.channelCreated( channel, creatingEvent );
-            }
-            else
+            @Override
+            protected ChannelFuture connectChannel( Bootstrap bootstrap )
             {
-                handler.channelFailedToCreate( this.id );
+                ListenerEvent creatingEvent = handler.channelCreating( id );
+                ChannelFuture channelFuture = connector.connect( address, bootstrap );
+                channelFuture.addListener( future -> {
+                    if ( future.isSuccess() )
+                    {
+                        // notify pool handler about a successful connection
+                        Channel channel = channelFuture.channel();
+                        setPoolId( channel, id );
+                        handler.channelCreated( channel, creatingEvent );
+                    }
+                    else
+                    {
+                        handler.channelFailedToCreate( id );
+                    }
+                } );
+                return channelFuture;
             }
-        } );
-        return channelFuture;
+        };
     }
 
     @Override
-    public CompletionStage<Void> repeatableCloseAsync()
+    public CompletionStage<Void> close()
     {
         if ( closed.compareAndSet( false, true ) )
         {
-            asCompletionStage( super.closeAsync(), closeFuture );
+            asCompletionStage( delegate.closeAsync(), closeFuture );
         }
         return closeFuture;
     }
 
+    @Override
+    public CompletionStage<Channel> acquire()
+    {
+        return asCompletionStage( delegate.acquire() );
+    }
+
+    @Override
+    public CompletionStage<Void> release( Channel channel )
+    {
+        return asCompletionStage( delegate.release( channel ) );
+    }
+
     @Override
     public boolean isClosed()
     {