Fix pipe task start failures: report, roll back and recover

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

@@ -358,6 +358,8 @@ public final class DataNodePipeMessages {
       "Successfully persisted all pipe's info to configNode.";
   public static final String THE_EXECUTOR_AND_HAS_BEEN_SUCCESSFULLY_SHUTDOWN =
       "The executor {} and {} has been successfully shutdown.";
+  public static final String LOG_THE_SHARED_CONSENSUS_EXECUTOR_HAS_BEEN_SHUT_DOWN_RECREATING_IT_C14080E3 =
+      "The shared consensus executor has been shut down. Recreating it.";
 
   // ===================== EVENT =====================
 

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

@@ -340,6 +340,8 @@ public final class DataNodePipeMessages {
       "成功将所有 Pipe 信息持久化到 ConfigNode。";
   public static final String THE_EXECUTOR_AND_HAS_BEEN_SUCCESSFULLY_SHUTDOWN =
       "执行器 {} 和 {} 已成功关闭。";
+  public static final String LOG_THE_SHARED_CONSENSUS_EXECUTOR_HAS_BEEN_SHUT_DOWN_RECREATING_IT_C14080E3 =
+      "共享的共识执行器已被关闭，正在重新创建该执行器。";
 
   // ===================== EVENT =====================
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/execution/PipeSubtaskExecutorManager.java

@@ -20,17 +20,25 @@
 package org.apache.iotdb.db.pipe.agent.task.execution;
 
 import org.apache.iotdb.commons.consensus.iotv2.container.IoTV2GlobalComponentContainer;
+import org.apache.iotdb.commons.pipe.agent.task.execution.PipeSubtaskExecutor;
 import org.apache.iotdb.commons.subscription.config.SubscriptionConfig;
+import org.apache.iotdb.db.i18n.DataNodePipeMessages;
 import org.apache.iotdb.db.pipe.consensus.IoTConsensusV2SubtaskExecutor;
 import org.apache.iotdb.db.subscription.task.execution.SubscriptionSubtaskExecutor;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import java.util.function.Supplier;
 
 /**
  * PipeTaskExecutor is responsible for executing the pipe tasks, and it is scheduled by the
  * PipeTaskScheduler. It is a singleton class.
  */
 public class PipeSubtaskExecutorManager {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(PipeSubtaskExecutorManager.class);
+
   private final PipeProcessorSubtaskExecutor processorExecutor;
   private final Supplier<PipeSinkSubtaskExecutor> connectorExecutorSupplier;
   private volatile SubscriptionSubtaskExecutor subscriptionExecutor;
@@ -43,9 +51,22 @@ public Supplier<PipeSinkSubtaskExecutor> getConnectorExecutorSupplier() {
     return connectorExecutorSupplier;
   }
 
-  public IoTConsensusV2SubtaskExecutor getConsensusExecutor() {
-    return (IoTConsensusV2SubtaskExecutor)
-        IoTV2GlobalComponentContainer.getInstance().getConsensusExecutor();
+  public synchronized IoTConsensusV2SubtaskExecutor getConsensusExecutor() {
+    final IoTV2GlobalComponentContainer container = IoTV2GlobalComponentContainer.getInstance();
+    final PipeSubtaskExecutor consensusExecutor = container.getConsensusExecutor();
+    if (consensusExecutor instanceof IoTConsensusV2SubtaskExecutor
+        && !consensusExecutor.isShutdown()) {
+      return (IoTConsensusV2SubtaskExecutor) consensusExecutor;
+    }
+
+    // The consensus executor is shared by all consensus pipes. Recreate it if it has been shut
+    // down, otherwise the following consensus pipes can never start their sink subtasks.
+    LOGGER.warn(
+        DataNodePipeMessages
+            .LOG_THE_SHARED_CONSENSUS_EXECUTOR_HAS_BEEN_SHUT_DOWN_RECREATING_IT_C14080E3);
+    final IoTConsensusV2SubtaskExecutor newConsensusExecutor = new IoTConsensusV2SubtaskExecutor();
+    container.setConsensusExecutor(newConsensusExecutor);
+    return newConsensusExecutor;
   }
 
   public SubscriptionSubtaskExecutor getSubscriptionExecutor() {

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/PipeDataNodeTaskAgentTest.java

@@ -25,6 +25,7 @@
 import org.apache.iotdb.commons.consensus.index.impl.SimpleProgressIndex;
 import org.apache.iotdb.commons.exception.pipe.PipeRuntimeCriticalException;
 import org.apache.iotdb.commons.exception.pipe.PipeRuntimeSinkCriticalException;
+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.PipeTaskManager;
 import org.apache.iotdb.commons.pipe.agent.task.meta.PipeMeta;
@@ -41,9 +42,11 @@
 import org.awaitility.Awaitility;
 import org.junit.Assert;
 import org.junit.Test;
+import org.mockito.Mockito;
 
 import java.io.IOException;
 import java.lang.reflect.Field;
+import java.lang.reflect.Method;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
@@ -576,6 +579,43 @@ public void testCarryOverRequiresStableLeaderAndLocalTask() {
             .getProgressIndex());
   }
 
+  /**
+   * A task that fails to start is dropped and unregistered, so the following meta sync creates it
+   * again.
+   */
+  @Test
+  public void testFailedTaskStartIsRolledBackForRetry() throws Exception {
+    final PipeDataNodeTaskAgent agent = newActiveTaskAgent();
+    final PipeMetaKeeper keeper = getField(agent, "pipeMetaKeeper");
+    final PipeTaskManager manager = getField(agent, "pipeTaskManager");
+
+    final int localNodeId = IoTDBDescriptor.getInstance().getConfig().getDataNodeId();
+    final PipeTaskMeta taskMeta = new PipeTaskMeta(MinimumProgressIndex.INSTANCE, localNodeId);
+    final ConcurrentMap<Integer, PipeTaskMeta> taskMetaMap = new ConcurrentHashMap<>();
+    taskMetaMap.put(REGION_ID, taskMeta);
+    final PipeMeta pipeMeta =
+        new PipeMeta(
+            new PipeStaticMeta(
+                "failedStartPipe", 1L, new HashMap<>(), new HashMap<>(), new HashMap<>()),
+            new PipeRuntimeMeta(taskMetaMap));
+    pipeMeta.getRuntimeMeta().getStatus().set(PipeStatus.RUNNING);
+    keeper.addPipeMeta(pipeMeta);
+
+    final PipeTask failedTask = Mockito.mock(PipeTask.class);
+    Mockito.doThrow(new PipeException("failed to start the sink")).when(failedTask).start();
+    manager.addPipeTask(pipeMeta.getStaticMeta(), REGION_ID, failedTask);
+
+    final Method startPipeTask =
+        PipeTaskAgent.class.getDeclaredMethod("startPipeTask", int.class, PipeStaticMeta.class);
+    startPipeTask.setAccessible(true);
+    startPipeTask.invoke(agent, REGION_ID, pipeMeta.getStaticMeta());
+
+    Mockito.verify(failedTask).drop();
+    Assert.assertNull(manager.getPipeTask(pipeMeta.getStaticMeta(), REGION_ID));
+    Assert.assertNull(pipeMeta.getRuntimeMeta().getConsensusGroupId2TaskMetaMap().get(REGION_ID));
+    Assert.assertEquals(PipeStatus.RUNNING, pipeMeta.getRuntimeMeta().getStatus().get());
+  }
+
   private PipeStaticMeta createStaticMeta(final long creationTime, final boolean historyEnabled) {
     final Map<String, String> sourceAttributes = new HashMap<>();
     sourceAttributes.put(

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/PipeSinkSubtaskExecutorTest.java

@@ -20,13 +20,24 @@
 package org.apache.iotdb.db.pipe.agent.task;
 
 import org.apache.iotdb.commons.pipe.agent.task.connection.UnboundedBlockingPendingQueue;
+import org.apache.iotdb.commons.pipe.agent.task.subtask.PipeSubtask;
+import org.apache.iotdb.commons.pipe.sink.protocol.PipeSinkWithSchedulingDelay;
 import org.apache.iotdb.db.pipe.agent.task.execution.PipeSinkSubtaskExecutor;
 import org.apache.iotdb.db.pipe.agent.task.subtask.sink.PipeSinkSubtask;
 import org.apache.iotdb.pipe.api.PipeConnector;
 
+import org.junit.Assert;
 import org.junit.Before;
+import org.junit.Test;
 import org.mockito.Mockito;
 
+import java.lang.reflect.Field;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+
 import static org.mockito.Mockito.mock;
 
 public class PipeSinkSubtaskExecutorTest extends PipeSubtaskExecutorTest {
@@ -46,4 +57,175 @@ public void setUp() throws Exception {
                 mock(UnboundedBlockingPendingQueue.class),
                 mock(PipeConnector.class)));
   }
+
+  /**
+   * A shut-down executor must reject {@code start()} and must not leave the subtask in the
+   * "submitting" state.
+   */
+  @Test
+  public void testStartAfterShutdown() {
+    executor.register(subtask);
+    executor.shutdown();
+    Assert.assertTrue(executor.isShutdown());
+
+    for (int i = 0; i < 2; ++i) {
+      try {
+        executor.start(subtask.getTaskID());
+        Assert.fail("The start of a subtask should fail if the executor is shut down");
+      } catch (final RejectedExecutionException e) {
+        // expected
+      }
+      Assert.assertFalse(subtask.isSubmittingSelf());
+    }
+  }
+
+  /** A rejected self-submission must be reported instead of being swallowed. */
+  @Test
+  public void testSubmitSelfAfterShutdown() throws Exception {
+    executor.register(subtask);
+    executor.shutdown();
+    allowSubmittingSelfAgain(subtask);
+
+    for (int i = 0; i < 2; ++i) {
+      try {
+        subtask.submitSelf();
+        Assert.fail("The submission of a subtask should fail if the executor is shut down");
+      } catch (final RejectedExecutionException e) {
+        // expected
+      }
+    }
+  }
+
+  /**
+   * A rejected delayed submission must clear the submitted flag, otherwise the following
+   * submissions would be skipped.
+   */
+  @Test
+  public void testDelayedSubmitSelfAfterShutdown() throws Exception {
+    final SinkWithSchedulingDelay sink = mock(SinkWithSchedulingDelay.class);
+    Mockito.when(sink.consumeSchedulingDelayMs()).thenReturn(1000L);
+
+    final PipeSubtask delayedSubtask =
+        Mockito.spy(
+            new PipeSinkSubtask(
+                "PipeConnectorSubtaskExecutorTest",
+                System.currentTimeMillis(),
+                "TestAttributeSortedString",
+                "TestAttributeSortedString",
+                0,
+                mock(UnboundedBlockingPendingQueue.class),
+                sink));
+
+    executor.register(delayedSubtask);
+    executor.shutdown();
+    allowSubmittingSelfAgain(delayedSubtask);
+
+    for (int i = 0; i < 2; ++i) {
+      try {
+        delayedSubtask.submitSelf();
+        Assert.fail("The submission of a subtask should fail if the executor is shut down");
+      } catch (final RejectedExecutionException e) {
+        // expected
+      }
+    }
+  }
+
+  /**
+   * Stopping and restarting a subtask while its previous call is still running must not start a
+   * second concurrent call.
+   */
+  @Test
+  public void testStartWhileThePreviousCallIsStillRunning() throws Exception {
+    final AtomicInteger concurrentCallCount = new AtomicInteger(0);
+    final AtomicInteger maxConcurrentCallCount = new AtomicInteger(0);
+    final CountDownLatch firstCallEntered = new CountDownLatch(1);
+    final CountDownLatch firstCallReleased = new CountDownLatch(1);
+    final CountDownLatch secondCallEntered = new CountDownLatch(1);
+
+    final PipeSubtask blockedSubtask =
+        new PipeSinkSubtask(
+            "PipeConnectorSubtaskExecutorTest",
+            System.currentTimeMillis(),
+            "TestAttributeSortedString",
+            "TestAttributeSortedString",
+            0,
+            mock(UnboundedBlockingPendingQueue.class),
+            mock(PipeConnector.class)) {
+          @Override
+          protected boolean executeOnce() {
+            final int current = concurrentCallCount.incrementAndGet();
+            maxConcurrentCallCount.accumulateAndGet(current, Math::max);
+            if (current == 1) {
+              firstCallEntered.countDown();
+              try {
+                firstCallReleased.await(5, TimeUnit.SECONDS);
+              } catch (final InterruptedException e) {
+                Thread.currentThread().interrupt();
+              }
+            } else {
+              secondCallEntered.countDown();
+            }
+            concurrentCallCount.decrementAndGet();
+            return false;
+          }
+        };
+
+    executor.register(blockedSubtask);
+    executor.start(blockedSubtask.getTaskID());
+    Assert.assertTrue(firstCallEntered.await(5, TimeUnit.SECONDS));
+
+    executor.stop(blockedSubtask.getTaskID());
+    executor.start(blockedSubtask.getTaskID());
+
+    try {
+      Assert.assertFalse(
+          "The subtask must not be executed in parallel with itself",
+          secondCallEntered.await(1, TimeUnit.SECONDS));
+      Assert.assertEquals(1, maxConcurrentCallCount.get());
+    } finally {
+      executor.stop(blockedSubtask.getTaskID());
+      firstCallReleased.countDown();
+    }
+  }
+
+  /**
+   * {@code start()} must fail after shutdown even if a pending delayed submission already marked
+   * the subtask as submitted.
+   */
+  @Test
+  public void testStartAfterShutdownWithPendingDelayedSubmission() {
+    final SinkWithSchedulingDelay sink = mock(SinkWithSchedulingDelay.class);
+    Mockito.when(sink.consumeSchedulingDelayMs()).thenReturn(60_000L);
+
+    final PipeSubtask delayedSubtask =
+        new PipeSinkSubtask(
+            "PipeConnectorSubtaskExecutorTest",
+            System.currentTimeMillis(),
+            "TestAttributeSortedString",
+            "TestAttributeSortedString",
+            0,
+            mock(UnboundedBlockingPendingQueue.class),
+            sink);
+
+    executor.register(delayedSubtask);
+    executor.start(delayedSubtask.getTaskID());
+
+    executor.shutdown();
+    try {
+      executor.start(delayedSubtask.getTaskID());
+      Assert.fail("The start of a subtask should fail if the executor is shut down");
+    } catch (final RejectedExecutionException e) {
+      // expected
+    }
+  }
+
+  /** Clears the stop flag to simulate a subtask that has not been told to stop yet. */
+  private void allowSubmittingSelfAgain(final PipeSubtask subtask) throws Exception {
+    final Field shouldStopSubmittingSelfField =
+        PipeSubtask.class.getDeclaredField("shouldStopSubmittingSelf");
+    shouldStopSubmittingSelfField.setAccessible(true);
+    ((AtomicBoolean) shouldStopSubmittingSelfField.get(subtask)).set(false);
+  }
+
+  private interface SinkWithSchedulingDelay extends PipeConnector, PipeSinkWithSchedulingDelay {}
 }

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/PipeSubtaskExecutorTest.java

@@ -25,6 +25,9 @@
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Test;
+import org.mockito.Mockito;
+
+import java.util.concurrent.RejectedExecutionException;
 
 import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.times;
@@ -115,6 +118,27 @@ public void testStop() {
     Assert.assertFalse(subtask.isSubmittingSelf());
   }
 
+  /**
+   * A rejected submission must be rolled back, otherwise the following start() calls would take the
+   * "already running" branch and report a false success.
+   */
+  @Test
+  public void testStartWithRejectedSubmission() {
+    executor.register(subtask);
+    Mockito.doThrow(new RejectedExecutionException()).when(subtask).submitSelf();
+
+    try {
+      executor.start(subtask.getTaskID());
+      Assert.fail("The start of a subtask should fail if the submission is rejected");
+    } catch (final RejectedExecutionException e) {
+      // expected
+    }
+
+    verify(subtask).allowSubmittingSelf();
+    verify(subtask).disallowSubmittingSelf();
+    Assert.assertFalse(executor.isShutdown());
+  }
+
   @Test
   public void testDeregister() {
     // test unregister a subtask which is not in the map