This is an automated email from the ASF dual-hosted git repository. jt2594838 pushed a commit to branch support_wal_in_import_data in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 2cd41b697d462a1cb9b48e2dfe7150f0d7a14904 Author: Tian Jiang <[email protected]> AuthorDate: Mon Aug 24 11:18:24 2026 +0800 add thread_num --- .../apache/iotdb/db/i18n/ImportWALMessages.java | 8 + .../apache/iotdb/db/i18n/ImportWALMessages.java | 8 + .../java/org/apache/iotdb/db/tools/ImportWAL.java | 437 ++++++++++++++++----- .../org/apache/iotdb/db/tools/ImportWALTest.java | 112 +++++- 4 files changed, 459 insertions(+), 106 deletions(-) diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java index 0dd487bfe8b..3061a1845be 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/ImportWALMessages.java @@ -41,6 +41,9 @@ public final class ImportWALMessages { public static final String MESSAGE_WHEN_ALL_WAL_FILES_ARE_REPLAYED_SUCCESSFULLY_DO_OPERATION_ON_SOURCE_WAL_FILES_OPTIONAL_PARAMETERS_ARE_NONE_DEFAULT_AND_DELETE_41963A66 = "When all WAL files are replayed successfully, do operation on source WAL files. Optional parameters are none (default) and delete."; + public static final String + MESSAGE_NUMBER_OF_THREADS_USED_TO_REPLAY_WAL_DIRECTORIES_IN_PARALLEL_DEFAULT_1_6AEF4F50 = + "Number of threads used to replay WAL directories in parallel. Default: 1."; public static final String MESSAGE_PRINT_THIS_HELP_MESSAGE_E800AF7A = "Print this help message."; public static final String MESSAGE_ARGUMENT_ERROR_ARG_A9767F62 = "Argument error: %s"; @@ -52,6 +55,11 @@ public final class ImportWALMessages { public static final String EXCEPTION_NO_WAL_FILES_FOUND_UNDER_ARG_45F7FA22 = "No WAL files found under: %s"; public static final String EXCEPTION_INVALID_PORT_ARG_A7CDD5AC = "Invalid port: %s"; + public static final String + EXCEPTION_INVALID_THREAD_COUNT_ARG_EXPECTED_A_POSITIVE_INTEGER_F3AE2CFD = + "Invalid thread count: %s. Expected a positive integer."; + public static final String EXCEPTION_WAL_REPLAY_WAS_INTERRUPTED_770BA8AD = + "WAL replay was interrupted."; public static final String MESSAGE_REPLAYED_ARG_OPERATIONS_FROM_ARG_WAL_FILES_SKIPPED_ARG_ENTRIES_F0D37E3A = "Replayed %d operations from %d WAL files; skipped %d entries."; diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java index 57052048e63..120c3265999 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/ImportWALMessages.java @@ -41,6 +41,9 @@ public final class ImportWALMessages { public static final String MESSAGE_WHEN_ALL_WAL_FILES_ARE_REPLAYED_SUCCESSFULLY_DO_OPERATION_ON_SOURCE_WAL_FILES_OPTIONAL_PARAMETERS_ARE_NONE_DEFAULT_AND_DELETE_41963A66 = "所有 WAL 文件成功重放后，对源 WAL 文件执行操作。可选参数为 none（默认）和 delete。"; + public static final String + MESSAGE_NUMBER_OF_THREADS_USED_TO_REPLAY_WAL_DIRECTORIES_IN_PARALLEL_DEFAULT_1_6AEF4F50 = + "并行重放 WAL 目录所用的线程数。默认：1。"; public static final String MESSAGE_PRINT_THIS_HELP_MESSAGE_E800AF7A = "打印帮助信息。"; public static final String MESSAGE_ARGUMENT_ERROR_ARG_A9767F62 = "参数错误：%s"; public static final String MESSAGE_WAL_IMPORT_FAILED_ARG_55C014BA = "WAL 导入失败：%s"; @@ -51,6 +54,11 @@ public final class ImportWALMessages { public static final String EXCEPTION_NO_WAL_FILES_FOUND_UNDER_ARG_45F7FA22 = "路径下未找到 WAL 文件：%s"; public static final String EXCEPTION_INVALID_PORT_ARG_A7CDD5AC = "无效端口：%s"; + public static final String + EXCEPTION_INVALID_THREAD_COUNT_ARG_EXPECTED_A_POSITIVE_INTEGER_F3AE2CFD = + "无效线程数：%s。应为正整数。"; + public static final String EXCEPTION_WAL_REPLAY_WAS_INTERRUPTED_770BA8AD = + "WAL 重放被中断。"; public static final String MESSAGE_REPLAYED_ARG_OPERATIONS_FROM_ARG_WAL_FILES_SKIPPED_ARG_ENTRIES_F0D37E3A = "已重放 %d 个操作（来自 %d 个 WAL 文件）；跳过 %d 个条目。"; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java index c1281c00ee3..2214b69db92 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/ImportWAL.java @@ -75,10 +75,17 @@ import java.util.ArrayList; import java.util.Collections; import java.util.Comparator; import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Locale; import java.util.Map; import java.util.Objects; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -89,7 +96,9 @@ public class ImportWAL { private static final String DEFAULT_HOST = "127.0.0.1"; private static final int DEFAULT_PORT = 6667; private static final String DEFAULT_USER = "root"; + private static final int DEFAULT_THREAD_NUM = 1; private static final int SNAPSHOT_TABLET_ROW_LIMIT = 1024; + private static final Object CONSOLE_PROMPT_LOCK = new Object(); private ImportWAL() {} @@ -120,53 +129,38 @@ public class ImportWAL { final String database = commandLine.getOptionValue("database"); final boolean deleteSource = shouldDeleteSource(commandLine.getOptionValue("on_success", "none")); + final int threadNum = + parseThreadNum( + commandLine.getOptionValue("thread_num", String.valueOf(DEFAULT_THREAD_NUM))); final String password = getPassword(commandLine); - final Session treeSession = - createSession( - commandLine.getOptionValue("host", DEFAULT_HOST), - parsePort(commandLine.getOptionValue("port", String.valueOf(DEFAULT_PORT))), - commandLine.getOptionValue("username", DEFAULT_USER), - password, - null); - final Session tableSession = - database == null - ? null - : createSession( - commandLine.getOptionValue("host", DEFAULT_HOST), - parsePort(commandLine.getOptionValue("port", String.valueOf(DEFAULT_PORT))), - commandLine.getOptionValue("username", DEFAULT_USER), - password, - database); - try { - treeSession.open(false); - if (tableSession != null) { - tableSession.open(false); - } - final ReplayStatistics statistics = - replayWALFiles( - walFiles, new WALReplayer(treeSession, tableSession, database), out, deleteSource); + final String host = commandLine.getOptionValue("host", DEFAULT_HOST); + final int port = parsePort(commandLine.getOptionValue("port", String.valueOf(DEFAULT_PORT))); + final String username = commandLine.getOptionValue("username", DEFAULT_USER); + final ReplayStatistics statistics = + replayWALDirectories( + walFiles, + threadNum, + () -> createWALReplayWorker(host, port, username, password, database), + out, + deleteSource); + out.printf( + ImportWALMessages + .MESSAGE_REPLAYED_ARG_OPERATIONS_FROM_ARG_WAL_FILES_SKIPPED_ARG_ENTRIES_F0D37E3A, + statistics.replayedOperationCount, + walFiles.size(), + statistics.skippedEntryCount); + out.println(); + out.printf( + ImportWALMessages + .MESSAGE_IMPORT_DURATION_ARG_SECONDS_TOTAL_SIZE_ARG_BYTES_AVERAGE_RATE_ARG_MB_PER_SECOND_4B4EA58D, + statistics.getElapsedSeconds(), + statistics.getTotalBytes(), + statistics.getAverageRateMbPerSecond()); + out.println(); + if (deleteSource) { out.printf( - ImportWALMessages - .MESSAGE_REPLAYED_ARG_OPERATIONS_FROM_ARG_WAL_FILES_SKIPPED_ARG_ENTRIES_F0D37E3A, - statistics.replayedOperationCount, - walFiles.size(), - statistics.skippedEntryCount); + ImportWALMessages.MESSAGE_DELETED_ARG_SOURCE_WAL_FILES_C7A5AA1B, walFiles.size()); out.println(); - out.printf( - ImportWALMessages - .MESSAGE_IMPORT_DURATION_ARG_SECONDS_TOTAL_SIZE_ARG_BYTES_AVERAGE_RATE_ARG_MB_PER_SECOND_4B4EA58D, - statistics.getElapsedSeconds(), - statistics.getTotalBytes(), - statistics.getAverageRateMbPerSecond()); - out.println(); - if (deleteSource) { - out.printf( - ImportWALMessages.MESSAGE_DELETED_ARG_SOURCE_WAL_FILES_C7A5AA1B, walFiles.size()); - out.println(); - } - } finally { - closeSession(tableSession); - closeSession(treeSession); } return CODE_OK; } catch (final Exception e) { @@ -228,6 +222,15 @@ public class ImportWAL { ImportWALMessages .MESSAGE_WHEN_ALL_WAL_FILES_ARE_REPLAYED_SUCCESSFULLY_DO_OPERATION_ON_SOURCE_WAL_FILES_OPTIONAL_PARAMETERS_ARE_NONE_DEFAULT_AND_DELETE_41963A66) .build()); + options.addOption( + Option.builder("tn") + .longOpt("thread_num") + .argName("thread_num") + .hasArg() + .desc( + ImportWALMessages + .MESSAGE_NUMBER_OF_THREADS_USED_TO_REPLAY_WAL_DIRECTORIES_IN_PARALLEL_DEFAULT_1_6AEF4F50) + .build()); options.addOption( Option.builder() .longOpt("help") @@ -314,6 +317,23 @@ public class ImportWAL { } } + static int parseThreadNum(final String threadNum) { + try { + final int value = Integer.parseInt(threadNum); + if (value <= 0) { + throw new NumberFormatException(threadNum); + } + return value; + } catch (final NumberFormatException e) { + throw new IllegalArgumentException( + String.format( + ImportWALMessages + .EXCEPTION_INVALID_THREAD_COUNT_ARG_EXPECTED_A_POSITIVE_INTEGER_F3AE2CFD, + threadNum), + e); + } + } + private static Session createSession( final String host, final int port, @@ -328,6 +348,29 @@ public class ImportWAL { return builder.build(); } + private static WALReplayWorker createWALReplayWorker( + final String host, + final int port, + final String username, + final String password, + final String database) + throws IOException { + final Session treeSession = createSession(host, port, username, password, null); + final Session tableSession = + database == null ? null : createSession(host, port, username, password, database); + try { + treeSession.open(false); + if (tableSession != null) { + tableSession.open(false); + } + return new SessionWALReplayer(treeSession, tableSession, database); + } catch (final IoTDBConnectionException e) { + closeSession(tableSession); + closeSession(treeSession); + throw new IOException(e.getMessage(), e); + } + } + private static void closeSession(final Session session) { if (session == null) { return; @@ -409,78 +452,219 @@ public class ImportWAL { final PrintStream progressStream, final boolean deleteSource) throws IOException { - final ReplayStatistics statistics = new ReplayStatistics(); final long startNanos = System.nanoTime(); + final ReplayStatistics statistics = createReplayStatistics(walFiles); for (final Path walFile : walFiles) { - statistics.totalBytes += Files.size(walFile); + recordCompletedFile( + statistics, + replayWALFile(walFile, replayer), + walFiles.size(), + startNanos, + progressStream); + } + if (deleteSource) { + deleteSourceWALFiles(walFiles); + } + return statistics; + } + + static ReplayStatistics replayWALDirectories( + final List<Path> walFiles, + final int threadNum, + final WALReplayWorkerFactory workerFactory, + final PrintStream progressStream, + final boolean deleteSource) + throws IOException { + if (threadNum <= 0) { + throw new IllegalArgumentException( + String.format( + ImportWALMessages + .EXCEPTION_INVALID_THREAD_COUNT_ARG_EXPECTED_A_POSITIVE_INTEGER_F3AE2CFD, + threadNum)); + } + final long startNanos = System.nanoTime(); + final ReplayStatistics statistics = createReplayStatistics(walFiles); + final List<List<Path>> walDirectories = groupWALFilesByDirectory(walFiles); + if (walDirectories.isEmpty()) { + statistics.elapsedNanos = System.nanoTime() - startNanos; + return statistics; + } + + final int workerCount = Math.min(threadNum, walDirectories.size()); + final AtomicInteger nextDirectoryIndex = new AtomicInteger(); + final AtomicReference<Throwable> replayFailure = new AtomicReference<>(); + final ExecutorService executor = Executors.newFixedThreadPool(workerCount); + final List<Future<?>> futures = new ArrayList<>(workerCount); + for (int i = 0; i < workerCount; i++) { + futures.add( + executor.submit( + () -> { + try (WALReplayWorker worker = workerFactory.create()) { + int directoryIndex; + // A worker owns one directory at a time so WAL files from that directory remain + // ordered, while independent directories can make progress concurrently. + while (replayFailure.get() == null + && (directoryIndex = nextDirectoryIndex.getAndIncrement()) + < walDirectories.size()) { + for (final Path walFile : walDirectories.get(directoryIndex)) { + recordCompletedFile( + statistics, + replayWALFile(walFile, worker), + walFiles.size(), + startNanos, + progressStream); + } + } + } catch (final Exception e) { + replayFailure.compareAndSet(null, e); + } + })); + } + executor.shutdown(); + + boolean interrupted = false; + for (final Future<?> future : futures) { + boolean completed = false; + // Wait for every worker even after a failure so sessions are closed before source deletion or + // the failure is reported to the caller. + while (!completed) { + try { + future.get(); + completed = true; + } catch (final InterruptedException e) { + interrupted = true; + replayFailure.compareAndSet( + null, + new IOException(ImportWALMessages.EXCEPTION_WAL_REPLAY_WAS_INTERRUPTED_770BA8AD, e)); + } catch (final ExecutionException e) { + replayFailure.compareAndSet(null, e.getCause()); + completed = true; + } + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + statistics.elapsedNanos = System.nanoTime() - startNanos; + rethrowReplayFailure(replayFailure.get()); + if (deleteSource) { + deleteSourceWALFiles(walFiles); + } + return statistics; + } + + static List<List<Path>> groupWALFilesByDirectory(final List<Path> walFiles) { + final List<Path> sortedWALFiles = new ArrayList<>(walFiles); + sortedWALFiles.sort(WAL_FILE_COMPARATOR); + final Map<Path, List<Path>> filesByDirectory = new LinkedHashMap<>(); + for (final Path walFile : sortedWALFiles) { + filesByDirectory + .computeIfAbsent(walFile.getParent(), ignored -> new ArrayList<>()) + .add(walFile); } + return new ArrayList<>(filesByDirectory.values()); + } + + private static ReplayStatistics createReplayStatistics(final List<Path> walFiles) + throws IOException { + final ReplayStatistics statistics = new ReplayStatistics(); for (final Path walFile : walFiles) { - try (WALReader reader = new WALReader(walFile.toFile())) { - long offset = reader.getWALCurrentReadOffset(); - while (reader.hasNext()) { - final WALEntry entry = reader.next(); - try { - if (replayer.replay(entry)) { - statistics.replayedOperationCount++; - } else { - statistics.skippedEntryCount++; - } - } catch (final IoTDBConnectionException | StatementExecutionException e) { - throw new WALReplayException( - String.format( - ImportWALMessages - .EXCEPTION_FAILED_TO_REPLAY_WAL_FILE_ARG_AT_OFFSET_ARG_ARG_FCFAF7F9, - walFile, - offset, - e.getMessage()), - e); + statistics.totalBytes += Files.size(walFile); + } + return statistics; + } + + private static ReplayStatistics replayWALFile(final Path walFile, final WALReplayWorker replayer) + throws IOException { + final ReplayStatistics statistics = new ReplayStatistics(); + try (WALReader reader = new WALReader(walFile.toFile())) { + long offset = reader.getWALCurrentReadOffset(); + while (reader.hasNext()) { + final WALEntry entry = reader.next(); + try { + if (replayer.replay(entry)) { + statistics.replayedOperationCount++; + } else { + statistics.skippedEntryCount++; } - offset = reader.getWALCurrentReadOffset(); - } - if (reader.isFileCorrupted()) { + } catch (final IoTDBConnectionException | StatementExecutionException e) { throw new WALReplayException( String.format( ImportWALMessages .EXCEPTION_FAILED_TO_REPLAY_WAL_FILE_ARG_AT_OFFSET_ARG_ARG_FCFAF7F9, walFile, - reader.getWALCurrentReadOffset(), - ImportWALMessages.EXCEPTION_THE_WAL_FILE_IS_TRUNCATED_OR_CORRUPTED_6B0734C5), - null); - } - statistics.completedFileCount++; - statistics.processedBytes += Files.size(walFile); - statistics.elapsedNanos = System.nanoTime() - startNanos; - if (progressStream != null) { - progressStream.printf( - ImportWALMessages - .MESSAGE_PROGRESS_ARG_COMPLETED_FILES_ARG_TOTAL_FILES_ARG_PROCESSED_BYTES_ARG_TOTAL_BYTES_ARG_PERCENT_ARG_ELAPSED_SECONDS_ARG_RATE_ARG_MB_PER_SECOND_F1C1356F, - statistics.completedFileCount, - walFiles.size(), - statistics.processedBytes, - statistics.totalBytes, - statistics.getProgressPercent(), - statistics.getElapsedSeconds(), - statistics.getAverageRateMbPerSecond()); - progressStream.println(); - } - } catch (final IOException e) { - if (e instanceof WALReplayException walReplayException) { - throw walReplayException; + offset, + e.getMessage()), + e); } + offset = reader.getWALCurrentReadOffset(); + } + if (reader.isFileCorrupted()) { throw new WALReplayException( String.format( ImportWALMessages .EXCEPTION_FAILED_TO_REPLAY_WAL_FILE_ARG_AT_OFFSET_ARG_ARG_FCFAF7F9, walFile, - 0, - e.getMessage()), - e); + reader.getWALCurrentReadOffset(), + ImportWALMessages.EXCEPTION_THE_WAL_FILE_IS_TRUNCATED_OR_CORRUPTED_6B0734C5), + null); } + statistics.completedFileCount++; + statistics.processedBytes += Files.size(walFile); + return statistics; + } catch (final IOException e) { + if (e instanceof WALReplayException walReplayException) { + throw walReplayException; + } + throw new WALReplayException( + String.format( + ImportWALMessages.EXCEPTION_FAILED_TO_REPLAY_WAL_FILE_ARG_AT_OFFSET_ARG_ARG_FCFAF7F9, + walFile, + 0, + e.getMessage()), + e); } - if (deleteSource) { - deleteSourceWALFiles(walFiles); + } + + private static void recordCompletedFile( + final ReplayStatistics statistics, + final ReplayStatistics completedFileStatistics, + final int totalFileCount, + final long startNanos, + final PrintStream progressStream) { + synchronized (statistics) { + statistics.add(completedFileStatistics); + statistics.elapsedNanos = System.nanoTime() - startNanos; + if (progressStream != null) { + progressStream.printf( + ImportWALMessages + .MESSAGE_PROGRESS_ARG_COMPLETED_FILES_ARG_TOTAL_FILES_ARG_PROCESSED_BYTES_ARG_TOTAL_BYTES_ARG_PERCENT_ARG_ELAPSED_SECONDS_ARG_RATE_ARG_MB_PER_SECOND_F1C1356F, + statistics.completedFileCount, + totalFileCount, + statistics.processedBytes, + statistics.totalBytes, + statistics.getProgressPercent(), + statistics.getElapsedSeconds(), + statistics.getAverageRateMbPerSecond()); + progressStream.println(); + } } - return statistics; + } + + private static void rethrowReplayFailure(final Throwable failure) throws IOException { + if (failure == null) { + return; + } + if (failure instanceof IOException ioException) { + throw ioException; + } + if (failure instanceof RuntimeException runtimeException) { + throw runtimeException; + } + if (failure instanceof Error error) { + throw error; + } + throw new IOException(failure); } private static void deleteSourceWALFiles(final List<Path> walFiles) throws IOException { @@ -498,7 +682,40 @@ public class ImportWAL { } } - static class WALReplayer { + @FunctionalInterface + interface WALReplayWorkerFactory { + + WALReplayWorker create() throws Exception; + } + + interface WALReplayWorker extends AutoCloseable { + + boolean replay(WALEntry entry) throws IoTDBConnectionException, StatementExecutionException; + + @Override + default void close() {} + } + + private static class SessionWALReplayer extends WALReplayer { + + private final Session treeSession; + private final Session tableSession; + + private SessionWALReplayer( + final Session treeSession, final Session tableSession, final String tableDatabaseName) { + super(treeSession, tableSession, tableDatabaseName); + this.treeSession = treeSession; + this.tableSession = tableSession; + } + + @Override + public void close() { + closeSession(tableSession); + closeSession(treeSession); + } + } + + static class WALReplayer implements WALReplayWorker { private final Session treeSession; private final Session tableSession; @@ -528,7 +745,8 @@ public class ImportWAL { null, null, ColumnFilterMatcher.matchAll(), tableDatabaseName); } - boolean replay(final WALEntry entry) + @Override + public boolean replay(final WALEntry entry) throws IoTDBConnectionException, StatementExecutionException { if (entry.getType() == WALEntryType.MEMORY_TABLE_SNAPSHOT || entry.getType() == WALEntryType.OLD_MEMORY_TABLE_SNAPSHOT) { @@ -558,12 +776,14 @@ public class ImportWAL { return null; } return entry -> { - final String answer = - console.readLine( - ImportWALMessages - .MESSAGE_UNSUPPORTED_WAL_OPERATION_ARG_SKIP_THIS_ENTRY_Y_N_DAFBE650, - entry.getType()); - return isSkipConfirmation(answer); + synchronized (CONSOLE_PROMPT_LOCK) { + final String answer = + console.readLine( + ImportWALMessages + .MESSAGE_UNSUPPORTED_WAL_OPERATION_ARG_SKIP_THIS_ENTRY_Y_N_DAFBE650, + entry.getType()); + return isSkipConfirmation(answer); + } }; } @@ -998,6 +1218,13 @@ public class ImportWAL { private long completedFileCount; private long elapsedNanos; + private void add(final ReplayStatistics statistics) { + replayedOperationCount += statistics.replayedOperationCount; + skippedEntryCount += statistics.skippedEntryCount; + processedBytes += statistics.processedBytes; + completedFileCount += statistics.completedFileCount; + } + long getReplayedOperationCount() { return replayedOperationCount; } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java index c8efb24ef9f..a40f020c18c 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/tools/ImportWALTest.java @@ -62,6 +62,12 @@ import java.nio.file.Path; import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; @@ -70,6 +76,7 @@ import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; @@ -104,7 +111,7 @@ public class ImportWALTest { files); } - /** Covers CLI discovery of the source-file operation without opening a Session. */ + /** Covers CLI discovery of source-file and parallel replay options without opening a Session. */ @Test public void testHelpDescribesDeleteSourceOption() { final ByteArrayOutputStream output = new ByteArrayOutputStream(); @@ -114,6 +121,79 @@ public class ImportWALTest { assertEquals(0, exitCode); assertTrue(output.toString().contains("--on_success")); + assertTrue(output.toString().contains("--thread_num")); + } + + /** Covers valid thread counts and rejects zero, negative, and non-numeric values. */ + @Test + public void testParseThreadNumOption() { + assertEquals(1, ImportWAL.parseThreadNum("1")); + assertEquals(4, ImportWAL.parseThreadNum("4")); + assertThrows(IllegalArgumentException.class, () -> ImportWAL.parseThreadNum("0")); + assertThrows(IllegalArgumentException.class, () -> ImportWAL.parseThreadNum("-1")); + assertThrows(IllegalArgumentException.class, () -> ImportWAL.parseThreadNum("invalid")); + } + + /** + * Covers directory-level parallel replay with two WAL files per directory. Different directories + * must overlap, while versions within each directory must retain ascending replay order. + */ + @Test + public void testReplayWALDirectoriesInParallelAndPreserveDirectoryOrder() throws Exception { + final Path source = temporaryFolder.newFolder("parallel-wal-root").toPath(); + final Path nodeA = Files.createDirectory(source.resolve("node-a")); + final Path nodeB = Files.createDirectory(source.resolve("node-b")); + final Path a1 = createWALFile(nodeA, 1); + final Path a2 = createWALFile(nodeA, 2); + final Path b1 = createWALFile(nodeB, 1); + final Path b2 = createWALFile(nodeB, 2); + writeWAL(a1.toFile(), new WALInfoEntry(1, WALTestUtils.getInsertRowNode("root.sg.a", 1))); + writeWAL(a2.toFile(), new WALInfoEntry(2, WALTestUtils.getInsertRowNode("root.sg.a", 2))); + writeWAL(b1.toFile(), new WALInfoEntry(3, WALTestUtils.getInsertRowNode("root.sg.b", 1))); + writeWAL(b2.toFile(), new WALInfoEntry(4, WALTestUtils.getInsertRowNode("root.sg.b", 2))); + final List<Path> walFiles = ImportWAL.collectWALFiles(source); + final CyclicBarrier replayBarrier = new CyclicBarrier(2); + final AtomicInteger activeReplays = new AtomicInteger(); + final AtomicInteger maxActiveReplays = new AtomicInteger(); + final AtomicInteger createdWorkers = new AtomicInteger(); + final Map<String, List<Long>> replayedTimestamps = new ConcurrentHashMap<>(); + + final ImportWAL.ReplayStatistics statistics = + ImportWAL.replayWALDirectories( + walFiles, + 2, + () -> { + createdWorkers.incrementAndGet(); + final Session session = mock(Session.class); + doAnswer( + invocation -> { + final Tablet tablet = invocation.getArgument(0); + final int active = activeReplays.incrementAndGet(); + maxActiveReplays.accumulateAndGet(active, Math::max); + try { + replayBarrier.await(5, TimeUnit.SECONDS); + replayedTimestamps + .computeIfAbsent( + tablet.getDeviceId(), ignored -> new CopyOnWriteArrayList<>()) + .add(tablet.getTimestamp(0)); + } finally { + activeReplays.decrementAndGet(); + } + return null; + }) + .when(session) + .insertTablet(any(Tablet.class)); + return new ImportWAL.WALReplayer(session, null, null); + }, + null, + false); + + assertEquals(2, createdWorkers.get()); + assertTrue(maxActiveReplays.get() >= 2); + assertEquals(Arrays.asList(1L, 2L), replayedTimestamps.get("root.sg.a")); + assertEquals(Arrays.asList(1L, 2L), replayedTimestamps.get("root.sg.b")); + assertEquals(4, statistics.getReplayedOperationCount()); + assertEquals(4, statistics.getCompletedFileCount()); } /** Covers the default retention value, deletion value, normalization, and invalid input. */ @@ -191,6 +271,36 @@ public class ImportWALTest { assertTrue(corruptedWALFile.exists()); } + /** + * Covers all-or-nothing deletion during directory-level parallel replay. A corrupted directory + * must retain WAL files from both the failed directory and another concurrently replayed one. + */ + @Test + public void testParallelReplayRetainsAllSourceFilesWhenAnyDirectoryFails() throws Exception { + final Path source = temporaryFolder.newFolder("parallel-retain-wal-root").toPath(); + final Path validDirectory = Files.createDirectory(source.resolve("valid")); + final Path corruptedDirectory = Files.createDirectory(source.resolve("corrupted")); + final Path validWALFile = createWALFile(validDirectory, 0); + final Path corruptedWALFile = createWALFile(corruptedDirectory, 0); + writeWAL( + validWALFile.toFile(), + new WALInfoEntry(1, WALTestUtils.getInsertRowNode("root.sg.parallel.retain", 1))); + Files.write(corruptedWALFile, new byte[] {WALEntryType.INSERT_ROW_NODE.getCode()}); + + assertThrows( + IOException.class, + () -> + ImportWAL.replayWALDirectories( + ImportWAL.collectWALFiles(source), + 2, + () -> new ImportWAL.WALReplayer(mock(Session.class), null, null), + null, + true)); + + assertTrue(Files.exists(validWALFile)); + assertTrue(Files.exists(corruptedWALFile)); + } + @Test public void testReplayReportsProgressAndFileStatistics() throws Exception { final File walFile = createWALFile(0);
