Fix snapshot load region replacement (#17684)

Author: shuwenwei

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java

@@ -123,21 +123,25 @@ public boolean clearSnapshot() {
 
   @Override
   public void loadSnapshot(File latestSnapshotRootDir) {
-    DataRegion newRegion =
-        new SnapshotLoader(
-                latestSnapshotRootDir.getAbsolutePath(),
-                region.getDatabaseName(),
-                region.getDataRegionIdString())
-            .loadSnapshotForStateMachine();
-    if (newRegion == null) {
-      logger.error("Fail to load snapshot from {}", latestSnapshotRootDir);
-      return;
-    }
-    this.region = newRegion;
+    String databaseName = region.getDatabaseName();
+    String dataRegionIdString = region.getDataRegionIdString();
+    DataRegionId regionId = new DataRegionId(Integer.parseInt(dataRegionIdString));
     try {
-      StorageEngine.getInstance()
-          .setDataRegion(
-              new DataRegionId(Integer.parseInt(region.getDataRegionIdString())), region);
+      DataRegion newRegion =
+          StorageEngine.getInstance()
+              .setDataRegionForSnapshotLoad(
+                  regionId,
+                  () ->
+                      new SnapshotLoader(
+                              latestSnapshotRootDir.getAbsolutePath(),
+                              databaseName,
+                              dataRegionIdString)
+                          .loadSnapshotForStateMachine());
+      if (newRegion == null) {
+        logger.error("Fail to load snapshot from {}", latestSnapshotRootDir);
+        return;
+      }
+      this.region = newRegion;
       ChunkCache.getInstance().clear();
       TimeSeriesMetadataCache.getInstance().clear();
       BloomFilterCache.getInstance().clear();

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java

@@ -115,6 +115,7 @@
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -922,6 +923,33 @@ public int getDataRegionNumber() {
   }
 
   /** This method is not thread-safe */
+  public DataRegion setDataRegionForSnapshotLoad(
+      DataRegionId regionId, Supplier<DataRegion> newRegionSupplier) {
+    if (dataRegionMap.containsKey(regionId)) {
+      DataRegion oldRegion = dataRegionMap.get(regionId);
+      oldRegion.markDeleted();
+      oldRegion.abortCompaction();
+      oldRegion.syncCloseAllWorkingTsFileProcessors();
+      oldRegion.deleteFolder(systemDir);
+      WRITING_METRICS.removeDataRegionMemoryCostMetrics(regionId);
+      WRITING_METRICS.removeFlushingMemTableStatusMetrics(regionId);
+      WRITING_METRICS.removeActiveMemtableCounterMetrics(regionId);
+      FileMetrics.getInstance()
+          .deleteRegion(oldRegion.getDatabaseName(), oldRegion.getDataRegionIdString());
+    }
+
+    DataRegion newRegion = newRegionSupplier.get();
+    if (newRegion != null) {
+      WRITING_METRICS.createFlushingMemTableStatusMetrics(regionId);
+      WRITING_METRICS.createDataRegionMemoryCostMetrics(newRegion);
+      WRITING_METRICS.createActiveMemtableCounterMetrics(regionId);
+      dataRegionMap.put(regionId, newRegion);
+    }
+    return newRegion;
+  }
+
+  /** This method is not thread-safe */
+  @TestOnly
   public void setDataRegion(DataRegionId regionId, DataRegion newRegion) {
     if (dataRegionMap.containsKey(regionId)) {
       DataRegion oldRegion = dataRegionMap.get(regionId);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotTaker.java

@@ -271,7 +271,7 @@ private void createHardLink(File target, File source) throws IOException {
     if (!target.getParentFile().exists()) {
       LOGGER.error("Hard link target dir {} doesn't exist", target.getParentFile());
     }
-    if (!checkHardLinkSourceFile(source)) {
+    if (!checkHardLinkSourceFile(source, 10)) {
       return;
     }
     Files.deleteIfExists(target.toPath());
@@ -280,8 +280,7 @@ private void createHardLink(File target, File source) throws IOException {
   }
 
   /** For "source file not exists" problem (jira787) debugging */
-  private boolean checkHardLinkSourceFile(File source) {
-    int retry = 10;
+  private boolean checkHardLinkSourceFile(File source, int retry) {
     while (!source.exists() && retry > 0) {
       LOGGER.warn(
           "Hard link source file {} doesn't exist, will retry for {} times...", source, retry);