Pipe: Fix connection leak caused by clients not closed after task dropped (2 situations) (#15910)

Author: SteveYurongSu

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/connector/protocol/IoTDBConfigRegionConnector.java

@@ -149,7 +149,7 @@ private void doTransferWrapper(
 
   private void doTransfer(final PipeConfigRegionWritePlanEvent pipeConfigRegionWritePlanEvent)
       throws PipeException {
-    final Pair<IoTDBSyncClient, Boolean> clientAndStatus = clientManager.getClient();
+    final Pair<IoTDBSyncClient, Boolean> clientAndStatus = getClientManager().getClient();
 
     final TPipeTransferResp resp;
     try {
@@ -175,7 +175,7 @@ private void doTransfer(final PipeConfigRegionWritePlanEvent pipeConfigRegionWri
     final TSStatus status = resp.getStatus();
     // Send handshake req and then re-transfer the event
     if (status.getCode() == TSStatusCode.PIPE_CONFIG_RECEIVER_HANDSHAKE_NEEDED.getStatusCode()) {
-      clientManager.sendHandshakeReq(clientAndStatus);
+      getClientManager().sendHandshakeReq(clientAndStatus);
     }
     // Only handle the failed statuses to avoid string format performance overhead
     if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()
@@ -214,7 +214,7 @@ private void doTransfer(final PipeConfigRegionSnapshotEvent snapshotEvent)
     final long creationTime = snapshotEvent.getCreationTime();
     final File snapshotFile = snapshotEvent.getSnapshotFile();
     final File templateFile = snapshotEvent.getTemplateFile();
-    final Pair<IoTDBSyncClient, Boolean> clientAndStatus = clientManager.getClient();
+    final Pair<IoTDBSyncClient, Boolean> clientAndStatus = getClientManager().getClient();
 
     // 1. Transfer snapshotFile, and template File if exists
     transferFilePieces(
@@ -265,7 +265,7 @@ private void doTransfer(final PipeConfigRegionSnapshotEvent snapshotEvent)
     final TSStatus status = resp.getStatus();
     // Send handshake req and then re-transfer the event
     if (status.getCode() == TSStatusCode.PIPE_CONFIG_RECEIVER_HANDSHAKE_NEEDED.getStatusCode()) {
-      clientManager.sendHandshakeReq(clientAndStatus);
+      getClientManager().sendHandshakeReq(clientAndStatus);
     }
     // Only handle the failed statuses to avoid string format performance overhead
     if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/subtask/connector/PipeConnectorSubtask.java

@@ -164,6 +164,11 @@ protected boolean executeOnce() {
   }
 
   private void transferHeartbeatEvent(final PipeHeartbeatEvent event) {
+    // DO NOT call heartbeat or transfer after closed, or will cause connection leak
+    if (isClosed.get()) {
+      return;
+    }
+
     try {
       outputPipeConnector.heartbeat();
       outputPipeConnector.transfer(event);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/connector/client/IoTDBDataNodeAsyncClientManager.java

@@ -299,6 +299,18 @@ public void onError(final Exception e) {
       client.resetMethodStateIfStopped();
       throw e;
     } finally {
+      if (isClosed) {
+        try {
+          client.close();
+          client.invalidateAll();
+        } catch (final Exception e) {
+          LOGGER.warn(
+              "Failed to close client {}:{} after handshake failure when the manager is closed.",
+              targetNodeUrl.getIp(),
+              targetNodeUrl.getPort(),
+              e);
+        }
+      }
       client.setShouldReturnSelf(true);
       client.returnSelf();
     }
@@ -348,8 +360,14 @@ public void close() {
               if (clientManager != null) {
                 try {
                   clientManager.close();
+                  LOGGER.info(
+                      "Closed AsyncPipeDataTransferServiceClientManager for receiver attributes: {}",
+                      receiverAttributes);
                 } catch (final Exception e) {
-                  LOGGER.warn("Failed to close client manager.", e);
+                  LOGGER.warn(
+                      "Failed to close AsyncPipeDataTransferServiceClientManager for receiver attributes: {}",
+                      receiverAttributes,
+                      e);
                 }
               }
               return null;

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/connector/protocol/thrift/async/IoTDBDataRegionAsyncConnector.java

@@ -198,7 +198,9 @@ public synchronized void handshake() throws Exception {
 
   @Override
   public void heartbeat() throws Exception {
-    syncConnector.heartbeat();
+    if (!isClosed()) {
+      syncConnector.heartbeat();
+    }
   }
 
   @Override

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/connector/protocol/thrift/async/handler/PipeTransferTrackableHandler.java

@@ -82,6 +82,8 @@ protected boolean tryTransfer(
     if (connector.isClosed()) {
       clearEventsReferenceCount();
       connector.eliminateHandler(this);
+      client.setShouldReturnSelf(true);
+      client.returnSelf();
       return false;
     }
     doTransfer(client, req);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/connector/protocol/thrift/async/handler/PipeTransferTsFileHandler.java

@@ -411,11 +411,25 @@ protected void onErrorInternal(final Exception exception) {
   }
 
   private void returnClientIfNecessary() {
-    if (client != null) {
-      client.setShouldReturnSelf(true);
-      client.returnSelf();
-      client = null;
+    if (client == null) {
+      return;
+    }
+
+    if (connector.isClosed()) {
+      try {
+        client.close();
+        client.invalidateAll();
+      } catch (final Exception e) {
+        LOGGER.warn(
+            "Failed to close or invalidate client when connector is closed. Client: {}, Exception: {}",
+            client,
+            e.getMessage(),
+            e);
+      }
     }
+    client.setShouldReturnSelf(true);
+    client.returnSelf();
+    client = null;
   }
 
   @Override

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

@@ -64,15 +64,19 @@ public V borrowClient(K node) throws ClientManagerException {
    * return of a client is automatic whenever a particular client is used.
    */
   public void returnClient(K node, V client) {
-    Optional.ofNullable(node)
-        .ifPresent(
-            x -> {
-              try {
-                pool.returnObject(node, client);
-              } catch (Exception e) {
-                LOGGER.warn("Return client {} for node {} to pool failed.", client, node, e);
-              }
-            });
+    if (node != null) {
+      try {
+        pool.returnObject(node, client);
+      } catch (Exception e) {
+        LOGGER.warn("Return client {} for node {} to pool failed.", client, node, e);
+      }
+    } else if (client instanceof ThriftClient) {
+      ((ThriftClient) client).invalidateAll();
+      LOGGER.warn(
+          "Return client {} to pool failed because the node is null. "
+              + "This may cause resource leak, please check your code.",
+          client);
+    }
   }
 
   @Override

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/async/AsyncPipeDataTransferServiceClient.java

@@ -127,7 +127,7 @@ public void setTimeoutDynamically(final int timeout) {
     }
   }
 
-  private void close() {
+  public void close() {
     ___transport.close();
     ___currentMethod = null;
   }

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/protocol/IoTDBSslSyncConnector.java

@@ -66,7 +66,14 @@ public abstract class IoTDBSslSyncConnector extends IoTDBConnector {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(IoTDBSslSyncConnector.class);
 
-  protected IoTDBSyncClientManager clientManager;
+  private volatile IoTDBSyncClientManager clientManager;
+
+  protected IoTDBSyncClientManager getClientManager() {
+    if (clientManager == null) {
+      throw new IllegalStateException("IoTDB sync client manager has been closed");
+    }
+    return clientManager;
+  }
 
   @Override
   public void validate(final PipeParameterValidator validator) throws Exception {
@@ -153,7 +160,7 @@ protected abstract IoTDBSyncClientManager constructClient(
 
   @Override
   public void handshake() throws Exception {
-    clientManager.checkClientStatusAndTryReconstructIfNecessary();
+    getClientManager().checkClientStatusAndTryReconstructIfNecessary();
   }
 
   @Override
@@ -229,7 +236,7 @@ protected void transferFilePieces(
         // Send handshake req and then re-transfer the event
         if (status.getCode()
             == TSStatusCode.PIPE_CONFIG_RECEIVER_HANDSHAKE_NEEDED.getStatusCode()) {
-          clientManager.sendHandshakeReq(clientAndStatus);
+          getClientManager().sendHandshakeReq(clientAndStatus);
         }
         // Only handle the failed statuses to avoid string format performance overhead
         if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()
@@ -255,6 +262,7 @@ protected abstract PipeTransferFilePieceReq getTransferMultiFilePieceReq(
   public void close() {
     if (clientManager != null) {
       clientManager.close();
+      clientManager = null;
     }
 
     super.close();