Merge pull request #414 from lutovich/1.5-driver-close-async

lutovich · web-flow · commit 2aaa75bbb1f1 · 2017-10-05T13:44:32.000+02:00
Expose `Driver#closeAsync()`
diff --git a/driver/src/main/java/org/neo4j/driver/internal/DirectConnectionProvider.java b/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()
diff --git a/driver/src/main/java/org/neo4j/driver/internal/DriverFactory.java b/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 )
             {
diff --git a/driver/src/main/java/org/neo4j/driver/internal/InternalDriver.java b/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() )
diff --git a/driver/src/main/java/org/neo4j/driver/internal/SessionFactory.java b/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();
 }
diff --git a/driver/src/main/java/org/neo4j/driver/internal/SessionFactoryImpl.java b/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();
     }
 
     /**
diff --git a/driver/src/main/java/org/neo4j/driver/internal/async/pool/AsyncConnectionPool.java b/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();
 }
diff --git a/driver/src/main/java/org/neo4j/driver/internal/async/pool/AsyncConnectionPoolImpl.java b/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 )
diff --git a/driver/src/main/java/org/neo4j/driver/internal/cluster/loadbalancing/LoadBalancer.java b/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 )
diff --git a/driver/src/main/java/org/neo4j/driver/internal/spi/ConnectionProvider.java b/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();
 }
diff --git a/driver/src/main/java/org/neo4j/driver/v1/Driver.java b/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();
 }
diff --git a/driver/src/test/java/org/neo4j/driver/internal/InternalDriverTest.java b/driver/src/test/java/org/neo4j/driver/internal/InternalDriverTest.java
@@ -0,0 +1,69 @@
+/*
+ * Copyright (c) 2002-2017 "Neo Technology,"
+ * Network Engine for Objects in Lund AB [http://neotechnology.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.neo4j.driver.internal;
+
+import org.junit.Test;
+
+import org.neo4j.driver.internal.security.SecurityPlan;
+
+import static java.util.concurrent.CompletableFuture.completedFuture;
+import static org.junit.Assert.assertNull;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.neo4j.driver.internal.async.Futures.getBlocking;
+import static org.neo4j.driver.internal.logging.DevNullLogging.DEV_NULL_LOGGING;
+
+public class InternalDriverTest
+{
+    @Test
+    public void shouldCloseSessionFactory()
+    {
+        SessionFactory sessionFactory = sessionFactoryMock();
+        InternalDriver driver = newDriver( sessionFactory );
+
+        assertNull( getBlocking( driver.closeAsync() ) );
+        verify( sessionFactory ).close();
+    }
+
+    @Test
+    public void shouldNotCloseSessionFactoryMultipleTimes()
+    {
+        SessionFactory sessionFactory = sessionFactoryMock();
+        InternalDriver driver = newDriver( sessionFactory );
+
+        assertNull( getBlocking( driver.closeAsync() ) );
+        assertNull( getBlocking( driver.closeAsync() ) );
+        assertNull( getBlocking( driver.closeAsync() ) );
+
+        verify( sessionFactory ).close();
+    }
+
+    private static InternalDriver newDriver( SessionFactory sessionFactory )
+    {
+        return new InternalDriver( SecurityPlan.insecure(), sessionFactory, DEV_NULL_LOGGING );
+    }
+
+    private static SessionFactory sessionFactoryMock()
+    {
+        SessionFactory sessionFactory = mock( SessionFactory.class );
+        when( sessionFactory.close() ).thenReturn( completedFuture( null ) );
+        return sessionFactory;
+    }
+}
diff --git a/driver/src/test/java/org/neo4j/driver/internal/RoutingDriverTest.java b/driver/src/test/java/org/neo4j/driver/internal/RoutingDriverTest.java
@@ -57,6 +57,7 @@
 import org.neo4j.driver.v1.exceptions.ServiceUnavailableException;
 
 import static java.util.Arrays.asList;
+import static java.util.concurrent.CompletableFuture.completedFuture;
 import static junit.framework.TestCase.fail;
 import static org.hamcrest.Matchers.containsString;
 import static org.junit.Assert.assertEquals;
@@ -364,6 +365,7 @@ private Driver driverWithPool( ConnectionPool pool )
         Logging logging = DEV_NULL_LOGGING;
         RoutingSettings settings = new RoutingSettings( 10, 5_000, null );
         AsyncConnectionPool asyncConnectionPool = mock( AsyncConnectionPool.class );
+        when( asyncConnectionPool.close() ).thenReturn( completedFuture( null ) );
         LoadBalancingStrategy loadBalancingStrategy = new LeastConnectedLoadBalancingStrategy( pool,
                 asyncConnectionPool, logging );
         ConnectionProvider connectionProvider = new LoadBalancer( SEED, settings, pool, asyncConnectionPool,
diff --git a/driver/src/test/java/org/neo4j/driver/internal/async/pool/AsyncConnectionPoolImplTest.java b/driver/src/test/java/org/neo4j/driver/internal/async/pool/AsyncConnectionPoolImplTest.java
@@ -63,7 +63,7 @@ public void setUp() throws Exception
     @After
     public void tearDown() throws Exception
     {
-        pool.closeAsync();
+        pool.close();
     }
 
     @Test
@@ -104,7 +104,7 @@ public void shouldFailToAcquireWhenPoolClosed() throws Exception
     {
         AsyncConnection connection = await( pool.acquire( neo4j.address() ) );
         await( connection.forceRelease() );
-        await( pool.closeAsync() );
+        await( pool.close() );
 
         try
         {
@@ -159,8 +159,8 @@ public void shouldCheckIfPoolHasAddress()
     @Test
     public void shouldNotCloseWhenClosed()
     {
-        assertNull( await( pool.closeAsync() ) );
-        assertTrue( pool.closeAsync().toCompletableFuture().isDone() );
+        assertNull( await( pool.close() ) );
+        assertTrue( pool.close().toCompletableFuture().isDone() );
     }
 
     private AsyncConnectionPoolImpl newPool() throws Exception

Original file line number	Diff line number	Diff line change
`@@ -92,7 +92,7 @@ public final Driver newInstance( URI uri, AuthToken authToken, RoutingSettings r`
`92`	`92`	`try`
`93`	`93`	`{`
`94`	`94`	`connectionPool.close();`
`95`		`- Futures.getBlocking( asyncConnectionPool.closeAsync() );`
	`95`	`+ Futures.getBlocking( asyncConnectionPool.close() );`
`96`	`96`	`}`
`97`	`97`	`catch ( Throwable closeError )`
`98`	`98`	`{`
Original file line number	Diff line number	Diff line change
`@@ -18,10 +18,14 @@`
`18`	`18`	`*/`
`19`	`19`	`package org.neo4j.driver.internal;`
`20`	`20`
	`21`	`+import java.util.concurrent.CompletionStage;`
	`22`	`+`
`21`	`23`	`import org.neo4j.driver.v1.AccessMode;`
`22`	`24`	`import org.neo4j.driver.v1.Session;`
`23`	`25`
`24`		`-public interface SessionFactory extends AutoCloseable`
	`26`	`+public interface SessionFactory`
`25`	`27`	`{`
`26`	`28`	`Session newInstance( AccessMode mode, Bookmark bookmark );`
	`29`	`+`
	`30`	`+ CompletionStage<Void> close();`
`27`	`31`	`}`
Original file line number	Diff line number	Diff line change
`@@ -18,6 +18,8 @@`
`18`	`18`	`*/`
`19`	`19`	`package org.neo4j.driver.internal;`
`20`	`20`
	`21`	`+import java.util.concurrent.CompletionStage;`
	`22`	`+`
`21`	`23`	`import org.neo4j.driver.internal.retry.RetryLogic;`
`22`	`24`	`import org.neo4j.driver.internal.spi.ConnectionProvider;`
`23`	`25`	`import org.neo4j.driver.v1.AccessMode;`
`@@ -57,9 +59,9 @@ protected NetworkSession createSession( ConnectionProvider connectionProvider, R`
`57`	`59`	`}`
`58`	`60`
`59`	`61`	`@Override`
`60`		`- public final void close() throws Exception`
	`62`	`+ public final CompletionStage<Void> close()`
`61`	`63`	`{`
`62`		`- connectionProvider.close();`
	`64`	`+ return connectionProvider.close();`
`63`	`65`	`}`
`64`	`66`
`65`	`67`	`/**`
Original file line number	Diff line number	Diff line change
`@@ -33,5 +33,5 @@ public interface AsyncConnectionPool`
`33`	`33`
`34`	`34`	`int activeConnections( BoltServerAddress address );`
`35`	`35`
`36`		`- CompletionStage<?> closeAsync();`
	`36`	`+ CompletionStage<Void> close();`
`37`	`37`	`}`
Original file line number	Diff line number	Diff line change
`@@ -109,7 +109,7 @@ public int activeConnections( BoltServerAddress address )`
`109`	`109`	`}`
`110`	`110`
`111`	`111`	`@Override`
`112`		`- public CompletionStage<?> closeAsync()`
	`112`	`+ public CompletionStage<Void> close()`
`113`	`113`	`{`
`114`	`114`	`if ( closed.compareAndSet( false, true ) )`
`115`	`115`	`{`
`@@ -128,7 +128,8 @@ public CompletionStage<?> closeAsync()`
`128`	`128`	`eventLoopGroup().shutdownGracefully();`
`129`	`129`	`}`
`130`	`130`	`}`
`131`		`- return Futures.asCompletionStage( eventLoopGroup().terminationFuture() );`
	`131`	`+ return Futures.asCompletionStage( eventLoopGroup().terminationFuture() )`
	`132`	`+ .thenApply( ignore -> null );`
`132`	`133`	`}`
`133`	`134`
`134`	`135`	`private ChannelPool getOrCreatePool( BoltServerAddress address )`
Original file line number	Diff line number	Diff line change
`@@ -27,7 +27,7 @@`
`27`	`27`	`* Interface defines a layer used by the driver to obtain connections. It is meant to be the only component that`
`28`	`28`	`* differs between "direct" and "routing" driver.`
`29`	`29`	`*/`
`30`		`-public interface ConnectionProvider extends AutoCloseable`
	`30`	`+public interface ConnectionProvider`
`31`	`31`	`{`
`32`	`32`	`/**`
`33`	`33`	`* Acquire new {@link PooledConnection pooled connection} for the given {@link AccessMode mode}.`
`@@ -38,4 +38,6 @@ public interface ConnectionProvider extends AutoCloseable`
`38`	`38`	`PooledConnection acquireConnection( AccessMode mode );`
`39`	`39`
`40`	`40`	`CompletionStage<AsyncConnection> acquireAsyncConnection( AccessMode mode );`
	`41`	`+`
	`42`	`+ CompletionStage<Void> close();`
`41`	`43`	`}`