[Pipe] Fix processor worker starvation on pipe stop

iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java

@@ -2588,4 +2588,6 @@ private DataNodePipeMessages() {}
       "Topic config for %s is unavailable during consensus subscription setup";
   public static final String LOG_FAILED_TO_RELEASE_TSFILE_PARSER_MEMORY_FOR_PIPE_ARG_CREATION_TIME_ARG_IN_DATAREGION_ARG_BECAUSE_NO_RESERVATION_EXISTS_BB8321C0 =
       "Failed to release TsFile parser memory for Pipe {} (creation time {}) in DataRegion {} because no reservation exists.";
+  public static final String LOG_PIPE_PROCESSOR_WORKER_ARG_HAS_BEEN_PROCESSING_THE_SAME_EVENT_FOR_ARG_MS_PIPE_ARG_DATAREGION_ARG_SUBTASK_ARG_EVENT_ARG_THREAD_STATE_ARG_STACK_ARG_63B40775 =
+      "Pipe processor worker {} has been processing the same event for {} ms. Pipe: {}, DataRegion: {}, subtask: {}, event: {}, thread state: {}. Stack:{}";
 }

iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java

@@ -2416,4 +2416,6 @@ private DataNodePipeMessages() {}
       "共识订阅设置期间 topic %s 的配置不可用";
   public static final String LOG_FAILED_TO_RELEASE_TSFILE_PARSER_MEMORY_FOR_PIPE_ARG_CREATION_TIME_ARG_IN_DATAREGION_ARG_BECAUSE_NO_RESERVATION_EXISTS_BB8321C0 =
       "无法释放 Pipe {}（创建时间 {}）在 DataRegion {} 中的 TsFile 解析器内存，因为不存在对应的预留。";
+  public static final String LOG_PIPE_PROCESSOR_WORKER_ARG_HAS_BEEN_PROCESSING_THE_SAME_EVENT_FOR_ARG_MS_PIPE_ARG_DATAREGION_ARG_SUBTASK_ARG_EVENT_ARG_THREAD_STATE_ARG_STACK_ARG_63B40775 =
+      "Pipe processor worker {} 已连续处理同一 event {} ms。Pipe：{}，DataRegion：{}，subtask：{}，event：{}，线程状态：{}。栈：{}";
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/connection/PipeEventCollector.java

@@ -28,6 +28,8 @@
 import org.apache.iotdb.commons.pipe.event.ProgressReportEvent;
 import org.apache.iotdb.db.i18n.DataNodePipeMessages;
 import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent;
+import org.apache.iotdb.db.pipe.agent.task.subtask.processor.PipeProcessorSubtaskExecutionGuard;
+import org.apache.iotdb.db.pipe.agent.task.subtask.processor.PipeProcessorSubtaskYieldException;
 import org.apache.iotdb.db.pipe.event.common.deletion.PipeDeleteDataNodeEvent;
 import org.apache.iotdb.db.pipe.event.common.heartbeat.PipeHeartbeatEvent;
 import org.apache.iotdb.db.pipe.event.common.tablet.PipeInsertNodeTabletInsertionEvent;
