Fix stale IoTConsensus peers after restart during removal

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

@@ -847,6 +847,14 @@ public final class ProcedureMessages {
       "RemoveRegionLocation remove region {} from DataNode {}, result is {}";
   public static final String REMOVEREGIONPEER_STATE_FAILED = "RemoveRegionPeer state {} failed";
   public static final String REMOVEREGIONPEER_STATE_SUCCESS = "RemoveRegionPeer state {} success";
+  public static final String
+      LOG_CANNOT_RECONCILE_PEER_LIST_FOR_ARG_AFTER_REMOVING_DATANODE_ARG_THE_CURRENT_REPLICA_SET_STILL_CONTAINS_THE_REMOVED_DATANODE_DBBF45DC =
+          "Cannot reconcile peer list for {} after removing DataNode {}: the current replica set still contains the removed DataNode.";
+  public static final String
+      LOG_FAILED_TO_RESET_PEER_LIST_FOR_ARG_ON_DATANODE_ARG_TO_ARG_ARG_C8D872B4 =
+          "Failed to reset peer list for {} on DataNode {} to {}: {}";
+  public static final String LOG_PIDARG_REMOVEREGION_STATE_ARG_WILL_RETRY_IN_ARG_MS_080917A9 =
+      "[pid{}][RemoveRegion] state {} will retry in {} ms.";
   public static final String REMOVEREGION_RATIS_TRANSFER_LEADER_FAIL_BUT_PROCEDURE_WILL_CONTINUE =
       "[RemoveRegion] Ratis transfer leader fail, but procedure will continue.";
   public static final String REMOVE_CONFIG_NODE = "Remove Config Node";

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

@@ -806,6 +806,14 @@ public final class ProcedureMessages {
       "RemoveRegionLocation：从 DataNode {} 移除 region {}，结果为 {}";
   public static final String REMOVEREGIONPEER_STATE_FAILED = "RemoveRegionPeer 状态 {} 失败";
   public static final String REMOVEREGIONPEER_STATE_SUCCESS = "RemoveRegionPeer 状态 {} 成功";
+  public static final String
+      LOG_CANNOT_RECONCILE_PEER_LIST_FOR_ARG_AFTER_REMOVING_DATANODE_ARG_THE_CURRENT_REPLICA_SET_STILL_CONTAINS_THE_REMOVED_DATANODE_DBBF45DC =
+          "无法纠正 {} 在移除 DataNode {} 后的 peer 列表：当前副本集仍包含已移除的 DataNode。";
+  public static final String
+      LOG_FAILED_TO_RESET_PEER_LIST_FOR_ARG_ON_DATANODE_ARG_TO_ARG_ARG_C8D872B4 =
+          "将 {} 在 DataNode {} 上的 peer 列表重置为 {} 失败：{}";
+  public static final String LOG_PIDARG_REMOVEREGION_STATE_ARG_WILL_RETRY_IN_ARG_MS_080917A9 =
+      "[pid{}][RemoveRegion] 状态 {} 将在 {} 毫秒后重试。";
   public static final String REMOVEREGION_RATIS_TRANSFER_LEADER_FAIL_BUT_PROCEDURE_WILL_CONTINUE =
       "[RemoveRegion] Ratis 转换 leader 失败，但 procedure 将继续。";
   public static final String REMOVE_CONFIG_NODE = "移除 ConfigNode";

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

@@ -539,6 +539,11 @@ private void executeProcedure(RootProcedureStack rootProcStack, Procedure<Env> p
         return;
       }
 
+      if (subprocs == null && proc.getState() == ProcedureState.WAITING_TIMEOUT) {
+        // Persist the suspended state before a timeout can make the procedure runnable again.
+        timeoutExecutor.add(proc);
+      }
+
       if (proc.isRunnable() && proc.isYieldAfterExecution(this.environment)) {
         yieldProcedure(proc);
         return;

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

@@ -360,6 +360,58 @@ public Map<Integer, TSStatus> resetPeerList(
     return clientHandler.getResponseMap();
   }
 
+  public boolean isIoTConsensusDataRegion(TConsensusGroupId regionId) {
+    return TConsensusGroupType.DataRegion.equals(regionId.getType())
+        && IOT_CONSENSUS.equals(CONF.getDataRegionConsensusProtocolClass());
+  }
+
+  /**
+   * Reconcile retained IoTConsensus replicas after the removed location has been committed. A
+   * restarting replica may have received the previous peer list while the removal was in progress.
+   * The existing reset RPC waits for consensus startup before applying this final list.
+   */
+  public boolean resetPeerListAfterRemove(
+      TConsensusGroupId regionId, TDataNodeLocation removedDataNode) {
+    if (!isIoTConsensusDataRegion(regionId)) {
+      return true;
+    }
+    Optional<TRegionReplicaSet> replicaSet = getRegionReplicaSet(regionId);
+    if (!replicaSet.isPresent()) {
+      // Database deletion may have removed the entire Region from ConfigNode metadata. Its
+      // deletion tasks own the cleanup, and there are no retained replicas left to reconcile.
+      return true;
+    }
+    if (replicaSet.get().getDataNodeLocations().stream()
+        .anyMatch(location -> location.getDataNodeId() == removedDataNode.getDataNodeId())) {
+      LOGGER.warn(
+          ProcedureMessages
+              .LOG_CANNOT_RECONCILE_PEER_LIST_FOR_ARG_AFTER_REMOVING_DATANODE_ARG_THE_CURRENT_REPLICA_SET_STILL_CONTAINS_THE_REMOVED_DATANODE_DBBF45DC,
+          regionId,
+          removedDataNode.getDataNodeId());
+      return false;
+    }
+    List<TDataNodeLocation> correctLocations = replicaSet.get().getDataNodeLocations();
+    Map<Integer, TDataNodeLocation> retainedNodes = new HashMap<>();
+    correctLocations.forEach(location -> retainedNodes.put(location.getDataNodeId(), location));
+    // Do not filter offline replicas: a missed reset must remain pending until they restart.
+    Map<Integer, TSStatus> responses = resetPeerList(regionId, correctLocations, retainedNodes);
+    boolean allSucceeded = true;
+    for (TDataNodeLocation location : correctLocations) {
+      TSStatus response = responses == null ? null : responses.get(location.getDataNodeId());
+      if (response == null || response.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+        LOGGER.warn(
+            ProcedureMessages
+                .LOG_FAILED_TO_RESET_PEER_LIST_FOR_ARG_ON_DATANODE_ARG_TO_ARG_ARG_C8D872B4,
+            regionId,
+            location.getDataNodeId(),
+            correctLocations,
+            response);
+        allSucceeded = false;
+      }
+    }
+    return allSucceeded;
+  }
+
   // TODO: will use 'procedure yield' to refactor later
   public TRegionMigrateResult waitTaskFinish(long taskId, TDataNodeLocation dataNodeLocation) {
     return waitTaskFinish(taskId, dataNodeLocation, null);
@@ -449,7 +501,7 @@ public void forceUpdateRegionCache(
         .forceUpdateRegionCache(regionId, newLocation.getDataNodeId(), regionStatus);
   }
 
-  public void removeRegionLocation(
+  public TSStatus removeRegionLocation(
       TConsensusGroupId regionId, TDataNodeLocation deprecatedLocation) {
     RemoveRegionLocationPlan req = new RemoveRegionLocationPlan(regionId, deprecatedLocation);
     TSStatus status = configManager.getPartitionManager().removeRegionLocation(req);
@@ -458,8 +510,14 @@ public void removeRegionLocation(
         regionId,
         getIdWithRpcEndpoint(deprecatedLocation),
         status);
-    configManager.getLoadManager().removeRegionCache(regionId, deprecatedLocation.getDataNodeId());
-    configManager.getLoadManager().getRouteBalancer().balanceRegionLeaderAndPriority();
+    if (!isIoTConsensusDataRegion(regionId)
+        || (status != null && status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode())) {
+      configManager
+          .getLoadManager()
+          .removeRegionCache(regionId, deprecatedLocation.getDataNodeId());
+      configManager.getLoadManager().getRouteBalancer().balanceRegionLeaderAndPriority();
+    }
+    return status;
   }
 
   /**

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/region/RemoveRegionPeerProcedure.java

@@ -32,6 +32,7 @@
 import org.apache.iotdb.confignode.procedure.env.ConfigNodeProcedureEnv;
 import org.apache.iotdb.confignode.procedure.env.RegionMaintainHandler;
 import org.apache.iotdb.confignode.procedure.exception.ProcedureException;
+import org.apache.iotdb.confignode.procedure.state.ProcedureState;
 import org.apache.iotdb.confignode.procedure.state.RemoveRegionPeerState;
 import org.apache.iotdb.confignode.procedure.store.ProcedureType;
 import org.apache.iotdb.mpp.rpc.thrift.TRegionMigrateResult;
@@ -55,6 +56,7 @@ public class RemoveRegionPeerProcedure extends RegionOperationProcedure<RemoveRe
   private static final Logger LOGGER = LoggerFactory.getLogger(RemoveRegionPeerProcedure.class);
   private static final int MAX_DELETE_OLD_REGION_PEER_RETRY = 3;
   private static final long DELETE_OLD_REGION_PEER_RETRY_INTERVAL_MS = 5_000;
+  private static final long RECONCILE_PEER_LIST_RETRY_INTERVAL_MS = 5_000;
   private TDataNodeLocation coordinator;
   private TDataNodeLocation targetDataNode;
   private transient int deleteOldRegionPeerAttempted = 0;
@@ -91,6 +93,8 @@ protected Flow executeFromState(ConfigNodeProcedureEnv env, RemoveRegionPeerStat
     }
     TSStatus tsStatus;
     RegionMaintainHandler handler = env.getRegionMaintainHandler();
+    final boolean reconcilePeerList =
+        state == REMOVE_REGION_LOCATION_CACHE && handler.isIoTConsensusDataRegion(regionId);
     try {
       switch (state) {
         case TRANSFER_REGION_LEADER:
@@ -193,7 +197,13 @@ protected Flow executeFromState(ConfigNodeProcedureEnv env, RemoveRegionPeerStat
           setNextState(REMOVE_REGION_LOCATION_CACHE);
           break;
         case REMOVE_REGION_LOCATION_CACHE:
-          handler.removeRegionLocation(regionId, targetDataNode);
+          tsStatus = handler.removeRegionLocation(regionId, targetDataNode);
+          if (reconcilePeerList
+              && (tsStatus == null
+                  || tsStatus.getCode() != SUCCESS_STATUS.getStatusCode()
+                  || !handler.resetPeerListAfterRemove(regionId, targetDataNode))) {
+            return suspendFinalizationForRetry();
+          }
           setKillPoint(state);
           LOGGER.info(ProcedureMessages.REMOVEREGIONPEER_STATE_SUCCESS, state);
           LOGGER.info(
@@ -211,12 +221,40 @@ protected Flow executeFromState(ConfigNodeProcedureEnv env, RemoveRegionPeerStat
       }
     } catch (Exception e) {
       LOGGER.error(ProcedureMessages.REMOVEREGIONPEER_STATE_FAILED, state, e);
+      if (reconcilePeerList) {
+        return suspendFinalizationForRetry();
+      }
       return Flow.NO_MORE_STATE;
     }
     LOGGER.info(ProcedureMessages.PID_REMOVEREGION_STATE_SUCCESS, getProcId(), state);
     return Flow.HAS_MORE_STATE;
   }
 
+  private Flow suspendFinalizationForRetry() {
+    // Keep the existing persisted state until every retained replica acknowledges the final list.
+    // Replaying the metadata removal after a leader switch is idempotent.
+    setNextState(REMOVE_REGION_LOCATION_CACHE);
+    setTimeout(RECONCILE_PEER_LIST_RETRY_INTERVAL_MS);
+    setState(ProcedureState.WAITING_TIMEOUT);
+    LOGGER.warn(
+        ProcedureMessages.LOG_PIDARG_REMOVEREGION_STATE_ARG_WILL_RETRY_IN_ARG_MS_080917A9,
+        getProcId(),
+        REMOVE_REGION_LOCATION_CACHE,
+        RECONCILE_PEER_LIST_RETRY_INTERVAL_MS);
+    return Flow.HAS_MORE_STATE;
+  }
+
+  @Override
+  protected synchronized boolean setTimeoutFailure(ConfigNodeProcedureEnv env) {
+    if (getState() == ProcedureState.WAITING_TIMEOUT
+        && getCurrentState() == REMOVE_REGION_LOCATION_CACHE) {
+      setState(ProcedureState.RUNNABLE);
+      env.getScheduler().addBack(this);
+      return false;
+    }
+    return super.setTimeoutFailure(env);
+  }
+
   @Override
   protected void rollbackState(ConfigNodeProcedureEnv env, RemoveRegionPeerState state)
       throws IOException, InterruptedException, ProcedureException {}