diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotFileSet.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotFileSet.java index 2d8ef0233886e..441b0935557eb 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotFileSet.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotFileSet.java @@ -22,6 +22,7 @@ import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile; import org.apache.iotdb.db.storageengine.dataregion.modification.v1.ModificationFileV1; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.utils.ObjectTypeUtils; import org.apache.tsfile.common.constant.TsFileConstant; @@ -37,6 +38,9 @@ public class SnapshotFileSet { TsFileResource.RESOURCE_SUFFIX.replace(".", ""), ModificationFileV1.FILE_SUFFIX.replace(".", ""), ModificationFile.FILE_SUFFIX.replace(".", ""), + ObjectTypeUtils.OBJECT_FILE_SUFFIX.replace(".", ""), + ObjectTypeUtils.OBJECT_TEMP_FILE_SUFFIX.replace(".", ""), + ObjectTypeUtils.OBJECT_BACK_FILE_SUFFIX.replace(".", ""), }; private static final Set DATA_FILE_SUFFIX_SET = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotLoader.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotLoader.java index 59fa796724129..2b7e3481609d9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotLoader.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotLoader.java @@ -26,7 +26,9 @@ import org.apache.iotdb.db.storageengine.dataregion.DataRegion; import org.apache.iotdb.db.storageengine.dataregion.flush.CompressionRatio; import org.apache.iotdb.db.storageengine.rescon.disk.FolderManager; +import org.apache.iotdb.db.storageengine.rescon.disk.TierManager; import org.apache.iotdb.db.storageengine.rescon.disk.strategy.DirectoryStrategyType; +import org.apache.iotdb.db.utils.ObjectTypeUtils; import org.apache.tsfile.external.commons.io.FileUtils; import org.slf4j.Logger; @@ -38,6 +40,8 @@ import java.nio.file.FileVisitor; import java.nio.file.Files; import java.nio.file.Path; +import java.nio.file.Paths; +import java.nio.file.SimpleFileVisitor; import java.nio.file.attribute.BasicFileAttributes; import java.util.ArrayList; import java.util.Arrays; @@ -46,6 +50,7 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; public class SnapshotLoader { private Logger LOGGER = LoggerFactory.getLogger(SnapshotLoader.class); @@ -231,6 +236,15 @@ private void deleteAllFilesInDataDirs() throws IOException { timePartitions.addAll(Arrays.asList(files)); } } + + File objectRegionDir = + Paths.get(dataDirPath) + .resolve(IoTDBConstant.OBJECT_FOLDER_NAME) + .resolve(dataRegionId) + .toFile(); + if (objectRegionDir.exists()) { + timePartitions.add(objectRegionDir); + } } try { @@ -312,6 +326,78 @@ private void createLinksFromSnapshotDirToDataDirWithoutLog(File sourceDir) createLinksFromSnapshotToSourceDir(targetSuffix, files, folderManager); } } + + File snapshotObjectDir = new File(sourceDir, IoTDBConstant.OBJECT_FOLDER_NAME); + if (snapshotObjectDir.exists()) { + FolderManager objectFolderManager = + new FolderManager( + TierManager.getInstance().getAllObjectFileFolders(), + DirectoryStrategyType.SEQUENCE_STRATEGY); + linkObjectTreeFromSnapshotToObjectDirs(snapshotObjectDir, objectFolderManager); + } + } + + private void linkObjectTreeFromSnapshotToObjectDirs( + File sourceObjectRoot, FolderManager folderManager) + throws DiskSpaceInsufficientException, IOException { + Path sourceRootPath = sourceObjectRoot.toPath(); + // Process files during traversal to avoid loading all object file paths into memory. + Files.walkFileTree( + sourceRootPath, + new SimpleFileVisitor() { + @Override + public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) + throws IOException { + if (!isObjectSnapshotCandidate(file.getFileName().toString())) { + return FileVisitResult.CONTINUE; + } + final Path sourceFile = file; + final Path targetRelPath = sourceRootPath.relativize(file); + try { + folderManager.getNextWithRetry( + currentObjectDir -> { + File targetFile = + new File(currentObjectDir).toPath().resolve(targetRelPath).toFile(); + try { + if (!targetFile.getParentFile().exists() + && !targetFile.getParentFile().mkdirs()) { + throw new IOException( + String.format( + "Cannot create directory %s", + targetFile.getParentFile().getAbsolutePath())); + } + try { + Files.createLink(targetFile.toPath(), sourceFile); + LOGGER.debug("Created hard link from {} to {}", sourceFile, targetFile); + return targetFile; + } catch (IOException e) { + LOGGER.info( + "Cannot create link from {} to {}, fallback to copy", + sourceFile, + targetFile); + } + Files.copy(sourceFile, targetFile.toPath()); + return targetFile; + } catch (Exception e) { + LOGGER.warn( + "Failed to process file {} in dir {}: {}", + sourceFile.getFileName(), + currentObjectDir, + e.getMessage(), + e); + throw e; + } + }); + } catch (Exception e) { + throw new IOException( + String.format( + "Failed to process object file after retries. Source: %s", + sourceFile.toAbsolutePath()), + e); + } + return FileVisitResult.CONTINUE; + } + }); } private void createLinksFromSnapshotToSourceDir( @@ -470,9 +556,105 @@ private int takeHardLinksFromSnapshotToDataDir( } } + File objectSnapshotRoot = + new File( + snapshotFolder.getAbsolutePath() + File.separator + IoTDBConstant.OBJECT_FOLDER_NAME); + if (objectSnapshotRoot.exists()) { + cnt += linkObjectSnapshotTreeToDataDir(objectSnapshotRoot, fileInfoSet); + } + return cnt; } + private int linkObjectSnapshotTreeToDataDir(File objectSnapshotRoot, Set fileInfoSet) + throws IOException { + final FolderManager folderManager; + try { + folderManager = + new FolderManager( + TierManager.getInstance().getAllObjectFileFolders(), + DirectoryStrategyType.SEQUENCE_STRATEGY); + } catch (DiskSpaceInsufficientException e) { + throw new IOException("Failed to initialize object folder manager", e); + } + Path rootPath = objectSnapshotRoot.toPath(); + AtomicInteger cnt = new AtomicInteger(0); + // Process files during traversal to avoid loading all object file paths into memory. + Files.walkFileTree( + rootPath, new ObjectSnapshotLinkFileVisitor(rootPath, fileInfoSet, folderManager, cnt)); + + return cnt.get(); + } + + private final class ObjectSnapshotLinkFileVisitor extends SimpleFileVisitor { + private final Path rootPath; + private final Set fileInfoSet; + private final FolderManager folderManager; + private final AtomicInteger cnt; + + private ObjectSnapshotLinkFileVisitor( + Path rootPath, Set fileInfoSet, FolderManager folderManager, AtomicInteger cnt) { + this.rootPath = rootPath; + this.fileInfoSet = fileInfoSet; + this.folderManager = folderManager; + this.cnt = cnt; + } + + @Override + public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException { + if (!isObjectSnapshotCandidate(file.getFileName().toString())) { + return FileVisitResult.CONTINUE; + } + String infoStr = getFileInfoString(file.toFile()); + if (!fileInfoSet.contains(infoStr)) { + throw new IOException( + String.format("File %s is not in the log file list", file.toAbsolutePath())); + } + final Path sourceFile = file; + final Path targetRelPath = rootPath.relativize(file); + try { + folderManager.getNextWithRetry( + currentObjectDir -> { + File targetFile = new File(currentObjectDir).toPath().resolve(targetRelPath).toFile(); + try { + if (!targetFile.getParentFile().exists() && !targetFile.getParentFile().mkdirs()) { + throw new IOException( + String.format( + "Cannot create directory %s", + targetFile.getParentFile().getAbsolutePath())); + } + try { + Files.createLink(targetFile.toPath(), sourceFile); + LOGGER.debug("Created hard link from {} to {}", sourceFile, targetFile); + return targetFile; + } catch (IOException e) { + LOGGER.info( + "Cannot create link from {} to {}, fallback to copy", sourceFile, targetFile); + } + Files.copy(sourceFile, targetFile.toPath()); + return targetFile; + } catch (Exception e) { + LOGGER.warn( + "Failed to process file {} in dir {}: {}", + sourceFile.getFileName(), + currentObjectDir, + e.getMessage(), + e); + throw e; + } + }); + } catch (Exception e) { + throw new IOException( + String.format( + "Failed to process object snapshot file after retries. Source: %s", + sourceFile.toAbsolutePath()), + e); + } + cnt.incrementAndGet(); + return FileVisitResult.CONTINUE; + } + } + private void createLinksFromSourceToTarget(File targetDir, File[] files, Set fileInfoSet) throws IOException { for (File file : files) { @@ -492,6 +674,33 @@ private void createLinksFromSourceToTarget(File targetDir, File[] files, Set= 0 && objectDirIndex < nameCount - 1) { + Path relativeToObject = filePath.subpath(objectDirIndex + 1, nameCount); + String fileName = relativeToObject.getFileName().toString(); + Path parentPath = relativeToObject.getParent(); + String middlePath = ""; + if (parentPath != null) { + List pathElements = new ArrayList<>(); + for (Path element : parentPath) { + pathElements.add(element.toString()); + } + middlePath = String.join("/", pathElements); + } + return fileName + + SnapshotLogger.SPLIT_CHAR + + middlePath + + SnapshotLogger.SPLIT_CHAR + + "object"; + } String[] splittedStr = file.getAbsolutePath().split(File.separator.equals("\\") ? "\\\\" : "/"); int length = splittedStr.length; return splittedStr[length - SnapshotLogger.FILE_NAME_OFFSET] @@ -501,6 +710,12 @@ private String getFileInfoString(File file) { + splittedStr[length - SnapshotLogger.SEQUENCE_OFFSET]; } + private boolean isObjectSnapshotCandidate(String fileName) { + return fileName.endsWith(ObjectTypeUtils.OBJECT_FILE_SUFFIX) + || fileName.endsWith(ObjectTypeUtils.OBJECT_TEMP_FILE_SUFFIX) + || fileName.endsWith(ObjectTypeUtils.OBJECT_BACK_FILE_SUFFIX); + } + public List getSnapshotFileInfo() throws IOException { File snapshotLogFile = getSnapshotLogFile(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotLogger.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotLogger.java index 3c07e0926a6fd..89b8c5345f10c 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotLogger.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotLogger.java @@ -18,12 +18,17 @@ */ package org.apache.iotdb.db.storageengine.dataregion.snapshot; +import org.apache.iotdb.commons.conf.IoTDBConstant; + import java.io.BufferedOutputStream; import java.io.File; import java.io.FileOutputStream; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; public class SnapshotLogger implements AutoCloseable { public static final String SNAPSHOT_LOG_NAME = "snapshot.log"; @@ -56,6 +61,26 @@ public void close() throws Exception { os.close(); } + public void logObjectRelativePath(Path relativePathFromObjectRoot) throws IOException { + String fileName = relativePathFromObjectRoot.getFileName().toString(); + Path parentPath = relativePathFromObjectRoot.getParent(); + String middlePath = ""; + if (parentPath != null) { + List pathElements = new ArrayList<>(); + for (Path element : parentPath) { + pathElements.add(element.toString()); + } + middlePath = String.join("/", pathElements); + } + os.write(fileName.getBytes(StandardCharsets.UTF_8)); + os.write(SPLIT_CHAR.getBytes(StandardCharsets.UTF_8)); + os.write(middlePath.getBytes(StandardCharsets.UTF_8)); + os.write(SPLIT_CHAR.getBytes(StandardCharsets.UTF_8)); + os.write(IoTDBConstant.OBJECT_FOLDER_NAME.getBytes(StandardCharsets.UTF_8)); + os.write("\n".getBytes(StandardCharsets.UTF_8)); + os.flush(); + } + /** * Log the logical info for the link file, including its file name, time partition, data region * id, database name, sequence or not. diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotTaker.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotTaker.java index f4313827c9ab3..82f8d1cd34810 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotTaker.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotTaker.java @@ -28,17 +28,24 @@ import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileManager; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.storageengine.rescon.disk.TierManager; +import org.apache.iotdb.db.utils.ObjectTypeUtils; +import org.apache.tsfile.fileSystem.FSFactoryProducer; +import org.apache.tsfile.fileSystem.fsFactory.FSFactory; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.File; import java.io.IOException; import java.nio.file.Files; +import java.nio.file.Path; import java.util.Arrays; import java.util.List; import java.util.Objects; import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; +import java.util.stream.Stream; /** * SnapshotTaker takes data snapshot for a DataRegion in one time. It does so by creating hard link @@ -48,6 +55,7 @@ */ public class SnapshotTaker { private static final Logger LOGGER = LoggerFactory.getLogger(SnapshotTaker.class); + private static final FSFactory FS_FACTORY = FSFactoryProducer.getFSFactory(); private final DataRegion dataRegion; private SnapshotLogger snapshotLogger; private List seqFiles; @@ -101,6 +109,7 @@ public boolean takeFullSnapshot( } success = createSnapshot(seqFiles, tempSnapshotId); success = success && createSnapshot(unseqFiles, tempSnapshotId); + success = success && snapshotObjectFiles(tempSnapshotId); success = success && snapshotCompressionRatio(snapshotDirPath); } finally { readUnlockTheFile(); @@ -267,6 +276,87 @@ private boolean createSnapshot(List resources, String snapshotId } } + private boolean snapshotObjectFiles(String tempSnapshotId) { + try { + String dataRegionIdString = dataRegion.getDataRegionIdString(); + for (String objectFolder : TierManager.getInstance().getAllObjectFileFolders()) { + File objectRegionDir = FS_FACTORY.getFile(objectFolder, dataRegionIdString); + if (!objectRegionDir.exists()) { + continue; + } + if (!snapshotSingleObjectRegionDir(objectRegionDir, tempSnapshotId)) { + return false; + } + } + return true; + } catch (IOException e) { + LOGGER.error("Failed to snapshot object files", e); + return false; + } + } + + private boolean snapshotSingleObjectRegionDir(File objectRegionDir, String tempSnapshotId) + throws IOException { + Path regionRoot = objectRegionDir.toPath(); + Path objectRoot = regionRoot.getParent(); + if (objectRoot == null) { + return false; + } + + Path dataRoot = objectRoot.getParent(); + if (dataRoot == null) { + LOGGER.error("Cannot locate data root for object dir {}", objectRegionDir.getAbsolutePath()); + return false; + } + + Path snapshotBaseDir = + dataRoot + .resolve(IoTDBConstant.SNAPSHOT_FOLDER_NAME) + .resolve( + dataRegion.getDatabaseName() + + IoTDBConstant.FILE_NAME_SEPARATOR + + dataRegion.getDataRegionIdString()) + .resolve(tempSnapshotId) + .resolve(IoTDBConstant.OBJECT_FOLDER_NAME); + + try (Stream paths = Files.walk(regionRoot)) { + for (Path file : + paths + .filter(Files::isRegularFile) + .filter(path -> isObjectSnapshotCandidate(path.getFileName().toString())) + .collect(Collectors.toList())) { + Path relPath = objectRoot.relativize(file); + Path targetPath = snapshotBaseDir.resolve(relPath); + createObjectHardLink(targetPath.toFile(), file.toFile(), relPath); + } + } + return true; + } + + private boolean isObjectSnapshotCandidate(String fileName) { + return fileName.endsWith(ObjectTypeUtils.OBJECT_FILE_SUFFIX) + || fileName.endsWith(ObjectTypeUtils.OBJECT_TEMP_FILE_SUFFIX) + || fileName.endsWith(ObjectTypeUtils.OBJECT_BACK_FILE_SUFFIX); + } + + private void createObjectHardLink(File target, File source, Path relativePathForLog) + throws IOException { + File parentDir = target.getParentFile(); + if (parentDir != null && !parentDir.exists() && !parentDir.mkdirs()) { + LOGGER.error("Cannot create snapshot object dir {}", parentDir); + throw new IOException("Failed to create directory " + parentDir); + } + + if (!checkHardLinkSourceFile(source)) { + return; + } + + Files.deleteIfExists(target.toPath()); + Files.createLink(target.toPath(), source.toPath()); + + snapshotLogger.logObjectRelativePath(relativePathForLog); + } + private void createHardLink(File target, File source) throws IOException { if (!target.getParentFile().exists()) { LOGGER.error("Hard link target dir {} doesn't exist", target.getParentFile()); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/ObjectTypeUtils.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/ObjectTypeUtils.java index 6d61056ed6a05..58cf0eca87915 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/ObjectTypeUtils.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/ObjectTypeUtils.java @@ -48,6 +48,9 @@ public class ObjectTypeUtils { private static final Logger logger = LoggerFactory.getLogger(ObjectTypeUtils.class); private static final TierManager TIER_MANAGER = TierManager.getInstance(); + public static final String OBJECT_FILE_SUFFIX = ".bin"; + public static final String OBJECT_TEMP_FILE_SUFFIX = ".tmp"; + public static final String OBJECT_BACK_FILE_SUFFIX = ".back"; private ObjectTypeUtils() {} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotObjectFilesTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotObjectFilesTest.java new file mode 100644 index 0000000000000..b35f98e741070 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotObjectFilesTest.java @@ -0,0 +1,280 @@ +/* + * 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.db.storageengine.dataregion.snapshot; + +import org.apache.iotdb.commons.conf.IoTDBConstant; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.storageengine.dataregion.DataRegion; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResourceStatus; +import org.apache.iotdb.db.storageengine.rescon.disk.TierManager; +import org.apache.iotdb.db.utils.EnvironmentUtils; +import org.apache.iotdb.db.utils.ObjectTypeUtils; + +import org.apache.tsfile.exception.write.WriteProcessException; +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.utils.TsFileGeneratorUtils; +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.File; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +import static org.apache.tsfile.common.constant.TsFileConstant.PATH_SEPARATOR; +import static org.apache.tsfile.common.constant.TsFileConstant.TSFILE_SUFFIX; + +/** + * SnapshotObjectFilesTest verifies the integrated snapshot capabilities for both standard TsFiles + * and custom Object files within IoTDB Storage Engine. + */ +public class SnapshotObjectFilesTest { + + private static final Logger LOGGER = LoggerFactory.getLogger(SnapshotObjectFilesTest.class); + + @Rule public TemporaryFolder tempFolder = new TemporaryFolder(); + + private String[][] originDataDirs; + private Path dataRootPath; + private Path snapshotPath; + + private final String testSgName = "root.testsg"; + private final String testRegionId = "0"; + private final String timePartition = "0"; + + @Before + public void setUp() throws Exception { + EnvironmentUtils.envSetUp(); + originDataDirs = IoTDBDescriptor.getInstance().getConfig().getTierDataDirs(); + + // Initialize sandbox directories + dataRootPath = tempFolder.newFolder("data").toPath(); + String[][] testDataDirs = new String[][] {{dataRootPath.toAbsolutePath().toString()}}; + IoTDBDescriptor.getInstance().getConfig().setTierDataDirs(testDataDirs); + TierManager.getInstance().resetFolders(); + + // Target directory where SnapshotTaker will export data + snapshotPath = tempFolder.newFolder("snapshot_export").toPath(); + } + + @After + public void tearDown() throws Exception { + IoTDBDescriptor.getInstance().getConfig().setTierDataDirs(originDataDirs); + TierManager.getInstance().resetFolders(); + EnvironmentUtils.cleanEnv(); + } + + /** + * Test basic snapshot creation. Logic: Populates DataRegion with TsFiles and Object files, then + * triggers SnapshotTaker. Verification: Validates the physical existence of linked files in the + * snapshot directory and the integrity of the Snapshot Log. + */ + @Test + public void testCreateSnapshotWithMixedFiles() throws Exception { + DataRegion region = createAndPopulateDataRegion(); + Set expectedObjectFiles = prepareObjectFiles(testRegionId); + + SnapshotTaker snapshotTaker = new SnapshotTaker(region); + // Execute snapshot + boolean success = + snapshotTaker.takeFullSnapshot(snapshotPath.toAbsolutePath().toString(), true); + Assert.assertTrue("Snapshot execution failed", success); + + // 1. Verify Snapshot Log + File logFile = snapshotPath.resolve(SnapshotLogger.SNAPSHOT_LOG_NAME).toFile(); + Assert.assertTrue("Snapshot log must exist", logFile.exists()); + + SnapshotLogAnalyzer analyzer = new SnapshotLogAnalyzer(logFile); + Assert.assertTrue("Log must mark snapshot as complete", analyzer.isSnapshotComplete()); + + Path actualRoot = resolveActualSnapshotRoot(); + // 2. Verify TsFile physical structure in snapshot + validateTsFileSnapshotStructure(actualRoot); + + // 3. Verify Object file physical structure + validateObjectFileSnapshotStructure(actualRoot, expectedObjectFiles); + + LOGGER.info("testCreateSnapshotWithMixedFiles completed successfully."); + } + + private Path resolveActualSnapshotRoot() { + // Get the first data directory configured in the test + String dataDir = IoTDBDescriptor.getInstance().getConfig().getTierDataDirs()[0][0]; + + return Paths.get(dataDir) + .resolve(IoTDBConstant.SNAPSHOT_FOLDER_NAME) // "snapshot" + .resolve(testSgName + IoTDBConstant.FILE_NAME_SEPARATOR + testRegionId) // "root.testsg-0" + .resolve(snapshotPath.getFileName().toString()); // "snapshot_export" + } + + /** + * Test snapshot recovery (loading). Logic: Creates a snapshot, closes the current region, and + * uses SnapshotLoader to restore. Verification: Checks if DataRegion is re-instantiated and files + * are restored to storage tiers. + */ + @Test + public void testLoadSnapshotWithMixedFiles() throws Exception { + DataRegion originalRegion = createAndPopulateDataRegion(); + Set expectedObjectFiles = prepareObjectFiles(testRegionId); + + // Create snapshot + new SnapshotTaker(originalRegion) + .takeFullSnapshot(snapshotPath.toAbsolutePath().toString(), true); + + // Load snapshot + SnapshotLoader loader = + new SnapshotLoader(snapshotPath.toAbsolutePath().toString(), testSgName, testRegionId); + DataRegion restoredRegion = loader.loadSnapshotForStateMachine(); + + Assert.assertNotNull("Restored DataRegion should not be null", restoredRegion); + Assert.assertEquals( + "Restored TsFile count mismatch", 5, restoredRegion.getTsFileManager().size(true)); + + // Verify Object files are back in storage tiers + for (Path path : expectedObjectFiles) { + Assert.assertTrue("Object file not restored to tiers: " + path, existsInStorageTiers(path)); + } + + LOGGER.info("testLoadSnapshotWithMixedFiles completed successfully."); + } + + // --- Industrial Helper Methods --- + + /** + * Generates TsFiles following the exact IoTDB directory hierarchy: + * data/sequence/{database}/{regionId}/{timePartition}/{fileName} + */ + private List writeTsFiles() throws IOException, WriteProcessException { + List resources = new ArrayList<>(); + // Align with IoTDBSnapshotTest structure + Path dataRegionDir = + dataRootPath + .resolve(IoTDBConstant.SEQUENCE_FOLDER_NAME) + .resolve(testSgName) + .resolve(testRegionId) + .resolve(timePartition); + + Files.createDirectories(dataRegionDir); + + for (int i = 1; i <= 5; i++) { + String fileName = String.format("%d-%d-0-0.tsfile", i, i); + Path tsFilePath = dataRegionDir.resolve(fileName); + + // Use standard generator for valid TsFile content + TsFileGeneratorUtils.generateMixTsFile(tsFilePath.toString(), 2, 2, 10, 0, 100, 10, 10); + + TsFileResource resource = new TsFileResource(tsFilePath.toFile()); + // Resource must be serialized to satisfy SnapshotTaker + IDeviceID deviceID = + IDeviceID.Factory.DEFAULT_FACTORY.create(testSgName + PATH_SEPARATOR + "d1"); + resource.updateStartTime(deviceID, 0); + resource.updateEndTime(deviceID, 100); + resource.setStatusForTest(TsFileResourceStatus.NORMAL); + resource.serialize(); + + resources.add(resource); + } + return resources; + } + + private Set prepareObjectFiles(String regionId) throws IOException { + Path objectBaseFolder = Paths.get(TierManager.getInstance().getAllObjectFileFolders().get(0)); + Set absolutePaths = new HashSet<>(); + + List relativeLogicPaths = + Arrays.asList( + Paths.get(regionId, timePartition, "obj_a" + ObjectTypeUtils.OBJECT_FILE_SUFFIX), + Paths.get(regionId, timePartition, "obj_b" + ObjectTypeUtils.OBJECT_TEMP_FILE_SUFFIX)); + + for (Path rel : relativeLogicPaths) { + Path abs = objectBaseFolder.resolve(rel); + Files.createDirectories(abs.getParent()); + Files.write(abs, "mock-object-data".getBytes(StandardCharsets.UTF_8)); + absolutePaths.add(abs); + } + return absolutePaths; + } + + private void validateTsFileSnapshotStructure(Path actualRoot) { + Path tsSnapshotDir = + actualRoot + .resolve(IoTDBConstant.SEQUENCE_FOLDER_NAME) + .resolve(testSgName) + .resolve(testRegionId) + .resolve(timePartition); + + Assert.assertTrue("TsFile snapshot directory missing", Files.exists(tsSnapshotDir)); + File[] files = tsSnapshotDir.toFile().listFiles((dir, name) -> name.endsWith(TSFILE_SUFFIX)); + Assert.assertNotNull(files); + Assert.assertEquals("Snapshot TsFile count mismatch", 5, files.length); + + for (File f : files) { + Assert.assertTrue( + "Resource file missing for " + f.getName(), + new File(f.getAbsolutePath() + TsFileResource.RESOURCE_SUFFIX).exists()); + } + } + + private void validateObjectFileSnapshotStructure(Path actualRoot, Set expectedFiles) { + Path objectSnapshotDir = actualRoot.resolve(IoTDBConstant.OBJECT_FOLDER_NAME); + Assert.assertTrue( + "Object snapshot directory missing at: " + objectSnapshotDir, + Files.exists(objectSnapshotDir)); + + for (Path sourcePath : expectedFiles) { + // In snapshot, object files retain their relative structure under 'object/' folder + Path fileName = sourcePath.getFileName(); + Path relativeParent = Paths.get(testRegionId).resolve(timePartition); + Path targetPath = objectSnapshotDir.resolve(relativeParent).resolve(fileName); + Assert.assertTrue("Object file missing in snapshot: " + targetPath, Files.exists(targetPath)); + } + } + + private boolean existsInStorageTiers(Path originalAbsPath) { + // Check all configured tiers for the existence of the file relative to the tier root + Path relativePath = + Paths.get(TierManager.getInstance().getAllObjectFileFolders().get(0)) + .relativize(originalAbsPath); + return TierManager.getInstance().getAllObjectFileFolders().stream() + .map(folder -> Paths.get(folder).resolve(relativePath)) + .anyMatch(Files::exists); + } + + private DataRegion createAndPopulateDataRegion() throws IOException, WriteProcessException { + List resources = writeTsFiles(); + DataRegion region = new DataRegion(testSgName, testRegionId); + region.getTsFileManager().addAll(resources, true); + return region; + } +}