[rpc] Handle endpoint changes for cached server connections

fluss-rpc/src/main/java/org/apache/fluss/rpc/RpcClient.java

@@ -53,21 +53,31 @@ static RpcClient create(Configuration conf, ClientMetricGroup clientMetricGroup)
     boolean connect(ServerNode node);
 
     /**
-     * Disconnects the connection to the given server node, if there is one. Any in-flight/pending
-     * requests for this connection will receive disconnections.
+     * Disconnects the connection to the given server endpoint, if there is one. Any in-flight or
+     * pending requests for this connection will receive disconnections.
+     *
+     * @param node The server node to disconnect
+     * @return A future that is completed when the disconnection is complete
+     */
+    CompletableFuture<Void> disconnect(ServerNode node);
+
+    /**
+     * Disconnects all connections associated with the given logical server uid. Any in-flight or
+     * pending requests for these connections will receive disconnections.
      *
      * @param serverUid The uid of the server node
      * @return A future that is completed when the disconnection is complete
      */
     CompletableFuture<Void> disconnect(String serverUid);
 
     /**
-     * Check if we are currently ready to send another request to the given server but don't attempt
-     * to connect if we aren't.
+     * Check if we are currently ready to send another request to the given server endpoint but
+     * don't attempt to connect if we aren't.
      *
-     * @return true if the node is ready
+     * @param node The server node to check
+     * @return true if the connection to the node is ready
      */
-    boolean isReady(String serverUid);
+    boolean isReady(ServerNode node);
 
     /**
      * Send an RPC request to the given server and return a future for the response. If the

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

@@ -41,6 +41,7 @@
 import javax.annotation.concurrent.ThreadSafe;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
@@ -66,10 +67,10 @@ public final class NettyClient implements RpcClient {
     private final EventLoopGroup eventGroup;
 
     /**
-     * Managed connections to Netty servers. The key is the server uid (e.g., "cs-2", "ts-3"), the
-     * value is the connection.
+     * Managed connections to Netty servers. Connections are identified by server uid, host and
+     * port.
      */
-    private final Map<String, ServerConnection> connections;
+    private final Map<ConnectionKey, ServerConnection> connections;
 
     /** Metric groups for client. */
     private final ClientMetricGroup clientMetricGroup;
@@ -120,37 +121,68 @@ public boolean connect(ServerNode node) {
     }
 
     /**
-     * Disconnects the connection to the given server node, if there is one. Any inflight/pending
-     * requests for this connection will receive disconnections.
+     * Disconnects the connection to the given server endpoint, if there is one. Any
+     * inflight/pending requests for this connection will receive disconnections.
      *
-     * @param serverUid The uid of the server node
+     * @param node The server node to disconnect
      * @return A future that completes when the connection is fully closed
      */
     @Override
