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) {} }