diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/task/PipeConfigNodeSubtask.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/task/PipeConfigNodeSubtask.java index 95376af638b1..9162c792c5ad 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/task/PipeConfigNodeSubtask.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/task/PipeConfigNodeSubtask.java @@ -176,16 +176,22 @@ private void initSink(final Map sinkAttributes) throws Exception */ @Override protected boolean executeOnce() throws Exception { - if (isClosed.get()) { + if (isClosed.get() || shouldStopSubmittingSelf.get()) { return false; } final Event event = lastEvent != null ? lastEvent : source.supply(); - // Record the last event for retry when exception occurs + // Keep the event for retry or restart if supply() returns after the task is stopped. setLastEvent(event); + // close() may have released the last event before supply() returned. + if (isClosed.get()) { + clearReferenceCountAndReleaseLastEvent(event); + return false; + } + try { - if (event == null) { + if (event == null || shouldStopSubmittingSelf.get()) { return false; } diff --git a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/pipe/agent/task/PipeConfigNodeSubtaskTest.java b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/pipe/agent/task/PipeConfigNodeSubtaskTest.java new file mode 100644 index 000000000000..2ca2aca5302c --- /dev/null +++ b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/pipe/agent/task/PipeConfigNodeSubtaskTest.java @@ -0,0 +1,150 @@ +/* + * 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.confignode.manager.pipe.agent.task; + +import org.apache.iotdb.commons.consensus.index.impl.MinimumProgressIndex; +import org.apache.iotdb.commons.pipe.agent.plugin.builtin.BuiltinPipePlugin; +import org.apache.iotdb.commons.pipe.agent.task.meta.PipeTaskMeta; +import org.apache.iotdb.commons.pipe.agent.task.subtask.PipeAbstractSinkSubtask; +import org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant; +import org.apache.iotdb.commons.pipe.event.EnrichedEvent; +import org.apache.iotdb.pipe.api.PipeConnector; +import org.apache.iotdb.pipe.api.PipeExtractor; + +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import java.lang.reflect.Field; +import java.util.Collections; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyZeroInteractions; +import static org.mockito.Mockito.when; + +public class PipeConfigNodeSubtaskTest { + + private PipeConfigNodeSubtask subtask; + private PipeExtractor source; + private PipeConnector sink; + private boolean subtaskClosed; + + @Before + public void setUp() throws Exception { + source = mock(PipeExtractor.class); + sink = mock(PipeConnector.class); + subtask = + new PipeConfigNodeSubtask( + "PipeConfigNodeSubtaskTest", + System.currentTimeMillis(), + Collections.emptyMap(), + Collections.emptyMap(), + Collections.singletonMap( + PipeSinkConstant.CONNECTOR_KEY, + BuiltinPipePlugin.DO_NOTHING_CONNECTOR.getPipePluginName()), + new PipeTaskMeta(MinimumProgressIndex.INSTANCE, Integer.MIN_VALUE)); + final Field sourceField = PipeConfigNodeSubtask.class.getDeclaredField("source"); + sourceField.setAccessible(true); + sourceField.set(subtask, source); + final Field sinkField = PipeAbstractSinkSubtask.class.getDeclaredField("outputPipeSink"); + sinkField.setAccessible(true); + sinkField.set(subtask, sink); + } + + @After + public void tearDown() { + if (!subtaskClosed) { + subtask.close(); + } + } + + @Test + public void testStoppedSubtaskDoesNotPollSource() throws Exception { + assertFalse(subtask.executeOnce()); + verifyZeroInteractions(source, sink); + } + + @Test + public void testStopDuringSupplyRetainsEventUntilRestart() throws Exception { + assertStopDuringSupply(false); + } + + @Test + public void testDropDuringSupplyReleasesEvent() throws Exception { + assertStopDuringSupply(true); + } + + private void assertStopDuringSupply(final boolean drop) throws Exception { + final EnrichedEvent event = mock(EnrichedEvent.class); + final CountDownLatch supplyStarted = new CountDownLatch(1); + final CountDownLatch supplyReleased = new CountDownLatch(1); + final ExecutorService worker = Executors.newSingleThreadExecutor(); + when(source.supply()) + .thenAnswer( + invocation -> { + supplyStarted.countDown(); + assertTrue(supplyReleased.await(10, TimeUnit.SECONDS)); + return event; + }); + subtask.allowSubmittingSelf(); + try { + final Future execution = worker.submit(() -> subtask.executeOnce()); + assertTrue(supplyStarted.await(10, TimeUnit.SECONDS)); + assertTrue(subtask.disallowSubmittingSelf()); + if (drop) { + subtask.close(); + subtaskClosed = true; + } + supplyReleased.countDown(); + + assertFalse(execution.get(10, TimeUnit.SECONDS)); + assertFalse(subtask.executeOnce()); + verify(sink, never()).transfer(event); + verify(event, never()).decreaseReferenceCount(anyString(), anyBoolean()); + + if (drop) { + verify(event).clearReferenceCount(anyString()); + } else { + verifyZeroInteractions(sink); + subtask.allowSubmittingSelf(); + assertTrue(subtask.executeOnce()); + verify(source).supply(); + verify(sink).transfer(event); + verify(event).decreaseReferenceCount(anyString(), eq(true)); + } + } finally { + supplyReleased.countDown(); + worker.shutdownNow(); + assertTrue(worker.awaitTermination(10, TimeUnit.SECONDS)); + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTaskRetention.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTaskRetention.java index 828df844ed7a..6197bbccc989 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTaskRetention.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTaskRetention.java @@ -62,14 +62,7 @@ final class LoadTaskRetention { /** Checks whether the staged directory must be retained for follower WAL synchronization. */ boolean mustRetain(final long searchIndex) { - // IoTConsensusV2 replicates through pipe, so no replica reads these bytes back out of the WAL - // and there is no index to wait for: the staged directory can go as soon as the task ends. - if (LoadConsensusCompat.usesIoTConsensusV2()) { - return false; - } - final long safelyDeletedSearchIndex = currentSafelyDeletedSearchIndex(); - return safelyDeletedSearchIndex != ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX - && safelyDeletedSearchIndex < searchIndex; + return currentSafelyDeletedSearchIndex() < searchIndex; } /** Takes ownership of a staged directory, postponing deletion until WAL watermarks allow it. */ @@ -122,14 +115,14 @@ void release() { } final long safelyDeletedSearchIndex = currentSafelyDeletedSearchIndex(); - final long boundary = - safelyDeletedSearchIndex == ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX - ? Long.MAX_VALUE - : safelyDeletedSearchIndex; + // The default V1 watermark means no replica progress has been confirmed yet. + if (safelyDeletedSearchIndex == ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX) { + return; + } for (final Map.Entry entry : loadId2RetainedTaskDir.entrySet()) { final RetainedTaskDir retained = entry.getValue(); - if (retained.reachedSearchIndex() <= boundary + if (retained.reachedSearchIndex() <= safelyDeletedSearchIndex && loadId2RetainedTaskDir.remove(entry.getKey(), retained)) { LOGGER.info( StorageEngineMessages @@ -165,10 +158,15 @@ private void ensureWalListenerRegistered() { } private long currentSafelyDeletedSearchIndex() { + // V2 replicates through pipe. Other protocols without a V1 WAL (for example Ratis) do not + // expand staged references from this node's WAL either, so neither has a watermark to wait for. + if (LoadConsensusCompat.usesIoTConsensusV2()) { + return Long.MAX_VALUE; + } return dataRegion .getWALNode() .map(IWALNode::getSafelyDeletedSearchIndex) - .orElse(ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX); + .orElse(Long.MAX_VALUE); } // ------------------------------------------------------------------------- diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileCleaner.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileCleaner.java index ee440de7790f..93a537c1e4ac 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileCleaner.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileCleaner.java @@ -350,14 +350,22 @@ private static boolean isComplete(final File taskDir) { } private static boolean replicasReached(final DataRegion dataRegion, final long searchIndex) { - final long safelyDeletedSearchIndex = - dataRegion - .getWALNode() - .map(wal -> wal.getSafelyDeletedSearchIndex()) - .orElse(ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX); - - return safelyDeletedSearchIndex == ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX - || safelyDeletedSearchIndex >= searchIndex; + // IoTConsensusV2 and consensus implementations without a WAL replicate the staged payload + // directly, so no follower can read it back from this node's WAL. IoTConsensus V1, however, + // must report a real watermark before its staged references are safe to delete. + if (LoadConsensusCompat.usesIoTConsensusV2()) { + return true; + } + return dataRegion + .getWALNode() + .map( + wal -> { + final long safelyDeletedSearchIndex = wal.getSafelyDeletedSearchIndex(); + return safelyDeletedSearchIndex + != ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX + && safelyDeletedSearchIndex >= searchIndex; + }) + .orElse(true); } // ------------------------------------------------------------------------- diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileCleanerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileCleanerTest.java index ca25a96ad4bb..01482f2dfe04 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileCleanerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileCleanerTest.java @@ -22,6 +22,7 @@ import org.apache.iotdb.commons.conf.ConfigurationFileUtils; import org.apache.iotdb.commons.conf.TrimProperties; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.consensus.iot.log.ConsensusReqReader; import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; @@ -186,6 +187,47 @@ public void testAnIdleTaskIsReclaimedOnlyOnceItsPiecesAreReplicated() throws Exc } } + @Test + public void testAFinishedTaskWaitsForTheFirstWatermark() throws Exception { + final String loadId = "unreported-watermark-task"; + final File taskDir = new File(tempDir, loadId); + Assert.assertTrue(taskDir.mkdirs()); + stageOneChunk(taskDir, 5L); + final LoadTsFileCleaner cleaner = LoadTsFileCleaner.getInstance(); + cleaner.register(loadId, dataRegion, taskDir, 7L, LoadTsFileConsensusOp.ABORT); + + walWatermark.set(ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX); + cleaner.sweepRegisteredTasks(); + Assert.assertTrue("an unreported watermark does not confirm replication", taskDir.exists()); + + walWatermark.set(5L); + cleaner.sweepRegisteredTasks(); + Assert.assertTrue("the terminal command is not replicated yet", taskDir.exists()); + + walWatermark.set(7L); + cleaner.sweepRegisteredTasks(); + Assert.assertFalse("the terminal command is replicated", taskDir.exists()); + } + + @Test + public void testIoTConsensusV2DoesNotWaitForAWALWatermark() throws Exception { + final String loadId = "v2-finished-task"; + final File taskDir = new File(tempDir, loadId); + Assert.assertTrue(taskDir.mkdirs()); + stageOneChunk(taskDir, 5L); + final LoadTsFileCleaner cleaner = LoadTsFileCleaner.getInstance(); + cleaner.register(loadId, dataRegion, taskDir, 7L, LoadTsFileConsensusOp.ABORT); + walWatermark.set(ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX); + + LoadConsensusCompat.setIoTConsensusV2ForTest(Boolean.TRUE); + try { + cleaner.sweepRegisteredTasks(); + Assert.assertFalse("V2 does not expand staged references from the WAL", taskDir.exists()); + } finally { + LoadConsensusCompat.setIoTConsensusV2ForTest(null); + } + } + @Test public void testAFinishedCommittedTaskIsReleasedOnceItsStagedFilesAreComplete() throws Exception { final String loadId = "finished-task"; diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManagerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManagerTest.java index c98b3afa1958..360fa83c6afb 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManagerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManagerTest.java @@ -21,6 +21,7 @@ import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.consensus.iot.log.ConsensusReqReader; import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.exception.load.LoadFileException; @@ -634,10 +635,8 @@ private static NonAlignedChunkData[] twoChunksOfTheSamePartition() { } @Test - public void testAbortedConsensusTaskReleasesStagedFilesWhenNoWatermarkIsReported() - throws Exception { - // Without a consensus watermark (single replica) nothing can read the staged bytes back, so the - // directory of a finished task is deleted right away. + public void testAbortedConsensusTaskReleasesStagedFilesWithoutAWAL() throws Exception { + // Protocols without a V1 WAL do not expand staged references from the WAL. final String loadId = "aborted-load"; final LoadTsFileManager manager = new LoadTsFileManager(dataRegion); final List refs = @@ -660,8 +659,8 @@ public void testFinishedConsensusTaskKeepsStagedFilesUntilTheWatermarkAdvances() try (final WALNode walNode = new WALNode("load-retention-test", new File(tempDir, "wal").getAbsolutePath())) { Mockito.when(dataRegion.getWALNode()).thenReturn(Optional.of(walNode)); - // The watermark lags behind the command, so a follower may still need the staged bytes. - walNode.setSafelyDeletedSearchIndex(5L); + // Until V1 reports its first watermark, a follower may still need the staged bytes. + walNode.setSafelyDeletedSearchIndex(ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX); final LoadTsFileManager manager = new LoadTsFileManager(dataRegion); final String loadId = "retained-load"; @@ -682,6 +681,10 @@ public void testFinishedConsensusTaskKeepsStagedFilesUntilTheWatermarkAdvances() assertTrue(progress.isFile()); assertNotNull(LoadTsFileProgress.readTerminal(progress)); + // A reported watermark that still lags behind the command must keep the staged bytes too. + walNode.setSafelyDeletedSearchIndex(5L); + assertTrue(staged.isFile()); + // Once every follower got past the command, the directory is released. walNode.setSafelyDeletedSearchIndex(10L); assertFalse(staged.exists()); @@ -708,6 +711,27 @@ public void testStagedDirectoryOfFinishedTaskIsDeletedOnRecovery() throws Except assertFalse(manager.prepare(prepareNode("finished-load"), Collections.emptyMap())); } + @Test + public void testRecoveredFinishedTaskWaitsForTheFirstWatermark() throws Exception { + useWatermarkWALNode(ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX); + final File taskDir = new File(tempDir, "root.load_manager_test-0/recovered-finished-load"); + assertTrue(taskDir.mkdirs()); + final File stagedFile = new File(taskDir, "123.tsfile"); + Files.write(stagedFile.toPath(), new byte[] {1, 2, 3}); + new LoadTsFileProgress(stagedFile).recordTerminal(LoadTsFileConsensusOp.ABORT, 7L); + + final LoadTsFileManager manager = new LoadTsFileManager(dataRegion); + try { + assertTrue("a restart must keep bytes until replication is confirmed", stagedFile.isFile()); + assertFalse(manager.prepare(prepareNode(taskDir.getName()), Collections.emptyMap())); + + advanceWatermark(7L); + assertFalse(taskDir.exists()); + } finally { + manager.stop(); + } + } + @Test public void testIoTConsensusV2WaitsForNoIndexBeforeReleasing() { LoadConsensusCompat.setIoTConsensusV2ForTest(Boolean.TRUE); @@ -778,7 +802,17 @@ public void testStagedDirectoryWithoutTerminalRecordIsKeptOnRecovery() throws Ex public void testCommitKeepsTheStagedBytesWhileAReplicaMayStillNeedThem() throws Exception { // The watermark lags behind the COMMIT command, so a replica that catches up from this node's // WAL still has to expand the pieces of the task from the staging directory. - useWatermarkWALNode(5L); + assertCommitKeepsTheStagedBytes(5L); + } + + @Test + public void testCommitKeepsTheStagedBytesBeforeTheFirstWatermark() throws Exception { + assertCommitKeepsTheStagedBytes(ConsensusReqReader.DEFAULT_SAFELY_DELETED_SEARCH_INDEX); + } + + private void assertCommitKeepsTheStagedBytes(final long safelyDeletedSearchIndex) + throws Exception { + useWatermarkWALNode(safelyDeletedSearchIndex); emulateStagedFileImport(); final String loadId = "committed-load";