From 1c83e2829820ca5fa9f58d41249bc1e23ac7c469 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:13:07 +0800 Subject: [PATCH 1/2] Fence IoTConsensus applies before storage shutdown (#18841) * Fix queued IoTConsensus writes during shutdown * Test queued IoTConsensus writes across storage shutdown * Reproduce WAL checkpoint loss from writes after shutdown * Fence IoTConsensus applies before storage shutdown (cherry picked from commit 532d6c13addc04946535729db388fb4c18feb6b9) --- .../it/autocreate/IoTDBPipeClusterIT.java | 38 +- .../consensus/iot/IoTConsensusServerImpl.java | 9 +- .../iot/IoTConsensusServerImplTest.java | 135 ++++ .../dataregion/DataRegionStateMachine.java | 6 +- .../IoTConsensusDataRegionStateMachine.java | 7 + .../db/service/DataNodeShutdownHook.java | 3 + .../iotdb/db/storageengine/StorageEngine.java | 13 + .../storageengine/dataregion/DataRegion.java | 32 + .../dataregion/IoTConsensusShutdownTest.java | 589 ++++++++++++++++++ 9 files changed, 814 insertions(+), 18 deletions(-) create mode 100644 iotdb-core/consensus/src/test/java/org/apache/iotdb/consensus/iot/IoTConsensusServerImplTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/consensus/statemachine/dataregion/IoTConsensusShutdownTest.java diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/autocreate/IoTDBPipeClusterIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/autocreate/IoTDBPipeClusterIT.java index 16e21fdb2956d..43568e887d96a 100644 --- a/integration-test/src/test/java/org/apache/iotdb/pipe/it/autocreate/IoTDBPipeClusterIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/autocreate/IoTDBPipeClusterIT.java @@ -442,6 +442,8 @@ private void flushTreeDataRegionReplicasAfterReplicationComplete( private void waitForTreeDataRegionReplicationComplete( final BaseEnv env, final List databases) { + // IoTConsensus accepts writes on any replica, including a previous leader during failover. + // Check every replica's outgoing replication before flushing or restarting the cluster. await() .pollInterval(500, TimeUnit.MILLISECONDS) .atMost(2, TimeUnit.MINUTES) @@ -449,24 +451,23 @@ private void waitForTreeDataRegionReplicationComplete( () -> { try (final SyncConfigNodeIServiceClient client = (SyncConfigNodeIServiceClient) env.getLeaderConfigNodeConnection()) { - final List leaderRegionInfoList = - showTreeDataRegionLeaders(databases, client); + final List regionInfoList = showTreeDataRegions(databases, client); Assert.assertFalse( - "No tree DataRegion leader found for databases " + databases, - leaderRegionInfoList.isEmpty()); + "No tree DataRegion found for databases " + databases, + regionInfoList.isEmpty()); - for (final TRegionInfo regionInfo : leaderRegionInfoList) { - final DataNodeWrapper leaderNode = + for (final TRegionInfo regionInfo : regionInfoList) { + final DataNodeWrapper replicaNode = findDataNodeWrapperByPort(env, regionInfo.getClientRpcPort()); final String metricsUrl = "http://" - + leaderNode.getIp() + + replicaNode.getIp() + ":" - + leaderNode.getMetricPort() + + replicaNode.getMetricPort() + "/metrics"; final String metricsContent = env.getUrlContent(metricsUrl, null); Assert.assertNotNull( - "Failed to fetch metrics from leader DataNode at " + metricsUrl, + "Failed to fetch metrics from replica DataNode at " + metricsUrl, metricsContent); assertSyncLagIsZero(metricsContent, buildDataRegionTag(regionInfo), metricsUrl); } @@ -476,6 +477,17 @@ private void waitForTreeDataRegionReplicationComplete( private List showTreeDataRegionLeaders( final List databases, final SyncConfigNodeIServiceClient client) throws TException { + final List result = new ArrayList<>(); + for (final TRegionInfo regionInfo : showTreeDataRegions(databases, client)) { + if (RegionRoleType.Leader.getRoleType().equals(regionInfo.getRoleType())) { + result.add(regionInfo); + } + } + return result; + } + + private List showTreeDataRegions( + final List databases, final SyncConfigNodeIServiceClient client) throws TException { final TShowRegionResp showRegionResp = client.showRegion( new TShowRegionReq() @@ -483,13 +495,7 @@ private List showTreeDataRegionLeaders( .setDatabases(databases)); Assert.assertEquals( TSStatusCode.SUCCESS_STATUS.getStatusCode(), showRegionResp.getStatus().getCode()); - final List result = new ArrayList<>(); - for (final TRegionInfo regionInfo : showRegionResp.getRegionInfoList()) { - if (RegionRoleType.Leader.getRoleType().equals(regionInfo.getRoleType())) { - result.add(regionInfo); - } - } - return result; + return showRegionResp.getRegionInfoList(); } private DataNodeWrapper findDataNodeWrapperByPort(final BaseEnv env, final int port) { diff --git a/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/IoTConsensusServerImpl.java b/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/IoTConsensusServerImpl.java index 3002b018e3ea8..051bfc392cd50 100644 --- a/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/IoTConsensusServerImpl.java +++ b/iotdb-core/consensus/src/main/java/org/apache/iotdb/consensus/iot/IoTConsensusServerImpl.java @@ -31,6 +31,7 @@ import org.apache.iotdb.commons.utils.CommonDateTimeUtils; import org.apache.iotdb.commons.utils.KillPoint.DataNodeKillPoints; import org.apache.iotdb.commons.utils.KillPoint.KillPoint; +import org.apache.iotdb.commons.utils.StatusUtils; import org.apache.iotdb.consensus.IStateMachine; import org.apache.iotdb.consensus.common.DataSet; import org.apache.iotdb.consensus.common.Peer; @@ -1010,7 +1011,13 @@ private TSStatus cacheAndInsertLatestNode(DeserializedBatchIndexedConsensusReque List subStatus = new LinkedList<>(); for (IConsensusRequest insertNode : request.getInsertNodes()) { insertNode.markAsGeneratedByRemoteConsensusLeader(); - subStatus.add(stateMachine.write(insertNode)); + // Reject requests that became read-only while queued. This is an early rejection only: + // the state machine must also fence the full apply against storage shutdown, since the + // replica can become read-only after this check. + subStatus.add( + isReadOnly() + ? StatusUtils.getStatus(TSStatusCode.SYSTEM_READ_ONLY) + : stateMachine.write(insertNode)); } long applyTime = System.nanoTime(); ioTConsensusServerMetrics.recordApplyCost(applyTime - sortTime); diff --git a/iotdb-core/consensus/src/test/java/org/apache/iotdb/consensus/iot/IoTConsensusServerImplTest.java b/iotdb-core/consensus/src/test/java/org/apache/iotdb/consensus/iot/IoTConsensusServerImplTest.java new file mode 100644 index 0000000000000..64c2ea56a04d3 --- /dev/null +++ b/iotdb-core/consensus/src/test/java/org/apache/iotdb/consensus/iot/IoTConsensusServerImplTest.java @@ -0,0 +1,135 @@ +/* + * 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.consensus.iot; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.consensus.DataRegionId; +import org.apache.iotdb.commons.utils.StatusUtils; +import org.apache.iotdb.consensus.common.Peer; +import org.apache.iotdb.consensus.common.request.DeserializedBatchIndexedConsensusRequest; +import org.apache.iotdb.consensus.common.request.IndexedConsensusRequest; +import org.apache.iotdb.consensus.config.IoTConsensusConfig; +import org.apache.iotdb.consensus.iot.util.TestEntry; +import org.apache.iotdb.consensus.iot.util.TestStateMachine; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +import java.util.Collections; +import java.util.TreeSet; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +public class IoTConsensusServerImplTest { + @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder(); + + @Test + public void testSyncLogBecomesReadOnlyWhileWaiting() throws Exception { + final ScheduledExecutorService backgroundTaskService = + Executors.newSingleThreadScheduledExecutor(); + final AtomicBoolean readOnly = new AtomicBoolean(); + final TestStateMachine stateMachine = + new TestStateMachine() { + @Override + public boolean isReadOnly() { + return readOnly.get(); + } + }; + final Peer thisNode = new Peer(new DataRegionId(1), 0, new TEndPoint("127.0.0.1", 6667)); + final IoTConsensusServerImpl server = + new IoTConsensusServerImpl( + temporaryFolder.getRoot().getAbsolutePath(), + thisNode, + new TreeSet<>(Collections.singletonList(thisNode)), + stateMachine, + backgroundTaskService, + null, + null, + IoTConsensusConfig.newBuilder() + .setReplication( + IoTConsensusConfig.Replication.newBuilder() + .setMaxWaitingTimeForWaitBatchInMs(100) + .build()) + .build()); + final DeserializedBatchIndexedConsensusRequest request = + new DeserializedBatchIndexedConsensusRequest(1, 2, 2); + for (int index = 1; index <= 2; index++) { + request.add( + new IndexedConsensusRequest( + index, index, Collections.singletonList(new TestEntry(index, thisNode)))); + } + final AtomicReference status = new AtomicReference<>(); + final Thread syncLogThread = new Thread(() -> status.set(server.syncLog(1, request))); + try { + assertFalse(server.isReadOnly()); + syncLogThread.start(); + final long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (syncLogThread.getState() != Thread.State.TIMED_WAITING + && System.nanoTime() < deadline) { + Thread.sleep(10); + } + assertEquals(Thread.State.TIMED_WAITING, syncLogThread.getState()); + readOnly.set(true); + syncLogThread.join(TimeUnit.SECONDS.toMillis(15)); + assertFalse(syncLogThread.isAlive()); + assertEquals(2, status.get().getSubStatusSize()); + assertEquals( + TSStatusCode.SYSTEM_READ_ONLY.getStatusCode(), + status.get().getSubStatus().get(0).getCode()); + assertEquals( + StatusUtils.getStatus(TSStatusCode.SYSTEM_READ_ONLY).getMessage(), + status.get().getSubStatus().get(0).getMessage()); + assertEquals( + TSStatusCode.SYSTEM_READ_ONLY.getStatusCode(), + status.get().getSubStatus().get(1).getCode()); + assertTrue(stateMachine.getRequestSet().isEmpty()); + + // The sender must be able to retry the complete batch after the replica becomes writable. + readOnly.set(false); + server.reloadConsensusConfig( + IoTConsensusConfig.newBuilder() + .setReplication( + IoTConsensusConfig.Replication.newBuilder() + .setMaxPendingBatchesNum(1) + .setMaxWaitingTimeForWaitBatchInMs(100) + .build()) + .build()); + final TSStatus retryStatus = server.syncLog(1, request); + assertTrue( + retryStatus.getSubStatus().stream() + .allMatch(s -> s.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode())); + assertEquals(2, stateMachine.getRequestSet().size()); + } finally { + syncLogThread.interrupt(); + syncLogThread.join(TimeUnit.SECONDS.toMillis(5)); + backgroundTaskService.shutdownNow(); + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java index a76889d8b63d7..66c0692fcf59b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java @@ -214,7 +214,7 @@ protected TSStatus write(PlanNode planNode) { TSStatus result = null; int retryTime = 0; while (retryTime < MAX_WRITE_RETRY_TIMES) { - result = planNode.accept(new DataExecutionVisitor(), region); + result = executeWrite(planNode); if (needRetry(result.getCode())) { retryTime++; logger.debug( @@ -239,6 +239,10 @@ protected TSStatus write(PlanNode planNode) { return result; } + protected TSStatus executeWrite(PlanNode planNode) { + return planNode.accept(new DataExecutionVisitor(), region); + } + protected void waitBeforeNextWriteRetry() throws InterruptedException { Thread.sleep(WRITE_RETRY_WAIT_TIME_IN_MS); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/IoTConsensusDataRegionStateMachine.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/IoTConsensusDataRegionStateMachine.java index 240c1b1caa0fe..56ed7d775e792 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/IoTConsensusDataRegionStateMachine.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/IoTConsensusDataRegionStateMachine.java @@ -49,6 +49,13 @@ public IoTConsensusDataRegionStateMachine(DataRegion region) { super(region); } + @Override + protected TSStatus executeWrite(PlanNode planNode) { + final DataRegion targetRegion = region; + return targetRegion.executeIoTConsensusWrite( + () -> planNode.accept(new DataExecutionVisitor(), targetRegion)); + } + @Override public TSStatus write(IConsensusRequest request) { try { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeShutdownHook.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeShutdownHook.java index 76966820abb87..68b9fafdb9239 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeShutdownHook.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeShutdownHook.java @@ -95,6 +95,9 @@ public void run() { // Reject write operations to make sure all tsfiles will be sealed CommonDescriptor.getInstance().getConfig().setStopping(true); CommonDescriptor.getInstance().getConfig().setNodeStatus(NodeStatus.ReadOnly); + // A pre-apply read-only check can race with this transition. Drain admitted IoTConsensus + // applies, including their WAL separators, and permanently reject new ones before flushing. + StorageEngine.getInstance().stopIoTConsensusWrites(); // Wait all wal are flushed WALManager.getInstance().waitAllWALFlushed(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java index a9d6490d5f425..becd7404a56ae 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java @@ -398,6 +398,7 @@ private void asyncRecoverTsFileResource() { @Override public void stop() { + stopIoTConsensusWrites(); loadTsFileManager.stop(); for (DataRegion dataRegion : dataRegionMap.values()) { if (dataRegion != null) { @@ -417,6 +418,7 @@ public void stop() { @Override public void shutdown(long milliseconds) throws ShutdownException { + stopIoTConsensusWrites(); loadTsFileManager.stop(); try { for (DataRegion dataRegion : dataRegionMap.values()) { @@ -484,6 +486,17 @@ public synchronized void reset() { dataRegionMap.clear(); } + public void stopIoTConsensusWrites() { + if (!CONFIG.getDataRegionConsensusProtocolClass().equals(ConsensusFactory.IOT_CONSENSUS)) { + return; + } + for (DataRegion dataRegion : dataRegionMap.values()) { + if (dataRegion != null) { + dataRegion.stopIoTConsensusWrites(); + } + } + } + /** flush command Sync asyncCloseOneProcessor all file node processors. */ public void syncCloseAllProcessor() { LOGGER.info("Start closing all database processor"); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/DataRegion.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/DataRegion.java index ecb3e114bf280..f73bdc792cd54 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/DataRegion.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/DataRegion.java @@ -33,6 +33,7 @@ import org.apache.iotdb.commons.service.metric.PerformanceOverviewMetrics; import org.apache.iotdb.commons.utils.CommonDateTimeUtils; import org.apache.iotdb.commons.utils.RetryUtils; +import org.apache.iotdb.commons.utils.StatusUtils; import org.apache.iotdb.commons.utils.TestOnly; import org.apache.iotdb.commons.utils.TimePartitionUtils; import org.apache.iotdb.consensus.ConsensusFactory; @@ -175,6 +176,7 @@ import java.util.concurrent.locks.Condition; import java.util.concurrent.locks.ReadWriteLock; import java.util.concurrent.locks.ReentrantReadWriteLock; +import java.util.function.Supplier; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -227,6 +229,12 @@ public class DataRegion implements IDataRegionForQuery { */ private final ReadWriteLock insertLock = new ReentrantReadWriteLock(); + // Covers a complete IoTConsensus apply, including its WAL separator. Keep this separate from + // insertLock so memory backpressure and retry waits do not prevent flushing memtables. + private final ReadWriteLock ioTConsensusWriteLock = new ReentrantReadWriteLock(); + // Accessed only under ioTConsensusWriteLock; stopping a region permanently closes admission. + private boolean ioTConsensusWritesStopped; + /** Condition to safely delete data region. */ private final Condition deletedCondition = insertLock.writeLock().newCondition(); @@ -1861,6 +1869,30 @@ public void timedFlushUnseqMemTable() { WritingMetrics.getInstance().recordTimedFlushMemTableCount(count); } + public TSStatus executeIoTConsensusWrite(Supplier write) { + ioTConsensusWriteLock.readLock().lock(); + try { + if (ioTConsensusWritesStopped + || CommonDescriptor.getInstance().getConfig().isStopping() + || CommonDescriptor.getInstance().getConfig().isReadOnly()) { + return StatusUtils.getStatus(TSStatusCode.SYSTEM_READ_ONLY); + } + return write.get(); + } finally { + ioTConsensusWriteLock.readLock().unlock(); + } + } + + /** Reject new IoTConsensus applies and drain admitted applies before closing storage or WAL. */ + public void stopIoTConsensusWrites() { + ioTConsensusWriteLock.writeLock().lock(); + try { + ioTConsensusWritesStopped = true; + } finally { + ioTConsensusWriteLock.writeLock().unlock(); + } + } + /** This method will be blocked until all tsfile processors are closed. */ public void syncCloseAllWorkingTsFileProcessors() { try { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/consensus/statemachine/dataregion/IoTConsensusShutdownTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/consensus/statemachine/dataregion/IoTConsensusShutdownTest.java new file mode 100644 index 0000000000000..35949f765a35e --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/consensus/statemachine/dataregion/IoTConsensusShutdownTest.java @@ -0,0 +1,589 @@ +/* + * 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.consensus.statemachine.dataregion; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.cluster.NodeStatus; +import org.apache.iotdb.commons.concurrent.ExceptionalCountDownLatch; +import org.apache.iotdb.commons.conf.CommonConfig; +import org.apache.iotdb.commons.conf.CommonDescriptor; +import org.apache.iotdb.commons.consensus.DataRegionId; +import org.apache.iotdb.commons.path.PartialPath; +import org.apache.iotdb.consensus.ConsensusFactory; +import org.apache.iotdb.consensus.common.Peer; +import org.apache.iotdb.consensus.common.request.DeserializedBatchIndexedConsensusRequest; +import org.apache.iotdb.consensus.common.request.IConsensusRequest; +import org.apache.iotdb.consensus.config.IoTConsensusConfig; +import org.apache.iotdb.consensus.iot.IoTConsensusServerImpl; +import org.apache.iotdb.db.conf.IoTDBConfig; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; +import org.apache.iotdb.db.storageengine.StorageEngine; +import org.apache.iotdb.db.storageengine.dataregion.DataRegion; +import org.apache.iotdb.db.storageengine.dataregion.compaction.schedule.CompactionTaskManager; +import org.apache.iotdb.db.storageengine.dataregion.flush.TsFileFlushPolicy; +import org.apache.iotdb.db.storageengine.dataregion.memtable.TsFileProcessor; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.storageengine.dataregion.wal.WALManager; +import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALBuffer; +import org.apache.iotdb.db.storageengine.dataregion.wal.node.WALNode; +import org.apache.iotdb.db.storageengine.dataregion.wal.recover.CheckpointRecoverUtils; +import org.apache.iotdb.db.storageengine.dataregion.wal.recover.WALRecoverManager; +import org.apache.iotdb.db.storageengine.dataregion.wal.recover.file.UnsealedTsFileRecoverPerformer; +import org.apache.iotdb.db.storageengine.dataregion.wal.utils.WALMode; +import org.apache.iotdb.db.storageengine.dataregion.wal.utils.listener.WALRecoverListener; +import org.apache.iotdb.db.utils.EnvironmentUtils; +import org.apache.iotdb.db.utils.constant.TestConstant; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.read.TsFileReader; +import org.apache.tsfile.read.TsFileSequenceReader; +import org.apache.tsfile.read.common.Path; +import org.apache.tsfile.read.common.RowRecord; +import org.apache.tsfile.read.expression.QueryExpression; +import org.apache.tsfile.read.query.dataset.QueryDataSet; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.awaitility.Awaitility; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +import java.io.File; +import java.lang.reflect.Field; +import java.util.Collections; +import java.util.TreeSet; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotSame; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.spy; + +public class IoTConsensusShutdownTest { + + private static final String DATABASE = "root.shutdown"; + private static final String DEVICE = DATABASE + ".d1"; + private static final String MEASUREMENT = "s1"; + + @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder(); + + private final IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig(); + private final CommonConfig commonConfig = CommonDescriptor.getInstance().getConfig(); + private WALMode previousWalMode; + private int previousDataNodeId; + private String previousConsensusProtocol; + private NodeStatus previousNodeStatus; + private String previousStatusReason; + private boolean previousStopping; + private DataRegion dataRegion; + private IoTConsensusDataRegionStateMachine stateMachine; + + @Before + public void setUp() throws Exception { + previousWalMode = config.getWalMode(); + previousDataNodeId = config.getDataNodeId(); + previousConsensusProtocol = config.getDataRegionConsensusProtocolClass(); + previousNodeStatus = commonConfig.getNodeStatus(); + previousStatusReason = commonConfig.getStatusReason(); + previousStopping = commonConfig.isStopping(); + config.setWalMode(WALMode.ASYNC); + config.setDataNodeId(0); + config.setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS); + commonConfig.setNodeStatus(NodeStatus.Running); + commonConfig.setStopping(false); + EnvironmentUtils.envSetUp(); + dataRegion = + new DataRegion( + TestConstant.OUTPUT_DATA_DIR + "shutdown-info", + "0", + new TsFileFlushPolicy.DirectFlushPolicy(), + DATABASE); + StorageEngine.getInstance().setDataRegion(new DataRegionId(0), dataRegion); + stateMachine = new IoTConsensusDataRegionStateMachine(dataRegion); + } + + @After + public void tearDown() throws Exception { + try { + commonConfig.setStopping(false); + commonConfig.setNodeStatus(NodeStatus.Running); + if (dataRegion != null) { + dataRegion.syncCloseAllWorkingTsFileProcessors(); + } + CompactionTaskManager.getInstance().stop(); + EnvironmentUtils.cleanEnv(); + } finally { + config.setWalMode(previousWalMode); + config.setDataNodeId(previousDataNodeId); + config.setDataRegionConsensusProtocolClass(previousConsensusProtocol); + commonConfig.setNodeStatus(NodeStatus.Running); + commonConfig.setNodeStatus(previousNodeStatus); + commonConfig.setStatusReason(previousStatusReason); + commonConfig.setStopping(previousStopping); + } + } + + @Test + public void testQueuedSyncLogDoesNotReopenStorageAfterShutdown() throws Exception { + writeAndAssertSuccess(1); + final ScheduledExecutorService backgroundTaskService = + Executors.newSingleThreadScheduledExecutor(); + final IoTConsensusServerImpl server = newServer(stateMachine, backgroundTaskService, 2); + final DeserializedBatchIndexedConsensusRequest request = + new DeserializedBatchIndexedConsensusRequest(1, 2, 2); + request.add(insertRow(2)); + request.add(insertRow(3)); + final AtomicReference status = new AtomicReference<>(); + final Thread syncLogThread = new Thread(() -> status.set(server.syncLog(1, request))); + try { + assertFalse(server.isReadOnly()); + syncLogThread.start(); + Awaitility.await() + .atMost(10, TimeUnit.SECONDS) + .until(() -> syncLogThread.getState() == Thread.State.TIMED_WAITING); + + // Follow the production order: read-only, final flush, storage stop, then WAL stop. + commonConfig.setStopping(true); + commonConfig.setNodeStatus(NodeStatus.ReadOnly); + StorageEngine.getInstance().stopIoTConsensusWrites(); + WALManager.getInstance().waitAllWALFlushed(); + StorageEngine.getInstance().syncCloseAllProcessor(); + final File flushedFile = dataRegion.getSequenceFileList().get(0).getTsFile(); + StorageEngine.getInstance().stop(); + WALManager.getInstance().stop(); + assertEquals(0, WALManager.getInstance().getWALNodesNum()); + assertTrue(syncLogThread.isAlive()); + + // An empty preceding batch wakes the ordered request without another storage write. + server.syncLog(1, new DeserializedBatchIndexedConsensusRequest(0, 0, 0)); + syncLogThread.join(TimeUnit.SECONDS.toMillis(10)); + assertFalse(syncLogThread.isAlive()); + assertEquals(2, status.get().getSubStatusSize()); + assertEquals( + status.get() + + ", WAL nodes: " + + WALManager.getInstance().getWALNodesNum() + + ", working TsFile processors: " + + dataRegion.getWorkSequenceTsFileProcessors().size(), + TSStatusCode.SYSTEM_READ_ONLY.getStatusCode(), + status.get().getSubStatus().get(0).getCode()); + assertEquals( + TSStatusCode.SYSTEM_READ_ONLY.getStatusCode(), + status.get().getSubStatus().get(1).getCode()); + assertTrue(dataRegion.getWorkSequenceTsFileProcessors().isEmpty()); + assertTrue(dataRegion.getWorkUnsequenceTsFileProcessors().isEmpty()); + assertEquals(0, WALManager.getInstance().getWALNodesNum()); + assertSingleRow(flushedFile, 1); + } finally { + syncLogThread.interrupt(); + syncLogThread.join(TimeUnit.SECONDS.toMillis(10)); + backgroundTaskService.shutdownNow(); + } + } + + @Test + public void testSyncLogBecomesReadOnlyAfterPreApplyCheck() throws Exception { + writeAndAssertSuccess(1); + final CountDownLatch preApplyCheckPassed = new CountDownLatch(1); + final CountDownLatch resumeApply = new CountDownLatch(1); + final IoTConsensusDataRegionStateMachine delayedStateMachine = spy(stateMachine); + doAnswer( + invocation -> { + preApplyCheckPassed.countDown(); + assertTrue(resumeApply.await(30, TimeUnit.SECONDS)); + return invocation.callRealMethod(); + }) + .when(delayedStateMachine) + .write(any(IConsensusRequest.class)); + final ScheduledExecutorService backgroundTaskService = + Executors.newSingleThreadScheduledExecutor(); + final ExecutorService syncLogExecutor = Executors.newSingleThreadExecutor(); + final IoTConsensusServerImpl server = newServer(delayedStateMachine, backgroundTaskService, 1); + final DeserializedBatchIndexedConsensusRequest request = + new DeserializedBatchIndexedConsensusRequest(1, 2, 2); + request.add(insertRow(2)); + request.add(insertRow(3)); + final Future syncLog = syncLogExecutor.submit(() -> server.syncLog(1, request)); + try { + // The consensus-layer read-only check has returned false, but apply has not started. + assertTrue(preApplyCheckPassed.await(10, TimeUnit.SECONDS)); + commonConfig.setStopping(true); + commonConfig.setNodeStatus(NodeStatus.ReadOnly); + StorageEngine.getInstance().stopIoTConsensusWrites(); + WALManager.getInstance().waitAllWALFlushed(); + StorageEngine.getInstance().syncCloseAllProcessor(); + final File flushedFile = dataRegion.getSequenceFileList().get(0).getTsFile(); + StorageEngine.getInstance().stop(); + WALManager.getInstance().stop(); + assertEquals(0, WALManager.getInstance().getWALNodesNum()); + + resumeApply.countDown(); + final TSStatus status = syncLog.get(10, TimeUnit.SECONDS); + assertEquals(2, status.getSubStatusSize()); + assertEquals( + TSStatusCode.SYSTEM_READ_ONLY.getStatusCode(), status.getSubStatus().get(0).getCode()); + assertEquals( + TSStatusCode.SYSTEM_READ_ONLY.getStatusCode(), status.getSubStatus().get(1).getCode()); + assertEquals(0, WALManager.getInstance().getWALNodesNum()); + assertTrue(dataRegion.getWorkSequenceTsFileProcessors().isEmpty()); + assertTrue(dataRegion.getWorkUnsequenceTsFileProcessors().isEmpty()); + assertSingleRow(flushedFile, 1); + } finally { + resumeApply.countDown(); + syncLogExecutor.shutdownNow(); + assertTrue(syncLogExecutor.awaitTermination(10, TimeUnit.SECONDS)); + backgroundTaskService.shutdownNow(); + } + } + + @Test + public void testShutdownDrainsAdmittedApplyBeforeClosingStorage() throws Exception { + writeAndAssertSuccess(1); + final CountDownLatch admittedApply = new CountDownLatch(1); + final CountDownLatch resumeApply = new CountDownLatch(1); + final InsertRowNode pendingRow = spy(insertRow(2)); + doAnswer( + invocation -> { + admittedApply.countDown(); + assertTrue(resumeApply.await(30, TimeUnit.SECONDS)); + return invocation.callRealMethod(); + }) + .when(pendingRow) + .accept(any(PlanVisitor.class), any(DataRegion.class)); + final ExecutorService executor = Executors.newFixedThreadPool(2); + final CountDownLatch shutdownStarted = new CountDownLatch(1); + final CountDownLatch admissionClosed = new CountDownLatch(1); + final AtomicReference shutdownThread = new AtomicReference<>(); + final Future write = + executor.submit(() -> stateMachine.write((IConsensusRequest) pendingRow)); + try { + // Unlike the pre-apply-check test, this request has passed the check under the lifecycle + // lock. + assertTrue(admittedApply.await(10, TimeUnit.SECONDS)); + commonConfig.setStopping(true); + commonConfig.setNodeStatus(NodeStatus.ReadOnly); + final Future shutdown = + executor.submit( + () -> { + shutdownThread.set(Thread.currentThread()); + shutdownStarted.countDown(); + StorageEngine.getInstance().stopIoTConsensusWrites(); + admissionClosed.countDown(); + WALManager.getInstance().waitAllWALFlushed(); + StorageEngine.getInstance().syncCloseAllProcessor(); + StorageEngine.getInstance().stop(); + WALManager.getInstance().stop(); + }); + assertTrue(shutdownStarted.await(10, TimeUnit.SECONDS)); + Awaitility.await() + .atMost(10, TimeUnit.SECONDS) + .until( + () -> + admissionClosed.getCount() == 0 + || shutdownThread.get().getState() == Thread.State.WAITING); + assertEquals(1, admissionClosed.getCount()); + assertFalse(shutdown.isDone()); + assertFalse(dataRegion.getSequenceFileList().get(0).isClosed()); + + resumeApply.countDown(); + assertEquals( + TSStatusCode.SUCCESS_STATUS.getStatusCode(), write.get(10, TimeUnit.SECONDS).getCode()); + shutdown.get(10, TimeUnit.SECONDS); + assertRows(dataRegion.getSequenceFileList().get(0).getTsFile(), 1, 2); + assertTrue(dataRegion.getWorkSequenceTsFileProcessors().isEmpty()); + assertTrue(dataRegion.getWorkUnsequenceTsFileProcessors().isEmpty()); + assertEquals(0, WALManager.getInstance().getWALNodesNum()); + + // Stopping storage closes admission permanently, even if the node status changes again. + commonConfig.setStopping(false); + commonConfig.setNodeStatus(NodeStatus.Running); + assertEquals( + TSStatusCode.SYSTEM_READ_ONLY.getStatusCode(), + stateMachine.write((IConsensusRequest) insertRow(3)).getCode()); + assertEquals(0, WALManager.getInstance().getWALNodesNum()); + } finally { + resumeApply.countDown(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } + + @Test + public void testShutdownDuringWriteRetryDoesNotHoldAdmissionLock() throws Exception { + writeAndAssertSuccess(1); + final CountDownLatch retryWaiting = new CountDownLatch(1); + final CountDownLatch resumeRetry = new CountDownLatch(1); + stateMachine = + new IoTConsensusDataRegionStateMachine(dataRegion) { + @Override + protected void waitBeforeNextWriteRetry() throws InterruptedException { + retryWaiting.countDown(); + assertTrue(resumeRetry.await(30, TimeUnit.SECONDS)); + } + }; + final AtomicInteger applyAttempts = new AtomicInteger(); + final InsertRowNode pendingRow = spy(insertRow(2)); + doAnswer( + invocation -> + applyAttempts.incrementAndGet() == 1 + ? new TSStatus(TSStatusCode.WRITE_PROCESS_REJECT.getStatusCode()) + : invocation.callRealMethod()) + .when(pendingRow) + .accept(any(PlanVisitor.class), any(DataRegion.class)); + final ExecutorService executor = Executors.newSingleThreadExecutor(); + final Future write = + executor.submit(() -> stateMachine.write((IConsensusRequest) pendingRow)); + try { + assertTrue(retryWaiting.await(10, TimeUnit.SECONDS)); + commonConfig.setStopping(true); + commonConfig.setNodeStatus(NodeStatus.ReadOnly); + StorageEngine.getInstance().stopIoTConsensusWrites(); + WALManager.getInstance().waitAllWALFlushed(); + StorageEngine.getInstance().syncCloseAllProcessor(); + final File flushedFile = dataRegion.getSequenceFileList().get(0).getTsFile(); + StorageEngine.getInstance().stop(); + WALManager.getInstance().stop(); + assertFalse(write.isDone()); + + resumeRetry.countDown(); + assertEquals( + TSStatusCode.SYSTEM_READ_ONLY.getStatusCode(), write.get(10, TimeUnit.SECONDS).getCode()); + assertEquals(1, applyAttempts.get()); + assertTrue(dataRegion.getWorkSequenceTsFileProcessors().isEmpty()); + assertTrue(dataRegion.getWorkUnsequenceTsFileProcessors().isEmpty()); + assertEquals(0, WALManager.getInstance().getWALNodesNum()); + assertSingleRow(flushedFile, 1); + } finally { + resumeRetry.countDown(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } + + @Test + public void testReadOnlyCanBecomeWritableBeforeShutdown() throws Exception { + writeAndAssertSuccess(1); + commonConfig.setNodeStatus(NodeStatus.ReadOnly); + assertEquals( + TSStatusCode.SYSTEM_READ_ONLY.getStatusCode(), + stateMachine.write((IConsensusRequest) insertRow(2)).getCode()); + commonConfig.setNodeStatus(NodeStatus.Running); + writeAndAssertSuccess(2); + dataRegion.syncCloseAllWorkingTsFileProcessors(); + assertRows(dataRegion.getSequenceFileList().get(0).getTsFile(), 1, 2); + } + + private IoTConsensusServerImpl newServer( + IoTConsensusDataRegionStateMachine targetStateMachine, + ScheduledExecutorService backgroundTaskService, + int maxPendingBatchesNum) + throws Exception { + final Peer thisNode = new Peer(new DataRegionId(0), 0, new TEndPoint("127.0.0.1", 6667)); + return new IoTConsensusServerImpl( + temporaryFolder.getRoot().getAbsolutePath(), + thisNode, + new TreeSet<>(Collections.singletonList(thisNode)), + targetStateMachine, + backgroundTaskService, + null, + null, + IoTConsensusConfig.newBuilder() + .setReplication( + IoTConsensusConfig.Replication.newBuilder() + .setMaxPendingBatchesNum(maxPendingBatchesNum) + .setMaxWaitingTimeForWaitBatchInMs(TimeUnit.MINUTES.toMillis(5)) + .build()) + .build()); + } + + @Test + public void testUncheckedWriteAfterWalManagerStopLosesCheckpoint() throws Exception { + writeAndAssertSuccess(1); + commonConfig.setNodeStatus(NodeStatus.ReadOnly); + WALManager.getInstance().waitAllWALFlushed(); + StorageEngine.getInstance().syncCloseAllProcessor(); + final File flushedFile = dataRegion.getSequenceFileList().get(0).getTsFile(); + final WALNode originalWalNode = (WALNode) dataRegion.getWALNode().get(); + StorageEngine.getInstance().stop(); + WALManager.getInstance().stop(); + + // Bypass lifecycle admission with the real visitor to reproduce the old late-apply behavior. + // Unlike a write before WAL stop, this recreates the node in a directory containing closed WAL. + writeUncheckedAndAssertSuccess(2); + final WALNode lateWalNode = (WALNode) dataRegion.getWALNode().get(); + assertNotSame(originalWalNode, lateWalNode); + assertEquals(originalWalNode.getLogDirectory(), lateWalNode.getLogDirectory()); + final File lateFile = + dataRegion + .getWorkSequenceTsFileProcessors() + .iterator() + .next() + .getTsFileResource() + .getTsFile(); + WALManager.getInstance().waitAllWALFlushed(); + WALManager.getInstance().stop(); + // The new CheckpointManager appends another header to the existing _0.checkpoint file. Even + // with a drained and normally closed WAL, recovery cannot find the late memtable in that file. + assertTrue( + CheckpointRecoverUtils.recoverMemTableInfo(lateWalNode.getLogDirectory()) + .getMemTableId2Info() + .isEmpty()); + dataRegion.forceCloseAllWorkingTsFileProcessors(); + dataRegion.getWorkSequenceTsFileProcessors().clear(); + dataRegion.getWorkUnsequenceTsFileProcessors().clear(); + + final WALRecoverManager recoverManager = WALRecoverManager.getInstance(); + recoverManager.clear(); + final UnsealedTsFileRecoverPerformer performer = + new UnsealedTsFileRecoverPerformer( + new TsFileResource(lateFile), true, recovered -> assertFalse(recovered.canWrite())); + recoverManager.addRecoverPerformer(performer); + recoverManager.setAllDataRegionScannedLatch(new ExceptionalCountDownLatch(0)); + recoverManager.recover(); + assertEquals(WALRecoverListener.Status.SUCCESS, performer.getRecoverListener().waitForResult()); + assertSingleRow(flushedFile, 1); + assertTrue(performer.getTsFileResource().isEmpty()); + } + + @Test + public void testAsyncWriteAfterTsFileFlushIsRecoveredWhenWalCloses() throws Exception { + writeAndAssertSuccess(1); + commonConfig.setNodeStatus(NodeStatus.ReadOnly); + WALManager.getInstance().waitAllWALFlushed(); + StorageEngine.getInstance().syncCloseAllProcessor(); + assertTrue(dataRegion.getWorkSequenceTsFileProcessors().isEmpty()); + final File flushedFile = dataRegion.getSequenceFileList().get(0).getTsFile(); + final WALBuffer buffer = ((WALNode) dataRegion.getWALNode().get()).getWALBuffer(); + + // Hold the real WAL sync executor so the async ACK cannot imply that the WAL is durable. + final Field syncBufferThreadField = WALBuffer.class.getDeclaredField("syncBufferThread"); + syncBufferThreadField.setAccessible(true); + final ExecutorService syncExecutor = (ExecutorService) syncBufferThreadField.get(buffer); + final CountDownLatch syncBlocked = new CountDownLatch(1); + final CountDownLatch resumeSync = new CountDownLatch(1); + final Future blocker = + syncExecutor.submit( + () -> { + syncBlocked.countDown(); + assertTrue(resumeSync.await(30, TimeUnit.SECONDS)); + return null; + }); + final File lateFile; + try { + assertTrue(syncBlocked.await(10, TimeUnit.SECONDS)); + // Call the real visitor directly to model the old unchecked SyncLog apply. + writeUncheckedAndAssertSuccess(2); + assertFalse(buffer.isAllWALEntriesConsumed()); + assertEquals(1, dataRegion.getWorkSequenceTsFileProcessors().size()); + final TsFileProcessor processor = + dataRegion.getWorkSequenceTsFileProcessors().iterator().next(); + assertEquals(1, processor.getWorkMemTable().getTotalPointsNum()); + lateFile = processor.getTsFileResource().getTsFile(); + } finally { + resumeSync.countDown(); + blocker.get(10, TimeUnit.SECONDS); + } + + // Normal WAL close drains the late write, even though there is no further TsFile flush. + WALManager.getInstance().stop(); + dataRegion.forceCloseAllWorkingTsFileProcessors(); + dataRegion.getWorkSequenceTsFileProcessors().clear(); + dataRegion.getWorkUnsequenceTsFileProcessors().clear(); + assertFalse(new File(lateFile.getPath() + TsFileResource.RESOURCE_SUFFIX).exists()); + + // Discard the memtable and recover through the actual checkpoint/WAL recovery path. + final WALRecoverManager recoverManager = WALRecoverManager.getInstance(); + recoverManager.clear(); + final UnsealedTsFileRecoverPerformer performer = + new UnsealedTsFileRecoverPerformer( + new TsFileResource(lateFile), true, recovered -> assertFalse(recovered.canWrite())); + recoverManager.addRecoverPerformer(performer); + recoverManager.setAllDataRegionScannedLatch(new ExceptionalCountDownLatch(0)); + recoverManager.recover(); + assertEquals(WALRecoverListener.Status.SUCCESS, performer.getRecoverListener().waitForResult()); + assertSingleRow(flushedFile, 1); + assertSingleRow(lateFile, 2); + } + + private void writeAndAssertSuccess(long time) throws Exception { + assertEquals( + TSStatusCode.SUCCESS_STATUS.getStatusCode(), + stateMachine.write((IConsensusRequest) insertRow(time)).getCode()); + } + + private void writeUncheckedAndAssertSuccess(long time) throws Exception { + assertEquals( + TSStatusCode.SUCCESS_STATUS.getStatusCode(), + insertRow(time).accept(new DataExecutionVisitor(), dataRegion).getCode()); + } + + private static InsertRowNode insertRow(long time) throws Exception { + final InsertRowNode node = + new InsertRowNode( + new PlanNodeId("shutdown"), + new PartialPath(DEVICE), + false, + new String[] {MEASUREMENT}, + new TSDataType[] {TSDataType.INT32}, + time, + new Object[] {(int) time * 10}, + false); + node.setMeasurementSchemas( + new MeasurementSchema[] {new MeasurementSchema(MEASUREMENT, TSDataType.INT32)}); + return node; + } + + private static void assertSingleRow(File file, long time) throws Exception { + assertRows(file, time); + } + + private static void assertRows(File file, long... times) throws Exception { + try (TsFileSequenceReader sequenceReader = new TsFileSequenceReader(file.getPath()); + TsFileReader reader = new TsFileReader(sequenceReader)) { + final QueryDataSet dataSet = + reader.query( + QueryExpression.create( + Collections.singletonList(new Path(DEVICE, MEASUREMENT, false)), null)); + for (long time : times) { + assertTrue(dataSet.hasNext()); + final RowRecord row = dataSet.next(); + assertEquals(time, row.getTimestamp()); + assertEquals((int) time * 10, row.getFields().get(0).getIntV()); + } + assertFalse(dataSet.hasNext()); + } + } +} From b721fd154142445f64d3d35bcbeaaee0ca502619 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Thu, 15 Oct 2026 18:41:06 +0800 Subject: [PATCH 2/2] Reuse DataRegion insert lock for IoTConsensus shutdown fencing (#18893) (cherry picked from commit d2b97bc8455710a4540df03a1f15192d276629f2) --- .../dataregion/DataExecutionVisitor.java | 9 +- .../dataregion/DataRegionStateMachine.java | 6 +- .../IoTConsensusDataRegionStateMachine.java | 7 - .../storageengine/dataregion/DataRegion.java | 147 +++++--- .../dataregion/IoTConsensusShutdownTest.java | 321 ++++++++++++++++-- ...nerSeqCompactionWithFastPerformerTest.java | 3 +- ...qCompactionWithReadChunkPerformerTest.java | 7 +- 7 files changed, 407 insertions(+), 93 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java index cb634ad479bb6..fd78fec2d4397 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java @@ -60,7 +60,6 @@ public TSStatus visitPlan(PlanNode node, DataRegion context) { public TSStatus visitInsertRow(InsertRowNode node, DataRegion dataRegion) { try { dataRegion.insert(node); - dataRegion.insertSeparatorToWAL(); return StatusUtils.OK; } catch (OutOfTTLException e) { LOGGER.warn("Error in executing plan node: {}, caused by {}", node, e.getMessage()); @@ -78,7 +77,6 @@ public TSStatus visitInsertRow(InsertRowNode node, DataRegion dataRegion) { public TSStatus visitInsertTablet(InsertTabletNode node, DataRegion dataRegion) { try { dataRegion.insertTablet(node); - dataRegion.insertSeparatorToWAL(); return StatusUtils.OK; } catch (OutOfTTLException e) { LOGGER.warn("Error in executing plan node: {}, caused by {}", node, e.getMessage()); @@ -115,7 +113,6 @@ public TSStatus visitInsertTablet(InsertTabletNode node, DataRegion dataRegion) public TSStatus visitInsertRows(InsertRowsNode node, DataRegion dataRegion) { try { dataRegion.insert(node); - dataRegion.insertSeparatorToWAL(); return StatusUtils.OK; } catch (WriteProcessRejectException e) { LOGGER.warn("Reject in executing plan node: {}, caused by {}", node, e.getMessage()); @@ -149,8 +146,9 @@ public TSStatus visitInsertRows(InsertRowsNode node, DataRegion dataRegion) { public TSStatus visitInsertMultiTablets(InsertMultiTabletsNode node, DataRegion dataRegion) { try { dataRegion.insertTablets(node); - dataRegion.insertSeparatorToWAL(); return StatusUtils.OK; + } catch (WriteProcessRejectException e) { + return RpcUtils.getStatus(e.getErrorCode(), e.getMessage()); } catch (BatchProcessException e) { LOGGER.warn("Batch failure in executing a InsertMultiTabletsNode."); TSStatus firstStatus = null; @@ -181,7 +179,6 @@ public TSStatus visitInsertRowsOfOneDevice( InsertRowsOfOneDeviceNode node, DataRegion dataRegion) { try { dataRegion.insert(node); - dataRegion.insertSeparatorToWAL(); return StatusUtils.OK; } catch (WriteProcessRejectException e) { LOGGER.warn("Reject in executing plan node: {}, caused by {}", node, e.getMessage()); @@ -244,6 +241,8 @@ public TSStatus visitDeleteData(DeleteDataNode node, DataRegion dataRegion) { dataRegion.insertSeparatorToWAL(); PipeInsertionDataNodeListener.getInstance().listenToDeleteData(node); return StatusUtils.OK; + } catch (WriteProcessRejectException e) { + return RpcUtils.getStatus(e.getErrorCode(), e.getMessage()); } catch (IOException | IllegalPathException e) { LOGGER.error("Error in executing plan node: {}", node, e); return new TSStatus(TSStatusCode.WRITE_PROCESS_ERROR.getStatusCode()); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java index 66c0692fcf59b..a76889d8b63d7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java @@ -214,7 +214,7 @@ protected TSStatus write(PlanNode planNode) { TSStatus result = null; int retryTime = 0; while (retryTime < MAX_WRITE_RETRY_TIMES) { - result = executeWrite(planNode); + result = planNode.accept(new DataExecutionVisitor(), region); if (needRetry(result.getCode())) { retryTime++; logger.debug( @@ -239,10 +239,6 @@ protected TSStatus write(PlanNode planNode) { return result; } - protected TSStatus executeWrite(PlanNode planNode) { - return planNode.accept(new DataExecutionVisitor(), region); - } - protected void waitBeforeNextWriteRetry() throws InterruptedException { Thread.sleep(WRITE_RETRY_WAIT_TIME_IN_MS); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/IoTConsensusDataRegionStateMachine.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/IoTConsensusDataRegionStateMachine.java index 56ed7d775e792..240c1b1caa0fe 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/IoTConsensusDataRegionStateMachine.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/IoTConsensusDataRegionStateMachine.java @@ -49,13 +49,6 @@ public IoTConsensusDataRegionStateMachine(DataRegion region) { super(region); } - @Override - protected TSStatus executeWrite(PlanNode planNode) { - final DataRegion targetRegion = region; - return targetRegion.executeIoTConsensusWrite( - () -> planNode.accept(new DataExecutionVisitor(), targetRegion)); - } - @Override public TSStatus write(IConsensusRequest request) { try { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/DataRegion.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/DataRegion.java index f73bdc792cd54..0f0509067875a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/DataRegion.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/DataRegion.java @@ -176,7 +176,6 @@ import java.util.concurrent.locks.Condition; import java.util.concurrent.locks.ReadWriteLock; import java.util.concurrent.locks.ReentrantReadWriteLock; -import java.util.function.Supplier; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -229,10 +228,7 @@ public class DataRegion implements IDataRegionForQuery { */ private final ReadWriteLock insertLock = new ReentrantReadWriteLock(); - // Covers a complete IoTConsensus apply, including its WAL separator. Keep this separate from - // insertLock so memory backpressure and retry waits do not prevent flushing memtables. - private final ReadWriteLock ioTConsensusWriteLock = new ReentrantReadWriteLock(); - // Accessed only under ioTConsensusWriteLock; stopping a region permanently closes admission. + // Accessed only under insertLock's write lock; stopping permanently closes write admission. private boolean ioTConsensusWritesStopped; /** Condition to safely delete data region. */ @@ -1070,6 +1066,7 @@ public void insert(InsertRowNode insertRowNode) throws WriteProcessException { writeLock("InsertRow"); PERFORMANCE_OVERVIEW_METRICS.recordScheduleLockCost(System.nanoTime() - startTime); try { + checkIoTConsensusWriteAllowed(); if (deleted) { return; } @@ -1098,6 +1095,7 @@ public void insert(InsertRowNode insertRowNode) throws WriteProcessException { System.nanoTime() - startTime); } } + logSeparatorToWAL(); } finally { writeUnlock(); } @@ -1112,10 +1110,19 @@ public void insert(InsertRowNode insertRowNode) throws WriteProcessException { public void insertTablet(InsertTabletNode insertTabletNode) throws BatchProcessException, WriteProcessException { StorageEngine.blockInsertionIfReject(); + insertTablet(insertTabletNode, true); + } + + private void insertTablet(InsertTabletNode insertTabletNode, boolean logSeparator) + throws BatchProcessException, WriteProcessException { long startTime = System.nanoTime(); writeLock("insertTablet"); PERFORMANCE_OVERVIEW_METRICS.recordScheduleLockCost(System.nanoTime() - startTime); try { + // A multi-tablet batch has already passed admission under insertLock. + if (logSeparator) { + checkIoTConsensusWriteAllowed(); + } if (deleted) { logger.info( "Won't insert tablet {}, because region is deleted", insertTabletNode.getSearchIndex()); @@ -1207,6 +1214,9 @@ public void insertTablet(InsertTabletNode insertTabletNode) if (!noFailure) { throw new BatchProcessException(results); } + if (logSeparator) { + logSeparatorToWAL(); + } } finally { writeUnlock(); } @@ -1869,27 +1879,25 @@ public void timedFlushUnseqMemTable() { WritingMetrics.getInstance().recordTimedFlushMemTableCount(count); } - public TSStatus executeIoTConsensusWrite(Supplier write) { - ioTConsensusWriteLock.readLock().lock(); - try { - if (ioTConsensusWritesStopped - || CommonDescriptor.getInstance().getConfig().isStopping() - || CommonDescriptor.getInstance().getConfig().isReadOnly()) { - return StatusUtils.getStatus(TSStatusCode.SYSTEM_READ_ONLY); - } - return write.get(); - } finally { - ioTConsensusWriteLock.readLock().unlock(); + /** Called under insertLock's write lock, after any memory-backpressure wait. */ + protected void checkIoTConsensusWriteAllowed() throws WriteProcessRejectException { + if (config.getDataRegionConsensusProtocolClass().equals(ConsensusFactory.IOT_CONSENSUS) + && (ioTConsensusWritesStopped + || CommonDescriptor.getInstance().getConfig().isStopping() + || CommonDescriptor.getInstance().getConfig().isReadOnly())) { + throw new WriteProcessRejectException( + StatusUtils.getStatus(TSStatusCode.SYSTEM_READ_ONLY).getMessage(), + TSStatusCode.SYSTEM_READ_ONLY.getStatusCode()); } } /** Reject new IoTConsensus applies and drain admitted applies before closing storage or WAL. */ public void stopIoTConsensusWrites() { - ioTConsensusWriteLock.writeLock().lock(); + writeLock("stopIoTConsensusWrites"); try { ioTConsensusWritesStopped = true; } finally { - ioTConsensusWriteLock.writeLock().unlock(); + writeUnlock(); } } @@ -2502,7 +2510,7 @@ private void getTwoKindsOfTsFiles( * @throws IOException */ public void deleteByDevice(PartialPath pattern, long startTime, long endTime, long searchIndex) - throws IOException { + throws IOException, WriteProcessRejectException { if (SettleService.getINSTANCE().getFilesToBeSettledCount().get() != 0) { throw new IOException( "Delete failed. " + "Please do not delete until the old files settled."); @@ -2515,6 +2523,7 @@ public void deleteByDevice(PartialPath pattern, long startTime, long endTime, lo boolean hasReleasedLock = false; try { + checkIoTConsensusWriteAllowed(); if (deleted) { return; } @@ -2544,6 +2553,8 @@ public void deleteByDevice(PartialPath pattern, long startTime, long endTime, lo hasReleasedLock = true; deleteDataInFiles(sealedTsFileResource, deletion, devicePaths, deviceMatchInfo); + } catch (WriteProcessRejectException e) { + throw e; } catch (Exception e) { throw new IOException(e); } finally { @@ -2554,7 +2565,8 @@ public void deleteByDevice(PartialPath pattern, long startTime, long endTime, lo } public void deleteDataDirectly( - PartialPath pathToDelete, long startTime, long endTime, long searchIndex) throws IOException { + PartialPath pathToDelete, long startTime, long endTime, long searchIndex) + throws IOException, WriteProcessRejectException { logger.info( "{} will delete data files directly for deleting data between {} and {}", databaseName + "-" + dataRegionIdString, @@ -2565,6 +2577,7 @@ public void deleteDataDirectly( boolean releasedLock = false; try { + checkIoTConsensusWriteAllowed(); if (deleted) { return; } @@ -2586,6 +2599,8 @@ public void deleteDataDirectly( writeUnlock(); releasedLock = true; deleteDataDirectlyInFile(sealedTsFileResource, pathToDelete, startTime, endTime); + } catch (WriteProcessRejectException e) { + throw e; } catch (Exception e) { throw new IOException(e); } finally { @@ -2638,20 +2653,29 @@ private List logDeletionInWAL( public void insertSeparatorToWAL() { writeLock("insertSeparatorToWAL"); try { - if (deleted) { + if (deleted || ioTConsensusWritesStopped) { return; } - getWALNode() - .ifPresent( - walNode -> - walNode.log( - TsFileProcessor.MEMTABLE_NOT_EXIST, - new ContinuousSameSearchIndexSeparatorNode())); + logSeparatorToWAL(); } finally { writeUnlock(); } } + /** Called under insertLock's write lock as part of an admitted write. */ + private void logSeparatorToWAL() { + if (deleted + || !config.getDataRegionConsensusProtocolClass().equals(ConsensusFactory.IOT_CONSENSUS)) { + return; + } + getWALNode() + .ifPresent( + walNode -> + walNode.log( + TsFileProcessor.MEMTABLE_NOT_EXIST, + new ContinuousSameSearchIndexSeparatorNode())); + } + @SuppressWarnings("OptionalGetWithoutIsPresent") private boolean canSkipDelete( TsFileResource tsFileResource, @@ -3784,6 +3808,7 @@ public void insert(InsertRowsOfOneDeviceNode insertRowsOfOneDeviceNode) writeLock("InsertRowsOfOneDevice"); PERFORMANCE_OVERVIEW_METRICS.recordScheduleLockCost(System.nanoTime() - startTime); try { + checkIoTConsensusWriteAllowed(); if (deleted) { return; } @@ -3872,12 +3897,13 @@ public void insert(InsertRowsOfOneDeviceNode insertRowsOfOneDeviceNode) System.nanoTime() - startTime); } } + if (!insertRowsOfOneDeviceNode.getResults().isEmpty()) { + throw new BatchProcessException("Partial failed inserting rows of one device"); + } + logSeparatorToWAL(); } finally { writeUnlock(); } - if (!insertRowsOfOneDeviceNode.getResults().isEmpty()) { - throw new BatchProcessException("Partial failed inserting rows of one device"); - } } public void insert(InsertRowsNode insertRowsNode) @@ -3887,6 +3913,7 @@ public void insert(InsertRowsNode insertRowsNode) writeLock("InsertRows"); PERFORMANCE_OVERVIEW_METRICS.recordScheduleLockCost(System.nanoTime() - startTime); try { + checkIoTConsensusWriteAllowed(); if (deleted) { return; } @@ -3937,6 +3964,7 @@ public void insert(InsertRowsNode insertRowsNode) if (!insertRowsNode.getResults().isEmpty()) { throw new BatchProcessException("Partial failed inserting rows"); } + logSeparatorToWAL(); } finally { writeUnlock(); } @@ -3959,34 +3987,45 @@ private void markAllMeasurementsFailed(final InsertRowNode insertRowNode) { * @param insertMultiTabletsNode batch of tablets belongs to multiple devices */ public void insertTablets(InsertMultiTabletsNode insertMultiTabletsNode) - throws BatchProcessException { - for (int i = 0; i < insertMultiTabletsNode.getInsertTabletNodeList().size(); i++) { - InsertTabletNode insertTabletNode = insertMultiTabletsNode.getInsertTabletNodeList().get(i); - try { - insertTablet(insertTabletNode); - } catch (WriteProcessException e) { - insertMultiTabletsNode - .getResults() - .put(i, RpcUtils.getStatus(e.getErrorCode(), e.getMessage())); - } catch (BatchProcessException e) { - // for each error - TSStatus firstStatus = null; - for (TSStatus status : e.getFailingStatus()) { - if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { - firstStatus = status; - } - // return WRITE_PROCESS_REJECT directly for the consensus retry logic - if (status.getCode() == TSStatusCode.WRITE_PROCESS_REJECT.getStatusCode()) { - insertMultiTabletsNode.getResults().put(i, status); - throw new BatchProcessException("Rejected inserting multi tablets"); + throws BatchProcessException, WriteProcessRejectException { + StorageEngine.blockInsertionIfReject(); + writeLock("insertTablets"); + try { + checkIoTConsensusWriteAllowed(); + if (deleted) { + return; + } + for (int i = 0; i < insertMultiTabletsNode.getInsertTabletNodeList().size(); i++) { + InsertTabletNode insertTabletNode = insertMultiTabletsNode.getInsertTabletNodeList().get(i); + try { + insertTablet(insertTabletNode, false); + } catch (WriteProcessException e) { + insertMultiTabletsNode + .getResults() + .put(i, RpcUtils.getStatus(e.getErrorCode(), e.getMessage())); + } catch (BatchProcessException e) { + // for each error + TSStatus firstStatus = null; + for (TSStatus status : e.getFailingStatus()) { + if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + firstStatus = status; + } + // return WRITE_PROCESS_REJECT directly for the consensus retry logic + if (status.getCode() == TSStatusCode.WRITE_PROCESS_REJECT.getStatusCode()) { + insertMultiTabletsNode.getResults().put(i, status); + throw new BatchProcessException("Rejected inserting multi tablets"); + } } + insertMultiTabletsNode.getResults().put(i, firstStatus); } - insertMultiTabletsNode.getResults().put(i, firstStatus); } - } - if (!insertMultiTabletsNode.getResults().isEmpty()) { - throw new BatchProcessException("Partial failed inserting multi tablets"); + if (!insertMultiTabletsNode.getResults().isEmpty()) { + throw new BatchProcessException("Partial failed inserting multi tablets"); + } + logSeparatorToWAL(); + } finally { + writeUnlock(); } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/consensus/statemachine/dataregion/IoTConsensusShutdownTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/consensus/statemachine/dataregion/IoTConsensusShutdownTest.java index 35949f765a35e..1ac445c679e11 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/consensus/statemachine/dataregion/IoTConsensusShutdownTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/consensus/statemachine/dataregion/IoTConsensusShutdownTest.java @@ -35,9 +35,18 @@ import org.apache.iotdb.consensus.iot.IoTConsensusServerImpl; import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.exception.DataRegionException; +import org.apache.iotdb.db.exception.WriteProcessRejectException; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeId; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.ContinuousSameSearchIndexSeparatorNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertMultiTabletsNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowsNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowsOfOneDeviceNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertTabletNode; import org.apache.iotdb.db.storageengine.StorageEngine; import org.apache.iotdb.db.storageengine.dataregion.DataRegion; import org.apache.iotdb.db.storageengine.dataregion.compaction.schedule.CompactionTaskManager; @@ -46,12 +55,14 @@ import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; import org.apache.iotdb.db.storageengine.dataregion.wal.WALManager; import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALBuffer; +import org.apache.iotdb.db.storageengine.dataregion.wal.node.IWALNode; import org.apache.iotdb.db.storageengine.dataregion.wal.node.WALNode; import org.apache.iotdb.db.storageengine.dataregion.wal.recover.CheckpointRecoverUtils; import org.apache.iotdb.db.storageengine.dataregion.wal.recover.WALRecoverManager; import org.apache.iotdb.db.storageengine.dataregion.wal.recover.file.UnsealedTsFileRecoverPerformer; import org.apache.iotdb.db.storageengine.dataregion.wal.utils.WALMode; import org.apache.iotdb.db.storageengine.dataregion.wal.utils.listener.WALRecoverListener; +import org.apache.iotdb.db.storageengine.rescon.memory.SystemInfo; import org.apache.iotdb.db.utils.EnvironmentUtils; import org.apache.iotdb.db.utils.constant.TestConstant; import org.apache.iotdb.rpc.TSStatusCode; @@ -73,7 +84,10 @@ import java.io.File; import java.lang.reflect.Field; +import java.util.Arrays; import java.util.Collections; +import java.util.List; +import java.util.Optional; import java.util.TreeSet; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; @@ -89,8 +103,11 @@ import static org.junit.Assert.assertNotSame; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; public class IoTConsensusShutdownTest { @@ -108,7 +125,7 @@ public class IoTConsensusShutdownTest { private NodeStatus previousNodeStatus; private String previousStatusReason; private boolean previousStopping; - private DataRegion dataRegion; + private TestDataRegion dataRegion; private IoTConsensusDataRegionStateMachine stateMachine; @Before @@ -125,12 +142,7 @@ public void setUp() throws Exception { commonConfig.setNodeStatus(NodeStatus.Running); commonConfig.setStopping(false); EnvironmentUtils.envSetUp(); - dataRegion = - new DataRegion( - TestConstant.OUTPUT_DATA_DIR + "shutdown-info", - "0", - new TsFileFlushPolicy.DirectFlushPolicy(), - DATABASE); + dataRegion = new TestDataRegion(); StorageEngine.getInstance().setDataRegion(new DataRegionId(0), dataRegion); stateMachine = new IoTConsensusDataRegionStateMachine(dataRegion); } @@ -271,18 +283,103 @@ public void testSyncLogBecomesReadOnlyAfterPreApplyCheck() throws Exception { @Test public void testShutdownDrainsAdmittedApplyBeforeClosingStorage() throws Exception { + checkShutdownDrainsAdmittedApply(false); + } + + @Test + public void testShutdownDrainsWalSeparatorBeforeClosingStorage() throws Exception { + checkShutdownDrainsAdmittedApply(true); + } + + @Test + public void testShutdownDrainsAdmittedMultiTabletBatch() throws Exception { writeAndAssertSuccess(1); - final CountDownLatch admittedApply = new CountDownLatch(1); - final CountDownLatch resumeApply = new CountDownLatch(1); - final InsertRowNode pendingRow = spy(insertRow(2)); + final CountDownLatch admittedBatch = new CountDownLatch(1); + final CountDownLatch resumeBatch = new CountDownLatch(1); + final InsertTabletNode firstTablet = spy(insertTablet(2)); doAnswer( invocation -> { - admittedApply.countDown(); - assertTrue(resumeApply.await(30, TimeUnit.SECONDS)); + admittedBatch.countDown(); + assertTrue(resumeBatch.await(30, TimeUnit.SECONDS)); return invocation.callRealMethod(); }) - .when(pendingRow) - .accept(any(PlanVisitor.class), any(DataRegion.class)); + .when(firstTablet) + .getDeviceID(); + final WALNode walNode = spy((WALNode) dataRegion.getWALNode().get()); + dataRegion.walNodeForTest = walNode; + final InsertMultiTabletsNode batch = + new InsertMultiTabletsNode( + new PlanNodeId("shutdown"), + Arrays.asList(0, 1), + Arrays.asList(firstTablet, insertTablet(3))); + final ExecutorService executor = Executors.newFixedThreadPool(2); + final AtomicReference shutdownThread = new AtomicReference<>(); + final Future write = + executor.submit(() -> stateMachine.write((IConsensusRequest) batch)); + try { + assertTrue(admittedBatch.await(10, TimeUnit.SECONDS)); + commonConfig.setStopping(true); + commonConfig.setNodeStatus(NodeStatus.ReadOnly); + final Future shutdown = + executor.submit( + () -> { + shutdownThread.set(Thread.currentThread()); + StorageEngine.getInstance().stopIoTConsensusWrites(); + WALManager.getInstance().waitAllWALFlushed(); + StorageEngine.getInstance().syncCloseAllProcessor(); + StorageEngine.getInstance().stop(); + WALManager.getInstance().stop(); + }); + Awaitility.await() + .atMost(10, TimeUnit.SECONDS) + .until( + () -> + shutdownThread.get() != null + && shutdownThread.get().getState() == Thread.State.WAITING); + assertFalse(shutdown.isDone()); + + // 1.3 delegates a batch to individual tablets. Once admitted, the entire batch must finish + // with one separator even if the node becomes read-only between its tablets. + resumeBatch.countDown(); + assertEquals( + TSStatusCode.SUCCESS_STATUS.getStatusCode(), write.get(10, TimeUnit.SECONDS).getCode()); + shutdown.get(10, TimeUnit.SECONDS); + verify(walNode, times(1)).log(anyLong(), any(ContinuousSameSearchIndexSeparatorNode.class)); + assertRows(dataRegion.getSequenceFileList().get(0).getTsFile(), 1, 2, 3); + assertEquals(0, WALManager.getInstance().getWALNodesNum()); + } finally { + resumeBatch.countDown(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } + + private void checkShutdownDrainsAdmittedApply(final boolean delaySeparator) throws Exception { + writeAndAssertSuccess(1); + final CountDownLatch admittedApply = new CountDownLatch(1); + final CountDownLatch resumeApply = new CountDownLatch(1); + final InsertRowNode pendingRow = spy(insertRow(2)); + if (delaySeparator) { + final WALNode walNode = spy((WALNode) dataRegion.getWALNode().get()); + dataRegion.walNodeForTest = walNode; + doAnswer( + invocation -> { + admittedApply.countDown(); + assertTrue(resumeApply.await(30, TimeUnit.SECONDS)); + return invocation.callRealMethod(); + }) + .when(walNode) + .log(anyLong(), any(ContinuousSameSearchIndexSeparatorNode.class)); + } else { + doAnswer( + invocation -> { + admittedApply.countDown(); + assertTrue(resumeApply.await(30, TimeUnit.SECONDS)); + return invocation.callRealMethod(); + }) + .when(pendingRow) + .getDeviceID(); + } final ExecutorService executor = Executors.newFixedThreadPool(2); final CountDownLatch shutdownStarted = new CountDownLatch(1); final CountDownLatch admissionClosed = new CountDownLatch(1); @@ -290,8 +387,7 @@ public void testShutdownDrainsAdmittedApplyBeforeClosingStorage() throws Excepti final Future write = executor.submit(() -> stateMachine.write((IConsensusRequest) pendingRow)); try { - // Unlike the pre-apply-check test, this request has passed the check under the lifecycle - // lock. + // This request has passed the storage admission check and holds insertLock. assertTrue(admittedApply.await(10, TimeUnit.SECONDS)); commonConfig.setStopping(true); commonConfig.setNodeStatus(NodeStatus.ReadOnly); @@ -381,7 +477,9 @@ protected void waitBeforeNextWriteRetry() throws InterruptedException { resumeRetry.countDown(); assertEquals( TSStatusCode.SYSTEM_READ_ONLY.getStatusCode(), write.get(10, TimeUnit.SECONDS).getCode()); - assertEquals(1, applyAttempts.get()); + // The retry enters the visitor, but its storage admission check rejects it before any WAL + // or processor can be created. + assertEquals(2, applyAttempts.get()); assertTrue(dataRegion.getWorkSequenceTsFileProcessors().isEmpty()); assertTrue(dataRegion.getWorkUnsequenceTsFileProcessors().isEmpty()); assertEquals(0, WALManager.getInstance().getWALNodesNum()); @@ -406,6 +504,144 @@ public void testReadOnlyCanBecomeWritableBeforeShutdown() throws Exception { assertRows(dataRegion.getSequenceFileList().get(0).getTsFile(), 1, 2); } + @Test + public void testWriteWaitingForInsertLockIsRejectedAfterShutdown() throws Exception { + writeAndAssertSuccess(1); + final ExecutorService executor = Executors.newSingleThreadExecutor(); + final AtomicReference writerThread = new AtomicReference<>(); + final Future write; + dataRegion.writeLock("testShutdownWhileWaitingForInsertLock"); + try { + write = + executor.submit( + () -> { + writerThread.set(Thread.currentThread()); + return stateMachine.write((IConsensusRequest) insertRow(2)); + }); + Awaitility.await() + .atMost(10, TimeUnit.SECONDS) + .until( + () -> + writerThread.get() != null + && writerThread.get().getState() == Thread.State.WAITING); + // Close admission while the writer has already entered the visitor and is queued for + // insertLock. The stopping and read-only flags deliberately remain unset. + dataRegion.stopIoTConsensusWrites(); + } finally { + dataRegion.writeUnlock(); + executor.shutdown(); + } + try { + assertEquals( + TSStatusCode.SYSTEM_READ_ONLY.getStatusCode(), write.get(10, TimeUnit.SECONDS).getCode()); + dataRegion.syncCloseAllWorkingTsFileProcessors(); + assertRows(dataRegion.getSequenceFileList().get(0).getTsFile(), 1); + } finally { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } + + @Test + public void testMemoryBackpressureDoesNotBlockShutdownOrFlush() throws Exception { + writeAndAssertSuccess(1); + final Field rejectedField = SystemInfo.class.getDeclaredField("rejected"); + rejectedField.setAccessible(true); + final boolean previousRejected = rejectedField.getBoolean(SystemInfo.getInstance()); + final ExecutorService executor = Executors.newFixedThreadPool(2); + final AtomicReference writerThread = new AtomicReference<>(); + rejectedField.setBoolean(SystemInfo.getInstance(), true); + final Future write = + executor.submit( + () -> { + writerThread.set(Thread.currentThread()); + return stateMachine.write((IConsensusRequest) insertRow(2)); + }); + try { + Awaitility.await() + .atMost(10, TimeUnit.SECONDS) + .until( + () -> + writerThread.get() != null + && writerThread.get().getState() == Thread.State.TIMED_WAITING); + executor.submit(() -> dataRegion.stopIoTConsensusWrites()).get(10, TimeUnit.SECONDS); + assertFalse(write.isDone()); + // The same insertLock remains available to flush the memory that the writer is waiting for. + executor + .submit(() -> dataRegion.syncCloseAllWorkingTsFileProcessors()) + .get(10, TimeUnit.SECONDS); + assertEquals( + TSStatusCode.SYSTEM_READ_ONLY.getStatusCode(), write.get(10, TimeUnit.SECONDS).getCode()); + assertTrue(dataRegion.getWorkSequenceTsFileProcessors().isEmpty()); + assertRows(dataRegion.getSequenceFileList().get(0).getTsFile(), 1); + } finally { + rejectedField.setBoolean(SystemInfo.getInstance(), previousRejected); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } + + @Test + public void testAllWritePathsRejectAfterStorageStop() throws Exception { + writeAndAssertSuccess(1); + final PlanNodeId id = new PlanNodeId("shutdown"); + final List writes = + Arrays.asList( + insertRow(2), + insertTablet(2), + new InsertRowsNode( + id, Collections.singletonList(0), Collections.singletonList(insertRow(2))), + new InsertRowsOfOneDeviceNode( + id, Collections.singletonList(0), Collections.singletonList(insertRow(2))), + new InsertMultiTabletsNode( + id, Collections.singletonList(0), Collections.singletonList(insertTablet(2))), + new DeleteDataNode( + id, Collections.singletonList(new PartialPath(DEVICE + ".s1")), 0, 10), + new DeleteDataNode( + id, Collections.singletonList(new PartialPath(DATABASE + ".**")), 0, 10)); + StorageEngine.getInstance().stop(); + WALManager.getInstance().stop(); + assertFalse(commonConfig.isReadOnly()); + assertFalse(commonConfig.isStopping()); + for (final PlanNode node : writes) { + assertEquals( + node.getClass().getSimpleName(), + TSStatusCode.SYSTEM_READ_ONLY.getStatusCode(), + stateMachine.write((IConsensusRequest) node).getCode()); + } + dataRegion.insertSeparatorToWAL(); + assertEquals(0, WALManager.getInstance().getWALNodesNum()); + assertTrue(dataRegion.getWorkSequenceTsFileProcessors().isEmpty()); + assertTrue(dataRegion.getWorkUnsequenceTsFileProcessors().isEmpty()); + assertRows(dataRegion.getSequenceFileList().get(0).getTsFile(), 1); + } + + @Test + public void testOtherConsensusAppliesRemainWritable() throws Exception { + final DataRegionStateMachine otherStateMachine = new DataRegionStateMachine(dataRegion); + commonConfig.setNodeStatus(NodeStatus.ReadOnly); + commonConfig.setStopping(true); + try { + long time = 0; + for (final String protocol : + Arrays.asList( + ConsensusFactory.RATIS_CONSENSUS, + ConsensusFactory.SIMPLE_CONSENSUS, + ConsensusFactory.IOT_CONSENSUS_V2)) { + config.setDataRegionConsensusProtocolClass(protocol); + assertEquals( + TSStatusCode.SUCCESS_STATUS.getStatusCode(), + otherStateMachine.write((IConsensusRequest) insertRow(++time)).getCode()); + } + } finally { + config.setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS); + commonConfig.setNodeStatus(NodeStatus.Running); + commonConfig.setStopping(false); + } + dataRegion.syncCloseAllWorkingTsFileProcessors(); + assertRows(dataRegion.getSequenceFileList().get(0).getTsFile(), 1, 2, 3); + } + private IoTConsensusServerImpl newServer( IoTConsensusDataRegionStateMachine targetStateMachine, ScheduledExecutorService backgroundTaskService, @@ -545,9 +781,54 @@ private void writeAndAssertSuccess(long time) throws Exception { } private void writeUncheckedAndAssertSuccess(long time) throws Exception { - assertEquals( - TSStatusCode.SUCCESS_STATUS.getStatusCode(), - insertRow(time).accept(new DataExecutionVisitor(), dataRegion).getCode()); + dataRegion.bypassWriteChecks = true; + try { + assertEquals( + TSStatusCode.SUCCESS_STATUS.getStatusCode(), + insertRow(time).accept(new DataExecutionVisitor(), dataRegion).getCode()); + } finally { + dataRegion.bypassWriteChecks = false; + } + } + + private static class TestDataRegion extends DataRegion { + // Only the two fault-injection tests bypass admission to reproduce the old late-write bug. + private boolean bypassWriteChecks; + private IWALNode walNodeForTest; + + @Override + public Optional getWALNode() { + return walNodeForTest == null ? super.getWALNode() : Optional.of(walNodeForTest); + } + + private TestDataRegion() throws DataRegionException { + super( + TestConstant.OUTPUT_DATA_DIR + "shutdown-info", + "0", + new TsFileFlushPolicy.DirectFlushPolicy(), + DATABASE); + } + + @Override + protected void checkIoTConsensusWriteAllowed() throws WriteProcessRejectException { + if (!bypassWriteChecks) { + super.checkIoTConsensusWriteAllowed(); + } + } + } + + private static InsertTabletNode insertTablet(final long time) throws Exception { + return new InsertTabletNode( + new PlanNodeId("shutdown"), + new PartialPath(DEVICE), + false, + new String[] {MEASUREMENT}, + new TSDataType[] {TSDataType.INT32}, + new MeasurementSchema[] {new MeasurementSchema(MEASUREMENT, TSDataType.INT32)}, + new long[] {time}, + null, + new Object[] {new int[] {(int) time * 10}}, + 1); } private static InsertRowNode insertRow(long time) throws Exception { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/compaction/inner/InnerSeqCompactionWithFastPerformerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/compaction/inner/InnerSeqCompactionWithFastPerformerTest.java index e7b032da2f203..74519ad9573e8 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/compaction/inner/InnerSeqCompactionWithFastPerformerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/compaction/inner/InnerSeqCompactionWithFastPerformerTest.java @@ -24,6 +24,7 @@ import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.exception.DataRegionException; import org.apache.iotdb.db.exception.StorageEngineException; +import org.apache.iotdb.db.exception.WriteProcessRejectException; import org.apache.iotdb.db.storageengine.buffer.BloomFilterCache; import org.apache.iotdb.db.storageengine.buffer.ChunkCache; import org.apache.iotdb.db.storageengine.buffer.TimeSeriesMetadataCache; @@ -1110,7 +1111,7 @@ public void testAppendChunk() throws Exception { @Test public void testCompactionWithDeletionsDuringCompactions() - throws MetadataException, IOException, DataRegionException { + throws MetadataException, IOException, DataRegionException, WriteProcessRejectException { // create source seq files List sourceResources = new ArrayList<>(); List> chunkPagePointsNum = new ArrayList<>(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/compaction/inner/InnerSeqCompactionWithReadChunkPerformerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/compaction/inner/InnerSeqCompactionWithReadChunkPerformerTest.java index fef89b402a62f..69f61309c994b 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/compaction/inner/InnerSeqCompactionWithReadChunkPerformerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/compaction/inner/InnerSeqCompactionWithReadChunkPerformerTest.java @@ -24,6 +24,7 @@ import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.exception.DataRegionException; import org.apache.iotdb.db.exception.StorageEngineException; +import org.apache.iotdb.db.exception.WriteProcessRejectException; import org.apache.iotdb.db.storageengine.buffer.BloomFilterCache; import org.apache.iotdb.db.storageengine.buffer.ChunkCache; import org.apache.iotdb.db.storageengine.buffer.TimeSeriesMetadataCache; @@ -1072,7 +1073,11 @@ public void testAppendChunk() throws Exception { @Test public void testCompactionWithDeletionsDuringCompactions() - throws MetadataException, IOException, DataRegionException, InterruptedException { + throws MetadataException, + IOException, + DataRegionException, + InterruptedException, + WriteProcessRejectException { // create source seq files List sourceResources = new ArrayList<>(); List> chunkPagePointsNum = new ArrayList<>();