Persist parent and child procedures atomically before scheduling

iotdb-core/confignode/src/main/i18n/en/org/apache/iotdb/confignode/i18n/ProcedureMessages.java

@@ -1453,4 +1453,6 @@ private ProcedureMessages() {}
       "finish procedure without rollback";
   public static final String LOG_DELETE_CHILD_PROCEDURES_WITHOUT_ROLLBACK_6E62D1C1 =
       "delete child procedures without rollback";
+  public static final String LOG_PERSIST_PARENT_AND_CHILD_PROCEDURES_54E2D0D2 =
+      "persist parent and child procedures";
 }

iotdb-core/confignode/src/main/i18n/zh/org/apache/iotdb/confignode/i18n/ProcedureMessages.java

@@ -1536,4 +1536,5 @@ private ProcedureMessages() {}
   public static final String LOG_FINISH_PROCEDURE_WITHOUT_ROLLBACK_41F7A354 = "结束流程（不回滚）";
   public static final String LOG_DELETE_CHILD_PROCEDURES_WITHOUT_ROLLBACK_6E62D1C1 =
       "删除子流程（不回滚）";
+  public static final String LOG_PERSIST_PARENT_AND_CHILD_PROCEDURES_54E2D0D2 = "持久化父流程和子流程";
 }

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlan.java

@@ -84,6 +84,7 @@
 import org.apache.iotdb.confignode.consensus.request.write.pipe.task.SetPipeStatusWithStoppedByRuntimeExceptionPlanV2;
 import org.apache.iotdb.confignode.consensus.request.write.procedure.DeleteProcedurePlan;
 import org.apache.iotdb.confignode.consensus.request.write.procedure.UpdateProcedurePlan;
+import org.apache.iotdb.confignode.consensus.request.write.procedure.UpdateProceduresPlan;
 import org.apache.iotdb.confignode.consensus.request.write.quota.SetSpaceQuotaPlan;
 import org.apache.iotdb.confignode.consensus.request.write.quota.SetThrottleQuotaPlan;
 import org.apache.iotdb.confignode.consensus.request.write.region.CreateRegionGroupsPlan;
@@ -280,6 +281,9 @@ public static ConfigPhysicalPlan create(final ByteBuffer buffer) throws IOExcept
         case DeleteProcedure:
           plan = new DeleteProcedurePlan();
           break;
+        case UpdateProcedures:
+          plan = new UpdateProceduresPlan();
+          break;
         case UpdateProcedure:
           plan = new UpdateProcedurePlan();
           break;

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlanType.java

