Load: Implement region operations cache for load × region migration detection (apache#15210)

liyuheng55555 · SteveYurongSu · web-flow · commit 8889255c0c48 · 2025-03-28T23:18:16.000+08:00
Co-authored-by: Steve Yurong Su &lt;rong@apache.org&gt;
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ProcedureManager.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ProcedureManager.java
@@ -167,6 +167,7 @@
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 public class ProcedureManager {
   private static final Logger LOGGER = LoggerFactory.getLogger(ProcedureManager.class);
@@ -709,16 +710,19 @@ private TSStatus checkMigrateRegion(
       TDataNodeLocation originalDataNode,
       TDataNodeLocation destDataNode,
       TDataNodeLocation coordinatorForAddPeer) {
-    String failMessage =
-        regionOperationCommonCheck(
-            regionGroupId,
-            destDataNode,
-            Arrays.asList(
-                new Pair<>("Original DataNode", originalDataNode),
-                new Pair<>("Destination DataNode", destDataNode),
-                new Pair<>("Coordinator for add peer", coordinatorForAddPeer)),
-            migrateRegionReq.getModel());
-    if (configManager
+    String failMessage;
+    if ((failMessage =
+            regionOperationCommonCheck(
+                regionGroupId,
+                destDataNode,
+                Arrays.asList(
+                    new Pair<>("Original DataNode", originalDataNode),
+                    new Pair<>("Destination DataNode", destDataNode),
+                    new Pair<>("Coordinator for add peer", coordinatorForAddPeer)),
+                migrateRegionReq.getModel()))
+        != null) {
+      // do nothing
+    } else if (configManager
         .getPartitionManager()
         .getAllReplicaSets(originalDataNode.getDataNodeId())
         .stream()
@@ -955,10 +959,7 @@ private String checkRegionOperationWithRemoveDataNode(
 
   private String checkRegionOperationDuplication(TConsensusGroupId regionId) {
     List<? extends RegionOperationProcedure<?>> otherRegionMemberChangeProcedures =
-        getExecutor().getProcedures().values().stream()
-            .filter(procedure -> !procedure.isFinished())
-            .filter(procedure -> procedure instanceof RegionOperationProcedure)
-            .map(procedure -> (RegionOperationProcedure<?>) procedure)
+        getRegionOperationProcedures()
             .filter(
                 regionMemberChangeProcedure ->
                     regionId.equals(regionMemberChangeProcedure.getRegionId()))
@@ -971,6 +972,20 @@ private String checkRegionOperationDuplication(TConsensusGroupId regionId) {
     return null;
   }
 
+  public List<TConsensusGroupId> getRegionOperationConsensusIds() {
+    return getRegionOperationProcedures()
+        .map(RegionOperationProcedure::getRegionId)
+        .distinct()
+        .collect(Collectors.toList());
+  }
+
+  private Stream<RegionOperationProcedure<?>> getRegionOperationProcedures() {
+    return getExecutor().getProcedures().values().stream()
+        .filter(procedure -> !procedure.isFinished())
+        .filter(procedure -> procedure instanceof RegionOperationProcedure)
+        .map(procedure -> (RegionOperationProcedure<?>) procedure);
+  }
+
   private String checkRegionOperationModelCorrectness(TConsensusGroupId regionId, Model model) {
     String databaseName = configManager.getPartitionManager().getRegionDatabase(regionId);
     boolean isTreeModelDatabase = databaseName.startsWith(SqlConstant.TREE_MODEL_DATABASE_PREFIX);
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/service/HeartbeatService.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/service/HeartbeatService.java
@@ -41,6 +41,7 @@
 import org.apache.iotdb.confignode.manager.load.cache.node.ConfigNodeHeartbeatCache;
 import org.apache.iotdb.confignode.manager.node.NodeManager;
 import org.apache.iotdb.confignode.rpc.thrift.TConfigNodeHeartbeatReq;
+import org.apache.iotdb.db.protocol.client.ConfigNodeInfo;
 import org.apache.iotdb.mpp.rpc.thrift.TDataNodeHeartbeatReq;
 
 import org.apache.tsfile.utils.Pair;
@@ -141,6 +142,11 @@ private TDataNodeHeartbeatReq genHeartbeatReq() {
     /* Generate heartbeat request */
     TDataNodeHeartbeatReq heartbeatReq = new TDataNodeHeartbeatReq();
     heartbeatReq.setHeartbeatTimestamp(System.nanoTime());
+    heartbeatReq.setLogicalClock(
+        configManager
+            .getConsensusManager()
+            .getConsensusImpl()
+            .getLogicalClock(ConfigNodeInfo.CONFIG_REGION_ID));
     // Always sample RegionGroups' leadership as the Region heartbeat
     heartbeatReq.setNeedJudgeLeader(true);
     // We sample DataNode's load in every 10 heartbeat loop
@@ -169,6 +175,12 @@ private TDataNodeHeartbeatReq genHeartbeatReq() {
       heartbeatReq.setDataNodes(configManager.getNodeManager().getRegisteredDataNodeLocations());
     }
 
+    // We broadcast region operations list every 100 heartbeat loops
+    if (heartbeatCounter.get() % 100 == 0) {
+      heartbeatReq.setCurrentRegionOperations(
+          configManager.getProcedureManager().getRegionOperationConsensusIds());
+    }
+
     /* Update heartbeat counter */
     heartbeatCounter.getAndIncrement();
 
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/env/ConfigNodeProcedureEnv.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/env/ConfigNodeProcedureEnv.java
@@ -59,6 +59,7 @@
 import org.apache.iotdb.confignode.rpc.thrift.TAddConsensusGroupReq;
 import org.apache.iotdb.confignode.rpc.thrift.TNodeVersionInfo;
 import org.apache.iotdb.consensus.exception.ConsensusException;
+import org.apache.iotdb.db.protocol.client.ConfigNodeInfo;
 import org.apache.iotdb.mpp.rpc.thrift.TActiveTriggerInstanceReq;
 import org.apache.iotdb.mpp.rpc.thrift.TCreateDataRegionReq;
 import org.apache.iotdb.mpp.rpc.thrift.TCreatePipePluginInstanceReq;
@@ -486,7 +487,15 @@ public List<TSStatus> notifyRegionMigrationToAllDataNodes(
     final Map<Integer, TDataNodeLocation> dataNodeLocationMap =
         configManager.getNodeManager().getRegisteredDataNodeLocations();
     final TNotifyRegionMigrationReq request =
-        new TNotifyRegionMigrationReq(consensusGroupId, isStart);
+        new TNotifyRegionMigrationReq(
+            configManager
+                .getConsensusManager()
+                .getConsensusImpl()
+                .getLogicalClock(ConfigNodeInfo.CONFIG_REGION_ID),
+            System.nanoTime(),
+            configManager.getProcedureManager().getRegionOperationConsensusIds());
+    request.setRegionId(consensusGroupId);
+    request.setIsStart(isStart);
 
     final DataNodeAsyncRequestContext<TNotifyRegionMigrationReq, TSStatus> clientHandler =
         new DataNodeAsyncRequestContext<>(
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/region/RegionMigrateProcedure.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/region/RegionMigrateProcedure.java
@@ -81,8 +81,8 @@ protected Flow executeFromState(ConfigNodeProcedureEnv env, RegionTransitionStat
               "[pid{}][MigrateRegion] started, {} will be migrated from DataNode {} to {}.",
               getProcId(),
               regionId,
-              handler.simplifiedLocation(originalDataNode),
-              handler.simplifiedLocation(destDataNode));
+              RegionMaintainHandler.simplifiedLocation(originalDataNode),
+              RegionMaintainHandler.simplifiedLocation(destDataNode));
           addChildProcedure(new NotifyRegionMigrationProcedure(regionId, true));
           setNextState(RegionTransitionState.ADD_REGION_PEER);
           break;
@@ -120,8 +120,8 @@ protected Flow executeFromState(ConfigNodeProcedureEnv env, RegionTransitionStat
               getProcId(),
               cleanHint,
               regionId,
-              handler.simplifiedLocation(originalDataNode),
-              handler.simplifiedLocation(destDataNode),
+              RegionMaintainHandler.simplifiedLocation(originalDataNode),
+              RegionMaintainHandler.simplifiedLocation(destDataNode),
               CommonDateTimeUtils.convertMillisecondToDurationStr(
                   System.currentTimeMillis() - getSubmittedTime()),
               DateTimeUtils.convertLongToDate(getSubmittedTime(), "ms"));
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java
@@ -1949,6 +1949,15 @@ public TDataNodeHeartbeatResp getDataNodeHeartBeat(TDataNodeHeartbeatReq req) th
       clusterTopology.updateTopology(req.getDataNodes(), req.getTopology());
     }
 
+    if (req.isSetCurrentRegionOperations()) {
+      RegionMigrateService.getInstance()
+          .notifyRegionMigration(
+              new TNotifyRegionMigrationReq(
+                  req.getLogicalClock(),
+                  req.getHeartbeatTimestamp(),
+                  req.getCurrentRegionOperations()));
+    }
+
     return resp;
   }
 
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java
@@ -240,9 +240,10 @@ public void start() {
             }
           }
 
-          if (RegionMigrateService.getInstance().getLastNotifyTime() > startTimeMs) {
+          if (RegionMigrateService.getInstance().getLastNotifyMigratingTime() > startTimeMs
+              || RegionMigrateService.getInstance().mayHaveMigratingRegions()) {
             LOGGER.warn(
-                "LoadTsFileScheduler: Region migration started or ended during loading TsFile {}, will convert to insertion to avoid data loss",
+                "LoadTsFileScheduler: Region migration was detected during loading TsFile {}, will convert to insertion to avoid data loss",
                 filePath);
             isLoadSingleTsFileSuccess = false;
           }
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/RegionMigrateService.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/RegionMigrateService.java
@@ -49,12 +49,12 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
-import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
 
 public class RegionMigrateService implements IService {
@@ -75,27 +75,76 @@ public class RegionMigrateService implements IService {
   private static final ConcurrentHashMap<Long, TRegionMigrateResult> taskResultMap =
       new ConcurrentHashMap<>();
 
-  private static final AtomicLong lastNotifyTime = new AtomicLong(Long.MIN_VALUE);
-
   private static final TRegionMigrateResult unfinishedResult = new TRegionMigrateResult();
 
+  private static class RegionMigrationStatusCache {
+
+    private long logicalClock = -1;
+    private long timestamp = -1;
+
+    private List<TConsensusGroupId> currentMigratingConsensusGroupIds = Collections.emptyList();
+
+    private long lastNotifyMigratingTime = Long.MIN_VALUE;
+
+    public synchronized void update(
+        long newLogicalClock, long newTimestamp, List<TConsensusGroupId> currentRegionOperations) {
+      if (newLogicalClock < logicalClock
+          || (newLogicalClock == logicalClock && newTimestamp <= timestamp)) {
+        return;
+      }
+      logicalClock = newLogicalClock;
+      timestamp = newTimestamp;
+
+      currentMigratingConsensusGroupIds = currentRegionOperations;
+
+      if (currentRegionOperations != null && !currentRegionOperations.isEmpty()) {
+        lastNotifyMigratingTime = System.currentTimeMillis();
+      }
+    }
+
+    public synchronized void notifyMigrating() {
+      lastNotifyMigratingTime = System.currentTimeMillis();
+    }
+
+    public synchronized boolean hasMigratingRegions() {
+      return currentMigratingConsensusGroupIds != null
+          && !currentMigratingConsensusGroupIds.isEmpty();
+    }
+
+    public synchronized long getLastNotifyMigratingTime() {
+      return lastNotifyMigratingTime;
+    }
+  }
+
+  private final RegionMigrationStatusCache regionMigrationStatusCache =
+      new RegionMigrationStatusCache();
+
   private RegionMigrateService() {}
 
   public static RegionMigrateService getInstance() {
     return Holder.INSTANCE;
   }
 
   public void notifyRegionMigration(TNotifyRegionMigrationReq req) {
-    lastNotifyTime.set(System.currentTimeMillis());
-    if (req.isIsStart()) {
-      LOGGER.info("Region {} is notified to begin migrating", req.getRegionId());
-    } else {
-      LOGGER.info("Region {} is notified to finish migrating", req.getRegionId());
+    regionMigrationStatusCache.update(
+        req.getLogicalClock(), req.getTimestamp(), req.getCurrentRegionOperations());
+
+    if (req.isSetIsStart() && req.isSetRegionId()) {
+      regionMigrationStatusCache.notifyMigrating();
+      if (req.isIsStart()) {
+        LOGGER.info("Region {} is notified to begin migrating", req.getRegionId());
+      } else {
+        LOGGER.info("Region {} is notified to finish migrating", req.getRegionId());
+      }
     }
   }
 
-  public long getLastNotifyTime() {
-    return lastNotifyTime.get();
+  public long getLastNotifyMigratingTime() {
+    return regionMigrationStatusCache.getLastNotifyMigratingTime();
+  }
+
+  public boolean mayHaveMigratingRegions() {
+    return regionMigrationStatusCache.hasMigratingRegions();
   }
 
   /**
diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift
@@ -53,8 +53,11 @@ struct TRegionMigrateResult {
 }
 
 struct TNotifyRegionMigrationReq {
-  1: required common.TConsensusGroupId regionId
-  2: required bool isStart
+  1: required i64 logicalClock
+  2: required i64 timestamp
+  3: optional common.TConsensusGroupId regionId
+  4: optional bool isStart
+  5: required list<common.TConsensusGroupId> currentRegionOperations
 }
 
 struct TCreatePeerReq {
@@ -279,6 +282,8 @@ struct TDataNodeHeartbeatReq {
   11: optional set<common.TEndPoint> configNodeEndPoints
   12: optional map<i32, common.TDataNodeLocation> dataNodes
   13: optional map<i32, set<i32>> topology
+  14: required i64 logicalClock
+  15: optional list<common.TConsensusGroupId> currentRegionOperations
 }
 
 struct TDataNodeActivation {

Original file line number	Diff line number	Diff line change
`@@ -1949,6 +1949,15 @@ public TDataNodeHeartbeatResp getDataNodeHeartBeat(TDataNodeHeartbeatReq req) th`
`1949`	`1949`	`clusterTopology.updateTopology(req.getDataNodes(), req.getTopology());`
`1950`	`1950`	`}`
`1951`	`1951`
	`1952`	`+ if (req.isSetCurrentRegionOperations()) {`
	`1953`	`+ RegionMigrateService.getInstance()`
	`1954`	`+ .notifyRegionMigration(`
	`1955`	`+ new TNotifyRegionMigrationReq(`
	`1956`	`+ req.getLogicalClock(),`
	`1957`	`+ req.getHeartbeatTimestamp(),`
	`1958`	`+ req.getCurrentRegionOperations()));`
	`1959`	`+ }`
	`1960`	`+`
`1952`	`1961`	`return resp;`
`1953`	`1962`	`}`
`1954`	`1963`
Original file line number	Diff line number	Diff line change
`@@ -240,9 +240,10 @@ public void start() {`
`240`	`240`	`}`
`241`	`241`	`}`
`242`	`242`
`243`		`- if (RegionMigrateService.getInstance().getLastNotifyTime() > startTimeMs) {`
	`243`	`+ if (RegionMigrateService.getInstance().getLastNotifyMigratingTime() > startTimeMs`
	`244`	`+ \|\| RegionMigrateService.getInstance().mayHaveMigratingRegions()) {`
`244`	`245`	`LOGGER.warn(`
`245`		`- "LoadTsFileScheduler: Region migration started or ended during loading TsFile {}, will convert to insertion to avoid data loss",`
	`246`	`+ "LoadTsFileScheduler: Region migration was detected during loading TsFile {}, will convert to insertion to avoid data loss",`
`246`	`247`	`filePath);`
`247`	`248`	`isLoadSingleTsFileSuccess = false;`
`248`	`249`	`}`
Original file line number	Diff line number	Diff line change
`@@ -53,8 +53,11 @@ struct TRegionMigrateResult {`
`53`	`53`	`}`
`54`	`54`
`55`	`55`	`struct TNotifyRegionMigrationReq {`
`56`		`- 1: required common.TConsensusGroupId regionId`
`57`		`- 2: required bool isStart`
	`56`	`+ 1: required i64 logicalClock`
	`57`	`+ 2: required i64 timestamp`
	`58`	`+ 3: optional common.TConsensusGroupId regionId`
	`59`	`+ 4: optional bool isStart`
	`60`	`+ 5: required list<common.TConsensusGroupId> currentRegionOperations`
`58`	`61`	`}`
`59`	`62`
`60`	`63`	`struct TCreatePeerReq {`
`@@ -279,6 +282,8 @@ struct TDataNodeHeartbeatReq {`
`279`	`282`	`11: optional set<common.TEndPoint> configNodeEndPoints`
`280`	`283`	`12: optional map<i32, common.TDataNodeLocation> dataNodes`
`281`	`284`	`13: optional map<i32, set<i32>> topology`
	`285`	`+ 14: required i64 logicalClock`
	`286`	`+ 15: optional list<common.TConsensusGroupId> currentRegionOperations`
`282`	`287`	`}`
`283`	`288`
`284`	`289`	`struct TDataNodeActivation {`