Pipe: improve incremental memory estimation

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/builder/PipeDataNodeTaskBuilder.java

@@ -21,11 +21,13 @@
 
 import org.apache.iotdb.commons.consensus.index.impl.MinimumProgressIndex;
 import org.apache.iotdb.commons.exception.IllegalPathException;
+import org.apache.iotdb.commons.pipe.agent.plugin.builtin.BuiltinPipePlugin;
 import org.apache.iotdb.commons.pipe.agent.task.PipeTaskAgent;
 import org.apache.iotdb.commons.pipe.agent.task.meta.PipeStaticMeta;
 import org.apache.iotdb.commons.pipe.agent.task.meta.PipeTaskMeta;
 import org.apache.iotdb.commons.pipe.agent.task.meta.PipeType;
 import org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant;
+import org.apache.iotdb.commons.pipe.config.constant.PipeSourceConstant;
 import org.apache.iotdb.commons.pipe.config.constant.SystemConstant;
 import org.apache.iotdb.db.pipe.agent.task.PipeDataNodeTask;
 import org.apache.iotdb.db.pipe.agent.task.execution.PipeProcessorSubtaskExecutor;
@@ -64,25 +66,22 @@ public class PipeDataNodeTaskBuilder {
   private static final PipeProcessorSubtaskExecutor PROCESSOR_EXECUTOR =
       PipeSubtaskExecutorManager.getInstance().getProcessorExecutor();
 
-  protected final Map<String, String> systemParameters = new HashMap<>();
-
   public PipeDataNodeTaskBuilder(
       final PipeStaticMeta pipeStaticMeta, final int regionId, final PipeTaskMeta pipeTaskMeta) {
     this.pipeStaticMeta = pipeStaticMeta;
     this.regionId = regionId;
     this.pipeTaskMeta = pipeTaskMeta;
-    generateSystemParameters();
   }
 
   public PipeDataNodeTask build() {
     // Event flow: source -> processor -> sink
 
     // Analyzes the PipeParameters to identify potential conflicts.
     final PipeParameters sourceParameters =
-        blendUserAndSystemParameters(pipeStaticMeta.getExtractorParameters());
+        blendUserAndSystemParameters(pipeStaticMeta.getExtractorParameters(), pipeTaskMeta);
     final PipeParameters sinkParameters =
-        blendUserAndSystemParameters(pipeStaticMeta.getConnectorParameters());
-    checkConflict(sourceParameters, sinkParameters);
+        blendUserAndSystemParameters(pipeStaticMeta.getConnectorParameters(), pipeTaskMeta);
+    preprocessParameters(sourceParameters, sinkParameters);
 
     // We first build the source and sink, then build the processor.
     final PipeTaskSourceStage sourceStage =
@@ -121,7 +120,7 @@ public PipeDataNodeTask build() {
         new PipeTaskProcessorStage(
             pipeStaticMeta.getPipeName(),
             pipeStaticMeta.getCreationTime(),
-            blendUserAndSystemParameters(pipeStaticMeta.getProcessorParameters()),
+            blendUserAndSystemParameters(pipeStaticMeta.getProcessorParameters(), pipeTaskMeta),
             regionId,
             sourceStage.getEventSupplier(),
             sinkStage.getPipeSinkPendingQueue(),
@@ -139,22 +138,25 @@ public PipeDataNodeTask build() {
         pipeStaticMeta.getPipeName(), regionId, sourceStage, processorStage, sinkStage);
   }
 
-  private void generateSystemParameters() {
+  public static PipeParameters blendUserAndSystemParameters(
+      final PipeParameters userParameters, final PipeTaskMeta pipeTaskMeta) {
+    // Deep copy the user parameters to avoid modification of the original parameters.
+    // If the original parameters are modified, progress index report will be affected.
+    final Map<String, String> blendedParameters = new HashMap<>(userParameters.getAttribute());
     if (!(pipeTaskMeta.getProgressIndex() instanceof MinimumProgressIndex)
         || pipeTaskMeta.isNewlyAdded()) {
-      systemParameters.put(SystemConstant.RESTART_OR_NEWLY_ADDED_KEY, Boolean.TRUE.toString());
+      blendedParameters.put(SystemConstant.RESTART_OR_NEWLY_ADDED_KEY, Boolean.TRUE.toString());
     }
+    return new PipeParameters(blendedParameters);
   }
 
-  private PipeParameters blendUserAndSystemParameters(final PipeParameters userParameters) {
-    // Deep copy the user parameters to avoid modification of the original parameters.
-    // If the original parameters are modified, progress index report will be affected.
-    final Map<String, String> blendedParameters = new HashMap<>(userParameters.getAttribute());
-    blendedParameters.putAll(systemParameters);
-    return new PipeParameters(blendedParameters);
+  public static void preprocessParameters(
+      final PipeParameters sourceParameters, final PipeParameters sinkParameters) {
+    checkConflict(sourceParameters, sinkParameters);
+    injectParameters(sourceParameters, sinkParameters);
   }
 
-  private void checkConflict(
+  private static void checkConflict(
       final PipeParameters sourceParameters, final PipeParameters sinkParameters) {
     final Pair<Boolean, Boolean> insertionDeletionListeningOptionPair;
     final boolean shouldTerminatePipeOnAllHistoricalEventsConsumed;
@@ -220,4 +222,30 @@ private void checkConflict(
       }
     }
   }
+
+  private static void injectParameters(
+      final PipeParameters sourceParameters, final PipeParameters sinkParameters) {
+    final boolean isSourceExternal =
+        !BuiltinPipePlugin.BUILTIN_SOURCES.contains(
+            sourceParameters
+                .getStringOrDefault(
+                    Arrays.asList(PipeSourceConstant.EXTRACTOR_KEY, PipeSourceConstant.SOURCE_KEY),
+                    BuiltinPipePlugin.IOTDB_EXTRACTOR.getPipePluginName())
+                .toLowerCase());
+
+    final String sinkPluginName =
+        sinkParameters
+            .getStringOrDefault(
+                Arrays.asList(PipeSinkConstant.CONNECTOR_KEY, PipeSinkConstant.SINK_KEY),
+                BuiltinPipePlugin.IOTDB_THRIFT_SINK.getPipePluginName())
+            .toLowerCase();
+    final boolean isWriteBackSink =
+        BuiltinPipePlugin.WRITE_BACK_CONNECTOR.getPipePluginName().equals(sinkPluginName)
+            || BuiltinPipePlugin.WRITE_BACK_SINK.getPipePluginName().equals(sinkPluginName);
+
+    if (isSourceExternal && isWriteBackSink) {
+      sinkParameters.addAttribute(
+          PipeSinkConstant.CONNECTOR_USE_EVENT_USER_NAME_KEY, Boolean.TRUE.toString());
+    }
+  }
 }

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

@@ -21,6 +21,7 @@
 
 import org.apache.iotdb.commons.consensus.DataRegionId;
 import org.apache.iotdb.commons.pipe.agent.task.connection.UnboundedBlockingPendingQueue;
+import org.apache.iotdb.commons.pipe.agent.task.meta.PipeRuntimeMeta;
 import org.apache.iotdb.commons.pipe.agent.task.progress.CommitterKey;
 import org.apache.iotdb.commons.pipe.agent.task.progress.PipeEventCommitManager;
 import org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant;
@@ -64,56 +65,26 @@ public synchronized String register(
       final Supplier<? extends PipeSinkSubtaskExecutor> executorSupplier,
       final PipeParameters pipeSinkParameters,
       final PipeTaskSinkRuntimeEnvironment environment) {
-    final String connectorName =
-        PipeSinkConstant.getConnectorOrSinkNameWithDefault(pipeSinkParameters);
-    final String connectorKey =
-        connectorName
-            // Convert the value of `CONNECTOR_KEY` or `SINK_KEY` to lowercase
-            // for matching in `CONNECTOR_CONSTRUCTORS`
-            .toLowerCase();
+    final String connectorKey = getConnectorKey(pipeSinkParameters);
     PipeEventCommitManager.getInstance()
         .register(
             environment.getPipeName(),
             environment.getCreationTime(),
             environment.getRegionId(),
             connectorKey);
 
-    final boolean isDataSinkConnector =
-        StorageEngine.getInstance()
-            .getAllDataRegionIds()
-            .contains(new DataRegionId(environment.getRegionId()));
-
-    final int sinkNum;
+    final boolean isDataRegionSink = isDataRegionSink(environment.getRegionId());
+    final int sinkNum = calculateSinkSubtaskNum(pipeSinkParameters, isDataRegionSink, connectorKey);
     boolean realTimeFirst = false;
-    boolean serializeByRegion = false;
-    String attributeSortedString = generateAttributeSortedString(pipeSinkParameters);
-    if (isDataSinkConnector) {
-      serializeByRegion = PipeSinkConstant.isSerializeByRegionEnabled(pipeSinkParameters);
-      sinkNum =
-          serializeByRegion
-              ? 1
-              : pipeSinkParameters.getIntOrDefault(
-                  Arrays.asList(
-                      PipeSinkConstant.CONNECTOR_IOTDB_PARALLEL_TASKS_KEY,
-                      PipeSinkConstant.SINK_IOTDB_PARALLEL_TASKS_KEY),
-                  PipeSinkConstant.SINGLE_THREAD_DEFAULT_SINK.contains(connectorKey)
-                      ? 1
-                      : PipeSinkConstant.CONNECTOR_IOTDB_PARALLEL_TASKS_DEFAULT_VALUE);
+    final String attributeSortedString =
+        generateAttributeSortedString(pipeSinkParameters, environment.getRegionId());
+    if (isDataRegionSink) {
       realTimeFirst =
           pipeSinkParameters.getBooleanOrDefault(
               Arrays.asList(
                   PipeSinkConstant.CONNECTOR_REALTIME_FIRST_KEY,
                   PipeSinkConstant.SINK_REALTIME_FIRST_KEY),
               PipeSinkConstant.CONNECTOR_REALTIME_FIRST_DEFAULT_VALUE);
-      attributeSortedString =
-          serializeByRegion
-              ? "data_region_" + environment.getRegionId() + "_" + attributeSortedString
-              : "data_" + attributeSortedString;
-    } else {
-      // Do not allow parallel tasks for schema region connectors
-      // to avoid the potential disorder of the schema region data transfer
-      sinkNum = 1;
-      attributeSortedString = "schema_" + attributeSortedString;
     }
     environment.setAttributeSortedString(attributeSortedString);
 
@@ -135,7 +106,7 @@ public synchronized String register(
 
       for (int connectorIndex = 0; connectorIndex < sinkNum; connectorIndex++) {
         final PipeConnector pipeConnector =
-            isDataSinkConnector
+            isDataRegionSink
                 ? PipeDataNodeAgent.plugin().dataRegion().reflectSink(pipeSinkParameters)
                 : PipeDataNodeAgent.plugin().schemaRegion().reflectSink(pipeSinkParameters);
         // 1. Construct, validate and customize PipeConnector, and then handshake (create
@@ -262,7 +233,61 @@ public UnboundedBlockingPendingQueue<Event> getPipeSinkPendingQueue(
         .getPendingQueue();
   }
 
-  private String generateAttributeSortedString(final PipeParameters pipeConnectorParameters) {
+  public synchronized boolean hasRegisteredSubtasks(
+      final PipeParameters pipeSinkParameters, final int regionId) {
+    return attributeSortedString2SubtaskLifeCycleMap.containsKey(
+        generateAttributeSortedString(pipeSinkParameters, regionId));
+  }
+
+  public static int calculateSinkSubtaskNum(
+      final PipeParameters pipeSinkParameters, final int regionId) {
+    final String connectorKey = getConnectorKey(pipeSinkParameters);
+    return calculateSinkSubtaskNum(pipeSinkParameters, isDataRegionSink(regionId), connectorKey);
+  }
+
+  public static String generateAttributeSortedString(
+      final PipeParameters pipeSinkParameters, final int regionId) {
+    final String attributeSortedString = generateAttributeSortedString(pipeSinkParameters);
+    if (isDataRegionSink(regionId)) {
+      return PipeSinkConstant.isSerializeByRegionEnabled(pipeSinkParameters)
+          ? "data_region_" + regionId + "_" + attributeSortedString
+          : "data_" + attributeSortedString;
+    }
+    return "schema_" + attributeSortedString;
+  }
+
+  private static String getConnectorKey(final PipeParameters pipeSinkParameters) {
+    return PipeSinkConstant.getConnectorOrSinkNameWithDefault(pipeSinkParameters).toLowerCase();
+  }
+
+  private static boolean isDataRegionSink(final int regionId) {
+    return StorageEngine.getInstance().getAllDataRegionIds().contains(new DataRegionId(regionId))
+        || PipeRuntimeMeta.isSourceExternal(regionId);
+  }
+
+  private static int calculateSinkSubtaskNum(
+      final PipeParameters pipeSinkParameters,
+      final boolean isDataRegionSink,
+      final String connectorKey) {
+    if (!isDataRegionSink) {
+      // Do not allow parallel tasks for schema region connectors to avoid the potential disorder of
+      // the schema region data transfer.
+      return 1;
+    }
+    if (PipeSinkConstant.isSerializeByRegionEnabled(pipeSinkParameters)) {
+      return 1;
+    }
+    return pipeSinkParameters.getIntOrDefault(
+        Arrays.asList(
+            PipeSinkConstant.CONNECTOR_IOTDB_PARALLEL_TASKS_KEY,
+            PipeSinkConstant.SINK_IOTDB_PARALLEL_TASKS_KEY),
+        PipeSinkConstant.SINGLE_THREAD_DEFAULT_SINK.contains(connectorKey)
+            ? 1
+            : PipeSinkConstant.CONNECTOR_IOTDB_PARALLEL_TASKS_DEFAULT_VALUE);
+  }
+
+  private static String generateAttributeSortedString(
+      final PipeParameters pipeConnectorParameters) {
     final TreeMap<String, String> sortedStringSourceMap =
         new TreeMap<>(pipeConnectorParameters.getAttribute());
     sortedStringSourceMap.remove(SystemConstant.RESTART_OR_NEWLY_ADDED_KEY);