@@ -88,6 +88,7 @@ public enum ConfigPhysicalPlanType {
   /** Procedure. */
   UpdateProcedure((short) 500),
   DeleteProcedure((short) 501),
+  UpdateProcedures((short) 502),
 
   /** Authority. */
   Author((short) 600),

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/procedure/UpdateProceduresPlan.java

@@ -0,0 +1,93 @@
+/*
+ * 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.confignode.consensus.request.write.procedure;
+
+import org.apache.iotdb.confignode.consensus.request.ConfigPhysicalPlan;
+import org.apache.iotdb.confignode.consensus.request.ConfigPhysicalPlanType;
+import org.apache.iotdb.confignode.procedure.Procedure;
+import org.apache.iotdb.confignode.procedure.env.ConfigNodeProcedureEnv;
+
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+
+/** A single consensus entry for a parent procedure and its newly submitted children. */
+public class UpdateProceduresPlan extends ConfigPhysicalPlan {
+  private List<Procedure<ConfigNodeProcedureEnv>> procedures;
+
+  public UpdateProceduresPlan() {
+    super(ConfigPhysicalPlanType.UpdateProcedures);
+  }
+
+  public UpdateProceduresPlan(List<Procedure<ConfigNodeProcedureEnv>> procedures) {
+    this();
+    this.procedures = new ArrayList<>(procedures);
+  }
+
+  public List<Procedure<ConfigNodeProcedureEnv>> getProcedures() {
+    return procedures;
+  }
+
+  @Override
+  protected void serializeImpl(DataOutputStream stream) throws IOException {
+    stream.writeShort(getType().getPlanType());
+    stream.writeInt(procedures.size());
+    for (Procedure<ConfigNodeProcedureEnv> procedure : procedures) {
+      ByteBuffer buffer = new UpdateProcedurePlan(procedure).serializeToByteBuffer();
+      stream.writeInt(buffer.remaining());
+      stream.write(buffer.array(), buffer.position(), buffer.remaining());
+    }
+  }
+
+  @Override
+  protected void deserializeImpl(ByteBuffer buffer) throws IOException {
+    int size = buffer.getInt();
+    procedures = new ArrayList<>(size);
+    for (int i = 0; i < size; i++) {
+      int length = buffer.getInt();
+      // Some procedures read optional trailing fields. Isolate each record from the next one.
+      ByteBuffer procedureBuffer = buffer.slice();
+      procedureBuffer.limit(length);
+      procedures.add(
+          ((UpdateProcedurePlan) ConfigPhysicalPlan.Factory.create(procedureBuffer))
+              .getProcedure());
+      buffer.position(buffer.position() + length);
+    }
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (this == obj) {
+      return true;
+    }
+    if (obj == null || getClass() != obj.getClass()) {
+      return false;
+    }
+    return procedures.equals(((UpdateProceduresPlan) obj).procedures);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(procedures);
+  }
+}

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/ProcedureInfo.java

@@ -27,6 +27,7 @@
 import org.apache.iotdb.commons.utils.TestOnly;
 import org.apache.iotdb.confignode.consensus.request.write.procedure.DeleteProcedurePlan;
 import org.apache.iotdb.confignode.consensus.request.write.procedure.UpdateProcedurePlan;
+import org.apache.iotdb.confignode.consensus.request.write.procedure.UpdateProceduresPlan;
 import org.apache.iotdb.confignode.i18n.ConfigNodeMessages;
 import org.apache.iotdb.confignode.manager.ConfigManager;
 import org.apache.iotdb.confignode.procedure.Procedure;
@@ -142,6 +143,15 @@ public TSStatus updateProcedure(UpdateProcedurePlan updateProcedurePlan) {
     return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
   }
 
+  // Recovery and snapshots must not observe only part of a parent/children batch.
+  public synchronized TSStatus updateProcedures(UpdateProceduresPlan plan) {
+    for (Procedure<ConfigNodeProcedureEnv> procedure : plan.getProcedures()) {
+      procedureMap.put(procedure.getProcId(), procedure);
+      lastProcId.updateAndGet(id -> Math.max(id, procedure.getProcId()));
+    }
+    return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
+  }
+
   @TestOnly
   public TSStatus oldUpdateProcedure(UpdateProcedurePlan updateProcedurePlan) {
     Procedure procedure = updateProcedurePlan.getProcedure();
@@ -194,7 +204,7 @@ private static Optional<Procedure> loadProcedure(Path procedureFilePath) {
   }
 
   @Override
-  public boolean processTakeSnapshot(File snapshotDir) throws TException, IOException {
+  public synchronized boolean processTakeSnapshot(File snapshotDir) throws TException, IOException {
     File procedureSnapshotDir = new File(snapshotDir, PROCEDURE_SNAPSHOT_DIR);
     if (procedureSnapshotDir.exists()) {
       LOGGER.error(
@@ -275,7 +285,7 @@ public void processLoadSnapshot(File snapshotDir) throws TException, IOException
             });
   }
 
-  public List<Procedure<ConfigNodeProcedureEnv>> getProcedures() {
+  public synchronized List<Procedure<ConfigNodeProcedureEnv>> getProcedures() {
     return new ArrayList<>(procedureMap.values());
   }
 

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/executor/ConfigPlanExecutor.java

@@ -110,6 +110,7 @@
 import org.apache.iotdb.confignode.consensus.request.write.pipe.task.SetPipeStatusWithStoppedByRuntimeExceptionPlanV2;
 import org.apache.iotdb.confignode.consensus.request.write.procedure.DeleteProcedurePlan;
 import org.apache.iotdb.confignode.consensus.request.write.procedure.UpdateProcedurePlan;
+import org.apache.iotdb.confignode.consensus.request.write.procedure.UpdateProceduresPlan;
 import org.apache.iotdb.confignode.consensus.request.write.quota.SetSpaceQuotaPlan;
 import org.apache.iotdb.confignode.consensus.request.write.quota.SetThrottleQuotaPlan;
 import org.apache.iotdb.confignode.consensus.request.write.region.CreateRegionGroupsPlan;
@@ -451,6 +452,8 @@ public TSStatus executeNonQueryPlan(ConfigPhysicalPlan physicalPlan)
         return partitionInfo.createDataPartition((CreateDataPartitionPlan) physicalPlan);
       case AutoCleanPartitionTable:
         return partitionInfo.autoCleanPartitionTable((AutoCleanPartitionTablePlan) physicalPlan);
+      case UpdateProcedures:
+        return procedureInfo.updateProcedures((UpdateProceduresPlan) physicalPlan);
       case UpdateProcedure:
         return procedureInfo.updateProcedure((UpdateProcedurePlan) physicalPlan);
       case DeleteProcedure:

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/ProcedureExecutor.java

@@ -537,7 +537,7 @@ private void executeProcedure(RootProcedureStack rootProcStack, Procedure<Env> p
       // state (e.g. AddRegionPeerProcedure parking at DO_ADD_REGION_PEER after waitTaskFinish() is
       // interrupted) would re-execute forever here and join() would hang. The persisted state lets
       // the next leader resume from where it stopped.
-      if (!isRunning() || !store.isRunning()) {
+      if (!isRunning() || !store.isRunning() || Thread.currentThread().isInterrupted()) {
         return;
       }
 
@@ -603,13 +603,14 @@ private void updateStoreOnExecution(
       if (LOG.isDebugEnabled()) {
         LOG.debug(ProcedureMessages.STORED_CHILDREN, proc, Arrays.toString(subprocs));
       }
-      try {
-        store.update(subprocs);
-      } catch (Exception e) {
-        // Do nothing since this step is idempotent. New CN leader can converge to the correct
-        // state when restore this procedure.
-        LOG.warn(ProcedureMessages.FAILED_TO_UPDATE_SUBPROCS_ON_EXECUTION, e);
-      }
+      // Persist the parent's next state and its children in one consensus entry before any child
+      // can run. Otherwise recovery can replay the parent's submission and create duplicate
+      // children.
+      Procedure<Env>[] proceduresToStore = Arrays.copyOf(subprocs, subprocs.length + 1);
+      proceduresToStore[subprocs.length] = proc;
+      RetryUtils.executeWithEndlessBackoffRetry(
+          () -> store.update(proceduresToStore),
+          ProcedureMessages.LOG_PERSIST_PARENT_AND_CHILD_PROCEDURES_54E2D0D2);
     } else {
       LOG.debug(ProcedureMessages.STORE_UPDATE, proc);
       if (proc.isFinished() && !proc.hasParent()) {

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/store/ConfigProcedureStore.java

@@ -19,23 +19,27 @@
 
 package org.apache.iotdb.confignode.procedure.store;
 
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
 import org.apache.iotdb.commons.conf.CommonDescriptor;
 import org.apache.iotdb.commons.utils.FileUtils;
 import org.apache.iotdb.commons.utils.TestOnly;
 import org.apache.iotdb.confignode.consensus.request.write.procedure.DeleteProcedurePlan;
 import org.apache.iotdb.confignode.consensus.request.write.procedure.UpdateProcedurePlan;
+import org.apache.iotdb.confignode.consensus.request.write.procedure.UpdateProceduresPlan;
 import org.apache.iotdb.confignode.i18n.ConfigNodeMessages;
 import org.apache.iotdb.confignode.manager.ConfigManager;
 import org.apache.iotdb.confignode.persistence.ProcedureInfo;
 import org.apache.iotdb.confignode.procedure.Procedure;
 import org.apache.iotdb.confignode.procedure.env.ConfigNodeProcedureEnv;
 import org.apache.iotdb.consensus.exception.ConsensusException;
+import org.apache.iotdb.rpc.TSStatusCode;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.File;
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 
@@ -109,9 +113,18 @@ public void update(Procedure<ConfigNodeProcedureEnv> procedure) throws Exception
   }
 
   @Override
-  public void update(Procedure[] subprocs) throws Exception {
-    for (Procedure subproc : subprocs) {
-      update(subproc);
+  public void update(Procedure[] procedures) throws Exception {
+    for (Procedure procedure : procedures) {
+      Objects.requireNonNull(
+          ProcedureFactory.getProcedureType(procedure),
+          ConfigNodeMessages.EXCEPTION_PROCEDURE_TYPE_IS_NULL_93147BD3);
+    }
+    TSStatus status =
+        configManager
+            .getConsensusManager()
+            .write(new UpdateProceduresPlan(Arrays.asList(procedures)));
+    if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+      throw new ConsensusException(status.toString());
     }
   }
 

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/store/IProcedureStore.java

@@ -40,7 +40,8 @@ public interface IProcedureStore<Env> {
 
   void update(Procedure<Env> procedure) throws Exception;
 
-  void update(Procedure<Env>[] subprocs) throws Exception;
+  /** Persist the entire batch atomically, including a parent and its newly created children. */
+  void update(Procedure<Env>[] procedures) throws Exception;
 
   void delete(long procId) throws Exception;
 

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/consensus/request/UpdateProceduresPlanTest.java

@@ -0,0 +1,114 @@
+/*
+ * 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.confignode.consensus.request;
+
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.confignode.consensus.request.write.procedure.UpdateProcedurePlan;
+import org.apache.iotdb.confignode.consensus.request.write.procedure.UpdateProceduresPlan;
+import org.apache.iotdb.confignode.manager.ConfigManager;
+import org.apache.iotdb.confignode.manager.consensus.ConsensusManager;
+import org.apache.iotdb.confignode.persistence.ProcedureInfo;
+import org.apache.iotdb.confignode.procedure.Procedure;
+import org.apache.iotdb.confignode.procedure.env.ConfigNodeProcedureEnv;
+import org.apache.iotdb.confignode.procedure.impl.testonly.NeverFinishProcedure;
+import org.apache.iotdb.confignode.procedure.store.ConfigProcedureStore;
+import org.apache.iotdb.consensus.exception.ConsensusException;
+import org.apache.iotdb.rpc.TSStatusCode;
+
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.mockito.ArgumentCaptor;
+
+import java.io.File;
+import java.util.Arrays;
+import java.util.List;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class UpdateProceduresPlanTest {
+  @Rule public TemporaryFolder folder = new TemporaryFolder();
+
+  @Test
+  public void batchRoundTripAndSnapshotRetainAllProcedures() throws Exception {
+    List<Procedure<ConfigNodeProcedureEnv>> procedures =
+        Arrays.asList(
+            new NeverFinishProcedure(7), new NeverFinishProcedure(8), new NeverFinishProcedure(4));
+    UpdateProceduresPlan batch = new UpdateProceduresPlan(procedures);
+    UpdateProceduresPlan restored =
+        (UpdateProceduresPlan) ConfigPhysicalPlan.Factory.create(batch.serializeToByteBuffer());
+    assertEquals(batch, restored);
+    ProcedureInfo info = new ProcedureInfo(null);
+    info.updateProcedures(restored);
+    File snapshot = folder.newFolder();
+    assertTrue(info.processTakeSnapshot(snapshot));
+    ProcedureInfo recovered = new ProcedureInfo(null);
+    recovered.processLoadSnapshot(snapshot);
+    assertEquals(info, recovered);
+    assertEquals(3, recovered.getProcedures().size());
+    assertEquals(9, recovered.getNextProcId());
+    // Existing single-procedure log entries retain their format and reader.
+    UpdateProcedurePlan single = new UpdateProcedurePlan(procedures.get(0));
+    assertEquals(single, ConfigPhysicalPlan.Factory.create(single.serializeToByteBuffer()));
+  }
+
+  @Test
+  public void storeWritesOneConsensusEntryForWholeBatch() throws Exception {
+    ConfigManager manager = mock(ConfigManager.class);
+    ConsensusManager consensus = mock(ConsensusManager.class);
+    when(manager.getConsensusManager()).thenReturn(consensus);
+    when(consensus.write(any()))
+        .thenReturn(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()));
+    ConfigProcedureStore store = new ConfigProcedureStore(manager, new ProcedureInfo(null));
+    Procedure[] procedures = {
+      new NeverFinishProcedure(7), new NeverFinishProcedure(8), new NeverFinishProcedure(4)
+    };
+    store.update(procedures);
+    ArgumentCaptor<ConfigPhysicalPlan> plan = ArgumentCaptor.forClass(ConfigPhysicalPlan.class);
+    verify(consensus).write(plan.capture());
+    assertEquals(
+        Arrays.asList(procedures), ((UpdateProceduresPlan) plan.getValue()).getProcedures());
+  }
+
+  @Test
+  public void rejectedBatchIsNotReportedAsPersisted() throws Exception {
+    ConfigManager manager = mock(ConfigManager.class);
+    ConsensusManager consensus = mock(ConsensusManager.class);
+    when(manager.getConsensusManager()).thenReturn(consensus);
+    when(consensus.write(any()))
+        .thenReturn(new TSStatus(TSStatusCode.INTERNAL_SERVER_ERROR.getStatusCode()));
+    ConfigProcedureStore store = new ConfigProcedureStore(manager, new ProcedureInfo(null));
+    try {
+      store.update(new Procedure[] {new NeverFinishProcedure(7)});
+      fail("Rejected consensus write must be retried by the executor");
+    } catch (ConsensusException expected) {
+      assertTrue(
+          expected
+              .getMessage()
+              .contains(String.valueOf(TSStatusCode.INTERNAL_SERVER_ERROR.getStatusCode())));
+    }
+  }
+}