Pipe: add IoTDBConnector sender port configuration

integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/treemodel/auto/basic/IoTDBPipeConnectorCustomPortIT.java

@@ -0,0 +1,166 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.pipe.it.dual.treemodel.auto.basic;
+
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.commons.client.sync.SyncConfigNodeIServiceClient;
+import org.apache.iotdb.confignode.rpc.thrift.TCreatePipeReq;
+import org.apache.iotdb.consensus.ConsensusFactory;
+import org.apache.iotdb.db.it.utils.TestUtils;
+import org.apache.iotdb.it.env.MultiEnvFactory;
+import org.apache.iotdb.it.framework.IoTDBTestRunner;
+import org.apache.iotdb.itbase.category.MultiClusterIT2DualTreeAutoBasic;
+import org.apache.iotdb.pipe.it.dual.treemodel.auto.AbstractPipeDualTreeModelAutoIT;
+import org.apache.iotdb.rpc.TSStatusCode;
+
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+@RunWith(IoTDBTestRunner.class)
+@Category({MultiClusterIT2DualTreeAutoBasic.class})
+public class IoTDBPipeConnectorCustomPortIT extends AbstractPipeDualTreeModelAutoIT {
+
+  @Override
+  @Before
+  public void setUp() {
+    // Override to enable air-gap
+    MultiEnvFactory.createEnv(2);
+    senderEnv = MultiEnvFactory.getEnv(0);
+    receiverEnv = MultiEnvFactory.getEnv(1);
+
+    senderEnv
+        .getConfig()
+        .getCommonConfig()
+        .setAutoCreateSchemaEnabled(true)
+        .setConfigNodeConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS)
+        .setSchemaRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS);
+    receiverEnv
+        .getConfig()
+        .getCommonConfig()
+        .setAutoCreateSchemaEnabled(true)
+        .setPipeAirGapReceiverEnabled(true)
+        .setConfigNodeConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS)
+        .setSchemaRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS);
+
+    senderEnv.initClusterEnvironment();
+    receiverEnv.initClusterEnvironment();
+  }
+
+  @Test
+  public void testAsyncPortRange() {
+    doTest(
+        "iotdb-thrift-async-connector", receiverEnv.getIP() + ":" + receiverEnv.getPort(), "range");
+  }
+
+  @Test
+  public void testAsyncPortCandidate() {
+    doTest(
+        "iotdb-thrift-async-connector",
+        receiverEnv.getIP() + ":" + receiverEnv.getPort(),
+        "candidate");
+  }
+
+  @Test
+  public void testSyncThriftPortRange() {
+    doTest(
+        "iotdb-thrift-sync-connector", receiverEnv.getIP() + ":" + receiverEnv.getPort(), "range");
+  }
+
+  @Test
+  public void testSyncThriftPortCandidate() {
+    doTest(
+        "iotdb-thrift-sync-connector",
+        receiverEnv.getIP() + ":" + receiverEnv.getPort(),
+        "candidate");
+  }
+
+  @Test
+  public void testAirGapPortRange() {
+    doTest(
+        "iotdb-air-gap-connector",
+        receiverEnv.getIP() + ":" + receiverEnv.getDataNodeWrapper(0).getPipeAirGapReceiverPort(),
+        "range");
+  }
+
+  @Test
+  public void testAirGapPortCandidate() {
+    doTest(
+        "iotdb-air-gap-connector",
+        receiverEnv.getIP() + ":" + receiverEnv.getDataNodeWrapper(0).getPipeAirGapReceiverPort(),
+        "candidate");
+  }
+
+  private void doTest(final String connector, final String urls, final String strategy) {
+    try (final SyncConfigNodeIServiceClient client =
+        (SyncConfigNodeIServiceClient) senderEnv.getLeaderConfigNodeConnection()) {
+      final Map<String, String> extractorAttributes = new HashMap<>();
+      final Map<String, String> processorAttributes = new HashMap<>();
+      final Map<String, String> connectorAttributes = new HashMap<>();
+      extractorAttributes.put("realtime.mode", "forced-log");
+      connectorAttributes.put("connector", connector);
+      connectorAttributes.put("batch.enable", "false");
+      connectorAttributes.put("node-urls", urls);
+      connectorAttributes.put("send-port.restriction-strategy", strategy);
+      if (strategy.equals("range")) {
+        connectorAttributes.put("send-port.range.min", "1024");
+        connectorAttributes.put("send-port.range.max", "1524");
+      } else {
+        StringBuilder candidateBuilder = new StringBuilder();
+        for (int i = 0; i < 30; i++) {
+          candidateBuilder.append(1024 + i).append(",");
+        }
+        candidateBuilder.append(1024 + 30);
+        connectorAttributes.put("send-port.candidate", candidateBuilder.toString());
+      }
+
+      final TSStatus status =
+          client.createPipe(
+              new TCreatePipeReq("p1", connectorAttributes)
+                  .setExtractorAttributes(extractorAttributes)
+                  .setProcessorAttributes(processorAttributes));
+
+      Assert.assertEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), status.getCode());
+
+      Thread.sleep(2000);
+      TestUtils.tryExecuteNonQueriesWithRetry(
+          senderEnv,
+          Arrays.asList(
+              "insert into root.db.d1(time, s1) values (2010-01-01T10:00:00+08:00, 1)",
+              "insert into root.db.d1(time, s1) values (2010-01-02T10:00:00+08:00, 2)",
+              "flush"));
+      TestUtils.assertDataEventuallyOnEnv(
+          receiverEnv,
+          "select count(*) from root.**",
+          "count(root.db.d1.s1),",
+          Collections.singleton("2,"));
+
+    } catch (Exception e) {
+      Assert.fail();
+    }
+  }
+}

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/sink/client/IoTDBConfigNodeSyncClientManager.java

