Fix Bolt handshake write handling and timeout management

Author: injectives

driver/src/main/java/org/neo4j/driver/internal/async/connection/ChannelConnectedListener.java

@@ -17,7 +17,6 @@
 package org.neo4j.driver.internal.async.connection;
 
 import static java.lang.String.format;
-import static org.neo4j.driver.internal.async.connection.BoltProtocolUtil.handshakeBuf;
 import static org.neo4j.driver.internal.async.connection.BoltProtocolUtil.handshakeString;
 
 import io.netty.channel.ChannelFuture;
@@ -57,7 +56,12 @@ public void operationComplete(ChannelFuture future) {
             var pipeline = channel.pipeline();
             pipeline.addLast(new HandshakeHandler(pipelineBuilder, handshakeCompletedPromise, logging));
             log.debug("C: [Bolt Handshake] %s", handshakeString());
-            channel.writeAndFlush(handshakeBuf(), channel.voidPromise());
+            channel.writeAndFlush(BoltProtocolUtil.handshakeBuf()).addListener(f -> {
+                if (!f.isSuccess()) {
+                    this.handshakeCompletedPromise.setFailure(new ServiceUnavailableException(
+                            String.format("Unable to write Bolt handshake to %s.", this.address), f.cause()));
+                }
+            });
         } else {
             handshakeCompletedPromise.setFailure(databaseUnavailableError(address, future.cause()));
         }

driver/src/main/java/org/neo4j/driver/internal/async/connection/ChannelConnectorImpl.java

@@ -142,7 +142,11 @@ private void installHandshakeCompletedListeners(
 
         // remove timeout handler from the pipeline once TLS and Bolt handshakes are completed. regular protocol
         // messages will flow next and we do not want to have read timeout for them
-        handshakeCompleted.addListener(future -> pipeline.remove(ConnectTimeoutHandler.class));
+        handshakeCompleted.addListener(future -> {
+            if (future.isSuccess()) {
+                pipeline.remove(ConnectTimeoutHandler.class);
+            }
+        });
 
         // add listener that sends an INIT message. connection is now fully established. channel pipeline if fully
         // set to send/receive messages for a selected protocol version

driver/src/test/java/org/neo4j/driver/internal/async/connection/ChannelConnectedListenerTest.java

@@ -17,6 +17,8 @@
 package org.neo4j.driver.internal.async.connection;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -27,7 +29,9 @@
 
 import io.netty.channel.ChannelPromise;
 import io.netty.channel.embedded.EmbeddedChannel;
+import io.netty.util.concurrent.Future;
 import java.io.IOException;
+import java.util.concurrent.CompletableFuture;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 import org.neo4j.driver.exceptions.ServiceUnavailableException;
@@ -70,6 +74,25 @@ void shouldWriteHandshakeWhenChannelConnected() {
         assertEquals(handshakeBuf(), channel.readOutbound());
     }
 
+    @Test
+    void shouldCompleteHandshakePromiseExceptionallyOnWriteFailure() {
+        var handshakeCompletedPromise = channel.newPromise();
+        var listener = newListener(handshakeCompletedPromise);
+        var channelConnectedPromise = channel.newPromise();
+        channelConnectedPromise.setSuccess();
+        channel.close();
+
+        listener.operationComplete(channelConnectedPromise);
+
+        assertTrue(handshakeCompletedPromise.isDone());
+        var future = new CompletableFuture<Future<?>>();
+        handshakeCompletedPromise.addListener(future::complete);
+        var handshakeFuture = future.join();
+        assertTrue(handshakeFuture.isDone());
+        assertFalse(handshakeFuture.isSuccess());
+        assertInstanceOf(ServiceUnavailableException.class, handshakeFuture.cause());
+    }
+
     private static ChannelConnectedListener newListener(ChannelPromise handshakeCompletedPromise) {
         return new ChannelConnectedListener(
                 LOCAL_DEFAULT, new ChannelPipelineBuilderImpl(), handshakeCompletedPromise, DEV_NULL_LOGGING);