From 49923f7e074e3601c364133f48ccb777abaeadb5 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Fri, 3 Jul 2026 15:43:53 +0800 Subject: [PATCH] Pipe: improve incremental memory estimation (#18049) * Pipe: improve incremental memory estimation * Pipe: add memory estimation unit tests * Pipe: localize memory check messages --- .../agent/task/PipeDataNodeTaskAgent.java | 306 +++++++++++++++--- .../task/builder/PipeDataNodeTaskBuilder.java | 60 +++- .../subtask/sink/PipeSinkSubtaskManager.java | 101 +++--- .../agent/task/PipeDataNodeTaskAgentTest.java | 159 +++++++++ .../builder/PipeDataNodeTaskBuilderTest.java | 111 +++++++ .../sink/PipeSinkSubtaskManagerTest.java | 95 ++++++ .../plugin/builtin/BuiltinPipePlugin.java | 10 + .../pipe/agent/task/PipeTaskAgent.java | 15 +- .../pipe/agent/task/meta/PipeRuntimeMeta.java | 9 + .../config/constant/PipeSinkConstant.java | 3 + 10 files changed, 771 insertions(+), 98 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/PipeDataNodeTaskAgentTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/builder/PipeDataNodeTaskBuilderTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtaskManagerTest.java 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 fba4b7d3ccfd1..5f4fc9b4cf044 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.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.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.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 @@ protected void createPipeTask( && !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 Map getAllConsensusPipe() { } @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> 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 @@ protected void calculateMemoryUsage( } } + private List> collectPipeTasksToBeCreated( + final PipeMeta pipeMetaFromCoordinator) throws IllegalPathException { + final PipeStaticMeta pipeStaticMeta = pipeMetaFromCoordinator.getStaticMeta(); + final PipeParameters sourceParameters = pipeStaticMeta.getExtractorParameters(); + final Set dataRegionIds = + new HashSet<>(StorageEngine.getInstance().getAllDataRegionIds()); + final Set schemaRegionIds = + new HashSet<>(SchemaEngine.getInstance().getAllSchemaRegionIds()); + final List> pipeTasksToBeCreated = new ArrayList<>(); + + for (final Map.Entry 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> pipeTasksToBeCreated, + final boolean ignoreRegisteredSinkSubtasks) { + long needMemory = 0; + int dataRegionTaskCount = 0; + final Set sinkSubtasksToBeCreated = new HashSet<>(); + + for (final Pair 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 @@ private boolean isInnerSource(final PipeParameters sourceParameters) { || 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 @@ private void calculateInsertNodeQueueMemory(final PipeParameters sourceParameter 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 @@ private long calculateTsFileParserMemory( 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 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 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 @@ private long calculateSendTsFileReadBufferMemory( 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 8d227d7c27210..cc4456b721b13 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 @@ 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 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 PipeDataNodeTask build() { // 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 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 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 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()); + } + } } 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 4385cc59504ce..6307dbba648b8 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 @@ 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 synchronized String register( final Supplier 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 synchronized String register( 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 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 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 0000000000000..0582a94170ddf --- /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 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 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 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 sourceAttributes = new HashMap<>(); + sourceAttributes.put(PipeSourceConstant.EXTRACTOR_HISTORY_ENABLE_KEY, Boolean.FALSE.toString()); + + final Map 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 0000000000000..e00f9500bbef7 --- /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 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 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 sourceAttributes = new HashMap<>(); + sourceAttributes.put(PipeSourceConstant.EXTRACTOR_KEY, "external-source"); + + final Map 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 0000000000000..9545d71d96569 --- /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 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 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 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 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 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 155e5c7ea054a..e5c041e8ec683 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 String getClassName() { return className; } + // used to distinguish between builtin and external sources + public static final Set 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 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 fcff317b3f337..642e90d14f257 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 @@ -502,11 +502,7 @@ private boolean createPipe(final PipeMeta pipeMetaFromCoordinator) throws Illega 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) { @@ -548,6 +544,15 @@ private boolean createPipe(final PipeMeta pipeMetaFromCoordinator) throws Illega 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 4724c619f9683..d9a579fb8176a 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 @@ -136,6 +136,15 @@ public void setIsStoppedByRuntimeException(boolean isStoppedByRuntimeException) 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 d19d6d2dc1653..083d4bfbf73f4 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 @@ private static String getDefaultConnectorOrSinkName(final PipeParameters paramet 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");