Use async API in InternalStreamConnection

katcharov · katcharov · commit 35061d0cc481 · 2023-10-26T15:30:00.000-06:00
diff --git a/driver-core/src/main/com/mongodb/connection/AsyncCompletionHandler.java b/driver-core/src/main/com/mongodb/connection/AsyncCompletionHandler.java
@@ -16,6 +16,7 @@
 
 package com.mongodb.connection;
 
+import com.mongodb.internal.async.SingleResultCallback;
 import com.mongodb.lang.Nullable;
 
 /**
@@ -38,4 +39,14 @@ public interface AsyncCompletionHandler<T> {
      * @param t the exception that describes the failure
      */
     void failed(Throwable t);
+
+    default SingleResultCallback<T> asCallback() {
+        return (r, t) -> {
+            if (t != null) {
+                failed(t);
+            } else {
+                completed(r);
+            }
+        };
+    }
 }
diff --git a/driver-core/src/main/com/mongodb/connection/Stream.java b/driver-core/src/main/com/mongodb/connection/Stream.java
@@ -17,6 +17,7 @@
 package com.mongodb.connection;
 
 import com.mongodb.ServerAddress;
+import com.mongodb.internal.async.SingleResultCallback;
 import org.bson.ByteBuf;
 
 import java.io.IOException;