+    public CompletableFuture<Void> disconnect(ServerNode node) {
+        LOG.debug("Disconnecting from server {}.", node);
+        checkArgument(!isClosed, "Netty client is closed.");
+
+        ServerConnection connection = connections.remove(ConnectionKey.from(node));
+
+        if (connection == null) {
+            return FutureUtils.completedVoidFuture();
+        }
+
+        return connection.close();
+    }
+
+    /**
+     * Disconnects all connections associated with the given logical server uid. Any
+     * inflight/pending requests for these connections will receive disconnections.
+     *
+     * @param serverUid The uid of the server node
+     * @return A future that completes when all associated connections are fully closed
+     */
+    @Override
     public CompletableFuture<Void> disconnect(String serverUid) {
         LOG.debug("Disconnecting from server {}.", serverUid);
         checkArgument(!isClosed, "Netty client is closed.");
-        ServerConnection connection = connections.remove(serverUid);
-        if (connection != null) {
-            return connection.close();
+
+        List<CompletableFuture<Void>> shutdownFutures = new ArrayList<>();
+
+        for (Map.Entry<ConnectionKey, ServerConnection> entry : connections.entrySet()) {
+            if (entry.getKey().belongsTo(serverUid)
+                    && connections.remove(entry.getKey(), entry.getValue())) {
+                shutdownFutures.add(entry.getValue().close());
+            }
+        }
+
+        if (shutdownFutures.isEmpty()) {
+            return FutureUtils.completedVoidFuture();
         }
-        return FutureUtils.completedVoidFuture();
+
+        return CompletableFuture.allOf(shutdownFutures.toArray(new CompletableFuture<?>[0]));
     }
 
     /**
-     * Check if we are currently ready to send another request to the given server but don't attempt
-     * to connect if we aren't.
+     * Check if we are currently ready to send another request to the given server endpoint but
+     * don't attempt to connect if we aren't.
      *
-     * @return true if the node is ready
+     * @param node The server node to check
+     * @return true if the connection to the node is ready
      */
     @Override
-    public boolean isReady(String serverUid) {
+    public boolean isReady(ServerNode node) {
         checkArgument(!isClosed, "Netty client is closed.");
-        ServerConnection connection = connections.get(serverUid);
-        if (connection == null) {
-            return false;
-        }
-        return connection.isReady();
+
+        ServerConnection connection = connections.get(ConnectionKey.from(node));
+
+        return connection != null && connection.isReady();
     }
 
     /** Send an RPC request to the given server and return a future for the response. */
@@ -166,7 +198,7 @@ public void close() throws Exception {
         try {
             isClosed = true;
             final List<CompletableFuture<Void>> shutdownFutures = new ArrayList<>();
-            for (Map.Entry<String, ServerConnection> conn : connections.entrySet()) {
+            for (Map.Entry<ConnectionKey, ServerConnection> conn : connections.entrySet()) {
                 if (connections.remove(conn.getKey(), conn.getValue())) {
                     shutdownFutures.add(conn.getValue().close());
                 }
@@ -181,22 +213,65 @@ public void close() throws Exception {
     }
 
     private ServerConnection getOrCreateConnection(ServerNode node) {
-        String serverId = node.uid();
+        ConnectionKey connectionKey = ConnectionKey.from(node);
         return connections.computeIfAbsent(
-                serverId,
+                connectionKey,
                 ignored -> {
                     LOG.debug("Creating connection to server {}.", node);
                     return new ServerConnection(
                             bootstrap,
                             node,
                             clientMetricGroup,
                             authenticatorSupplier.get(),
-                            (con, ignore) -> connections.remove(serverId, con));
+                            (con, ignore) -> connections.remove(connectionKey, con));
                 });
     }
 
     @VisibleForTesting
-    Map<String, ServerConnection> connections() {
-        return connections;
+    Collection<ServerConnection> connections() {
+        return connections.values();
+    }
+
+    private static final class ConnectionKey {
+
+        private final String serverUid;
+        private final String host;
+        private final int port;
+
+        private ConnectionKey(String serverUid, String host, int port) {
+            this.serverUid = serverUid;
+            this.host = host;
+            this.port = port;
+        }
+
+        private static ConnectionKey from(ServerNode node) {
+            return new ConnectionKey(node.uid(), node.host(), node.port());
+        }
+
+        private boolean belongsTo(String serverUid) {
+            return this.serverUid.equals(serverUid);
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) {
+                return true;
+            }
+
+            if (!(o instanceof ConnectionKey)) {
+                return false;
+            }
+
+            ConnectionKey that = (ConnectionKey) o;
+            return port == that.port && serverUid.equals(that.serverUid) && host.equals(that.host);
+        }
+
+        @Override
+        public int hashCode() {
+            int result = serverUid.hashCode();
+            result = 31 * result + host.hashCode();
+            result = 31 * result + port;
+            return result;
+        }
     }
 }

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

@@ -115,7 +115,9 @@ final class ServerConnection {
             BiConsumer<ServerConnection, Throwable> closeCallback) {
         this.node = node;
         this.state = ConnectionState.CONNECTING;
-        this.connectionMetrics = clientMetricGroup.createConnectionMetricGroup(node.uid());
+        this.connectionMetrics =
+                clientMetricGroup.createConnectionMetricGroup(
+                        node.uid() + "@" + node.host() + ":" + node.port());
         this.authenticator = authenticator;
         this.backoff = new ExponentialBackoff(100L, 2, 5000L, 0.2);
         whenClose(closeCallback);