[AINode] Revert ClientPoolFactory

Author: Alchuang22-dev

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

@@ -23,7 +23,7 @@
 import org.apache.iotdb.common.rpc.thrift.TEndPoint;
 import org.apache.iotdb.commons.client.IClientManager;
 import org.apache.iotdb.confignode.client.async.handlers.heartbeat.AINodeHeartbeatHandler;
-import org.apache.iotdb.db.protocol.client.ConfigNodeClientManager;
+import org.apache.iotdb.db.protocol.client.AINodeClientFactory;
 import org.apache.iotdb.db.protocol.client.ainode.AsyncAINodeServiceClient;
 
 public class AsyncAINodeHeartbeatClientPool {
@@ -33,9 +33,7 @@ public class AsyncAINodeHeartbeatClientPool {
   private AsyncAINodeHeartbeatClientPool() {
     clientManager =
         new IClientManager.Factory<TEndPoint, AsyncAINodeServiceClient>()
-            .createClientManager(
-                new ConfigNodeClientManager.ClientPoolFactory
-                    .AsyncAINodeHeartbeatServiceClientPoolFactory());
+            .createClientManager(new AINodeClientFactory.AINodeHeartbeatClientPoolFactory());
   }
 
   public void getAINodeHeartBeat(

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

@@ -20,11 +20,11 @@
 package org.apache.iotdb.confignode.client.async;
 
 import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.commons.client.ClientPoolFactory;
 import org.apache.iotdb.commons.client.IClientManager;
 import org.apache.iotdb.commons.client.async.AsyncConfigNodeInternalServiceClient;
 import org.apache.iotdb.confignode.client.async.handlers.heartbeat.ConfigNodeHeartbeatHandler;
 import org.apache.iotdb.confignode.rpc.thrift.TConfigNodeHeartbeatReq;
-import org.apache.iotdb.db.protocol.client.ConfigNodeClientManager;
 
 public class AsyncConfigNodeHeartbeatClientPool {
 
@@ -34,8 +34,7 @@ private AsyncConfigNodeHeartbeatClientPool() {
     clientManager =
         new IClientManager.Factory<TEndPoint, AsyncConfigNodeInternalServiceClient>()
             .createClientManager(
-                new ConfigNodeClientManager.ClientPoolFactory
-                    .AsyncConfigNodeHeartbeatServiceClientPoolFactory());
+                new ClientPoolFactory.AsyncConfigNodeHeartbeatServiceClientPoolFactory());
   }
 
   /**

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

@@ -20,11 +20,11 @@
 package org.apache.iotdb.confignode.client.async;
 
 import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.commons.client.ClientPoolFactory;
 import org.apache.iotdb.commons.client.IClientManager;
 import org.apache.iotdb.commons.client.async.AsyncDataNodeInternalServiceClient;
 import org.apache.iotdb.confignode.client.async.handlers.audit.DataNodeWriteAuditLogHandler;
 import org.apache.iotdb.confignode.client.async.handlers.heartbeat.DataNodeHeartbeatHandler;
-import org.apache.iotdb.db.protocol.client.ConfigNodeClientManager;
 import org.apache.iotdb.mpp.rpc.thrift.TAuditLogReq;
 import org.apache.iotdb.mpp.rpc.thrift.TDataNodeHeartbeatReq;
 
@@ -37,8 +37,7 @@ private AsyncDataNodeHeartbeatClientPool() {
     clientManager =
         new IClientManager.Factory<TEndPoint, AsyncDataNodeInternalServiceClient>()
             .createClientManager(
-                new ConfigNodeClientManager.ClientPoolFactory
-                    .AsyncDataNodeHeartbeatServiceClientPoolFactory());
+                new ClientPoolFactory.AsyncDataNodeHeartbeatServiceClientPoolFactory());
   }
 
   /**

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

@@ -24,13 +24,13 @@
 import org.apache.iotdb.common.rpc.thrift.TNodeLocations;
 import org.apache.iotdb.common.rpc.thrift.TSStatus;
 import org.apache.iotdb.common.rpc.thrift.TSetConfigurationReq;
+import org.apache.iotdb.commons.client.ClientPoolFactory;
 import org.apache.iotdb.commons.client.IClientManager;
 import org.apache.iotdb.commons.client.exception.ClientManagerException;
 import org.apache.iotdb.commons.client.sync.SyncConfigNodeIServiceClient;
 import org.apache.iotdb.confignode.client.CnToCnNodeRequestType;
 import org.apache.iotdb.confignode.rpc.thrift.TAddConsensusGroupReq;
 import org.apache.iotdb.confignode.rpc.thrift.TConfigNodeRegisterReq;
-import org.apache.iotdb.db.protocol.client.ConfigNodeClientManager;
 import org.apache.iotdb.rpc.RpcUtils;
 import org.apache.iotdb.rpc.TSStatusCode;
 
@@ -54,9 +54,7 @@ public class SyncConfigNodeClientPool {
   private SyncConfigNodeClientPool() {
     clientManager =
         new IClientManager.Factory<TEndPoint, SyncConfigNodeIServiceClient>()
-            .createClientManager(
-                new ConfigNodeClientManager.ClientPoolFactory
-                    .SyncConfigNodeIServiceClientPoolFactory());
+            .createClientManager(new ClientPoolFactory.SyncConfigNodeIServiceClientPoolFactory());
     configNodeLeader = new TEndPoint();
   }
 

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

@@ -23,11 +23,11 @@
 import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation;
 import org.apache.iotdb.common.rpc.thrift.TEndPoint;
 import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.commons.client.ClientPoolFactory;
 import org.apache.iotdb.commons.client.IClientManager;
 import org.apache.iotdb.commons.client.exception.ClientManagerException;
 import org.apache.iotdb.commons.client.sync.SyncDataNodeInternalServiceClient;
 import org.apache.iotdb.commons.exception.UncheckedStartupException;
-import org.apache.iotdb.db.protocol.client.ConfigNodeClientManager;
 import org.apache.iotdb.mpp.rpc.thrift.TCleanDataNodeCacheReq;
 import org.apache.iotdb.mpp.rpc.thrift.TCreateDataRegionReq;
 import org.apache.iotdb.mpp.rpc.thrift.TCreatePeerReq;
@@ -73,8 +73,7 @@ private SyncDataNodeClientPool() {
     clientManager =
         new IClientManager.Factory<TEndPoint, SyncDataNodeInternalServiceClient>()
             .createClientManager(
-                new ConfigNodeClientManager.ClientPoolFactory
-                    .SyncDataNodeInternalServiceClientPoolFactory());
+                new ClientPoolFactory.SyncDataNodeInternalServiceClientPoolFactory());
     buildActionMap();
     checkActionMapCompleteness();
   }

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/env/RegionMaintainHandler.java

@@ -27,6 +27,7 @@
 import org.apache.iotdb.common.rpc.thrift.TRegionMigrateFailedType;
 import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet;
 import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.commons.client.ClientPoolFactory;
 import org.apache.iotdb.commons.client.IClientManager;
 import org.apache.iotdb.commons.client.sync.SyncDataNodeInternalServiceClient;
 import org.apache.iotdb.commons.cluster.NodeStatus;
@@ -44,7 +45,6 @@
 import org.apache.iotdb.confignode.manager.ConfigManager;
 import org.apache.iotdb.confignode.manager.load.cache.consensus.ConsensusGroupHeartbeatSample;
 import org.apache.iotdb.confignode.procedure.exception.ProcedureException;
-import org.apache.iotdb.db.protocol.client.ConfigNodeClientManager;
 import org.apache.iotdb.mpp.rpc.thrift.TCreatePeerReq;
 import org.apache.iotdb.mpp.rpc.thrift.TMaintainPeerReq;
 import org.apache.iotdb.mpp.rpc.thrift.TRegionLeaderChangeResp;
@@ -84,8 +84,7 @@ public RegionMaintainHandler(ConfigManager configManager) {
     dataNodeClientManager =
         new IClientManager.Factory<TEndPoint, SyncDataNodeInternalServiceClient>()
             .createClientManager(
-                new ConfigNodeClientManager.ClientPoolFactory
-                    .SyncDataNodeInternalServiceClientPoolFactory());
+                new ClientPoolFactory.SyncDataNodeInternalServiceClientPoolFactory());
   }
 
   public static String getIdWithRpcEndpoint(TDataNodeLocation location) {

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

@@ -21,6 +21,7 @@
 
 import org.apache.iotdb.common.rpc.thrift.TEndPoint;
 import org.apache.iotdb.commons.audit.UserEntity;
+import org.apache.iotdb.commons.client.ClientPoolFactory;
 import org.apache.iotdb.commons.client.IClientManager;
 import org.apache.iotdb.commons.client.ThriftClient;
 import org.apache.iotdb.commons.client.async.AsyncPipeDataTransferServiceClient;
@@ -34,7 +35,6 @@
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
 import org.apache.iotdb.db.pipe.sink.payload.evolvable.request.PipeTransferDataNodeHandshakeV1Req;
 import org.apache.iotdb.db.pipe.sink.payload.evolvable.request.PipeTransferDataNodeHandshakeV2Req;
-import org.apache.iotdb.db.protocol.client.ConfigNodeClientManager;
 import org.apache.iotdb.pipe.api.exception.PipeConnectionException;
 import org.apache.iotdb.pipe.api.exception.PipeException;
 import org.apache.iotdb.rpc.TSStatusCode;
@@ -134,10 +134,9 @@ public IoTDBDataNodeAsyncClientManager(
             new IClientManager.Factory<TEndPoint, AsyncPipeDataTransferServiceClient>()
                 .createClientManager(
                     isTSFileUsed
-                        ? new ConfigNodeClientManager.ClientPoolFactory
+                        ? new ClientPoolFactory
                             .AsyncPipeTsFileDataTransferServiceClientPoolFactory()
-                        : new ConfigNodeClientManager.ClientPoolFactory
-                            .AsyncPipeDataTransferServiceClientPoolFactory()));
+                        : new ClientPoolFactory.AsyncPipeDataTransferServiceClientPoolFactory()));
       }
       endPoint2Client = ASYNC_PIPE_DATA_TRANSFER_CLIENT_MANAGER_HOLDER.get(receiverAttributes);
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/AINodeClientFactory.java

@@ -21,25 +21,32 @@
 
 import org.apache.iotdb.common.rpc.thrift.TEndPoint;
 import org.apache.iotdb.commons.client.ClientManager;
+import org.apache.iotdb.commons.client.ClientManagerMetrics;
+import org.apache.iotdb.commons.client.IClientPoolFactory;
 import org.apache.iotdb.commons.client.factory.ThriftClientFactory;
+import org.apache.iotdb.commons.client.property.ClientPoolProperty;
 import org.apache.iotdb.commons.client.property.ThriftClientProperty;
+import org.apache.iotdb.commons.concurrent.ThreadName;
+import org.apache.iotdb.commons.conf.CommonConfig;
+import org.apache.iotdb.commons.conf.CommonDescriptor;
 import org.apache.iotdb.db.protocol.client.ainode.AINodeClient;
+import org.apache.iotdb.db.protocol.client.ainode.AsyncAINodeServiceClient;
 
 import org.apache.commons.pool2.PooledObject;
 import org.apache.commons.pool2.impl.DefaultPooledObject;
+import org.apache.commons.pool2.impl.GenericKeyedObjectPool;
 
 import java.util.Optional;
 
-/**
- * A dedicated Thrift client factory for AINodeClient. This removes AINode pooling/creation from the
- * generic ClientPoolFactory to avoid cross-node client mixing.
- */
+/** Dedicated factory for AINodeClient + AINodeClientPoolFactory. */
 public class AINodeClientFactory extends ThriftClientFactory<TEndPoint, AINodeClient> {
 
+  private static final int connectionTimeout =
+      CommonDescriptor.getInstance().getConfig().getDnConnectionTimeoutInMS();
+
   public AINodeClientFactory(
-      ClientManager<TEndPoint, AINodeClient> clientManager,
-      ThriftClientProperty thriftClientProperty) {
-    super(clientManager, thriftClientProperty);
+      ClientManager<TEndPoint, AINodeClient> manager, ThriftClientProperty thriftProperty) {
+    super(manager, thriftProperty);
   }
 
   @Override
@@ -59,4 +66,68 @@ public boolean validateObject(TEndPoint key, PooledObject<AINodeClient> pooledOb
         .map(org.apache.thrift.transport.TTransport::isOpen)
         .orElse(false);
   }
+
+  /** The PoolFactory originally inside ClientPoolFactory — now moved here. */
+  public static class AINodeClientPoolFactory
+      implements IClientPoolFactory<TEndPoint, AINodeClient> {
+
+    @Override
+    public GenericKeyedObjectPool<TEndPoint, AINodeClient> createClientPool(
+        ClientManager<TEndPoint, AINodeClient> manager) {
+
+      // Build thrift client properties
+      ThriftClientProperty thriftProperty =
+          new ThriftClientProperty.Builder()
+              .setConnectionTimeoutMs(connectionTimeout)
+              .setRpcThriftCompressionEnabled(
+                  CommonDescriptor.getInstance().getConfig().isRpcThriftCompressionEnabled())
+              .build();
+
+      GenericKeyedObjectPool<TEndPoint, AINodeClient> pool =
+          new GenericKeyedObjectPool<>(
+              new AINodeClientFactory(manager, thriftProperty),
+              new ClientPoolProperty.Builder<AINodeClient>()
+                  .setMaxClientNumForEachNode(
+                      CommonDescriptor.getInstance().getConfig().getMaxClientNumForEachNode())
+                  .build()
+                  .getConfig());
+
+      ClientManagerMetrics.getInstance()
+          .registerClientManager(this.getClass().getSimpleName(), pool);
+
+      return pool;
+    }
+  }
+
+  public static class AINodeHeartbeatClientPoolFactory
+      implements IClientPoolFactory<TEndPoint, AsyncAINodeServiceClient> {
+
+    @Override
+    public GenericKeyedObjectPool<TEndPoint, AsyncAINodeServiceClient> createClientPool(
+        ClientManager<TEndPoint, AsyncAINodeServiceClient> manager) {
+
+      final CommonConfig conf = CommonDescriptor.getInstance().getConfig();
+
+      GenericKeyedObjectPool<TEndPoint, AsyncAINodeServiceClient> clientPool =
+          new GenericKeyedObjectPool<>(
+              new AsyncAINodeServiceClient.Factory(
+                  manager,
+                  new ThriftClientProperty.Builder()
+                      .setConnectionTimeoutMs(conf.getCnConnectionTimeoutInMS())
+                      .setRpcThriftCompressionEnabled(conf.isRpcThriftCompressionEnabled())
+                      .setSelectorNumOfAsyncClientManager(conf.getSelectorNumOfClientManager())
+                      .setPrintLogWhenEncounterException(false)
+                      .build(),
+                  ThreadName.ASYNC_DATANODE_HEARTBEAT_CLIENT_POOL.getName()),
+              new ClientPoolProperty.Builder<AsyncAINodeServiceClient>()
+                  .setMaxClientNumForEachNode(conf.getMaxClientNumForEachNode())
+                  .build()
+                  .getConfig());
+
+      ClientManagerMetrics.getInstance()
+          .registerClientManager(this.getClass().getSimpleName(), clientPool);
+
+      return clientPool;
+    }
+  }
 }