Fix missing finish bug, refactor

Author: katcharov

driver-core/src/main/com/mongodb/connection/AsyncCompletionHandler.java

@@ -40,6 +40,9 @@ public interface AsyncCompletionHandler<T> {
      */
     void failed(Throwable t);
 
+    /**
+     * @return this handler as a callback
+     */
     default SingleResultCallback<T> asCallback() {
         return (r, t) -> {
             if (t != null) {

driver-core/src/main/com/mongodb/connection/TlsChannelStreamFactoryFactory.java

@@ -203,7 +203,7 @@ public boolean supportsAdditionalTimeout() {
         @SuppressWarnings("deprecation")
         @Override
         public void openAsync(final SingleResultCallback<Void> callback) {
-            AsyncCompletionHandler<Void> handler = callback.toHandler();
+            AsyncCompletionHandler<Void> handler = callback.asHandler();
             isTrue("unopened", getChannel() == null);
             try {
                 SocketChannel socketChannel = SocketChannel.open();

driver-core/src/main/com/mongodb/connection/netty/NettyStream.java

@@ -166,7 +166,7 @@ public void open() throws IOException {
     @SuppressWarnings("deprecation")
     @Override
     public void openAsync(final SingleResultCallback<Void> callback) {
-        AsyncCompletionHandler<Void> handler = callback.toHandler();
+        AsyncCompletionHandler<Void> handler = callback.asHandler();
         Queue<SocketAddress> socketAddressQueue;
 
         try {

driver-core/src/main/com/mongodb/internal/async/SingleResultCallback.java

@@ -37,7 +37,10 @@ public interface SingleResultCallback<T> {
      */
     void onResult(@Nullable T result, @Nullable Throwable t);
 
-    default AsyncCompletionHandler<T> toHandler() {
+    /**
+     * @return this callback as a handler
+     */
+    default AsyncCompletionHandler<T> asHandler() {
         return new AsyncCompletionHandler<T>() {
             @Override
             public void completed(@Nullable final T result) {

driver-core/src/main/com/mongodb/internal/connection/AsynchronousSocketChannelStream.java

@@ -58,7 +58,7 @@ public AsynchronousSocketChannelStream(final ServerAddress serverAddress, final
     @SuppressWarnings("deprecation")
     @Override
     public void openAsync(final SingleResultCallback<Void> callback) {
-        AsyncCompletionHandler<Void> handler = callback.toHandler();
+        AsyncCompletionHandler<Void> handler = callback.asHandler();
 
         isTrue("unopened", getChannel() == null);
         Queue<SocketAddress> socketAddressQueue;

driver-core/src/main/com/mongodb/internal/connection/InternalStreamConnectionInitializer.java

@@ -109,10 +109,10 @@ public void startHandshakeAsync(final InternalConnection internalConnection,
                 executeCommandAsync("admin", helloCommandDocument, clusterConnectionMode, serverApi, internalConnection, c2);
             }).onErrorIf(e -> e instanceof MongoException, (t, c2) -> {
                 throw mapHelloException((MongoException) t);
-            }).thenApply((helloResult, c2) -> {
+            }).<InternalConnectionInitializationDescription>thenApply((helloResult, c2) -> {
                 setSpeculativeAuthenticateResponse(helloResult);
                 c2.complete(createInitializationDescription(helloResult, internalConnection, startTime));
-            });
+            }).finish(c);
         }).finish(callback);
     }
 

driver-core/src/test/unit/com/mongodb/internal/async/AsyncFunctionsTest.java

@@ -122,8 +122,8 @@ below, SHOULD be used (for consistency, and ease of comparison/review).
         1. Is everything inside the boilerplate?
         2. Is "callback" supplied to "finish"?
         3. In each block and nested block, is that same block's "c" always passed/completed at the end of execution?
-        4. Is any c.complete followed by a return, to end execution?
-        5. Do any sync methods still need to be converted to async?
+        4. Is every c.complete followed by a return, to end execution?
+        5. Have all sync method calls been converted to async, where needed?
         */
     }
 

driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/ClientSessionBinding.java

@@ -104,27 +104,24 @@ public void getWriteConnectionSource(final SingleResultCallback<AsyncConnectionS
 
     private void getConnectionSource(final AsyncSupplier<AsyncConnectionSource> connectionSourceSupplier,
             final SingleResultCallback<AsyncConnectionSource> callback) {
-        // wrapper applied at end
         beginAsync().<AsyncConnectionSource>thenSupply(c -> {
             if (!session.hasActiveTransaction()) {
                 connectionSourceSupplier.getAsync(c);
-                return;
-            }
-            if (TransactionContext.get(session) != null) {
+            } else if (TransactionContext.get(session) != null) {
                 wrapped.getConnectionSource(assertNotNull(session.getPinnedServerAddress()), c);
-                return;
+            } else {
+                beginAsync().<AsyncConnectionSource>thenSupply(c2 -> {
+                    connectionSourceSupplier.getAsync(c2);
+                }).<AsyncConnectionSource>thenApply((source, c2) -> {
+                    ClusterType clusterType = assertNotNull(source).getServerDescription().getClusterType();
+                    if (clusterType == SHARDED || clusterType == LOAD_BALANCED) {
+                        TransactionContext<AsyncConnection> transactionContext = new TransactionContext<>(clusterType);
+                        session.setTransactionContext(source.getServerDescription().getAddress(), transactionContext);
+                        transactionContext.release();  // The session is responsible for retaining a reference to the context
+                    }
+                    c2.complete(source);
+                }).finish(c);
             }
-            beginAsync().<AsyncConnectionSource>thenSupply(c2 -> {
-                connectionSourceSupplier.getAsync(c2);
-            }).<AsyncConnectionSource>thenApply((source, c2) -> {
-                ClusterType clusterType = assertNotNull(source).getServerDescription().getClusterType();
-                if (clusterType == SHARDED || clusterType == LOAD_BALANCED) {
-                    TransactionContext<AsyncConnection> transactionContext = new TransactionContext<>(clusterType);
-                    session.setTransactionContext(source.getServerDescription().getAddress(), transactionContext);
-                    transactionContext.release();  // The session is responsible for retaining a reference to the context
-                } //
-                c2.complete(source);
-            }).finish(c);
         }).<AsyncConnectionSource>thenApply((source, c) -> {
             c.complete(new SessionBindingAsyncConnectionSource(source));
         }).finish(callback);
@@ -194,12 +191,12 @@ public void getConnection(final SingleResultCallback<AsyncConnection> callback)
                 if (transactionContext == null || !transactionContext.isConnectionPinningRequired()) {
                     wrapped.getConnection(c);
                     return;
-                } //
+                }
                 AsyncConnection pinnedAsyncConnection = transactionContext.getPinnedConnection();
                 if (pinnedAsyncConnection != null) {
                     c.complete(pinnedAsyncConnection.retain());
                     return;
-                } //
+                }
                 beginAsync().<AsyncConnection>thenSupply(c2 -> {
                     wrapped.getConnection(c2);
                 }).<AsyncConnection>thenApply((connection, c2) -> {

driver-sync/src/main/com/mongodb/client/internal/ClientSessionBinding.java

@@ -34,7 +34,6 @@
 import com.mongodb.internal.session.SessionContext;
 import com.mongodb.lang.Nullable;
 
-import java.util.function.Function;
 import java.util.function.Supplier;
 
 import static com.mongodb.connection.ClusterType.LOAD_BALANCED;
@@ -126,23 +125,21 @@ public OperationContext getOperationContext() {
     }
 
     private ConnectionSource getConnectionSource(final Supplier<ConnectionSource> connectionSourceSupplier) {
-        Function<ConnectionSource, ConnectionSource> wrapper = c -> new SessionBindingConnectionSource(c);
-
+        ConnectionSource source;
         if (!session.hasActiveTransaction()) {
-            return wrapper.apply(connectionSourceSupplier.get());
-        }
-        if (TransactionContext.get(session) != null) {
-            return wrapper.apply(
-                    wrapped.getConnectionSource(assertNotNull(session.getPinnedServerAddress())));
-        }
-        ConnectionSource source = connectionSourceSupplier.get();
-        ClusterType clusterType = source.getServerDescription().getClusterType();
-        if (clusterType == SHARDED || clusterType == LOAD_BALANCED) {
-            TransactionContext<Connection> transactionContext = new TransactionContext<>(clusterType);
-            session.setTransactionContext(source.getServerDescription().getAddress(), transactionContext);
-            transactionContext.release();  // The session is responsible for retaining a reference to the context
+            source = connectionSourceSupplier.get();
+        } else if (TransactionContext.get(session) != null) {
+            source = wrapped.getConnectionSource(assertNotNull(session.getPinnedServerAddress()));
+        } else {
+            source = connectionSourceSupplier.get();
+            ClusterType clusterType = source.getServerDescription().getClusterType();
+            if (clusterType == SHARDED || clusterType == LOAD_BALANCED) {
+                TransactionContext<Connection> transactionContext = new TransactionContext<>(clusterType);
+                session.setTransactionContext(source.getServerDescription().getAddress(), transactionContext);
+                transactionContext.release();  // The session is responsible for retaining a reference to the context
+            }
         }
-        return wrapper.apply(source);
+        return new SessionBindingConnectionSource(source);
     }
 
     private class SessionBindingConnectionSource implements ConnectionSource {