@@ -63,6 +65,8 @@ public class PipeEventCollector implements EventCollector {
   private final boolean skipParsing;
 
   private final boolean isUsedForConsensusPipe;
+  private PipeProcessorSubtaskExecutionGuard processorExecutionGuard =
+      PipeProcessorSubtaskExecutionGuard.disabled();
 
   private final AtomicInteger collectInvocationCount = new AtomicInteger(0);
   private boolean hasNoGeneratedEvent = true;
@@ -83,6 +87,11 @@ public PipeEventCollector(
     this.isUsedForConsensusPipe = isUsedInConsensusPipe;
   }
 
+  public void setProcessorExecutionGuard(
+      final PipeProcessorSubtaskExecutionGuard processorExecutionGuard) {
+    this.processorExecutionGuard = processorExecutionGuard;
+  }
+
   @Override
   public void collect(final Event event) {
     try {
@@ -97,6 +106,8 @@ public void collect(final Event event) {
       } else if (!(event instanceof ProgressReportEvent)) {
         collectEvent(event);
       }
+    } catch (final PipeProcessorSubtaskYieldException e) {
+      throw e;
     } catch (final PipeException e) {
       throw e;
     } catch (final Exception e) {
@@ -131,7 +142,7 @@ private void parseAndCollectEvent(final PipeRawTabletInsertionEvent sourceEvent)
   }
 
   private void parseAndCollectEvent(final PipeTsFileInsertionEvent sourceEvent) throws Exception {
-    if (!sourceEvent.waitForTsFileClose()) {
+    if (!sourceEvent.waitForTsFileClose(processorExecutionGuard)) {
       LOGGER.warn(
           DataNodePipeMessages.PIPE_SKIPPING_TEMPORARY_TSFILE_WHICH_SHOULDN_T,
           sourceEvent.getTsFile());
@@ -148,7 +159,9 @@ private void parseAndCollectEvent(final PipeTsFileInsertionEvent sourceEvent) th
     }
 
     sourceEvent.consumeTabletInsertionEventsWithRetry(
-        this::collectParsedRawTableEvent, "PipeEventCollector::parseAndCollectEvent");
+        this::collectParsedRawTableEvent,
+        "PipeEventCollector::parseAndCollectEvent",
+        processorExecutionGuard);
     sourceEvent.close();
     if (sourceEvent.isGeneratedByHistoricalExtractor()) {
       PipeTerminateEvent.markHistoricalTsFileSplit(

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

@@ -57,6 +57,7 @@
 
 import java.util.Objects;
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
 public class PipeProcessorSubtask extends PipeReportableSubtask {
@@ -74,6 +75,11 @@ public class PipeProcessorSubtask extends PipeReportableSubtask {
   private final EventSupplier inputEventSupplier;
   private final PipeProcessor pipeProcessor;
   private final PipeEventCollector outputEventCollector;
+  private final PipeProcessorSubtaskExecutionGuard executionGuard =
+      new PipeProcessorSubtaskExecutionGuard();
+  private final AtomicBoolean isResumingFromYield = new AtomicBoolean(false);
+  private final AtomicReference<EventProcessingContext> eventProcessingContext =
+      new AtomicReference<>();
 
   // This variable is used to distinguish between old and new subtasks before and after stuck
   // restart.
@@ -94,6 +100,7 @@ public PipeProcessorSubtask(
     this.inputEventSupplier = inputEventSupplier;
     this.pipeProcessor = pipeProcessor;
     this.outputEventCollector = outputEventCollector;
+    this.outputEventCollector.setProcessorExecutionGuard(executionGuard);
     this.subtaskCreationTime = System.currentTimeMillis();
 
     // Only register dataRegions
@@ -106,7 +113,7 @@ public PipeProcessorSubtask(
   @Override
   public void bindExecutors(
       final ListeningExecutorService subtaskWorkerThreadPoolExecutor,
-      final ListeningScheduledExecutorService ignoredScheduledExecutor,
+      final ListeningScheduledExecutorService subtaskWorkerScheduledExecutor,
       final ExecutorService ignored,
       final PipeSubtaskScheduler subtaskScheduler) {
     this.subtaskWorkerThreadPoolExecutor = subtaskWorkerThreadPoolExecutor;
@@ -117,19 +124,31 @@ public void bindExecutors(
       synchronized (PipeProcessorSubtaskWorkerManager.class) {
         if (subtaskWorkerManager.get() == null) {
           subtaskWorkerManager.set(
-              new PipeProcessorSubtaskWorkerManager(subtaskWorkerThreadPoolExecutor));
+              new PipeProcessorSubtaskWorkerManager(
+                  subtaskWorkerThreadPoolExecutor, subtaskWorkerScheduledExecutor));
         }
       }
     }
     subtaskWorkerManager.get().schedule(this);
   }
 
+  @Override
+  public Boolean call() throws Exception {
+    executionGuard.enter();
+    try {
+      return super.call();
+    } finally {
+      executionGuard.exit();
+    }
+  }
+
   @Override
   protected boolean executeOnce() throws Exception {
     if (isClosed.get()) {
       return false;
     }
 
+    executionGuard.check();
     final Event event =
         lastEvent != null
             ? lastEvent
@@ -141,7 +160,13 @@ protected boolean executeOnce() throws Exception {
       return false;
     }
 
-    outputEventCollector.resetFlags();
+    executionGuard.check();
+    if (!isResumingFromYield.getAndSet(false)) {
+      outputEventCollector.resetFlags();
+    }
+    final EventProcessingContext currentEventProcessingContext =
+        new EventProcessingContext(event, System.nanoTime());
+    eventProcessingContext.set(currentEventProcessingContext);
     try {
       if (event instanceof EnrichedEvent) {
         ((EnrichedEvent) event).throwIfNoPrivilege();
@@ -178,13 +203,16 @@ protected boolean executeOnce() throws Exception {
                 event1 -> {
                   try {
                     pipeProcessor.process(event1, outputEventCollector);
+                  } catch (PipeProcessorSubtaskYieldException e) {
+                    throw e;
                   } catch (PipeRuntimeOutOfMemoryCriticalException e) {
                     throw e;
                   } catch (Exception e) {
                     throw new PipeException(e.getMessage(), e);
                   }
                 },
-                "PipeProcessorSubtask::executeOnce");
+                "PipeProcessorSubtask::executeOnce",
+                executionGuard);
             tsFileInsertionEvent.close();
             if (tsFileInsertionEvent.isGeneratedByHistoricalExtractor()) {
               PipeTerminateEvent.markHistoricalTsFileSplit(
@@ -242,6 +270,9 @@ protected boolean executeOnce() throws Exception {
             .enrichWithCommitterKeyAndCommitId((EnrichedEvent) event, creationTime, regionId);
       }
       decreaseReferenceCountAndReleaseLastEvent(event, shouldReport);
+    } catch (final PipeProcessorSubtaskYieldException e) {
+      isResumingFromYield.set(true);
+      throw e;
     } catch (final PipeRuntimeOutOfMemoryCriticalException e) {
       recordResourceFailure(event, PipeResourceFailureType.MEMORY_TIMEOUT);
       PipeLogger.log(
@@ -250,7 +281,12 @@ protected boolean executeOnce() throws Exception {
           e.getMessage());
       return false;
     } catch (final Exception e) {
-      if (ExceptionUtils.getRootCause(e) instanceof PipeRuntimeOutOfMemoryCriticalException) {
+      final Throwable rootCause = ExceptionUtils.getRootCause(e);
+      if (rootCause instanceof PipeProcessorSubtaskYieldException) {
+        isResumingFromYield.set(true);
+        throw (PipeProcessorSubtaskYieldException) rootCause;
+      }
+      if (rootCause instanceof PipeRuntimeOutOfMemoryCriticalException) {
         recordResourceFailure(event, PipeResourceFailureType.MEMORY_TIMEOUT);
         PipeLogger.log(
             LOGGER::info,
@@ -275,6 +311,8 @@ protected boolean executeOnce() throws Exception {
             e.getMessage() != null ? " Message: " + e.getMessage() : "");
         clearReferenceCountAndReleaseLastEvent(event);
       }
+    } finally {
+      eventProcessingContext.compareAndSet(currentEventProcessingContext, null);
     }
 
     return true;
@@ -287,6 +325,20 @@ public void submitSelf() {
     // and the worker will be submitted to the executor
   }
 
+  @Override
+  protected void onAllowSubmittingSelf() {
+    executionGuard.start();
+  }
+
+  @Override
+  protected void onDisallowSubmittingSelf() {
+    executionGuard.stop();
+    final Event event = lastEvent;
+    if (event instanceof PipeTsFileInsertionEvent) {
+      ((PipeTsFileInsertionEvent) event).cancelTsFileParserMemoryReservationIfPending();
+    }
+  }
+
   public boolean isStoppedByException() {
     return lastEvent instanceof EnrichedEvent && retryCount.get() > MAX_RETRY_TIMES;
   }
@@ -316,6 +368,29 @@ boolean isClosed() {
     return isClosed.get();
   }
 
+  EventProcessingContext getEventProcessingContext() {
+    return eventProcessingContext.get();
+  }
+
+  static final class EventProcessingContext {
+
+    private final Event event;
+    private final long startTimeInNanos;
+
+    EventProcessingContext(final Event event, final long startTimeInNanos) {
+      this.event = event;
+      this.startTimeInNanos = startTimeInNanos;
+    }
+
+    Event getEvent() {
+      return event;
+    }
+
+    long getStartTimeInNanos() {
+      return startTimeInNanos;
+    }
+  }
+
   @Override
   public boolean equals(final Object obj) {
     if (this == obj) {

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

@@ -0,0 +1,110 @@
+/*
+ * 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.processor;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
+
+/**
+ * Guards one processor subtask invocation against concurrent STOP/START operations.
+ *
+ * <p>An invocation captures the current execution epoch. STOP invalidates that epoch before START
+ * can enable a new one, so an invocation started before STOP must yield even if the pipe is started
+ * again immediately.
+ */
+public class PipeProcessorSubtaskExecutionGuard {
+
+  private static final PipeProcessorSubtaskExecutionGuard DISABLED_GUARD =
+      new PipeProcessorSubtaskExecutionGuard(false);
+
+  private final boolean enabled;
+  private final AtomicBoolean isRunning = new AtomicBoolean(false);
+  private final AtomicLong executionEpoch = new AtomicLong(0);
+  private final ThreadLocal<Long> invocationEpoch = new ThreadLocal<>();
+
+  public PipeProcessorSubtaskExecutionGuard() {
+    this(true);
+  }
+
+  private PipeProcessorSubtaskExecutionGuard(final boolean enabled) {
+    this.enabled = enabled;
+  }
+
+  public static PipeProcessorSubtaskExecutionGuard disabled() {
+    return DISABLED_GUARD;
+  }
+
+  public boolean isEnabled() {
+    return enabled;
+  }
+
+  void start() {
+    if (enabled) {
+      isRunning.set(true);
+    }
+  }
+
+  void stop() {
+    if (enabled) {
+      isRunning.set(false);
+      executionEpoch.incrementAndGet();
+    }
+  }
+
+  void enter() {
+    if (!enabled) {
+      return;
+    }
+
+    final long currentEpoch = executionEpoch.get();
+    invocationEpoch.set(currentEpoch);
+    if (!isRunning.get() || currentEpoch != executionEpoch.get()) {
+      invocationEpoch.remove();
+      throw PipeProcessorSubtaskYieldException.pauseRequested();
+    }
+  }
+
+  void exit() {
+    if (enabled) {
+      invocationEpoch.remove();
+    }
+  }
+
+  public void check() {
+    if (!isCurrentInvocationValid()) {
+      throw PipeProcessorSubtaskYieldException.pauseRequested();
+    }
+  }
+
+  public boolean isCurrentInvocationValid() {
+    if (!enabled) {
+      return true;
+    }
+
+    final Long currentInvocationEpoch = invocationEpoch.get();
+    return currentInvocationEpoch != null
+        && isRunning.get()
+        && currentInvocationEpoch == executionEpoch.get();
+  }
+
+  public void yieldIfParserNotAdmitted() {
+    throw PipeProcessorSubtaskYieldException.parserNotAdmitted();
+  }
+}