[client] Fix exception occurs when initializing server connection. (#1740)

Author: loserwang1024

fluss-rpc/src/main/java/org/apache/fluss/rpc/netty/client/NettyClient.java

@@ -192,15 +192,13 @@ private ServerConnection getOrCreateConnection(ServerNode node) {
                 serverId,
                 ignored -> {
                     LOG.debug("Creating connection to server {}.", node);
-                    ServerConnection connection =
-                            new ServerConnection(
-                                    bootstrap,
-                                    node,
-                                    clientMetricGroup,
-                                    authenticatorSupplier.get(),
-                                    isInnerClient);
-                    connection.whenClose(ignore -> connections.remove(serverId, connection));
-                    return connection;
+                    return new ServerConnection(
+                            bootstrap,
+                            node,
+                            clientMetricGroup,
+                            authenticatorSupplier.get(),
+                            (con, ignore) -> connections.remove(serverId, con),
+                            isInnerClient);
                 });
     }
 

fluss-rpc/src/main/java/org/apache/fluss/rpc/netty/client/ServerConnection.java

@@ -58,7 +58,7 @@
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
-import java.util.function.Consumer;
+import java.util.function.BiConsumer;
 
 import static org.apache.fluss.utils.IOUtils.closeQuietly;
 
@@ -104,15 +104,20 @@ final class ServerConnection {
             ServerNode node,
             ClientMetricGroup clientMetricGroup,
             ClientAuthenticator authenticator,
+            BiConsumer<ServerConnection, Throwable> closeCallback,
             boolean isInnerClient) {
         this.node = node;
         this.state = ConnectionState.CONNECTING;
         this.connectionMetricGroup = clientMetricGroup.createConnectionMetricGroup(node.uid());
+        this.authenticator = authenticator;
+        this.backoff = new ExponentialBackoff(100L, 2, 5000L, 0.2);
+        whenClose(closeCallback);
+
+        // connect and handle should be last in case of other variables are nullable and close
+        // callback is not registered when connection established.
         bootstrap
                 .connect(node.host(), node.port())
                 .addListener(future -> establishConnection((ChannelFuture) future, isInnerClient));
-        this.authenticator = authenticator;
-        this.backoff = new ExponentialBackoff(100L, 2, 5000L, 0.2);
     }
 
     public ServerNode getServerNode() {
@@ -131,8 +136,8 @@ public CompletableFuture<ApiMessage> send(ApiKeys apikey, ApiMessage request) {
     }
 
     /** Register a callback to be called when the connection is closed. */
-    public void whenClose(Consumer<Throwable> closeCallback) {
-        closeFuture.whenComplete((v, throwable) -> closeCallback.accept(throwable));
+    private void whenClose(BiConsumer<ServerConnection, Throwable> closeCallback) {
+        closeFuture.whenComplete((v, throwable) -> closeCallback.accept(this, throwable));
     }
 
     /** Close the connection. */

fluss-rpc/src/test/java/org/apache/fluss/rpc/netty/client/NettyClientTest.java

@@ -234,6 +234,25 @@ void testMultipleEndpoint() throws Exception {
         }
     }
 
+    @Test
+    void testExceptionWhenInitializeServerConnection() throws Exception {
+        ApiVersionsRequest request =
+                new ApiVersionsRequest()
+                        .setClientSoftwareName("testing_client_100")
+                        .setClientSoftwareVersion("1.0");
+        // close the netty server.
+        nettyServer.close();
+
+        // send request and create server connection.
+        assertThatThrownBy(
+                        () ->
+                                nettyClient
+                                        .sendRequest(serverNode, ApiKeys.API_VERSIONS, request)
+                                        .get())
+                .hasMessageContaining("Disconnected from node");
+        assertThat(nettyClient.connections()).isEmpty();
+    }
+
     private void buildNettyServer(int serverId) throws Exception {
         try (NetUtils.Port availablePort = getAvailablePort()) {
             serverNode =

fluss-rpc/src/test/java/org/apache/fluss/rpc/netty/client/ServerConnectionTest.java

@@ -96,6 +96,7 @@ void testConnectionClose() {
                         serverNode,
                         TestingClientMetricGroup.newInstance(),
                         clientAuthenticator,
+                        (con, ignore) -> {},
                         false);
         ConnectionState connectionState = connection.getConnectionState();
         assertThat(connectionState).isEqualTo(ConnectionState.CONNECTING);