@@ -48,6 +48,10 @@ public IoTDBConfigNodeSyncClientManager(
       String loadTsFileStrategy,
       boolean validateTsFile,
       boolean shouldMarkAsPipeRequest,
+      String customSendPortStrategy,
+      int minSendPortRange,
+      int maxSendPortRange,
+      List<Integer> candidatePorts,
       final boolean skipIfNoPrivileges) {
     super(
         endPoints,
@@ -62,6 +66,10 @@ public IoTDBConfigNodeSyncClientManager(
         loadTsFileStrategy,
         validateTsFile,
         shouldMarkAsPipeRequest,
+        customSendPortStrategy,
+        minSendPortRange,
+        maxSendPortRange,
+        candidatePorts,
         skipIfNoPrivileges);
   }
 

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/sink/protocol/IoTDBConfigRegionSink.java

@@ -77,6 +77,10 @@ protected IoTDBSyncClientManager constructClient(
       final String loadTsFileStrategy,
       final boolean validateTsFile,
       final boolean shouldMarkAsPipeRequest,
+      String customSendPortStrategy,
+      int minSendPortRange,
+      int maxSendPortRange,
+      List<Integer> candidatePorts,
       final boolean skipIfNoPrivileges) {
     return new IoTDBConfigNodeSyncClientManager(
         nodeUrls,
@@ -90,6 +94,10 @@ protected IoTDBSyncClientManager constructClient(
         loadTsFileStrategy,
         validateTsFile,
         shouldMarkAsPipeRequest,
+        customSendPortStrategy,
+        minSendPortRange,
+        maxSendPortRange,
+        candidatePorts,
         skipIfNoPrivileges);
   }
 

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

@@ -103,6 +103,10 @@ public IoTDBDataNodeAsyncClientManager(
       final boolean validateTsFile,
       final boolean shouldMarkAsPipeRequest,
       final boolean isTSFileUsed,
+      final String customSendPortStrategy,
+      final int minSendPortRange,
+      final int maxSendPortRange,
+      List<Integer> candidatePorts,
       final boolean skipIfNoPrivileges) {
     super(
         endPoints,
@@ -113,6 +117,10 @@ public IoTDBDataNodeAsyncClientManager(
         loadTsFileStrategy,
         validateTsFile,
         shouldMarkAsPipeRequest,
+        customSendPortStrategy,
+        minSendPortRange,
+        maxSendPortRange,
+        candidatePorts,
         skipIfNoPrivileges);
 
     endPointSet = new HashSet<>(endPoints);
@@ -134,9 +142,16 @@ public IoTDBDataNodeAsyncClientManager(
             new IClientManager.Factory<TEndPoint, AsyncPipeDataTransferServiceClient>()
                 .createClientManager(
                     isTSFileUsed
-                        ? new ClientPoolFactory
-                            .AsyncPipeTsFileDataTransferServiceClientPoolFactory()
-                        : new ClientPoolFactory.AsyncPipeDataTransferServiceClientPoolFactory()));
+                        ? new ClientPoolFactory.AsyncPipeTsFileDataTransferServiceClientPoolFactory(
+                            customSendPortStrategy,
+                            minSendPortRange,
+                            maxSendPortRange,
+                            candidatePorts)
+                        : new ClientPoolFactory.AsyncPipeDataTransferServiceClientPoolFactory(
+                            customSendPortStrategy,
+                            minSendPortRange,
+                            maxSendPortRange,
+                            candidatePorts)));
       }
       endPoint2Client = ASYNC_PIPE_DATA_TRANSFER_CLIENT_MANAGER_HOLDER.get(receiverAttributes);
 

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

@@ -58,6 +58,10 @@ public IoTDBDataNodeSyncClientManager(
       final String loadTsFileStrategy,
       final boolean validateTsFile,
       final boolean shouldMarkAsPipeRequest,
+      final String customSendPortStrategy,
+      final int minSendPortRange,
+      final int maxSendPortRange,
+      final List<Integer> candidatePorts,
       final boolean skipIfNoPrivileges) {
     super(
         endPoints,
@@ -72,6 +76,10 @@ public IoTDBDataNodeSyncClientManager(
         loadTsFileStrategy,
         validateTsFile,
         shouldMarkAsPipeRequest,
+        customSendPortStrategy,
+        minSendPortRange,
+        maxSendPortRange,
+        candidatePorts,
         skipIfNoPrivileges);
   }
 

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

@@ -167,6 +167,10 @@ public void customize(
             loadTsFileValidation,
             shouldMarkAsPipeRequest,
             false,
+            customSendPortStrategy,
+            minSendPortRange,
+            maxSendPortRange,
+            candidatePorts,
             skipIfNoPrivileges);
 
     transferTsFileClientManager =
@@ -183,6 +187,10 @@ public void customize(
             loadTsFileValidation,
             shouldMarkAsPipeRequest,
             isSplitTSFileBatchModeEnabled,
+            customSendPortStrategy,
+            minSendPortRange,
+            maxSendPortRange,
+            candidatePorts,
             skipIfNoPrivileges);
 
     if (isTabletBatchModeEnabled) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/sync/IoTDBDataNodeSyncSink.java

@@ -53,6 +53,10 @@ protected IoTDBSyncClientManager constructClient(
       final String loadTsFileStrategy,
       final boolean validateTsFile,
       final boolean shouldMarkAsPipeRequest,
+      final String customSendPortStrategy,
+      final int minSendPortRange,
+      final int maxSendPortRange,
+      List<Integer> candidatePorts,
       final boolean skipIfNoPrivileges) {
     clientManager =
         new IoTDBDataNodeSyncClientManager(
@@ -68,6 +72,10 @@ protected IoTDBSyncClientManager constructClient(
             loadTsFileStrategy,
             validateTsFile,
             shouldMarkAsPipeRequest,
+            customSendPortStrategy,
+            minSendPortRange,
+            maxSendPortRange,
+            candidatePorts,
             skipIfNoPrivileges);
     return clientManager;
   }