From d8c59cd5b895bbd3e68992d20db5724f45a6cadc Mon Sep 17 00:00:00 2001 From: luoluoyuyu Date: Fri, 16 Oct 2026 15:05:56 +0800 Subject: [PATCH] LOAD TsFile: use the Pair the resolver takes and cover the HA cases The router and the resolver each wrapped a (device, time partition) pair into a private record only to unwrap it again on the way in or out. Pair already carries the value equality the deduplication maps need, so keep the pair itself and drop both SlotKey records, the record and its conversion. Cover the two cluster cases the consensus LOAD path has to survive: - a LOAD that runs while one DataNode is down: restarting that DataNode must replay the LOAD commands it missed and serve the loaded file. - a LOAD whose region is migrated mid-flight: the LOAD has to follow the refreshed route or fall back to the tablet path, but every point of the file must be imported and the migration must finish. --- .../iotdb/db/it/IoTDBLoadTsFileClusterIT.java | 71 +++- .../it/IoTDBLoadTsFileRegionMigrationIT.java | 337 ++++++++++++++++++ .../scheduler/load/DataPartitionResolver.java | 10 +- .../scheduler/load/DataPartitionRouter.java | 33 +- 4 files changed, 419 insertions(+), 32 deletions(-) create mode 100644 integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileRegionMigrationIT.java diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileClusterIT.java b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileClusterIT.java index a3d39c6fcd65..3603c3dd09ba 100644 --- a/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileClusterIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileClusterIT.java @@ -21,6 +21,7 @@ import org.apache.iotdb.consensus.ConsensusFactory; import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper; import org.apache.iotdb.it.framework.IoTDBTestRunner; import org.apache.iotdb.it.utils.TsFileGenerator; import org.apache.iotdb.itbase.category.ClusterIT; @@ -260,6 +261,74 @@ public void testLoadAfterOneDataNodeDown() throws Exception { } } + /** + * A LOAD that runs while one DataNode is down is replicated by consensus to the surviving + * replicas. When the stopped DataNode restarts, its replicas replay the LOAD commands they + * missed, so the restarted DataNode ends up serving the very same file. + */ + @Test + public void testRestartedDataNodeIsSyncedWithTheLoadedFile() throws Exception { + final long firstPointCount = generateTsFile(tmpDir, "sync-first.tsfile", 5_000, 1L); + + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + createDatabaseAndSchema(statement); + statement.execute("load \"" + tmpDir.getAbsolutePath() + "\""); + awaitPointCount(statement, firstPointCount); + } + + // LOAD the second file while DataNode 0 is down. With three replicas it holds a copy of every + // region, so it is the DataNode that has to catch up with the LOAD when it comes back. + final DataNodeWrapper restartedDataNode = EnvFactory.getEnv().getDataNodeWrapper(0); + restartedDataNode.stop(); + Assert.assertFalse(restartedDataNode.isAlive()); + + final long secondPointCount = generateTsFile(tmpDir, "sync-second.tsfile", 3_000, 100_000L); + final long expectedPointCount = firstPointCount + secondPointCount; + try (final Connection connection = + DriverManager.getConnection( + Config.IOTDB_URL_PREFIX + + EnvFactory.getEnv().getDataNodeWrapper(1).getIpAndPortString(), + "root", + "root"); + final Statement statement = connection.createStatement()) { + statement.execute("load \"" + tmpDir.getAbsolutePath() + "\""); + awaitPointCount(statement, expectedPointCount); + } + + restartedDataNode.start(); + Assert.assertTrue(restartedDataNode.isAlive()); + + // The connection is pinned to the restarted DataNode, so this waits for the replicas it + // replayed the LOAD into and fails if the restart left them behind. + awaitPointCount(restartedDataNode, expectedPointCount); + } + + /** Waits until the given statement sees the expected number of imported points. */ + private static void awaitPointCount(final Statement statement, final long expectedPointCount) { + Awaitility.await() + .atMost(2, TimeUnit.MINUTES) + .pollInterval(1, TimeUnit.SECONDS) + .ignoreExceptions() + .untilAsserted(() -> Assert.assertEquals(expectedPointCount, queryPointCount(statement))); + } + + /** Waits until the points of every loaded file are queryable through the given DataNode. */ + private static void awaitPointCount( + final DataNodeWrapper dataNode, final long expectedPointCount) { + Awaitility.await() + .atMost(2, TimeUnit.MINUTES) + .pollInterval(1, TimeUnit.SECONDS) + .ignoreExceptions() + .untilAsserted( + () -> { + try (final Connection connection = EnvFactory.getEnv().getConnection(dataNode); + final Statement statement = connection.createStatement()) { + Assert.assertEquals(expectedPointCount, queryPointCount(statement)); + } + }); + } + /** With two of three DataNodes down the LOAD must fail with a message. */ @Test public void testLoadWithTwoDataNodesDownFailsWithAMeaningfulMessage() throws Exception { @@ -365,7 +434,7 @@ private long generateTsFile( } } - private long queryPointCount(final Statement statement) throws Exception { + private static long queryPointCount(final Statement statement) throws Exception { long actualPointCount = 0; final List allDevices = new ArrayList<>(DEVICES); allDevices.addAll(ALIGNED_DEVICES); diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileRegionMigrationIT.java b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileRegionMigrationIT.java new file mode 100644 index 000000000000..f45a42df17da --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileRegionMigrationIT.java @@ -0,0 +1,337 @@ +/* + * 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.it; + +import org.apache.iotdb.commons.schema.column.ColumnHeaderConstant; +import org.apache.iotdb.consensus.ConsensusFactory; +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.it.utils.TsFileGenerator; +import org.apache.iotdb.itbase.category.ClusterIT; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.apache.tsfile.write.schema.IMeasurementSchema; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.awaitility.Awaitility; +import org.junit.After; +import org.junit.AfterClass; +import org.junit.Assert; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.io.File; +import java.nio.file.Files; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +/** + * LOAD TsFile has to keep working when the topology changes while a file is being imported: the + * route of a region can be re-resolved after a piece was rejected, and a LOAD that notices a + * migration falls back to inserting the file through the ordinary write path. Either way every + * point of the file has to end up in the cluster and the migration has to finish. + * + *

