(iotdb) 03/03: borrowClient(String deviceId)

[ro...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

rong pushed a commit to branch pipe-cache-leader
in repository https://gitbox.apache.org/repos/asf/iotdb.git

commit 95d03a602a5cbb66c82f98c5ad62e9598b30a2cf
Author: Steve Yurong Su <ro...@apache.org>
AuthorDate: Fri Dec 29 01:31:59 2023 +0800

    borrowClient(String deviceId)
---
 .../async/IoTDBThriftAsyncClientManager.java       | 26 ++++++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/connector/protocol/thrift/async/IoTDBThriftAsyncClientManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/connector/protocol/thrift/async/IoTDBThriftAsyncClientManager.java
index 385bc050a83..477165d2bd8 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/connector/protocol/thrift/async/IoTDBThriftAsyncClientManager.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/connector/protocol/thrift/async/IoTDBThriftAsyncClientManager.java
@@ -81,8 +81,30 @@ public class IoTDBThriftAsyncClientManager extends IoTDBThriftClientManager {
     }
   }
 
-  public AsyncPipeDataTransferServiceClient borrowClient(String deviceId) {
-    return null;
+  public AsyncPipeDataTransferServiceClient borrowClient(String deviceId) throws Exception {
+    if (!useLeaderCache) {
+      return borrowClient();
+    }
+
+    final TEndPoint endPoint = leaderCacheManager.getLeaderEndPoint(deviceId);
+    if (endPoint == null) {
+      return borrowClient();
+    }
+
+    try {
+      final AsyncPipeDataTransferServiceClient client = endPoint2Client.borrowClient(endPoint);
+      if (handshakeIfNecessary(endPoint, client)) {
+        return client;
+      }
+    } catch (Exception e) {
+      LOGGER.warn(
+          "failed to borrow client {}:{} for cached leader.",
+          endPoint.getIp(),
+          endPoint.getPort(),
+          e);
+    }
+
+    return borrowClient();
   }
 
   /**