This is an automated email from the ASF dual-hosted git repository.
jt2594838 pushed a commit to branch dev/1.3
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/dev/1.3 by this push:
new c8e0935d5d7 Pipe: improve incremental memory estimation (#18049)
(#18113)
c8e0935d5d7 is described below
commit c8e0935d5d708e8d1641f4fe67428ee8a1ab3843
Author: Caideyipi <[email protected]>
AuthorDate: Mon Jul 6 12:17:10 2026 +0800
Pipe: improve incremental memory estimation (#18049) (#18113)
* Pipe: improve incremental memory estimation
* Pipe: add memory estimation unit tests
* Pipe: localize memory check messages
---
.../db/pipe/agent/task/PipeDataNodeTaskAgent.java | 306 ++++++++++++++++++---
.../task/builder/PipeDataNodeTaskBuilder.java | 60 ++--
.../task/subtask/sink/PipeSinkSubtaskManager.java | 101 ++++---
.../pipe/agent/task/PipeDataNodeTaskAgentTest.java | 159 +++++++++++
.../task/builder/PipeDataNodeTaskBuilderTest.java | 111 ++++++++
.../subtask/sink/PipeSinkSubtaskManagerTest.java | 95 +++++++
.../agent/plugin/builtin/BuiltinPipePlugin.java | 10 +
.../commons/pipe/agent/task/PipeTaskAgent.java | 15 +-
.../pipe/agent/task/meta/PipeRuntimeMeta.java | 9 +
.../pipe/config/constant/PipeSinkConstant.java | 3 +
10 files changed, 771 insertions(+), 98 deletions(-)
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/PipeDataNodeTaskAgent.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/PipeDataNodeTaskAgent.java
index fba4b7d3ccf..5f4fc9b4cf0 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/PipeDataNodeTaskAgent.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/PipeDataNodeTaskAgent.java
@@ -19,6 +19,7 @@
package org.apache.iotdb.db.pipe.agent.task;
+import org.apache.iotdb.common.rpc.thrift.TEndPoint;
import org.apache.iotdb.common.rpc.thrift.TPipeHeartbeatResp;
import org.apache.iotdb.common.rpc.thrift.TSStatus;
import org.apache.iotdb.commons.concurrent.IoTThreadFactory;
@@ -35,6 +36,7 @@ import
org.apache.iotdb.commons.pipe.agent.plugin.builtin.BuiltinPipePlugin;
import org.apache.iotdb.commons.pipe.agent.task.PipeTask;
import org.apache.iotdb.commons.pipe.agent.task.PipeTaskAgent;
import org.apache.iotdb.commons.pipe.agent.task.meta.PipeMeta;
+import org.apache.iotdb.commons.pipe.agent.task.meta.PipeRuntimeMeta;
import org.apache.iotdb.commons.pipe.agent.task.meta.PipeStaticMeta;
import org.apache.iotdb.commons.pipe.agent.task.meta.PipeStatus;
import org.apache.iotdb.commons.pipe.agent.task.meta.PipeTaskMeta;
@@ -44,6 +46,7 @@ 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.commons.pipe.resource.log.PipeLogger;
+import org.apache.iotdb.commons.utils.NodeUrlUtils;
import org.apache.iotdb.consensus.exception.ConsensusException;
import org.apache.iotdb.consensus.pipe.consensuspipe.ConsensusPipeName;
import org.apache.iotdb.db.conf.IoTDBConfig;
@@ -52,6 +55,7 @@ import
org.apache.iotdb.db.consensus.SchemaRegionConsensusImpl;
import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent;
import org.apache.iotdb.db.pipe.agent.task.builder.PipeDataNodeBuilder;
import org.apache.iotdb.db.pipe.agent.task.builder.PipeDataNodeTaskBuilder;
+import org.apache.iotdb.db.pipe.agent.task.subtask.sink.PipeSinkSubtaskManager;
import org.apache.iotdb.db.pipe.metric.overview.PipeDataNodeSinglePipeMetrics;
import org.apache.iotdb.db.pipe.metric.overview.PipeTsFileToTabletsMetrics;
import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager;
@@ -189,7 +193,14 @@ public class PipeDataNodeTaskAgent extends PipeTaskAgent {
&&
!SchemaRegionListeningFilter.parseListeningPlanTypeSet(sourceParameters).isEmpty();
// Advance the extractor parameters parsing logic to avoid creating
un-relevant pipeTasks
- if (needConstructDataRegionTask || needConstructSchemaRegionTask) {
+ if (PipeRuntimeMeta.isSourceExternal(consensusGroupId)
+ || needConstructDataRegionTask
+ || needConstructSchemaRegionTask) {
+ calculateMemoryUsage(
+ pipeStaticMeta,
+ Collections.singletonList(new Pair<>(consensusGroupId,
pipeTaskMeta)),
+ false);
+
final PipeDataNodeTask pipeTask =
new PipeDataNodeTaskBuilder(pipeStaticMeta, consensusGroupId,
pipeTaskMeta).build();
pipeTask.create();
@@ -801,25 +812,49 @@ public class PipeDataNodeTaskAgent extends PipeTaskAgent {
}
@Override
- protected void calculateMemoryUsage(
+ protected void calculateMemoryUsage(final PipeMeta pipeMetaFromCoordinator)
+ throws IllegalPathException {
+ final PipeStaticMeta staticMeta = pipeMetaFromCoordinator.getStaticMeta();
+ if (!PipeConfig.getInstance().isPipeEnableMemoryCheck()
+ || !isInnerSource(staticMeta.getExtractorParameters())
+ || !PipeType.USER.equals(staticMeta.getPipeType())) {
+ return;
+ }
+
+ final PipeMeta pipeMetaInAgent =
pipeMetaKeeper.getPipeMeta(staticMeta.getPipeName());
+ final boolean ignoreRegisteredSinkSubtasks =
+ Objects.nonNull(pipeMetaInAgent)
+ && (!pipeMetaInAgent.getStaticMeta().equals(staticMeta)
+ ||
PipeStatus.DROPPED.equals(pipeMetaInAgent.getRuntimeMeta().getStatus().get()));
+ calculateMemoryUsage(
+ staticMeta,
+ collectPipeTasksToBeCreated(pipeMetaFromCoordinator),
+ ignoreRegisteredSinkSubtasks);
+ }
+
+ private void calculateMemoryUsage(
final PipeStaticMeta staticMeta,
- final PipeParameters sourceParameters,
- final PipeParameters processorParameters,
- final PipeParameters sinkParameters) {
+ final List<Pair<Integer, PipeTaskMeta>> pipeTasksToBeCreated,
+ final boolean ignoreRegisteredSinkSubtasks)
+ throws IllegalPathException {
if (!PipeConfig.getInstance().isPipeEnableMemoryCheck()
- || !isInnerSource(sourceParameters)
+ || !isInnerSource(staticMeta.getExtractorParameters())
|| !PipeType.USER.equals(staticMeta.getPipeType())) {
return;
}
- calculateInsertNodeQueueMemory(sourceParameters);
+ if (pipeTasksToBeCreated.isEmpty()) {
+ calculateInsertNodeQueueMemory(staticMeta.getExtractorParameters(), 1);
+ return;
+ }
- long needMemory = 0;
+ final MemoryEstimation memoryEstimation =
+ calculateIncrementalMemoryUsage(
+ staticMeta, pipeTasksToBeCreated, ignoreRegisteredSinkSubtasks);
+ calculateInsertNodeQueueMemory(
+ staticMeta.getExtractorParameters(),
memoryEstimation.dataRegionTaskCount);
- needMemory += calculateTsFileParserMemory(sourceParameters,
sinkParameters);
- needMemory += calculateSinkBatchMemory(sinkParameters);
- needMemory += calculateSendTsFileReadBufferMemory(sourceParameters,
sinkParameters);
- needMemory += calculateAssignerMemory(sourceParameters);
+ final long needMemory = memoryEstimation.nonFloatingMemoryInBytes;
PipeMemoryManager pipeMemoryManager = PipeDataNodeResourceManager.memory();
final long freeMemorySizeInBytes =
pipeMemoryManager.getFreeMemorySizeInBytes();
@@ -841,6 +876,109 @@ public class PipeDataNodeTaskAgent extends PipeTaskAgent {
}
}
+ private List<Pair<Integer, PipeTaskMeta>> collectPipeTasksToBeCreated(
+ final PipeMeta pipeMetaFromCoordinator) throws IllegalPathException {
+ final PipeStaticMeta pipeStaticMeta =
pipeMetaFromCoordinator.getStaticMeta();
+ final PipeParameters sourceParameters =
pipeStaticMeta.getExtractorParameters();
+ final Set<DataRegionId> dataRegionIds =
+ new HashSet<>(StorageEngine.getInstance().getAllDataRegionIds());
+ final Set<SchemaRegionId> schemaRegionIds =
+ new HashSet<>(SchemaEngine.getInstance().getAllSchemaRegionIds());
+ final List<Pair<Integer, PipeTaskMeta>> pipeTasksToBeCreated = new
ArrayList<>();
+
+ for (final Map.Entry<Integer, PipeTaskMeta> consensusGroupIdToPipeTaskMeta
:
+
pipeMetaFromCoordinator.getRuntimeMeta().getConsensusGroupId2TaskMetaMap().entrySet())
{
+ final int consensusGroupId = consensusGroupIdToPipeTaskMeta.getKey();
+ final PipeTaskMeta pipeTaskMeta =
consensusGroupIdToPipeTaskMeta.getValue();
+ if (pipeTaskMeta.getLeaderNodeId() != CONFIG.getDataNodeId()) {
+ continue;
+ }
+
+ final boolean needConstructTask;
+ if (isSourceExternal(sourceParameters)) {
+ needConstructTask = true;
+ } else {
+ final DataRegionId dataRegionId = new DataRegionId(consensusGroupId);
+ final boolean needConstructDataRegionTask =
+ dataRegionIds.contains(dataRegionId)
+ && DataRegionListeningFilter.shouldDataRegionBeListened(
+ sourceParameters, dataRegionId);
+ final boolean needConstructSchemaRegionTask =
+ schemaRegionIds.contains(new SchemaRegionId(consensusGroupId))
+ &&
!SchemaRegionListeningFilter.parseListeningPlanTypeSet(sourceParameters)
+ .isEmpty();
+ needConstructTask = needConstructDataRegionTask ||
needConstructSchemaRegionTask;
+ }
+
+ if (needConstructTask) {
+ pipeTasksToBeCreated.add(new Pair<>(consensusGroupId, pipeTaskMeta));
+ }
+ }
+ return pipeTasksToBeCreated;
+ }
+
+ private MemoryEstimation calculateIncrementalMemoryUsage(
+ final PipeStaticMeta staticMeta,
+ final List<Pair<Integer, PipeTaskMeta>> pipeTasksToBeCreated,
+ final boolean ignoreRegisteredSinkSubtasks) {
+ long needMemory = 0;
+ int dataRegionTaskCount = 0;
+ final Set<String> sinkSubtasksToBeCreated = new HashSet<>();
+
+ for (final Pair<Integer, PipeTaskMeta> regionIdAndTaskMeta :
pipeTasksToBeCreated) {
+ final int regionId = regionIdAndTaskMeta.getLeft();
+ final PipeTaskMeta pipeTaskMeta = regionIdAndTaskMeta.getRight();
+ final PipeParameters sourceParameters =
+ PipeDataNodeTaskBuilder.blendUserAndSystemParameters(
+ staticMeta.getExtractorParameters(), pipeTaskMeta);
+ final PipeParameters sinkParameters =
+ PipeDataNodeTaskBuilder.blendUserAndSystemParameters(
+ staticMeta.getConnectorParameters(), pipeTaskMeta);
+ PipeDataNodeTaskBuilder.preprocessParameters(sourceParameters,
sinkParameters);
+
+ final boolean isDataRegionTask = isDataRegionTask(regionId);
+ if (isDataRegionTask) {
+ dataRegionTaskCount++;
+ needMemory += calculateTsFileParserMemory(sourceParameters,
sinkParameters);
+ }
+
+ final String sinkSubtaskId =
+ PipeSinkSubtaskManager.generateAttributeSortedString(sinkParameters,
regionId);
+ if (isDataRegionTask
+ && !sinkSubtasksToBeCreated.contains(sinkSubtaskId)
+ && (ignoreRegisteredSinkSubtasks
+ || !PipeSinkSubtaskManager.instance()
+ .hasRegisteredSubtasks(sinkParameters, regionId))) {
+ sinkSubtasksToBeCreated.add(sinkSubtaskId);
+ final int sinkSubtaskNum =
+ PipeSinkSubtaskManager.calculateSinkSubtaskNum(sinkParameters,
regionId);
+ needMemory += calculateSinkBatchMemory(sinkParameters) *
sinkSubtaskNum;
+ needMemory +=
+ calculateSendTsFileReadBufferMemory(sourceParameters,
sinkParameters) * sinkSubtaskNum;
+ }
+ }
+
+ if (dataRegionTaskCount > 0) {
+ needMemory +=
calculateAssignerMemory(staticMeta.getExtractorParameters());
+ }
+ return new MemoryEstimation(needMemory, dataRegionTaskCount);
+ }
+
+ private boolean isDataRegionTask(final int regionId) {
+ return StorageEngine.getInstance().getAllDataRegionIds().contains(new
DataRegionId(regionId))
+ || PipeRuntimeMeta.isSourceExternal(regionId);
+ }
+
+ private static class MemoryEstimation {
+ private final long nonFloatingMemoryInBytes;
+ private final int dataRegionTaskCount;
+
+ private MemoryEstimation(final long nonFloatingMemoryInBytes, final int
dataRegionTaskCount) {
+ this.nonFloatingMemoryInBytes = nonFloatingMemoryInBytes;
+ this.dataRegionTaskCount = dataRegionTaskCount;
+ }
+ }
+
private boolean isInnerSource(final PipeParameters sourceParameters) {
final String pluginName =
sourceParameters
@@ -853,7 +991,20 @@ public class PipeDataNodeTaskAgent extends PipeTaskAgent {
||
pluginName.equals(BuiltinPipePlugin.IOTDB_SOURCE.getPipePluginName());
}
- private void calculateInsertNodeQueueMemory(final PipeParameters
sourceParameters) {
+ private boolean isSourceExternal(final PipeParameters sourceParameters) {
+ return !BuiltinPipePlugin.BUILTIN_SOURCES.contains(
+ sourceParameters
+ .getStringOrDefault(
+ Arrays.asList(PipeSourceConstant.EXTRACTOR_KEY,
PipeSourceConstant.SOURCE_KEY),
+ BuiltinPipePlugin.IOTDB_EXTRACTOR.getPipePluginName())
+ .toLowerCase());
+ }
+
+ private void calculateInsertNodeQueueMemory(
+ final PipeParameters sourceParameters, final int dataRegionTaskCount) {
+ if (dataRegionTaskCount <= 0) {
+ return;
+ }
// Realtime source is enabled by default, so we only need to check the
source realtime
if (!sourceParameters.getBooleanOrDefault(
@@ -872,16 +1023,17 @@ public class PipeDataNodeTaskAgent extends PipeTaskAgent
{
return;
}
+ final long needFloatingMemory =
+ PipeConfig.getInstance().getPipeInsertNodeQueueMemory() *
dataRegionTaskCount;
final long allocatedMemorySizeInBytes =
this.getAllFloatingMemoryUsageInByte();
final long remainingMemory =
- PipeMemoryManager.getTotalFloatingMemorySizeInBytes() -
allocatedMemorySizeInBytes;
- if (remainingMemory <
PipeConfig.getInstance().getPipeInsertNodeQueueMemory()) {
+
PipeDataNodeResourceManager.memory().getTotalFloatingMemorySizeInBytes()
+ - allocatedMemorySizeInBytes;
+ if (remainingMemory < needFloatingMemory) {
final String message =
String.format(
- "%s Need Floating memory: %d bytes, free Floating memory: %d
bytes",
- MESSAGE_PIPE_NOT_ENOUGH_MEMORY,
- PipeConfig.getInstance().getPipeInsertNodeQueueMemory(),
- remainingMemory);
+ "%s Need Floating memory: %d bytes, free Floating memory: %d
bytes",
+ MESSAGE_PIPE_NOT_ENOUGH_MEMORY, needFloatingMemory,
remainingMemory);
LOGGER.warn(message);
throw new PipeException(message);
}
@@ -957,36 +1109,112 @@ public class PipeDataNodeTaskAgent extends
PipeTaskAgent {
return PipeConfig.getInstance().getTsFileParserMemory();
}
- private long calculateSinkBatchMemory(final PipeParameters sinkParameters) {
+ private static long calculateSinkBatchMemory(final PipeParameters
sinkParameters) {
+ final String format =
+ sinkParameters.getStringOrDefault(
+ Arrays.asList(PipeSinkConstant.CONNECTOR_FORMAT_KEY,
PipeSinkConstant.SINK_FORMAT_KEY),
+ PipeSinkConstant.CONNECTOR_FORMAT_HYBRID_VALUE);
+ final boolean usingTsFileBatch =
PipeSinkConstant.CONNECTOR_FORMAT_TS_FILE_VALUE.equals(format);
- // If the sink format is tsfile , we need to use batch
- boolean needUseBatch =
- PipeSinkConstant.CONNECTOR_FORMAT_TS_FILE_VALUE.equals(
- sinkParameters.getStringOrDefault(
+ // TsFile format always uses a batch. Other formats only use a batch when
batch mode is enabled.
+ final boolean needUseBatch =
+ usingTsFileBatch
+ || sinkParameters.getBooleanOrDefault(
Arrays.asList(
- PipeSinkConstant.CONNECTOR_FORMAT_KEY,
PipeSinkConstant.SINK_FORMAT_KEY),
- PipeSinkConstant.CONNECTOR_FORMAT_HYBRID_VALUE));
+ PipeSinkConstant.CONNECTOR_IOTDB_BATCH_MODE_ENABLE_KEY,
+ PipeSinkConstant.SINK_IOTDB_BATCH_MODE_ENABLE_KEY),
+
PipeSinkConstant.CONNECTOR_IOTDB_BATCH_MODE_ENABLE_DEFAULT_VALUE);
- if (needUseBatch) {
- return PipeConfig.getInstance().getSinkBatchMemoryTsFile();
+ if (!needUseBatch) {
+ return 0;
}
- // If the sink is batch mode, we need to use batch
- needUseBatch =
- sinkParameters.getBooleanOrDefault(
+ final long batchSizeInBytes =
+ sinkParameters.getLongOrDefault(
Arrays.asList(
- PipeSinkConstant.CONNECTOR_IOTDB_BATCH_MODE_ENABLE_KEY,
- PipeSinkConstant.SINK_IOTDB_BATCH_MODE_ENABLE_KEY),
- PipeSinkConstant.CONNECTOR_IOTDB_BATCH_MODE_ENABLE_DEFAULT_VALUE);
+ PipeSinkConstant.CONNECTOR_IOTDB_BATCH_SIZE_KEY,
+ PipeSinkConstant.SINK_IOTDB_BATCH_SIZE_KEY),
+ usingTsFileBatch
+ ?
PipeSinkConstant.CONNECTOR_IOTDB_TS_FILE_BATCH_SIZE_DEFAULT_VALUE
+ :
PipeSinkConstant.CONNECTOR_IOTDB_PLAIN_BATCH_SIZE_DEFAULT_VALUE);
- if (!needUseBatch) {
- return 0;
+ return batchSizeInBytes * calculateBatchShardCount(sinkParameters,
usingTsFileBatch);
+ }
+
+ private static long calculateBatchShardCount(
+ final PipeParameters sinkParameters, final boolean usingTsFileBatch) {
+ if (usingTsFileBatch
+ || !sinkParameters.getBooleanOrDefault(
+ Arrays.asList(
+ PipeSinkConstant.CONNECTOR_LEADER_CACHE_ENABLE_KEY,
+ PipeSinkConstant.SINK_LEADER_CACHE_ENABLE_KEY),
+ PipeSinkConstant.CONNECTOR_LEADER_CACHE_ENABLE_DEFAULT_VALUE)) {
+ return 1;
}
- return PipeConfig.getInstance().getSinkBatchMemoryInsertNode();
+ // Plain batches always allocate the default batch and may lazily allocate
one batch per target
+ // endpoint when leader cache splits events by endpoint.
+ return 1L + calculateTargetEndPointCount(sinkParameters);
+ }
+
+ private static int calculateTargetEndPointCount(final PipeParameters
sinkParameters) {
+ final Set<TEndPoint> targetEndPoints = new HashSet<>();
+ try {
+ addTargetEndPoint(
+ targetEndPoints,
+ sinkParameters,
+ PipeSinkConstant.CONNECTOR_IOTDB_IP_KEY,
+ PipeSinkConstant.CONNECTOR_IOTDB_HOST_KEY,
+ PipeSinkConstant.CONNECTOR_IOTDB_PORT_KEY);
+ addTargetEndPoint(
+ targetEndPoints,
+ sinkParameters,
+ PipeSinkConstant.SINK_IOTDB_IP_KEY,
+ PipeSinkConstant.SINK_IOTDB_HOST_KEY,
+ PipeSinkConstant.SINK_IOTDB_PORT_KEY);
+ if
(sinkParameters.hasAttribute(PipeSinkConstant.CONNECTOR_IOTDB_NODE_URLS_KEY)) {
+ targetEndPoints.addAll(
+ NodeUrlUtils.parseTEndPointUrls(
+ Arrays.asList(
+ sinkParameters
+
.getStringByKeys(PipeSinkConstant.CONNECTOR_IOTDB_NODE_URLS_KEY)
+ .replace(" ", "")
+ .split(","))));
+ }
+ if
(sinkParameters.hasAttribute(PipeSinkConstant.SINK_IOTDB_NODE_URLS_KEY)) {
+ targetEndPoints.addAll(
+ NodeUrlUtils.parseTEndPointUrls(
+ Arrays.asList(
+ sinkParameters
+
.getStringByKeys(PipeSinkConstant.SINK_IOTDB_NODE_URLS_KEY)
+ .replace(" ", "")
+ .split(","))));
+ }
+ } catch (final Exception ignored) {
+ return 1;
+ }
+ return Math.max(1, targetEndPoints.size());
+ }
+
+ private static void addTargetEndPoint(
+ final Set<TEndPoint> targetEndPoints,
+ final PipeParameters sinkParameters,
+ final String ipKey,
+ final String hostKey,
+ final String portKey) {
+ if (sinkParameters.hasAttribute(ipKey) &&
sinkParameters.hasAttribute(portKey)) {
+ targetEndPoints.add(
+ new TEndPoint(
+ sinkParameters.getStringByKeys(ipKey),
sinkParameters.getIntByKeys(portKey)));
+ }
+ if (sinkParameters.hasAttribute(hostKey) &&
sinkParameters.hasAttribute(portKey)) {
+ targetEndPoints.add(
+ new TEndPoint(
+ sinkParameters.getStringByKeys(hostKey),
sinkParameters.getIntByKeys(portKey)));
+ }
}
- private long calculateSendTsFileReadBufferMemory(
+ private static long calculateSendTsFileReadBufferMemory(
final PipeParameters sourceParameters, final PipeParameters
sinkParameters) {
// If the source is history enable, we need to transfer tsfile
boolean needTransferTsFile =
@@ -1012,7 +1240,7 @@ public class PipeDataNodeTaskAgent extends PipeTaskAgent {
return 0;
}
- return PipeConfig.getInstance().getSendTsFileReadBuffer();
+ return PipeConfig.getInstance().getPipeSinkReadFileBufferSize();
}
private long calculateAssignerMemory(final PipeParameters sourceParameters) {
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/builder/PipeDataNodeTaskBuilder.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/builder/PipeDataNodeTaskBuilder.java
index 8d227d7c272..cc4456b721b 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/builder/PipeDataNodeTaskBuilder.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/builder/PipeDataNodeTaskBuilder.java
@@ -21,11 +21,13 @@ package org.apache.iotdb.db.pipe.agent.task.builder;
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,14 +66,11 @@ 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() {
@@ -79,10 +78,10 @@ public class PipeDataNodeTaskBuilder {
// 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 class PipeDataNodeTaskBuilder {
new PipeTaskProcessorStage(
pipeStaticMeta.getPipeName(),
pipeStaticMeta.getCreationTime(),
-
blendUserAndSystemParameters(pipeStaticMeta.getProcessorParameters()),
+
blendUserAndSystemParameters(pipeStaticMeta.getProcessorParameters(),
pipeTaskMeta),
regionId,
sourceStage.getEventSupplier(),
sinkStage.getPipeSinkPendingQueue(),
@@ -139,22 +138,25 @@ public class PipeDataNodeTaskBuilder {
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 @@ public class PipeDataNodeTaskBuilder {
}
}
}
+
+ 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());
+ }
+ }
}
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtaskManager.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtaskManager.java
index 4385cc59504..6307dbba648 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtaskManager.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtaskManager.java
@@ -21,6 +21,7 @@ package org.apache.iotdb.db.pipe.agent.task.subtask.sink;
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,13 +65,7 @@ public class PipeSinkSubtaskManager {
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(),
@@ -78,42 +73,18 @@ public class PipeSinkSubtaskManager {
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 class PipeSinkSubtaskManager {
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 class PipeSinkSubtaskManager {
.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);
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/PipeDataNodeTaskAgentTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/PipeDataNodeTaskAgentTest.java
new file mode 100644
index 00000000000..0582a94170d
--- /dev/null
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/PipeDataNodeTaskAgentTest.java
@@ -0,0 +1,159 @@
+/*
+ * 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.db.pipe.agent.task;
+
+import org.apache.iotdb.commons.conf.CommonDescriptor;
+import org.apache.iotdb.commons.pipe.agent.task.meta.PipeMeta;
+import org.apache.iotdb.commons.pipe.agent.task.meta.PipeRuntimeMeta;
+import org.apache.iotdb.commons.pipe.agent.task.meta.PipeStaticMeta;
+import org.apache.iotdb.commons.pipe.config.PipeConfig;
+import org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant;
+import org.apache.iotdb.commons.pipe.config.constant.PipeSourceConstant;
+import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent;
+import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameters;
+import org.apache.iotdb.pipe.api.exception.PipeException;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.lang.reflect.Method;
+import java.util.HashMap;
+import java.util.Map;
+
+public class PipeDataNodeTaskAgentTest {
+
+ @Test
+ public void testCreateMemoryCheckStillRunsWhenNoPipeTasksNeedToBeCreated()
throws Exception {
+ final boolean originalPipeEnableMemoryCheck =
+ CommonDescriptor.getInstance().getConfig().isPipeEnableMemoryChecked();
+ final long originalPipeInsertNodeQueueMemory =
+
CommonDescriptor.getInstance().getConfig().getPipeInsertNodeQueueMemory();
+ final double originalPipeTotalFloatingMemoryProportion =
+
CommonDescriptor.getInstance().getConfig().getPipeTotalFloatingMemoryProportion();
+
+ try {
+
CommonDescriptor.getInstance().getConfig().setIsPipeEnableMemoryChecked(true);
+
CommonDescriptor.getInstance().getConfig().setPipeInsertNodeQueueMemory(1);
+
CommonDescriptor.getInstance().getConfig().setPipeTotalFloatingMemoryProportion(0);
+
+ Assert.assertThrows(
+ PipeException.class,
+ () ->
+ PipeDataNodeAgent.task()
+ .calculateMemoryUsage(
+ new PipeMeta(
+ new PipeStaticMeta(
+ "p", 1L, new HashMap<>(), new HashMap<>(), new
HashMap<>()),
+ new PipeRuntimeMeta())));
+ } finally {
+ CommonDescriptor.getInstance()
+ .getConfig()
+ .setIsPipeEnableMemoryChecked(originalPipeEnableMemoryCheck);
+ CommonDescriptor.getInstance()
+ .getConfig()
+ .setPipeInsertNodeQueueMemory(originalPipeInsertNodeQueueMemory);
+ CommonDescriptor.getInstance()
+ .getConfig()
+
.setPipeTotalFloatingMemoryProportion(originalPipeTotalFloatingMemoryProportion);
+ }
+ }
+
+ @Test
+ public void testPlainBatchMemoryIncludesLeaderCacheEndpointShards() throws
Exception {
+ final Map<String, String> sinkAttributes = new HashMap<>();
+ sinkAttributes.put(
+ PipeSinkConstant.CONNECTOR_FORMAT_KEY,
PipeSinkConstant.CONNECTOR_FORMAT_TABLET_VALUE);
+ sinkAttributes.put(PipeSinkConstant.CONNECTOR_IOTDB_BATCH_SIZE_KEY,
"1024");
+ sinkAttributes.put(
+ PipeSinkConstant.CONNECTOR_IOTDB_NODE_URLS_KEY, "127.0.0.1:6667,
127.0.0.2:6667");
+ sinkAttributes.put(PipeSinkConstant.CONNECTOR_IOTDB_IP_KEY, "127.0.0.3");
+ sinkAttributes.put(PipeSinkConstant.CONNECTOR_IOTDB_PORT_KEY, "6667");
+
+ Assert.assertEquals(
+ 4 * 1024L, invokeCalculateSinkBatchMemory(new
PipeParameters(sinkAttributes)));
+
+ sinkAttributes.put(
+ PipeSinkConstant.CONNECTOR_LEADER_CACHE_ENABLE_KEY,
Boolean.FALSE.toString());
+ Assert.assertEquals(1024L, invokeCalculateSinkBatchMemory(new
PipeParameters(sinkAttributes)));
+ }
+
+ @Test
+ public void testTsFileBatchMemoryIgnoresLeaderCacheEndpointShards() throws
Exception {
+ final Map<String, String> sinkAttributes = new HashMap<>();
+ sinkAttributes.put(
+ PipeSinkConstant.CONNECTOR_FORMAT_KEY,
PipeSinkConstant.CONNECTOR_FORMAT_TS_FILE_VALUE);
+ sinkAttributes.put(PipeSinkConstant.CONNECTOR_IOTDB_BATCH_SIZE_KEY,
"2048");
+ sinkAttributes.put(
+ PipeSinkConstant.CONNECTOR_IOTDB_NODE_URLS_KEY,
"127.0.0.1:6667,127.0.0.2:6667");
+
+ Assert.assertEquals(2048L, invokeCalculateSinkBatchMemory(new
PipeParameters(sinkAttributes)));
+ }
+
+ @Test
+ public void testPlainBatchMemoryReturnsZeroWhenBatchModeIsDisabled() throws
Exception {
+ final Map<String, String> sinkAttributes = new HashMap<>();
+ sinkAttributes.put(
+ PipeSinkConstant.CONNECTOR_FORMAT_KEY,
PipeSinkConstant.CONNECTOR_FORMAT_TABLET_VALUE);
+ sinkAttributes.put(
+ PipeSinkConstant.CONNECTOR_IOTDB_BATCH_MODE_ENABLE_KEY,
Boolean.FALSE.toString());
+ sinkAttributes.put(PipeSinkConstant.CONNECTOR_IOTDB_BATCH_SIZE_KEY,
"1024");
+
+ Assert.assertEquals(0L, invokeCalculateSinkBatchMemory(new
PipeParameters(sinkAttributes)));
+ }
+
+ @Test
+ public void testSendTsFileReadBufferMemoryUsesSinkReadFileBufferSize()
throws Exception {
+ final Map<String, String> sourceAttributes = new HashMap<>();
+ sourceAttributes.put(PipeSourceConstant.EXTRACTOR_HISTORY_ENABLE_KEY,
Boolean.FALSE.toString());
+
+ final Map<String, String> sinkAttributes = new HashMap<>();
+ sinkAttributes.put(
+ PipeSinkConstant.CONNECTOR_FORMAT_KEY,
PipeSinkConstant.CONNECTOR_FORMAT_TABLET_VALUE);
+ Assert.assertEquals(
+ 0L,
+ invokeCalculateSendTsFileReadBufferMemory(
+ new PipeParameters(sourceAttributes), new
PipeParameters(sinkAttributes)));
+
+ sinkAttributes.put(
+ PipeSinkConstant.CONNECTOR_FORMAT_KEY,
PipeSinkConstant.CONNECTOR_FORMAT_HYBRID_VALUE);
+ Assert.assertEquals(
+ PipeConfig.getInstance().getPipeSinkReadFileBufferSize(),
+ invokeCalculateSendTsFileReadBufferMemory(
+ new PipeParameters(sourceAttributes), new
PipeParameters(sinkAttributes)));
+ }
+
+ private long invokeCalculateSinkBatchMemory(final PipeParameters
sinkParameters)
+ throws Exception {
+ final Method method =
+ PipeDataNodeTaskAgent.class.getDeclaredMethod(
+ "calculateSinkBatchMemory", PipeParameters.class);
+ method.setAccessible(true);
+ return (long) method.invoke(null, sinkParameters);
+ }
+
+ private long invokeCalculateSendTsFileReadBufferMemory(
+ final PipeParameters sourceParameters, final PipeParameters
sinkParameters) throws Exception {
+ final Method method =
+ PipeDataNodeTaskAgent.class.getDeclaredMethod(
+ "calculateSendTsFileReadBufferMemory", PipeParameters.class,
PipeParameters.class);
+ method.setAccessible(true);
+ return (long) method.invoke(null, sourceParameters, sinkParameters);
+ }
+}
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/builder/PipeDataNodeTaskBuilderTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/builder/PipeDataNodeTaskBuilderTest.java
new file mode 100644
index 00000000000..e00f9500bbe
--- /dev/null
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/builder/PipeDataNodeTaskBuilderTest.java
@@ -0,0 +1,111 @@
+/*
+ * 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.db.pipe.agent.task.builder;
+
+import org.apache.iotdb.commons.consensus.index.impl.MinimumProgressIndex;
+import org.apache.iotdb.commons.consensus.index.impl.SimpleProgressIndex;
+import org.apache.iotdb.commons.pipe.agent.plugin.builtin.BuiltinPipePlugin;
+import org.apache.iotdb.commons.pipe.agent.task.meta.PipeTaskMeta;
+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.pipe.api.customizer.parameter.PipeParameters;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class PipeDataNodeTaskBuilderTest {
+
+ @Test
+ public void testBlendUserAndSystemParametersDoesNotMutateOriginal() {
+ final Map<String, String> attributes = new HashMap<>();
+ attributes.put(PipeSourceConstant.EXTRACTOR_PATTERN_KEY, "root.sg.**");
+ final PipeParameters userParameters = new PipeParameters(attributes);
+
+ final PipeParameters blendedParameters =
+ PipeDataNodeTaskBuilder.blendUserAndSystemParameters(
+ userParameters, new PipeTaskMeta(MinimumProgressIndex.INSTANCE,
1));
+
+ Assert.assertEquals(
+ "root.sg.**",
blendedParameters.getStringByKeys(PipeSourceConstant.EXTRACTOR_PATTERN_KEY));
+
Assert.assertFalse(blendedParameters.hasAttribute(SystemConstant.RESTART_OR_NEWLY_ADDED_KEY));
+
Assert.assertFalse(userParameters.hasAttribute(SystemConstant.RESTART_OR_NEWLY_ADDED_KEY));
+ }
+
+ @Test
+ public void testBlendUserAndSystemParametersMarksRestartOrNewlyAddedTask() {
+ final PipeParameters restartedParameters =
+ PipeDataNodeTaskBuilder.blendUserAndSystemParameters(
+ new PipeParameters(new HashMap<>()),
+ new PipeTaskMeta(new SimpleProgressIndex(1, 2L), 1));
+
+ Assert.assertEquals(
+ Boolean.TRUE.toString(),
+
restartedParameters.getStringByKeys(SystemConstant.RESTART_OR_NEWLY_ADDED_KEY));
+
+ final PipeTaskMeta newlyAddedTaskMeta =
+ new PipeTaskMeta(MinimumProgressIndex.INSTANCE, 1).markAsNewlyAdded();
+ final PipeParameters newlyAddedParameters =
+ PipeDataNodeTaskBuilder.blendUserAndSystemParameters(
+ new PipeParameters(new HashMap<>()), newlyAddedTaskMeta);
+
+ Assert.assertEquals(
+ Boolean.TRUE.toString(),
+
newlyAddedParameters.getStringByKeys(SystemConstant.RESTART_OR_NEWLY_ADDED_KEY));
+ }
+
+ @Test
+ public void testPreprocessParametersInjectsRuntimeDefaults() {
+ final Map<String, String> sourceAttributes = new HashMap<>();
+ sourceAttributes.put(PipeSourceConstant.EXTRACTOR_INCLUSION_KEY,
"data.delete");
+ final PipeParameters sourceParameters = new
PipeParameters(sourceAttributes);
+ final PipeParameters sinkParameters = new PipeParameters(new HashMap<>());
+
+ PipeDataNodeTaskBuilder.preprocessParameters(sourceParameters,
sinkParameters);
+
+ Assert.assertEquals(
+ Boolean.FALSE.toString(),
+
sinkParameters.getStringByKeys(PipeSinkConstant.CONNECTOR_REALTIME_FIRST_KEY));
+ Assert.assertEquals(
+ Boolean.TRUE.toString(),
+
sinkParameters.getStringByKeys(PipeSinkConstant.SINK_ENABLE_SEND_TSFILE_LIMIT));
+ }
+
+ @Test
+ public void
testPreprocessParametersInjectsEventUserForExternalWriteBackSink() {
+ final Map<String, String> sourceAttributes = new HashMap<>();
+ sourceAttributes.put(PipeSourceConstant.EXTRACTOR_KEY, "external-source");
+
+ final Map<String, String> sinkAttributes = new HashMap<>();
+ sinkAttributes.put(
+ PipeSinkConstant.CONNECTOR_KEY,
BuiltinPipePlugin.WRITE_BACK_CONNECTOR.getPipePluginName());
+ final PipeParameters sinkParameters = new PipeParameters(sinkAttributes);
+
+ PipeDataNodeTaskBuilder.preprocessParameters(
+ new PipeParameters(sourceAttributes), sinkParameters);
+
+ Assert.assertEquals(
+ Boolean.TRUE.toString(),
+
sinkParameters.getStringByKeys(PipeSinkConstant.CONNECTOR_USE_EVENT_USER_NAME_KEY));
+ }
+}
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtaskManagerTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtaskManagerTest.java
new file mode 100644
index 00000000000..9545d71d965
--- /dev/null
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtaskManagerTest.java
@@ -0,0 +1,95 @@
+/*
+ * 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.db.pipe.agent.task.subtask.sink;
+
+import org.apache.iotdb.commons.pipe.agent.plugin.builtin.BuiltinPipePlugin;
+import org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant;
+import org.apache.iotdb.commons.pipe.config.constant.SystemConstant;
+import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameters;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class PipeSinkSubtaskManagerTest {
+
+ @Test
+ public void
testGenerateAttributeSortedStringAddsRegionPrefixAndIgnoresRestartFlag() {
+ final Map<String, String> attributes = new HashMap<>();
+ attributes.put("z", "1");
+ attributes.put("a", "2");
+ attributes.put(SystemConstant.RESTART_OR_NEWLY_ADDED_KEY,
Boolean.TRUE.toString());
+
+ Assert.assertEquals(
+ "data_region_-1_{a=2, z=1}",
+ PipeSinkSubtaskManager.generateAttributeSortedString(
+ new PipeParameters(new HashMap<>(attributes)), -1));
+
+ attributes.put(PipeSinkConstant.CONNECTOR_SERIALIZE_BY_REGION_KEY,
Boolean.FALSE.toString());
+ Assert.assertEquals(
+ "data_{a=2, connector.serialize-by-region=false, z=1}",
+ PipeSinkSubtaskManager.generateAttributeSortedString(
+ new PipeParameters(new HashMap<>(attributes)), -1));
+ }
+
+ @Test
+ public void testCalculateSinkSubtaskNumForDataRegionSink() {
+ final Map<String, String> parallelAttributes = new HashMap<>();
+ parallelAttributes.put(
+ PipeSinkConstant.CONNECTOR_SERIALIZE_BY_REGION_KEY,
Boolean.FALSE.toString());
+
parallelAttributes.put(PipeSinkConstant.CONNECTOR_IOTDB_PARALLEL_TASKS_KEY,
"3");
+ Assert.assertEquals(
+ 3,
+ PipeSinkSubtaskManager.calculateSinkSubtaskNum(new
PipeParameters(parallelAttributes), -1));
+
+ final Map<String, String> serializedAttributes = new HashMap<>();
+ serializedAttributes.put(
+ PipeSinkConstant.CONNECTOR_SERIALIZE_BY_REGION_KEY,
Boolean.TRUE.toString());
+
serializedAttributes.put(PipeSinkConstant.CONNECTOR_IOTDB_PARALLEL_TASKS_KEY,
"3");
+ Assert.assertEquals(
+ 1,
+ PipeSinkSubtaskManager.calculateSinkSubtaskNum(
+ new PipeParameters(serializedAttributes), -1));
+ }
+
+ @Test
+ public void
testCalculateSinkSubtaskNumUsesSingleThreadDefaultSinkAndSchemaRegionLimit() {
+ final Map<String, String> singleThreadAttributes = new HashMap<>();
+ singleThreadAttributes.put(
+ PipeSinkConstant.CONNECTOR_SERIALIZE_BY_REGION_KEY,
Boolean.FALSE.toString());
+ singleThreadAttributes.put(
+ PipeSinkConstant.CONNECTOR_KEY,
BuiltinPipePlugin.OPC_UA_SINK.getPipePluginName());
+ Assert.assertEquals(
+ 1,
+ PipeSinkSubtaskManager.calculateSinkSubtaskNum(
+ new PipeParameters(singleThreadAttributes), -1));
+
+ final Map<String, String> schemaRegionAttributes = new HashMap<>();
+ schemaRegionAttributes.put(
+ PipeSinkConstant.CONNECTOR_SERIALIZE_BY_REGION_KEY,
Boolean.FALSE.toString());
+
schemaRegionAttributes.put(PipeSinkConstant.CONNECTOR_IOTDB_PARALLEL_TASKS_KEY,
"3");
+ Assert.assertEquals(
+ 1,
+ PipeSinkSubtaskManager.calculateSinkSubtaskNum(
+ new PipeParameters(schemaRegionAttributes), Integer.MAX_VALUE));
+ }
+}
diff --git
a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/plugin/builtin/BuiltinPipePlugin.java
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/plugin/builtin/BuiltinPipePlugin.java
index 155e5c7ea05..e5c041e8ec6 100644
---
a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/plugin/builtin/BuiltinPipePlugin.java
+++
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/plugin/builtin/BuiltinPipePlugin.java
@@ -126,6 +126,16 @@ public enum BuiltinPipePlugin {
return className;
}
+ // used to distinguish between builtin and external sources
+ public static final Set<String> BUILTIN_SOURCES =
+ Collections.unmodifiableSet(
+ new HashSet<>(
+ Arrays.asList(
+ DO_NOTHING_EXTRACTOR.getPipePluginName().toLowerCase(),
+ IOTDB_EXTRACTOR.getPipePluginName().toLowerCase(),
+ DO_NOTHING_SOURCE.getPipePluginName().toLowerCase(),
+ IOTDB_SOURCE.getPipePluginName().toLowerCase())));
+
public static final Set<String> SHOW_PIPE_PLUGINS_BLACKLIST =
Collections.unmodifiableSet(
new HashSet<>(
diff --git
a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/task/PipeTaskAgent.java
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/task/PipeTaskAgent.java
index 66ed052ea6f..ee629b065df 100644
---
a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/task/PipeTaskAgent.java
+++
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/task/PipeTaskAgent.java
@@ -510,11 +510,7 @@ public abstract class PipeTaskAgent {
final String pipeName =
pipeMetaFromCoordinator.getStaticMeta().getPipeName();
final long creationTime =
pipeMetaFromCoordinator.getStaticMeta().getCreationTime();
- calculateMemoryUsage(
- pipeMetaFromCoordinator.getStaticMeta(),
- pipeMetaFromCoordinator.getStaticMeta().getExtractorParameters(),
- pipeMetaFromCoordinator.getStaticMeta().getProcessorParameters(),
- pipeMetaFromCoordinator.getStaticMeta().getConnectorParameters());
+ calculateMemoryUsage(pipeMetaFromCoordinator);
final PipeMeta existedPipeMeta = pipeMetaKeeper.getPipeMeta(pipeName);
if (existedPipeMeta != null) {
@@ -561,6 +557,15 @@ public abstract class PipeTaskAgent {
return needToStartPipe;
}
+ protected void calculateMemoryUsage(final PipeMeta pipeMetaFromCoordinator)
+ throws IllegalPathException {
+ calculateMemoryUsage(
+ pipeMetaFromCoordinator.getStaticMeta(),
+ pipeMetaFromCoordinator.getStaticMeta().getExtractorParameters(),
+ pipeMetaFromCoordinator.getStaticMeta().getProcessorParameters(),
+ pipeMetaFromCoordinator.getStaticMeta().getConnectorParameters());
+ }
+
protected void calculateMemoryUsage(
final PipeStaticMeta staticMeta,
final PipeParameters extractorParameters,
diff --git
a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/task/meta/PipeRuntimeMeta.java
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/task/meta/PipeRuntimeMeta.java
index 0e52297edff..63aa7b86c43 100644
---
a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/task/meta/PipeRuntimeMeta.java
+++
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/task/meta/PipeRuntimeMeta.java
@@ -145,6 +145,15 @@ public class PipeRuntimeMeta {
this.isStoppedByRuntimeException.set(isStoppedByRuntimeException);
}
+ /**
+ * We use negative regionId to identify the external pipe source, which is
not a consensus group
+ * id. Then we can reuse the regionId to schedule the external pipe source
and store the progress
+ * information.
+ */
+ public static boolean isSourceExternal(int regionId) {
+ return regionId < 0;
+ }
+
public ByteBuffer serialize() throws IOException {
PublicBAOS byteArrayOutputStream = new PublicBAOS();
DataOutputStream outputStream = new
DataOutputStream(byteArrayOutputStream);
diff --git
a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeSinkConstant.java
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeSinkConstant.java
index d19d6d2dc16..083d4bfbf73 100644
---
a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeSinkConstant.java
+++
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeSinkConstant.java
@@ -359,6 +359,9 @@ public class PipeSinkConstant {
public static final String CONNECTOR_OPC_DA_PROGID_KEY =
"connector.opcda.progid";
public static final String SINK_OPC_DA_PROGID_KEY = "sink.opcda.progid";
+ public static final String CONNECTOR_USE_EVENT_USER_NAME_KEY =
"connector.use-event-user-name";
+ public static final String SINK_USE_EVENT_USER_NAME_KEY =
"sink.use-event-user-name";
+ public static final boolean CONNECTOR_USE_EVENT_USER_NAME_DEFAULT_VALUE =
false;
private PipeSinkConstant() {
throw new IllegalStateException("Utility class");
