Expose Driver#closeAsync()

To allow disposing of the driver and it's resources (including network
connections, Netty event loop threads, etc.) in an async manner.

Author: lutovich

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

@@ -21,7 +21,6 @@
 import java.util.concurrent.CompletionStage;
 
 import org.neo4j.driver.internal.async.AsyncConnection;
-import org.neo4j.driver.internal.async.Futures;
 import org.neo4j.driver.internal.async.pool.AsyncConnectionPool;
 import org.neo4j.driver.internal.net.BoltServerAddress;
 import org.neo4j.driver.internal.spi.ConnectionPool;
@@ -61,10 +60,18 @@ public CompletionStage<AsyncConnection> acquireAsyncConnection( AccessMode mode
     }
 
     @Override
-    public void close() throws Exception
+    public CompletionStage<Void> close()
     {
-        pool.close();
-        Futures.getBlocking( asyncPool.closeAsync() );
+        // todo: remove this try-catch when blocking API works on top of async
+        try
+        {
+            pool.close();
+        }
+        catch ( Exception e )
+        {
+            throw new RuntimeException( e );
+        }
+        return asyncPool.close();
     }
 
     public BoltServerAddress getAddress()

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

@@ -92,7 +92,7 @@ public final Driver newInstance( URI uri, AuthToken authToken, RoutingSettings r
             try
             {
                 connectionPool.close();
-                Futures.getBlocking( asyncConnectionPool.closeAsync() );
+                Futures.getBlocking( asyncConnectionPool.close() );
             }
             catch ( Throwable closeError )
             {

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

@@ -18,16 +18,18 @@
  */
 package org.neo4j.driver.internal;
 
+import java.util.concurrent.CompletionStage;
 import java.util.concurrent.atomic.AtomicBoolean;
 
+import org.neo4j.driver.internal.async.Futures;
 import org.neo4j.driver.internal.security.SecurityPlan;
 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.lang.String.format;
+import static java.util.concurrent.CompletableFuture.completedFuture;
 
 public class InternalDriver implements Driver
 {
@@ -95,23 +97,26 @@ private Session newSession( AccessMode mode, Bookmark bookmark )
         Session session = sessionFactory.newInstance( mode, bookmark );
         if ( closed.get() )
         {
-            // the driver is already closed and we either 1. obtain this session from the old session pool
-            // or 2. we obtain this session from a new session pool
-            // For 1. this closeResources will take no effect as everything is already closed.
-            // For 2. this closeResources will close the new connection pool just created to ensure no resource leak.
-            closeResources();
+            // session does not immediately acquire connection, it is fine to just throw
             throw driverCloseException();
         }
         return session;
     }
 
     @Override
     public final void close()
+    {
+        Futures.getBlocking( closeAsync() );
+    }
+
+    @Override
+    public CompletionStage<Void> closeAsync()
     {
         if ( closed.compareAndSet( false, true ) )
         {
-            closeResources();
+            return sessionFactory.close();
         }
+        return completedFuture( null );
     }
 
     /**
@@ -126,18 +131,6 @@ public final SessionFactory getSessionFactory()
         return sessionFactory;
     }
 
-    private void closeResources()
-    {
-        try
-        {
-            sessionFactory.close();
-        }
-        catch ( Exception ex )
-        {
-            log.error( format( "~~ [ERROR] %s", ex.getMessage() ), ex );
-        }
-    }
-
     private void assertOpen()
     {
         if ( closed.get() )

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

@@ -18,10 +18,14 @@
  */
 package org.neo4j.driver.internal;
 
+import java.util.concurrent.CompletionStage;
+
 import org.neo4j.driver.v1.AccessMode;
 import org.neo4j.driver.v1.Session;
 
-public interface SessionFactory extends AutoCloseable
+public interface SessionFactory
 {
     Session newInstance( AccessMode mode, Bookmark bookmark );
+
+    CompletionStage<Void> close();
 }

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

@@ -18,6 +18,8 @@
  */
 package org.neo4j.driver.internal;
 
+import java.util.concurrent.CompletionStage;
+
 import org.neo4j.driver.internal.retry.RetryLogic;
 import org.neo4j.driver.internal.spi.ConnectionProvider;
 import org.neo4j.driver.v1.AccessMode;
@@ -57,9 +59,9 @@ protected NetworkSession createSession( ConnectionProvider connectionProvider, R
     }
 
     @Override
-    public final void close() throws Exception
+    public final CompletionStage<Void> close()
     {
-        connectionProvider.close();
+        return connectionProvider.close();
     }
 
     /**

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

@@ -33,5 +33,5 @@ public interface AsyncConnectionPool
 
     int activeConnections( BoltServerAddress address );
 
-    CompletionStage<?> closeAsync();
+    CompletionStage<Void> close();
 }

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

@@ -109,7 +109,7 @@ public int activeConnections( BoltServerAddress address )
     }
 
     @Override
-    public CompletionStage<?> closeAsync()
+    public CompletionStage<Void> close()
     {
         if ( closed.compareAndSet( false, true ) )
         {
@@ -128,7 +128,8 @@ public CompletionStage<?> closeAsync()
                 eventLoopGroup().shutdownGracefully();
             }
         }
-        return Futures.asCompletionStage( eventLoopGroup().terminationFuture() );
+        return Futures.asCompletionStage( eventLoopGroup().terminationFuture() )
+                .thenApply( ignore -> null );
     }
 
     private ChannelPool getOrCreatePool( BoltServerAddress address )

driver/src/main/java/org/neo4j/driver/internal/cluster/loadbalancing/LoadBalancer.java

@@ -26,7 +26,6 @@
 
 import org.neo4j.driver.internal.RoutingErrorHandler;
 import org.neo4j.driver.internal.async.AsyncConnection;
-import org.neo4j.driver.internal.async.Futures;
 import org.neo4j.driver.internal.async.RoutingAsyncConnection;
 import org.neo4j.driver.internal.async.pool.AsyncConnectionPool;
 import org.neo4j.driver.internal.cluster.AddressSet;
@@ -52,7 +51,7 @@
 
 import static java.util.concurrent.CompletableFuture.completedFuture;
 
-public class LoadBalancer implements ConnectionProvider, RoutingErrorHandler, AutoCloseable
+public class LoadBalancer implements ConnectionProvider, RoutingErrorHandler
 {
     private static final String LOAD_BALANCER_LOG_NAME = "LoadBalancer";
 
@@ -131,10 +130,18 @@ public void onWriteFailure( BoltServerAddress address )
     }
 
     @Override
-    public void close() throws Exception
+    public CompletionStage<Void> close()
     {
-        connections.close();
-        Futures.getBlocking( asyncConnectionPool.closeAsync() );
+        try
+        {
+            connections.close();
+        }
+        catch ( Exception e )
+        {
+            throw new RuntimeException( e );
+        }
+
+        return asyncConnectionPool.close();
     }
 
     private PooledConnection acquireConnection( AccessMode mode, AddressSet servers )

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

@@ -27,7 +27,7 @@
  * Interface defines a layer used by the driver to obtain connections. It is meant to be the only component that
  * differs between "direct" and "routing" driver.
  */
-public interface ConnectionProvider extends AutoCloseable
+public interface ConnectionProvider
 {
     /**
      * Acquire new {@link PooledConnection pooled connection} for the given {@link AccessMode mode}.
@@ -38,4 +38,6 @@ public interface ConnectionProvider extends AutoCloseable
     PooledConnection acquireConnection( AccessMode mode );
 
     CompletionStage<AsyncConnection> acquireAsyncConnection( AccessMode mode );
+
+    CompletionStage<Void> close();
 }

driver/src/main/java/org/neo4j/driver/v1/Driver.java

@@ -18,6 +18,8 @@
  */
 package org.neo4j.driver.v1;
 
+import java.util.concurrent.CompletionStage;
+
 /**
  * Accessor for a specific Neo4j graph database.
  * <p>
@@ -140,4 +142,6 @@ public interface Driver extends AutoCloseable
      * Close all the resources assigned to this driver, including any open connections.
      */
     void close();
+
+    CompletionStage<Void> closeAsync();
 }