The data replication factor is 2 with three DataNodes, so a region always has exactly one + * DataNode that can be its migration destination. + */ +@RunWith(IoTDBTestRunner.class) +@Category({ClusterIT.class}) +public class IoTDBLoadTsFileRegionMigrationIT { + + private static final long PARTITION_INTERVAL = 100_000_000L; + private static final String DATABASE = "root.load_migrate"; + private static final List DEVICES = + Arrays.asList( + DATABASE + ".d1", DATABASE + ".d2", DATABASE + ".d3", DATABASE + ".d4", DATABASE + ".d5"); + private static final List MEASUREMENTS = Arrays.asList("s1", "s2", "s3", "s4"); + private static final int POINT_COUNT_PER_DEVICE = 20_000; + private static final int DATA_REPLICATION_FACTOR = 2; + private static final int SCHEMA_REPLICATION_FACTOR = 3; + private static final int CONFIG_NODE_NUM = 1; + private static final int DATA_NODE_NUM = 3; + + private File tmpDir; + + @BeforeClass + public static void setUpClass() { + EnvFactory.getEnv().getConfig().getCommonConfig().setTimePartitionInterval(PARTITION_INTERVAL); + EnvFactory.getEnv().getConfig().getCommonConfig().setEnforceStrongPassword(false); + EnvFactory.getEnv().getConfig().getCommonConfig().setPipeMemoryManagementEnabled(false); + EnvFactory.getEnv().getConfig().getCommonConfig().setDatanodeMemoryProportion("1:10:1:1:1:0"); + EnvFactory.getEnv().getConfig().getCommonConfig().setTargetChunkPointNum(1_000); + EnvFactory.getEnv().getConfig().getCommonConfig().setMaxNumberOfPointsInPage(500); + EnvFactory.getEnv() + .getConfig() + .getCommonConfig() + .setConfigNodeConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS) + .setSchemaRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS) + .setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS) + .setSchemaReplicationFactor(SCHEMA_REPLICATION_FACTOR) + .setDataReplicationFactor(DATA_REPLICATION_FACTOR); + EnvFactory.getEnv().initClusterEnvironment(CONFIG_NODE_NUM, DATA_NODE_NUM); + } + + @AfterClass + public static void tearDownClass() { + EnvFactory.getEnv().cleanClusterEnvironment(); + } + + @Before + public void setUp() throws Exception { + tmpDir = new File(Files.createTempDirectory("load-migrate-it").toUri()); + } + + @After + public void tearDown() throws Exception { + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + statement.execute("delete database " + DATABASE); + } catch (final Exception ignored) { + } + + final File[] files = tmpDir.listFiles(); + if (files != null) { + for (final File file : files) { + Assert.assertTrue(deleteRecursively(file)); + } + } + Assert.assertTrue(tmpDir.delete()); + } + + /** + * Starts a LOAD of a file large enough to still be dispatching pieces when a region of its + * database is migrated away from a DataNode it was routed to, then asserts that both the LOAD and + * the migration finish with every point of the file imported. + */ + @Test + public void testLoadWhileARegionMigrates() throws Exception { + final long writtenPointCount = generateTsFile(new File(tmpDir, "load-migrating.tsfile")); + + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + createDatabaseAndSchema(statement); + } + + final ExecutorService loader = Executors.newSingleThreadExecutor(); + try { + final Future load = + loader.submit( + () -> { + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + statement.execute("load \"" + tmpDir.getAbsolutePath() + "\""); + return null; + } catch (final Exception e) { + return e; + } + }); + + // Wait until the LOAD allocated the data region of its points, then migrate that region while + // the LOAD is still splitting and dispatching the pieces of the file. + final int[] migration = awaitMigratableRegion(); + Assert.assertFalse("the LOAD finished before the migration could overlap it", load.isDone()); + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + statement.execute( + "migrate region " + migration[0] + " from " + migration[1] + " to " + migration[2]); + } + + // The LOAD may follow the migration with a refreshed route or fall back to the tablet path, + // but it must not report a failure to the client. + final Exception loadFailure = load.get(15, TimeUnit.MINUTES); + Assert.assertNull("the LOAD must survive a region migration: " + loadFailure, loadFailure); + + awaitRegionMigrated(migration[0], migration[1], migration[2]); + + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + awaitPointCount(statement, writtenPointCount); + } + } finally { + loader.shutdownNow(); + } + } + + /** + * Waits until the LOAD has allocated a data region of the database, and returns the region, one + * of the DataNodes that replicates it and a DataNode that does not, in that order. + */ + private static int[] awaitMigratableRegion() { + final AtomicReference migration = new AtomicReference<>(); + Awaitility.await() + .atMost(5, TimeUnit.MINUTES) + .pollInterval(200, TimeUnit.MILLISECONDS) + .ignoreExceptions() + .untilAsserted( + () -> { + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + final Set allDataNodes = allDataNodeIds(statement); + final Map> regionMap = dataRegionMap(statement); + Assert.assertFalse("the LOAD must allocate a data region", regionMap.isEmpty()); + + for (final Map.Entry> entry : regionMap.entrySet()) { + final Set replicas = entry.getValue(); + final int destination = + allDataNodes.stream() + .filter(dataNodeId -> !replicas.contains(dataNodeId)) + .findAny() + .orElse(-1); + if (destination < 0) { + continue; + } + migration.set( + new int[] {entry.getKey(), replicas.iterator().next(), destination}); + return; + } + Assert.fail("no data region has a free migration destination yet"); + } + }); + return migration.get(); + } + + /** Waits until the migrated region left the source DataNode and reached the destination. */ + private static void awaitRegionMigrated( + final int regionId, final int sourceDataNode, final int destinationDataNode) { + Awaitility.await() + .atMost(5, TimeUnit.MINUTES) + .pollInterval(1, TimeUnit.SECONDS) + .ignoreExceptions() + .untilAsserted( + () -> { + try (final Connection connection = EnvFactory.getEnv().getConnection(); + final Statement statement = connection.createStatement()) { + final Set replicas = dataRegionMap(statement).get(regionId); + Assert.assertNotNull("the migrated region must still exist", replicas); + Assert.assertFalse( + "the migrated region must leave DataNode " + sourceDataNode, + replicas.contains(sourceDataNode)); + Assert.assertTrue( + "the migrated region must reach DataNode " + destinationDataNode, + replicas.contains(destinationDataNode)); + } + }); + } + + /** Waits until the given statement sees the expected number of imported points. */ + private static void awaitPointCount(final Statement statement, final long expectedPointCount) { + Awaitility.await() + .atMost(5, TimeUnit.MINUTES) + .pollInterval(1, TimeUnit.SECONDS) + .ignoreExceptions() + .untilAsserted(() -> Assert.assertEquals(expectedPointCount, queryPointCount(statement))); + } + + private void createDatabaseAndSchema(final Statement statement) throws SQLException { + statement.execute("create database " + DATABASE); + for (final String device : DEVICES) { + for (final String measurement : MEASUREMENTS) { + statement.execute( + "create timeseries " + device + "." + measurement + " " + TSDataType.INT64.name()); + } + } + } + + /** The DataRegion rows of `show regions`, mapped to the DataNodes that replicate each region. */ + private static Map> dataRegionMap(final Statement statement) + throws SQLException { + final Map> regionMap = new HashMap<>(); + try (final ResultSet resultSet = statement.executeQuery("show regions")) { + while (resultSet.next()) { + if (!"DataRegion".equals(resultSet.getString(ColumnHeaderConstant.TYPE)) + || !DATABASE.equals(resultSet.getString(ColumnHeaderConstant.DATABASE))) { + continue; + } + final int regionId = resultSet.getInt(ColumnHeaderConstant.REGION_ID); + final int dataNodeId = resultSet.getInt(ColumnHeaderConstant.DATA_NODE_ID); + regionMap.computeIfAbsent(regionId, id -> new HashSet<>()).add(dataNodeId); + } + } + return regionMap; + } + + private static Set allDataNodeIds(final Statement statement) throws SQLException { + final Set dataNodeIds = new HashSet<>(); + try (final ResultSet resultSet = statement.executeQuery("show datanodes")) { + while (resultSet.next()) { + dataNodeIds.add(resultSet.getInt(ColumnHeaderConstant.NODE_ID)); + } + } + return dataNodeIds; + } + + private long generateTsFile(final File tsFile) throws Exception { + final List schemas = + Arrays.asList( + new MeasurementSchema(MEASUREMENTS.get(0), TSDataType.INT64, TSEncoding.PLAIN), + new MeasurementSchema(MEASUREMENTS.get(1), TSDataType.INT64, TSEncoding.PLAIN), + new MeasurementSchema(MEASUREMENTS.get(2), TSDataType.INT64, TSEncoding.PLAIN), + new MeasurementSchema(MEASUREMENTS.get(3), TSDataType.INT64, TSEncoding.PLAIN)); + try (final TsFileGenerator generator = new TsFileGenerator(tsFile)) { + for (final String device : DEVICES) { + generator.registerTimeseries(device, schemas); + generator.generateData(device, POINT_COUNT_PER_DEVICE, 1L, false, 1L); + } + return generator.getTotalNumber(); + } + } + + private static long queryPointCount(final Statement statement) throws Exception { + long actualPointCount = 0; + for (final String device : DEVICES) { + for (final String measurement : MEASUREMENTS) { + try (final ResultSet resultSet = + statement.executeQuery("select count(" + measurement + ") from " + device)) { + Assert.assertTrue(resultSet.next()); + actualPointCount += resultSet.getLong(1); + } + } + } + return actualPointCount; + } + + private static boolean deleteRecursively(final File file) { + final File[] children = file.listFiles(); + if (children != null) { + for (final File child : children) { + if (!deleteRecursively(child)) { + return false; + } + } + } + return file.delete(); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionResolver.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionResolver.java index f20e74636ff2..3a9dbadf12f0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionResolver.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionResolver.java @@ -43,7 +43,8 @@ class DataPartitionResolver { private final String userName; /** The replica set resolved for every slot so far. */ - private final Map resolvedPartitions = new HashMap<>(); + private final Map, TRegionReplicaSet> resolvedPartitions = + new HashMap<>(); DataPartitionResolver(final DataPartitionBatchFetcher partitionFetcher, final String userName) { this.partitionFetcher = @@ -63,7 +64,7 @@ List resolve(final List> final List indexesToFetch = new ArrayList<>(); for (int i = 0; i < slotList.size(); i++) { final Pair slot = slotList.get(i); - final TRegionReplicaSet resolved = resolvedPartitions.get(new SlotKey(slot.left, slot.right)); + final TRegionReplicaSet resolved = resolvedPartitions.get(slot); replicaSets.add(resolved); if (resolved == null) { slotsToFetch.add(slot); @@ -88,13 +89,10 @@ List resolve(final List> final TRegionReplicaSet replicaSet = fetched.get(i); replicaSets.set(indexesToFetch.get(i), replicaSet); final Pair slot = slotsToFetch.get(i); - resolvedPartitions.put(new SlotKey(slot.left, slot.right), replicaSet); + resolvedPartitions.put(slot, replicaSet); } } return replicaSets; } - - /** One (device, time partition) pair, the key the partitions are resolved by. */ - private record SlotKey(IDeviceID device, TTimePartitionSlot slot) {} } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionRouter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionRouter.java index 0d7478b307fa..f9caf454f0b7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionRouter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionRouter.java @@ -67,8 +67,9 @@ List route(final List chunkDataList) { final int totalChunks = chunkDataList.size(); final int[] chunkPartitionIndexes = new int[totalChunks]; - final List distinctSlotList = new ArrayList<>(); - final Map slotToIndexMap = new HashMap<>(totalChunks); + final List> distinctSlotList = new ArrayList<>(); + final Map, Integer> slotToIndexMap = + new HashMap<>(totalChunks); // 1. Deduplicate (device, time-partition) pairs and record mapping index per chunk for (int i = 0; i < totalChunks; i++) { @@ -88,7 +89,7 @@ List route(final List chunkDataList) { chunkData.getTimePartitionSlot(), DataNodeQueryMessages.EXCEPTION_CHUNK_TIME_PARTITION_SLOT_CANNOT_BE_NULL_E5B04C6F); - final SlotKey slotKey = new SlotKey(device, slot); + final Pair slotKey = new Pair<>(device, slot); final Integer existingIndex = slotToIndexMap.get(slotKey); if (existingIndex == null) { @@ -102,8 +103,7 @@ List route(final List chunkDataList) { } // 2. Resolve the replica sets of the distinct slots, reusing the ones the judgement resolved. - final List replicaSets = - partitionResolver.resolve(toQuerySlots(distinctSlotList)); + final List replicaSets = partitionResolver.resolve(distinctSlotList); // 3. Map resolved replica sets back to original input order final List routedReplicaSets = new ArrayList<>(totalChunks); @@ -111,34 +111,17 @@ List route(final List chunkDataList) { final int partitionIndex = chunkPartitionIndexes[i]; final TRegionReplicaSet replicaSet = replicaSets.get(partitionIndex); if (replicaSet == null) { - final SlotKey slotKey = distinctSlotList.get(partitionIndex); + final Pair slotKey = distinctSlotList.get(partitionIndex); throw new IllegalStateException( String.format( DataNodeQueryMessages .EXCEPTION_NULL_REPLICA_SET_RESOLVED_FOR_DEVICE_ARG_AT_PARTITION_ARG_19CA94B2, - slotKey.device(), - slotKey.slot())); + slotKey.getLeft(), + slotKey.getRight())); } routedReplicaSets.add(replicaSet); } return routedReplicaSets; } - - // ------------------------------------------------------------------------- - // Internal Model - // ------------------------------------------------------------------------- - - /** The distinct slots in the pair form the resolver takes. */ - private static List> toQuerySlots( - final List distinctSlots) { - final List> slots = new ArrayList<>(distinctSlots.size()); - for (final SlotKey slot : distinctSlots) { - slots.add(new Pair<>(slot.device(), slot.slot())); - } - return slots; - } - - /** Composite key for slot deduplication; the only slot representation this class holds. */ - private record SlotKey(IDeviceID device, TTimePartitionSlot slot) {} }