@@ -43,7 +44,7 @@ public interface Stream extends BufferProvider{
      *
      * @param handler the completion handler for opening the stream
      */
-    void openAsync(AsyncCompletionHandler<Void> handler);
+    void openAsync(SingleResultCallback<Void> handler);
 
     /**
      * Write each buffer in the list to the stream in order, blocking until all are completely written.
diff --git a/driver-core/src/main/com/mongodb/connection/TlsChannelStreamFactoryFactory.java b/driver-core/src/main/com/mongodb/connection/TlsChannelStreamFactoryFactory.java
@@ -19,6 +19,7 @@
 import com.mongodb.MongoClientException;
 import com.mongodb.MongoSocketOpenException;
 import com.mongodb.ServerAddress;
+import com.mongodb.internal.async.SingleResultCallback;
 import com.mongodb.internal.connection.AsynchronousChannelStream;
 import com.mongodb.internal.connection.ExtendedAsynchronousByteChannel;
 import com.mongodb.internal.connection.PowerOfTwoBufferPool;
@@ -201,7 +202,8 @@ public boolean supportsAdditionalTimeout() {
 
         @SuppressWarnings("deprecation")
         @Override
-        public void openAsync(final AsyncCompletionHandler<Void> handler) {
+        public void openAsync(final SingleResultCallback<Void> callback) {
+            AsyncCompletionHandler<Void> handler = callback.toHandler();
             isTrue("unopened", getChannel() == null);
             try {
                 SocketChannel socketChannel = SocketChannel.open();
diff --git a/driver-core/src/main/com/mongodb/connection/netty/NettyStream.java b/driver-core/src/main/com/mongodb/connection/netty/NettyStream.java
@@ -28,6 +28,7 @@
 import com.mongodb.connection.SocketSettings;
 import com.mongodb.connection.SslSettings;
 import com.mongodb.connection.Stream;
+import com.mongodb.internal.async.SingleResultCallback;
 import com.mongodb.internal.connection.netty.NettyByteBuf;
 import com.mongodb.lang.Nullable;
 import io.netty.bootstrap.Bootstrap;
@@ -158,13 +159,14 @@ public ByteBuf getBuffer(final int size) {
     @Override
     public void open() throws IOException {
         FutureAsyncCompletionHandler<Void> handler = new FutureAsyncCompletionHandler<>();
-        openAsync(handler);
+        openAsync(handler.asCallback());
         handler.get();
     }
 
     @SuppressWarnings("deprecation")
     @Override
-    public void openAsync(final AsyncCompletionHandler<Void> handler) {
+    public void openAsync(final SingleResultCallback<Void> callback) {
+        AsyncCompletionHandler<Void> handler = callback.toHandler();
         Queue<SocketAddress> socketAddressQueue;
 
         try {
diff --git a/driver-core/src/main/com/mongodb/internal/connection/AsynchronousChannelStream.java b/driver-core/src/main/com/mongodb/internal/connection/AsynchronousChannelStream.java
@@ -127,7 +127,7 @@ private void readAsync(final int numBytes, final int additionalTimeout, final As
     @Override
     public void open() throws IOException {
         FutureAsyncCompletionHandler<Void> handler = new FutureAsyncCompletionHandler<>();
-        openAsync(handler);
+        openAsync(handler.asCallback());
         handler.getOpen();
     }
 
diff --git a/driver-core/src/main/com/mongodb/internal/connection/AsynchronousSocketChannelStream.java b/driver-core/src/main/com/mongodb/internal/connection/AsynchronousSocketChannelStream.java
@@ -21,6 +21,7 @@
 import com.mongodb.ServerAddress;
 import com.mongodb.connection.AsyncCompletionHandler;
 import com.mongodb.connection.SocketSettings;
+import com.mongodb.internal.async.SingleResultCallback;
 import com.mongodb.lang.Nullable;
 
 import java.io.IOException;
@@ -56,7 +57,9 @@ public AsynchronousSocketChannelStream(final ServerAddress serverAddress, final
 
     @SuppressWarnings("deprecation")
     @Override
-    public void openAsync(final AsyncCompletionHandler<Void> handler) {
+    public void openAsync(final SingleResultCallback<Void> callback) {
+        AsyncCompletionHandler<Void> handler = callback.toHandler();
+
         isTrue("unopened", getChannel() == null);
         Queue<SocketAddress> socketAddressQueue;
 
diff --git a/driver-core/src/main/com/mongodb/internal/connection/InternalStreamConnection.java b/driver-core/src/main/com/mongodb/internal/connection/InternalStreamConnection.java
@@ -72,6 +72,7 @@
 import static com.mongodb.assertions.Assertions.assertNotNull;
 import static com.mongodb.assertions.Assertions.isTrue;
 import static com.mongodb.assertions.Assertions.notNull;
+import static com.mongodb.internal.async.AsyncRunnable.beginAsync;
 import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback;
 import static com.mongodb.internal.connection.CommandHelper.HELLO;
 import static com.mongodb.internal.connection.CommandHelper.LEGACY_HELLO;
@@ -205,15 +206,16 @@ public int getGeneration() {
 
     @Override
     public void open() {
-        isTrue("Open already called", stream == null);
-        stream = streamFactory.create(getServerAddressWithResolver());
         try {
+            isTrue("Open already called", stream == null);
+            stream = streamFactory.create(getServerAddressWithResolver());
             stream.open();
 
             InternalConnectionInitializationDescription initializationDescription = connectionInitializer.startHandshake(this);
-            initAfterHandshakeStart(initializationDescription);
 
+            initAfterHandshakeStart(initializationDescription);
             initializationDescription = connectionInitializer.finishHandshake(this, initializationDescription);
+
             initAfterHandshakeFinish(initializationDescription);
         } catch (Throwable t) {
             close();
@@ -227,45 +229,27 @@ public void open() {
 
     @Override
     public void openAsync(final SingleResultCallback<Void> callback) {
-        isTrue("Open already called", stream == null, callback);
-        try {
-            stream = streamFactory.create(getServerAddressWithResolver());
-            stream.openAsync(new AsyncCompletionHandler<Void>() {
-                @Override
-                public void completed(@Nullable final Void aVoid) {
-                    connectionInitializer.startHandshakeAsync(InternalStreamConnection.this,
-                            (initialResult, initialException) -> {
-                                    if (initialException != null) {
-                                        close();
-                                        callback.onResult(null, initialException);
-                                    } else {
-                                        assertNotNull(initialResult);
-                                        initAfterHandshakeStart(initialResult);
-                                        connectionInitializer.finishHandshakeAsync(InternalStreamConnection.this,
-                                                initialResult, (completedResult, completedException) ->  {
-                                                        if (completedException != null) {
-                                                            close();
-                                                            callback.onResult(null, completedException);
-                                                        } else {
-                                                            assertNotNull(completedResult);
-                                                            initAfterHandshakeFinish(completedResult);
-                                                            callback.onResult(null, null);
-                                                        }
-                                                });
-                                    }
-                            });
-                }
-
-                @Override
-                public void failed(final Throwable t) {
-                    close();
-                    callback.onResult(null, t);
-                }
-            });
-        } catch (Throwable t) {
+        beginAsync().thenRun(
+            beginAsync().thenRun(c -> {
+                isTrue("Open already called", stream == null, callback);
+                stream = streamFactory.create(getServerAddressWithResolver());
+                stream.openAsync(c);
+            }).<InternalConnectionInitializationDescription>thenSupply(c -> {
+                connectionInitializer.startHandshakeAsync(this, c);
+            }).<InternalConnectionInitializationDescription>thenApply((initializationDescription, c)  -> {
+                initAfterHandshakeStart(assertNotNull(initializationDescription));
+                connectionInitializer.finishHandshakeAsync(this, initializationDescription, c);
+            }).thenConsume((initializationDescription, c)  -> {
+                initAfterHandshakeFinish(assertNotNull(initializationDescription));
+            })
+        ).onErrorIf(t -> true, (t, c) -> {
             close();
-            callback.onResult(null, t);
-        }
+            if (t instanceof MongoException) {
+                throw (MongoException) t;
+            } else {
+                throw new MongoException(assertNotNull(t).toString(), t);
+            }
+        }).finish(callback);
     }
 
     private ServerAddress getServerAddressWithResolver() {
@@ -336,7 +320,7 @@ private Compressor createCompressor(final MongoCompressor mongoCompressor) {
     public void close() {
         // All but the first call is a no-op
         if (!isClosed.getAndSet(true) && (stream != null)) {
-                stream.close();
+            stream.close();
         }
     }
 
@@ -352,8 +336,9 @@ public boolean isClosed() {
 
     @Nullable
     @Override
-    public <T> T sendAndReceive(final CommandMessage message, final Decoder<T> decoder, final SessionContext sessionContext,
-                                final RequestContext requestContext, final OperationContext operationContext) {
+    public <T> T sendAndReceive(final CommandMessage message,
+            final Decoder<T> decoder, final SessionContext sessionContext,
+            final RequestContext requestContext, final OperationContext operationContext) {
         CommandEventSender commandEventSender;
 
         try (ByteBufferBsonOutput bsonOutput = new ByteBufferBsonOutput(this)) {
@@ -476,8 +461,10 @@ private <T> T receiveCommandMessageResponse(final Decoder<T> decoder,
     }
 
     @Override
-    public <T> void sendAndReceiveAsync(final CommandMessage message, final Decoder<T> decoder, final SessionContext sessionContext,
-            final RequestContext requestContext, final OperationContext operationContext, final SingleResultCallback<T> callback) {
+    public <T> void sendAndReceiveAsync(final CommandMessage message,
+            final Decoder<T> decoder, final SessionContext sessionContext,
+            final RequestContext requestContext, final OperationContext operationContext,
+            final SingleResultCallback<T> callback) {
         notNull("stream is open", stream, callback);
 
         if (isClosed()) {
@@ -575,11 +562,9 @@ private <T> T getCommandResult(final Decoder<T> decoder, final ResponseBuffers r
     @Override
     public void sendMessage(final List<ByteBuf> byteBuffers, final int lastRequestId) {
         notNull("stream is open", stream);
-
         if (isClosed()) {
             throw new MongoSocketClosedException("Cannot write to a closed stream", getServerAddress());
         }
-
         try {
             stream.write(byteBuffers);
         } catch (Exception e) {
@@ -609,14 +594,14 @@ private ResponseBuffers receiveMessageWithAdditionalTimeout(final int additional
 
     @Override
     public void sendMessageAsync(final List<ByteBuf> byteBuffers, final int lastRequestId, final SingleResultCallback<Void> callback) {
-        notNull("stream is open", stream, callback);
-
-        if (isClosed()) {
-            callback.onResult(null, new MongoSocketClosedException("Can not read from a closed socket", getServerAddress()));
-            return;
-        }
+        beginAsync().thenRun(c -> {
+            notNull("stream is open", stream, callback);
+            if (isClosed()) {
+                throw new MongoSocketClosedException("Can not read from a closed socket", getServerAddress());
+            }
 
-        writeAsync(byteBuffers, errorHandlingCallback(callback, LOGGER));
+            writeAsync(byteBuffers, errorHandlingCallback(callback, LOGGER));
+        }).finish(callback);
     }
 
     private void writeAsync(final List<ByteBuf> byteBuffers, final SingleResultCallback<Void> callback) {
diff --git a/driver-core/src/main/com/mongodb/internal/connection/SocketStream.java b/driver-core/src/main/com/mongodb/internal/connection/SocketStream.java
@@ -26,6 +26,7 @@
 import com.mongodb.connection.SocketSettings;
 import com.mongodb.connection.SslSettings;
 import com.mongodb.connection.Stream;
+import com.mongodb.internal.async.SingleResultCallback;
 import org.bson.ByteBuf;
 
 import javax.net.SocketFactory;
@@ -207,7 +208,7 @@ public ByteBuf read(final int numBytes, final int additionalTimeout) throws IOEx
     }
 
     @Override
-    public void openAsync(final AsyncCompletionHandler<Void> handler) {
+    public void openAsync(final SingleResultCallback<Void> callback) {
         throw new UnsupportedOperationException(getClass() + " does not support asynchronous operations.");
     }
 
diff --git a/driver-core/src/test/functional/com/mongodb/connection/netty/NettyStreamSpecification.groovy b/driver-core/src/test/functional/com/mongodb/connection/netty/NettyStreamSpecification.groovy
@@ -76,7 +76,7 @@ class NettyStreamSpecification extends Specification {
 
         def stream = new NettyStreamFactory(SocketSettings.builder().connectTimeout(1000, TimeUnit.MILLISECONDS).build(),
                 SslSettings.builder().build()).create(serverAddress)
-        def callback = new CallbackErrorHolder()
+        def callback = new CallbackErrorHolder().asCallback()
 
         when:
         stream.openAsync(callback)
diff --git a/driver-core/src/test/functional/com/mongodb/internal/connection/AsyncSocketChannelStreamSpecification.groovy b/driver-core/src/test/functional/com/mongodb/internal/connection/AsyncSocketChannelStreamSpecification.groovy
@@ -83,7 +83,7 @@ class AsyncSocketChannelStreamSpecification extends Specification {
         def stream = new AsynchronousSocketChannelStream(serverAddress,
                 SocketSettings.builder().connectTimeout(100, MILLISECONDS).build(),
                 new PowerOfTwoBufferPool(), null)
-        def callback = new CallbackErrorHolder()
+        def callback = new CallbackErrorHolder().asCallback()
 
         when:
         stream.openAsync(callback)
diff --git a/driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/crypt/KeyManagementService.java b/driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/crypt/KeyManagementService.java
@@ -85,7 +85,7 @@ public void failed(final Throwable t) {
                     stream.close();
                     sink.error(t);
                 }
-            });
+            }.asCallback());
         }).onErrorMap(this::unWrapException);
     }
 

Original file line number	Diff line number	Diff line change
`@@ -127,7 +127,7 @@ private void readAsync(final int numBytes, final int additionalTimeout, final As`
`127`	`127`	`@Override`
`128`	`128`	`public void open() throws IOException {`
`129`	`129`	`FutureAsyncCompletionHandler<Void> handler = new FutureAsyncCompletionHandler<>();`
`130`		`- openAsync(handler);`
	`130`	`+ openAsync(handler.asCallback());`
`131`	`131`	`handler.getOpen();`
`132`	`132`	`}`
`133`	`133`
Original file line number	Diff line number	Diff line change
`@@ -26,6 +26,7 @@`
`26`	`26`	`import com.mongodb.connection.SocketSettings;`
`27`	`27`	`import com.mongodb.connection.SslSettings;`
`28`	`28`	`import com.mongodb.connection.Stream;`
	`29`	`+import com.mongodb.internal.async.SingleResultCallback;`
`29`	`30`	`import org.bson.ByteBuf;`
`30`	`31`
`31`	`32`	`import javax.net.SocketFactory;`
`@@ -207,7 +208,7 @@ public ByteBuf read(final int numBytes, final int additionalTimeout) throws IOEx`
`207`	`208`	`}`
`208`	`209`
`209`	`210`	`@Override`
`210`		`- public void openAsync(final AsyncCompletionHandler<Void> handler) {`
	`211`	`+ public void openAsync(final SingleResultCallback<Void> callback) {`
`211`	`212`	`throw new UnsupportedOperationException(getClass() + " does not support asynchronous operations.");`
`212`	`213`	`}`
`213`	`214`
Original file line number	Diff line number	Diff line change
`@@ -85,7 +85,7 @@ public void failed(final Throwable t) {`
`85`	`85`	`stream.close();`
`86`	`86`	`sink.error(t);`
`87`	`87`	`}`
`88`		`- });`
	`88`	`+ }.asCallback());`
`89`	`89`	`}).onErrorMap(this::unWrapException);`
`90`	`90`	`}`
`91`	`91`