Gracefully manage async client threads (#17683)

Author: CRZbulabula

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/AsyncAINodeHeartbeatClientPool.java

@@ -21,6 +21,7 @@
 
 import org.apache.iotdb.ainode.rpc.thrift.TAIHeartbeatReq;
 import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.commons.client.ClientManager;
 import org.apache.iotdb.commons.client.ClientPoolFactory;
 import org.apache.iotdb.commons.client.IClientManager;
 import org.apache.iotdb.commons.client.async.AsyncAINodeInternalServiceClient;
@@ -47,10 +48,22 @@ private AsyncAINodeHeartbeatClientPool() {
    */
   public void getAINodeHeartBeat(
       TEndPoint endPoint, TAIHeartbeatReq req, AINodeHeartbeatHandler handler) {
+    AsyncAINodeInternalServiceClient client = null;
+    boolean dispatched = false;
     try {
-      clientManager.borrowClient(endPoint).getAIHeartbeat(req, handler);
+      client = clientManager.borrowClient(endPoint);
+      client.getAIHeartbeat(req, handler);
+      dispatched = true;
     } catch (Exception ignore) {
       // Just ignore
+    } finally {
+      // After the async call is dispatched, the client's onComplete/onError callback is
+      // responsible for returning the client. If the RPC was not dispatched (exception
+      // before/during the call), the client must be returned here to prevent pool leakage.
+      if (!dispatched && client != null && clientManager instanceof ClientManager) {
+        ((ClientManager<TEndPoint, AsyncAINodeInternalServiceClient>) clientManager)
+            .returnClient(endPoint, client);
+      }
     }
   }
 

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/AsyncConfigNodeHeartbeatClientPool.java

@@ -20,6 +20,7 @@
 package org.apache.iotdb.confignode.client.async;
 
 import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.commons.client.ClientManager;
 import org.apache.iotdb.commons.client.ClientPoolFactory;
 import org.apache.iotdb.commons.client.IClientManager;
 import org.apache.iotdb.commons.client.async.AsyncConfigNodeInternalServiceClient;
@@ -48,10 +49,22 @@ public void getConfigNodeHeartBeat(
       TEndPoint endPoint,
       TConfigNodeHeartbeatReq heartbeatReq,
       ConfigNodeHeartbeatHandler handler) {
+    AsyncConfigNodeInternalServiceClient client = null;
+    boolean dispatched = false;
     try {
-      clientManager.borrowClient(endPoint).getConfigNodeHeartBeat(heartbeatReq, handler);
+      client = clientManager.borrowClient(endPoint);
+      client.getConfigNodeHeartBeat(heartbeatReq, handler);
+      dispatched = true;
     } catch (Exception ignore) {
       // Just ignore
+    } finally {
+      // After the async call is dispatched, the client's onComplete/onError callback is
+      // responsible for returning the client. If the RPC was not dispatched (exception
+      // before/during the call), the client must be returned here to prevent pool leakage.
+      if (!dispatched && client != null && clientManager instanceof ClientManager) {
+        ((ClientManager<TEndPoint, AsyncConfigNodeInternalServiceClient>) clientManager)
+            .returnClient(endPoint, client);
+      }
     }
   }
 

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/AsyncDataNodeHeartbeatClientPool.java

@@ -20,6 +20,7 @@
 package org.apache.iotdb.confignode.client.async;
 
 import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.commons.client.ClientManager;
 import org.apache.iotdb.commons.client.ClientPoolFactory;
 import org.apache.iotdb.commons.client.IClientManager;
 import org.apache.iotdb.commons.client.async.AsyncDataNodeInternalServiceClient;
@@ -49,19 +50,42 @@ private AsyncDataNodeHeartbeatClientPool() {
    */
   public void getDataNodeHeartBeat(
       TEndPoint endPoint, TDataNodeHeartbeatReq req, DataNodeHeartbeatHandler handler) {
+    AsyncDataNodeInternalServiceClient client = null;
+    boolean dispatched = false;
     try {
-      clientManager.borrowClient(endPoint).getDataNodeHeartBeat(req, handler);
+      client = clientManager.borrowClient(endPoint);
+      client.getDataNodeHeartBeat(req, handler);
+      dispatched = true;
     } catch (Exception ignore) {
       // Just ignore
+    } finally {
+      returnClientIfNotDispatched(endPoint, client, dispatched);
     }
   }
 
   public void writeAuditLog(
       TEndPoint endPoint, TAuditLogReq req, DataNodeWriteAuditLogHandler handler) {
+    AsyncDataNodeInternalServiceClient client = null;
+    boolean dispatched = false;
     try {
-      clientManager.borrowClient(endPoint).writeAuditLog(req, handler);
+      client = clientManager.borrowClient(endPoint);
+      client.writeAuditLog(req, handler);
+      dispatched = true;
     } catch (Exception e) {
       // Just ignore
+    } finally {
+      returnClientIfNotDispatched(endPoint, client, dispatched);
+    }
+  }
+
+  // After the async call is dispatched, the client's onComplete/onError callback is responsible
+  // for returning the client. If the RPC was not dispatched (exception before/during the call),
+  // the client must be returned here to prevent pool leakage.
+  private void returnClientIfNotDispatched(
+      TEndPoint endPoint, AsyncDataNodeInternalServiceClient client, boolean dispatched) {
+    if (!dispatched && client != null && clientManager instanceof ClientManager) {
+      ((ClientManager<TEndPoint, AsyncDataNodeInternalServiceClient>) clientManager)
+          .returnClient(endPoint, client);
     }
   }
 

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/sync/SyncConfigNodeClientPool.java

@@ -138,6 +138,11 @@ public TSStatus removeConfigNode(
     while (status.getCode() == TSStatusCode.REDIRECTION_RECOMMEND.getStatusCode()) {
       TimeUnit.MILLISECONDS.sleep(2000);
       updateConfigNodeLeader(status);
+      if (configNodeLeader == null) {
+        LOGGER.warn(
+            "Redirection recommended for removeConfigNode but no leader endpoint provided, abort retry.");
+        break;
+      }
       try (SyncConfigNodeIServiceClient clientLeader =
           clientManager.borrowClient(configNodeLeader)) {
         status = clientLeader.removeConfigNode(configNodeLocation);

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/request/AsyncRequestManager.java

@@ -20,6 +20,7 @@
 package org.apache.iotdb.commons.client.request;
 
 import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.commons.client.ClientManager;
 import org.apache.iotdb.commons.client.IClientManager;
 import org.apache.iotdb.commons.client.exception.ClientManagerException;
 import org.apache.iotdb.commons.i18n.ClientMessages;
@@ -178,27 +179,38 @@ protected void sendAsyncRequest(
       int requestId,
       NodeLocation targetNode,
       int retryCount) {
+    final TEndPoint endPoint = nodeLocationToEndPoint(targetNode);
+    Client client = null;
+    boolean dispatched = false;
     try {
       if (!actionMap.containsKey(requestContext.getRequestType())) {
         throw new UnsupportedOperationException(
             "unsupported request type "
                 + requestContext.getRequestType()
                 + ", please set it in AsyncRequestManager::initActionMapBuilder()");
       }
-      Client client = clientManager.borrowClient(nodeLocationToEndPoint(targetNode));
+      client = clientManager.borrowClient(endPoint);
       adjustClientTimeoutIfNecessary(requestContext.getRequestType(), client);
       Object req = requestContext.getRequest(requestId);
       AsyncRequestRPCHandler<?, RequestType, NodeLocation> handler =
           buildHandler(requestContext, requestId, targetNode);
       Objects.requireNonNull(actionMap.get(requestContext.getRequestType()))
           .accept(req, client, handler);
+      // After accept() returns, the async callback (onComplete/onError) takes over the
+      // responsibility of returning the client to the pool. Before this point, if any exception
+      // is thrown, the client must be returned/invalidated here to prevent pool leakage.
+      dispatched = true;
     } catch (Exception e) {
       LOGGER.warn(
           ClientMessages.ASYNC_REQUEST_FAILED_ON_NODE,
           requestContext.getRequestType(),
-          nodeLocationToEndPoint(targetNode),
+          endPoint,
           e.getMessage(),
           retryCount);
+    } finally {
+      if (!dispatched && client != null && clientManager instanceof ClientManager) {
+        ((ClientManager<TEndPoint, Client>) clientManager).returnClient(endPoint, client);
+      }
     }
   }