diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/treemodel/auto/enhanced/IoTDBPipeAutoDropIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/treemodel/auto/enhanced/IoTDBPipeAutoDropIT.java index a3e5d32cb72a..19b30f628d94 100644 --- a/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/treemodel/auto/enhanced/IoTDBPipeAutoDropIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/treemodel/auto/enhanced/IoTDBPipeAutoDropIT.java @@ -175,16 +175,8 @@ public void testAutoDropFinitePipesWithoutDataRegion() throws Exception { .setExtractorAttributes(realtimeSourceAttributes)) .getCode()); - Assert.assertEquals( - TSStatusCode.SUCCESS_STATUS.getStatusCode(), - client.startPipe("query_pipe_without_data_region").getCode()); - Assert.assertEquals( - TSStatusCode.SUCCESS_STATUS.getStatusCode(), - client.startPipe("history_pipe_without_data_region").getCode()); - Assert.assertEquals( - TSStatusCode.SUCCESS_STATUS.getStatusCode(), - client.startPipe("realtime_pipe_without_data_region").getCode()); - + // Pipes start automatically after creation. Explicitly starting finite pipes can race with + // their auto-drop when they have no DataRegions to process. await() .pollInSameThread() .pollInterval(1L, TimeUnit.SECONDS) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/PipeTsFileInsertionEvent.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/PipeTsFileInsertionEvent.java index 9690eed5d758..a83a027627b2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/PipeTsFileInsertionEvent.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/PipeTsFileInsertionEvent.java @@ -55,6 +55,7 @@ import org.apache.iotdb.db.pipe.resource.tsfile.PipeTsFileResourceManager; import org.apache.iotdb.db.pipe.source.dataregion.realtime.assigner.PipeTsFileEpochProgressIndexKeeper; import org.apache.iotdb.db.storageengine.dataregion.memtable.TsFileProcessor; +import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent; import org.apache.iotdb.pipe.api.event.dml.insertion.TsFileInsertionEvent; @@ -274,7 +275,9 @@ private PipeTsFileInsertionEvent( // hard-link it to each pipe dir this.tsFile = Objects.isNull(tsFile) ? resource.getTsFile() : tsFile; - this.isWithMod = isWithMod && resource.anyModFileExists(); + // Historical files and their mods have already been pinned. Compaction may have removed + // the original resource before this event is created, so inspect the pinned mods later. + this.isWithMod = isWithMod && (isGeneratedByHistoricalExtractor || resource.anyModFileExists()); this.modFile = this.isWithMod ? resource.getExclusiveModFile().getFile() : null; // TODO: process the shared mod file this.sharedModFile = @@ -463,9 +466,22 @@ public boolean internallyIncreaseResourceReferenceCount(final String holderMessa PipeTsFileResourceManager.getPipeTsFileResourcePipeName(pipeName, creationTime); final File originalTsFile = tsFile; final File originalModFile = modFile; + final boolean originalIsWithMod = isWithMod; File increasedTsFile = null; boolean increased = false; try { + if (isGeneratedByHistoricalExtractor) { + tsFile = + PipeTsFileResourceManager.getHardlinkOrCopiedFileInPipeDir( + tsFile, pipeTsFileResourcePipeName); + if (isWithMod) { + modFile = ModificationFile.getExclusiveMods(tsFile); + isWithMod = modFile.exists(); + if (!isWithMod) { + modFile = null; + } + } + } increasedTsFile = PipeDataNodeResourceManager.tsfile() .increaseFileReference(tsFile, true, pipeTsFileResourcePipeName); @@ -490,6 +506,7 @@ public boolean internallyIncreaseResourceReferenceCount(final String holderMessa } tsFile = originalTsFile; modFile = originalModFile; + isWithMod = originalIsWithMod; LOGGER.warn( String.format( DataNodePipeMessages.INCREASE_REFERENCE_COUNT_TSFILE_OR_MODFILE_ERROR_HOLDER_FMT, diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/tsfile/PipeTsFileResourceManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/tsfile/PipeTsFileResourceManager.java index db9816636812..02898d17b4b7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/tsfile/PipeTsFileResourceManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/tsfile/PipeTsFileResourceManager.java @@ -427,14 +427,17 @@ public Map getMeasurementDataTypeMapFromCache( : hardlinkOrCopiedFileToTsFilePublicResourceMap; } - public void pinTsFileResource( + /** Returns whether a companion mod-file reference was acquired. */ + public boolean pinTsFileResource( final TsFileResource resource, final boolean withMods, final @Nullable String pipeName) throws IOException { final File pinnedTsFile = increaseFileReference(resource.getTsFile(), true, pipeName); try { if (withMods && resource.getExclusiveModFile().exists()) { increaseFileReference(resource.getExclusiveModFile().getFile(), false, pipeName); + return true; } + return false; } catch (final IOException | RuntimeException e) { // Pinning is a two-file operation. Do not leave the TsFile pinned when the mod file cannot // be pinned (for example, when the pipe directory is temporarily unavailable). @@ -448,9 +451,7 @@ public void pinTsFileResource( } public void unpinTsFileResource( - final TsFileResource resource, - final boolean shouldTransferModFile, - final @Nullable String pipeName) + final TsFileResource resource, final @Nullable String pipeName, final boolean pinnedModFile) throws IOException { Exception firstException = null; try { @@ -462,7 +463,9 @@ public void unpinTsFileResource( // Always attempt the mod-file cleanup even when resolving/decreasing the TsFile fails. A // failed first cleanup must not strand the second reference. - if (shouldTransferModFile && resource.exclusiveModFileExists()) { + // Use the state captured when pinning. Compaction can remove the original mod file, and + // an event can acquire mod-file references even when no mod file was originally pinned. + if (pinnedModFile) { try { decreaseFileReference( getHardlinkOrCopiedFileInPipeDir(resource.getExclusiveModFile().getFile(), pipeName), diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/source/dataregion/historical/PipeHistoricalDataRegionTsFileAndDeletionSource.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/source/dataregion/historical/PipeHistoricalDataRegionTsFileAndDeletionSource.java index a28598ddec15..041d33c2faff 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/source/dataregion/historical/PipeHistoricalDataRegionTsFileAndDeletionSource.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/source/dataregion/historical/PipeHistoricalDataRegionTsFileAndDeletionSource.java @@ -175,6 +175,7 @@ public class PipeHistoricalDataRegionTsFileAndDeletionSource private Queue pendingQueue; private final Map> filteredTsFileResources2TableNames = new HashMap<>(); + private final Set resourcesWithPinnedModFile = new HashSet<>(); private final Map pendingResource2ReplicateIndexForIoTV2 = new HashMap<>(); private final Set historicalProgressReportResources = new HashSet<>(); @@ -800,8 +801,11 @@ private void extractTsFiles( // Pin the resource, in case the file is removed by compaction or anything. // Will unpin it after the PipeTsFileInsertionEvent is created and pinned. try { - PipeDataNodeResourceManager.tsfile() - .pinTsFileResource(resource, shouldTransferModFile, pipeNameWithCreationTime); + if (PipeDataNodeResourceManager.tsfile() + .pinTsFileResource( + resource, shouldTransferModFile, pipeNameWithCreationTime)) { + resourcesWithPinnedModFile.add(resource); + } return false; } catch (final IOException e) { ++statistics.pinFailedCount; @@ -1217,7 +1221,8 @@ protected boolean consumeSkippedHistoricalTsFileEventIfNecessary(final TsFileRes } finally { try { PipeDataNodeResourceManager.tsfile() - .unpinTsFileResource(resource, shouldTransferModFile, pipeNameWithCreationTime); + .unpinTsFileResource( + resource, pipeNameWithCreationTime, resourcesWithPinnedModFile.remove(resource)); } catch (final IOException e) { LOGGER.warn( DataNodePipeMessages.PIPE_FAILED_TO_UNPIN_SKIPPED_HISTORICAL_TSFILERESOURCE, @@ -1329,7 +1334,8 @@ protected Event supplyTsFileEvent(final TsFileResource resource) { if (shouldUnpinResource) { try { PipeDataNodeResourceManager.tsfile() - .unpinTsFileResource(resource, shouldTransferModFile, pipeNameWithCreationTime); + .unpinTsFileResource( + resource, pipeNameWithCreationTime, resourcesWithPinnedModFile.remove(resource)); } catch (final IOException e) { LOGGER.warn( DataNodePipeMessages.PIPE_FAILED_TO_UNPIN_TSFILERESOURCE_AFTER_CREATING, @@ -1454,7 +1460,9 @@ public synchronized void close() { try { PipeDataNodeResourceManager.tsfile() .unpinTsFileResource( - (TsFileResource) resource, shouldTransferModFile, pipeNameWithCreationTime); + (TsFileResource) resource, + pipeNameWithCreationTime, + resourcesWithPinnedModFile.remove(resource)); } catch (final IOException e) { LOGGER.warn( DataNodePipeMessages.PIPE_FAILED_TO_UNPIN_TSFILERESOURCE_AFTER_DROPPING, @@ -1468,6 +1476,7 @@ public synchronized void close() { pendingQueue = null; } pendingResource2ReplicateIndexForIoTV2.clear(); + resourcesWithPinnedModFile.clear(); historicalProgressReportResources.clear(); pendingHistoricalProgressIndexToReport = null; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TableConfigTaskVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TableConfigTaskVisitor.java index 2af8eccc1738..f91969ae32ad 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TableConfigTaskVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TableConfigTaskVisitor.java @@ -1370,6 +1370,15 @@ public static void checkAndEnrichSourceUser( .FAILED_TO_S_PIPE_S_IN_IOTDB_SOURCE_PASSWORD_MUST_BE_SET_WHEN_THE_USERNAME_IS_SPECIFIED, isAlter ? DataNodeQueryMessages.ALTER : DataNodeQueryMessages.CREATE, pipeName)); + } else { + updatePipeUserIdFromUsername( + replacedSourceAttributes, + sourceParameters, + PipeSourceConstant.SOURCE_IOTDB_USER_ID, + PipeSourceConstant.EXTRACTOR_IOTDB_USER_KEY, + PipeSourceConstant.SOURCE_IOTDB_USER_KEY, + PipeSourceConstant.EXTRACTOR_IOTDB_USERNAME_KEY, + PipeSourceConstant.SOURCE_IOTDB_USERNAME_KEY); } } @@ -1442,6 +1451,15 @@ public static void checkAndEnrichSinkUser( .FAILED_TO_S_PIPE_S_IN_WRITE_BACK_SINK_PASSWORD_MUST_BE_SET_WHEN_THE_USERNAME_IS, isAlter ? DataNodeQueryMessages.ALTER : DataNodeQueryMessages.CREATE, pipeName)); + } else { + updatePipeUserIdFromUsername( + connectorAttributes, + connectorParameters, + PipeSinkConstant.SINK_IOTDB_USER_ID, + PipeSinkConstant.CONNECTOR_IOTDB_USER_KEY, + PipeSinkConstant.SINK_IOTDB_USER_KEY, + PipeSinkConstant.CONNECTOR_IOTDB_USERNAME_KEY, + PipeSinkConstant.SINK_IOTDB_USERNAME_KEY); } } @@ -1510,6 +1528,14 @@ public static void markSourceAuthenticationAsExplicitIfNecessary( PipeSourceConstant.SOURCE_IOTDB_PASSWORD_KEY)) { sourceAttributes.put( SystemConstant.SOURCE_AUTHENTICATION_INJECTED_KEY, Boolean.FALSE.toString()); + updatePipeUserIdFromUsername( + sourceAttributes, + sourceParameters, + PipeSourceConstant.SOURCE_IOTDB_USER_ID, + PipeSourceConstant.EXTRACTOR_IOTDB_USER_KEY, + PipeSourceConstant.SOURCE_IOTDB_USER_KEY, + PipeSourceConstant.EXTRACTOR_IOTDB_USERNAME_KEY, + PipeSourceConstant.SOURCE_IOTDB_USERNAME_KEY); } } @@ -1524,6 +1550,26 @@ public static void markSinkAuthenticationAsExplicitIfNecessary( PipeSinkConstant.CONNECTOR_IOTDB_PASSWORD_KEY, PipeSinkConstant.SINK_IOTDB_PASSWORD_KEY)) { sinkAttributes.put(SystemConstant.SINK_AUTHENTICATION_INJECTED_KEY, Boolean.FALSE.toString()); + updatePipeUserIdFromUsername( + sinkAttributes, + sinkParameters, + PipeSinkConstant.SINK_IOTDB_USER_ID, + PipeSinkConstant.CONNECTOR_IOTDB_USER_KEY, + PipeSinkConstant.SINK_IOTDB_USER_KEY, + PipeSinkConstant.CONNECTOR_IOTDB_USERNAME_KEY, + PipeSinkConstant.SINK_IOTDB_USERNAME_KEY); + } + } + + private static void updatePipeUserIdFromUsername( + final Map attributes, + final PipeParameters parameters, + final String userIdAttributeKey, + final String... usernameAttributeKeys) { + final String username = parameters.getStringByKeys(usernameAttributeKeys); + if (username != null) { + attributes.put( + userIdAttributeKey, String.valueOf(AuthorityChecker.getUserId(username).orElse(-1L))); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java index 6ed389ddafbb..e8d8c8a83037 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java @@ -518,34 +518,39 @@ private boolean reject(final SubscriptionQueueRejectionReason reason) { Objects.isNull(reason) ? SubscriptionQueueRejectionReason.QUEUE_CAPACITY : reason; lastOfferRejectionReason.set(effectiveReason); rejectionListener.accept(effectiveReason); - return false; + return effectiveReason != SubscriptionQueueRejectionReason.INACTIVE_OR_CLOSED; } @Override public boolean offer(final IndexedConsensusRequest request) { final boolean offered; + boolean wakeupRequired = false; synchronized (this) { final SubscriptionQueueRejectionReason admissionReason = admissionSupplier.get(); if (admissionReason != SubscriptionQueueRejectionReason.NONE) { - return reject(admissionReason); - } - if (!requestMemoryManager.reserve(request)) { - return reject(SubscriptionQueueRejectionReason.CONSENSUS_REQUEST_MEMORY_LIMIT); - } - try { - offered = super.offer(request); - } catch (final Throwable t) { - requestMemoryManager.free(request); - throw t; - } - if (offered) { - retainedRequestBytes.addAndGet(request.getRetainedMemorySize()); + wakeupRequired = reject(admissionReason); + offered = false; + } else if (!requestMemoryManager.reserve(request)) { + wakeupRequired = reject(SubscriptionQueueRejectionReason.CONSENSUS_REQUEST_MEMORY_LIMIT); + offered = false; } else { - requestMemoryManager.free(request); - reject(SubscriptionQueueRejectionReason.QUEUE_CAPACITY); + try { + offered = super.offer(request); + } catch (final Throwable t) { + requestMemoryManager.free(request); + throw t; + } + if (offered) { + retainedRequestBytes.addAndGet(request.getRetainedMemorySize()); + } else { + requestMemoryManager.free(request); + wakeupRequired = reject(SubscriptionQueueRejectionReason.QUEUE_CAPACITY); + } } } - if (offered) { + // The rejected request remains available in WAL. Wake the replay worker so a queue that has + // already gone dormant does not wait for a later successful realtime offer to make progress. + if (offered || wakeupRequired) { wakeupHook.run(); } return offered; diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/PipeTsFileInsertionEventTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/PipeTsFileInsertionEventTest.java index 5d5be1dbadee..61e68055f904 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/PipeTsFileInsertionEventTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/PipeTsFileInsertionEventTest.java @@ -35,12 +35,18 @@ import org.apache.iotdb.db.auth.AuthorityChecker; import org.apache.iotdb.db.pipe.event.common.tsfile.PipeCompactedTsFileInsertionEvent; import org.apache.iotdb.db.pipe.event.common.tsfile.PipeTsFileInsertionEvent; +import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager; +import org.apache.iotdb.db.pipe.resource.tsfile.PipeTsFileResourceManager; import org.apache.iotdb.db.pipe.source.dataregion.realtime.assigner.PipeTsFileEpochProgressIndexKeeper; import org.apache.iotdb.db.queryengine.plan.relational.security.AccessControl; import org.apache.iotdb.db.queryengine.plan.relational.security.TreeAccessCheckContext; import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.RelationalAuthorStatement; import org.apache.iotdb.db.queryengine.plan.statement.Statement; import org.apache.iotdb.db.storageengine.dataregion.compaction.utils.CompactionTestFileWriter; +import org.apache.iotdb.db.storageengine.dataregion.modification.DeletionPredicate; +import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile; +import org.apache.iotdb.db.storageengine.dataregion.modification.TableDeletionEntry; +import org.apache.iotdb.db.storageengine.dataregion.modification.TagPredicate.FullExactMatch; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResourceStatus; import org.apache.iotdb.db.storageengine.dataregion.tsfile.generator.TsFileNameGenerator; @@ -72,6 +78,108 @@ import static org.apache.iotdb.db.auth.AuthorityChecker.NO_PERMISSION_PROMOTION; public class PipeTsFileInsertionEventTest { + @Test + public void testHistoricalModsSurviveOriginalFileRemoval() throws Exception { + assertHistoricalModsAfterOriginalFileRemoval(true, true); + } + + @Test + public void testHistoricalFileWithoutModsSurvivesOriginalFileRemoval() throws Exception { + assertHistoricalModsAfterOriginalFileRemoval(false, true); + } + + @Test + public void testHistoricalModsRemainDisabledAfterOriginalFileRemoval() throws Exception { + assertHistoricalModsAfterOriginalFileRemoval(true, false); + } + + private void assertHistoricalModsAfterOriginalFileRemoval( + final boolean hasMods, final boolean transferMods) throws Exception { + final File tempDir = Files.createTempDirectory("pipeHistoricalMods").toFile(); + final String pipeName = "historical_mods"; + final long creationTime = 1L; + final String resourcePipeName = + PipeTsFileResourceManager.getPipeTsFileResourcePipeName(pipeName, creationTime); + final PipeTsFileResourceManager manager = PipeDataNodeResourceManager.tsfile(); + PipeTsFileInsertionEvent event = null; + boolean increased = false; + File pinnedTsFile = null; + File pinnedModFile = null; + try { + final File dataDir = new File(tempDir, "unsequence/db1/0/0"); + Assert.assertTrue(dataDir.mkdirs()); + final TsFileResource resource = + createNonEmptyTsFileResource(dataDir, "1-1-0-0.tsfile", 1L, 0); + resource.setStatusForTest(TsFileResourceStatus.NORMAL); + final TableDeletionEntry deletion = + new TableDeletionEntry( + new DeletionPredicate( + "t1", + new FullExactMatch( + IDeviceID.Factory.DEFAULT_FACTORY.create(new String[] {"t1", "d1"}))), + new TimeRange(1L, 1L)); + if (hasMods) { + resource.getExclusiveModFile().write(deletion); + resource.getExclusiveModFile().close(); + } + manager.pinTsFileResource(resource, transferMods, resourcePipeName); + pinnedTsFile = + PipeTsFileResourceManager.getHardlinkOrCopiedFileInPipeDir( + resource.getTsFile(), resourcePipeName); + pinnedModFile = ModificationFile.getExclusiveMods(pinnedTsFile); + + // Compaction removes the original after historical selection has pinned both files. + Assert.assertTrue(resource.remove()); + Assert.assertFalse(resource.anyModFileExists()); + Assert.assertEquals(hasMods && transferMods, pinnedModFile.exists()); + event = + new PipeTsFileInsertionEvent( + true, + "db1", + resource, + null, + transferMods, + false, + true, + Collections.singleton("t1"), + pipeName, + creationTime, + null, + null, + new TablePattern(true, null, null), + null, + null, + null, + true, + Long.MIN_VALUE, + Long.MAX_VALUE); + increased = event.internallyIncreaseResourceReferenceCount(getClass().getName()); + Assert.assertTrue(increased); + Assert.assertEquals(hasMods && transferMods, event.isWithMod()); + if (hasMods && transferMods) { + Assert.assertEquals(pinnedModFile, event.getModFile()); + Assert.assertEquals( + Collections.singletonList(deletion), + ModificationFile.readAllModifications(event.getTsFile(), true)); + } else { + Assert.assertNull(event.getModFile()); + } + } finally { + if (increased) { + event.internallyDecreaseResourceReferenceCount(getClass().getName()); + } else if (event != null) { + event.close(); + } + if (pinnedModFile != null) { + manager.decreaseFileReference(pinnedModFile, resourcePipeName); + } + if (pinnedTsFile != null) { + manager.decreaseFileReference(pinnedTsFile, resourcePipeName); + } + FileUtils.deleteFileOrDirectory(tempDir); + } + } + @Test public void testAuthCheck() throws Exception { final AccessControl oldControl = AuthorityChecker.getAccessControl(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/PipeTsFileResourceManagerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/PipeTsFileResourceManagerTest.java index d4b8c45f0f21..b823f0fe2bf6 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/PipeTsFileResourceManagerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/PipeTsFileResourceManagerTest.java @@ -28,6 +28,7 @@ import org.apache.iotdb.db.storageengine.dataregion.modification.ModEntry; import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile; import org.apache.iotdb.db.storageengine.dataregion.modification.TreeDeletionEntry; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.metadata.enums.TSEncoding; @@ -246,6 +247,57 @@ public void testDecreaseTsFile() throws IOException { Assert.assertFalse(Files.exists(originModFile.toPath())); } + @Test + public void testUnpinAfterOriginalModsRemoved() throws Exception { + final TsFileResource resource = new TsFileResource(new File(TS_FILE_NAME)); + resource + .getExclusiveModFile() + .write(new TreeDeletionEntry(new MeasurementPath("root.lemming.**"), 0, Long.MAX_VALUE)); + resource.getExclusiveModFile().close(); + + final boolean pinnedModFile = + pipeTsFileResourceManager.pinTsFileResource(resource, true, PIPE_NAME); + Assert.assertTrue(pinnedModFile); + final File pipeTsFile = + PipeTsFileResourceManager.getHardlinkOrCopiedFileInPipeDir(resource.getTsFile(), PIPE_NAME); + final File pipeMods = ModificationFile.getExclusiveMods(pipeTsFile); + Assert.assertEquals(1, pipeTsFileResourceManager.getFileReferenceCount(pipeMods, PIPE_NAME)); + + Assert.assertTrue(resource.remove()); + Assert.assertFalse(resource.exclusiveModFileExists()); + Assert.assertTrue(pipeMods.exists()); + pipeTsFileResourceManager.unpinTsFileResource(resource, PIPE_NAME, pinnedModFile); + + Assert.assertEquals(0, pipeTsFileResourceManager.getFileReferenceCount(pipeTsFile, PIPE_NAME)); + Assert.assertEquals(0, pipeTsFileResourceManager.getFileReferenceCount(pipeMods, PIPE_NAME)); + Assert.assertEquals(0, pipeTsFileResourceManager.getFileReferenceCount(pipeMods, null)); + Assert.assertFalse(pipeMods.exists()); + } + + @Test + public void testUnpinDoesNotReleaseModsCreatedAfterPinning() throws Exception { + final TsFileResource resource = new TsFileResource(new File(TS_FILE_NAME)); + final boolean pinnedModFile = + pipeTsFileResourceManager.pinTsFileResource(resource, true, PIPE_NAME); + Assert.assertFalse(pinnedModFile); + + // An event can acquire mods created after historical selection pinned the TsFile alone. + resource + .getExclusiveModFile() + .write(new TreeDeletionEntry(new MeasurementPath("root.lemming.**"), 0, Long.MAX_VALUE)); + resource.getExclusiveModFile().close(); + final File eventMods = + pipeTsFileResourceManager.increaseFileReference( + resource.getExclusiveModFile().getFile(), false, PIPE_NAME); + pipeTsFileResourceManager.unpinTsFileResource(resource, PIPE_NAME, pinnedModFile); + + Assert.assertEquals(1, pipeTsFileResourceManager.getFileReferenceCount(eventMods, PIPE_NAME)); + Assert.assertTrue(eventMods.exists()); + pipeTsFileResourceManager.decreaseFileReference(eventMods, PIPE_NAME); + Assert.assertEquals(0, pipeTsFileResourceManager.getFileReferenceCount(eventMods, PIPE_NAME)); + Assert.assertFalse(eventMods.exists()); + } + @Test public void testConcurrentIncreaseTsFile() throws Exception { assertConcurrentIncreaseFileReference(new File(TS_FILE_NAME), true); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/config/PipeAuthenticationIdentityTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/config/PipeAuthenticationIdentityTest.java new file mode 100644 index 000000000000..741dd0171ba2 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/config/PipeAuthenticationIdentityTest.java @@ -0,0 +1,131 @@ +/* + * 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.queryengine.plan.execution.config; + +import org.apache.iotdb.commons.audit.UserEntity; +import org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant; +import org.apache.iotdb.commons.pipe.config.constant.PipeSourceConstant; +import org.apache.iotdb.commons.pipe.config.constant.SystemConstant; +import org.apache.iotdb.db.auth.AuthorityChecker; +import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameters; + +import org.junit.Test; +import org.junit.runner.RunWith; +import org.powermock.core.classloader.annotations.PowerMockIgnore; +import org.powermock.core.classloader.annotations.PrepareForTest; +import org.powermock.modules.junit4.PowerMockRunner; + +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; + +import static org.junit.Assert.assertEquals; +import static org.mockito.Mockito.when; +import static org.powermock.api.mockito.PowerMockito.mockStatic; + +@PowerMockIgnore({"com.sun.org.apache.xerces.*", "javax.xml.*", "org.xml.*", "javax.management.*"}) +@RunWith(PowerMockRunner.class) +@PrepareForTest(AuthorityChecker.class) +public class PipeAuthenticationIdentityTest { + + @Test + public void testModifiedCredentialsReplaceInjectedSuperUserId() { + mockStatic(AuthorityChecker.class); + when(AuthorityChecker.getUserId("thulab")).thenReturn(Optional.of(10001L)); + + final Map sourceAttributes = new HashMap<>(); + TableConfigTaskVisitor.checkAndEnrichSourceUser( + "pipe", sourceAttributes, new UserEntity(0L, "root", "host"), false); + + final Map sourceChanges = new HashMap<>(); + sourceChanges.put("username", "thulab"); + sourceChanges.put("password", "password"); + TableConfigTaskVisitor.markSourceAuthenticationAsExplicitIfNecessary(sourceChanges); + new PipeParameters(sourceAttributes) + .addOrReplaceEquivalentAttributes(new PipeParameters(sourceChanges)); + + final Map sinkAttributes = new HashMap<>(); + sinkAttributes.put(PipeSinkConstant.SINK_KEY, "write-back-sink"); + TableConfigTaskVisitor.checkAndEnrichSinkUser( + "pipe", sinkAttributes, new UserEntity(0L, "root", "host"), false); + + final Map sinkChanges = new HashMap<>(); + sinkChanges.put(PipeSinkConstant.CONNECTOR_IOTDB_USER_KEY, "thulab"); + sinkChanges.put(PipeSinkConstant.CONNECTOR_IOTDB_PASSWORD_KEY, "password"); + TableConfigTaskVisitor.markSinkAuthenticationAsExplicitIfNecessary(sinkChanges); + new PipeParameters(sinkAttributes) + .addOrReplaceEquivalentAttributes(new PipeParameters(sinkChanges)); + + assertEquals("10001", sourceAttributes.get(PipeSourceConstant.SOURCE_IOTDB_USER_ID)); + assertEquals( + Boolean.FALSE.toString(), + sourceAttributes.get(SystemConstant.SOURCE_AUTHENTICATION_INJECTED_KEY)); + assertEquals("10001", sinkAttributes.get(PipeSinkConstant.SINK_IOTDB_USER_ID)); + assertEquals( + Boolean.FALSE.toString(), + sinkAttributes.get(SystemConstant.SINK_AUTHENTICATION_INJECTED_KEY)); + } + + @Test + public void testExplicitCredentialsUseTheirOwnUserIdOnCreateAndReplace() { + mockStatic(AuthorityChecker.class); + when(AuthorityChecker.getUserId("thulab")).thenReturn(Optional.of(10001L)); + + for (final boolean isAlter : new boolean[] {false, true}) { + final Map sourceAttributes = new HashMap<>(); + sourceAttributes.put(PipeSourceConstant.EXTRACTOR_IOTDB_USER_KEY, "thulab"); + sourceAttributes.put(PipeSourceConstant.EXTRACTOR_IOTDB_PASSWORD_KEY, "password"); + TableConfigTaskVisitor.checkAndEnrichSourceUser( + "pipe", sourceAttributes, new UserEntity(0L, "root", "host"), isAlter); + + final Map sinkAttributes = new HashMap<>(); + sinkAttributes.put(PipeSinkConstant.SINK_KEY, "write-back-sink"); + sinkAttributes.put(PipeSinkConstant.SINK_IOTDB_USERNAME_KEY, "thulab"); + sinkAttributes.put(PipeSinkConstant.SINK_IOTDB_PASSWORD_KEY, "password"); + TableConfigTaskVisitor.checkAndEnrichSinkUser( + "pipe", sinkAttributes, new UserEntity(0L, "root", "host"), isAlter); + + assertEquals("10001", sourceAttributes.get(PipeSourceConstant.SOURCE_IOTDB_USER_ID)); + assertEquals("10001", sinkAttributes.get(PipeSinkConstant.SINK_IOTDB_USER_ID)); + } + } + + @Test + public void testPasswordOnlyModificationPreservesExistingUserId() { + final Map sourceAttributes = new HashMap<>(); + sourceAttributes.put(PipeSourceConstant.SOURCE_IOTDB_USER_ID, "10001"); + final Map sourceChanges = new HashMap<>(); + sourceChanges.put("password", "new-password"); + TableConfigTaskVisitor.markSourceAuthenticationAsExplicitIfNecessary(sourceChanges); + new PipeParameters(sourceAttributes) + .addOrReplaceEquivalentAttributes(new PipeParameters(sourceChanges)); + + final Map sinkAttributes = new HashMap<>(); + sinkAttributes.put(PipeSinkConstant.SINK_IOTDB_USER_ID, "10001"); + final Map sinkChanges = new HashMap<>(); + sinkChanges.put("password", "new-password"); + TableConfigTaskVisitor.markSinkAuthenticationAsExplicitIfNecessary(sinkChanges); + new PipeParameters(sinkAttributes) + .addOrReplaceEquivalentAttributes(new PipeParameters(sinkChanges)); + + assertEquals("10001", sourceAttributes.get(PipeSourceConstant.SOURCE_IOTDB_USER_ID)); + assertEquals("10001", sinkAttributes.get(PipeSinkConstant.SINK_IOTDB_USER_ID)); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueTest.java index 614934751159..e3de8b312799 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueTest.java @@ -1197,11 +1197,12 @@ public void testPendingQueueUsesSharedConsensusMemoryBudgetAndClearReleasesIt() final double previousQueueRatio = memoryManager.getMaxMemoryRatioForQueue(); final AtomicLongMemoryBlock testMemoryBlock = new AtomicLongMemoryBlock("SubscriptionPendingTest", null, 100L); + final AtomicInteger firstQueueWakeups = new AtomicInteger(); final BlockingQueue firstQueue = (BlockingQueue) constructor.newInstance( 8, - (Runnable) () -> {}, + (Runnable) () -> firstQueueWakeups.incrementAndGet(), (Supplier) () -> SubscriptionQueueRejectionReason.NONE, (Consumer) reason -> {}); @@ -1218,12 +1219,14 @@ public void testPendingQueueUsesSharedConsensusMemoryBudgetAndClearReleasesIt() try { final IndexedConsensusRequest sharedRequest = createSizedRequest(1L, 10L, 30); assertTrue(firstQueue.offer(sharedRequest)); + assertEquals(1, firstQueueWakeups.get()); assertTrue(secondQueue.offer(sharedRequest)); assertEquals(40L, testMemoryBlock.getUsedMemoryInBytes()); assertEquals(40L, retainedBytesMethod.invoke(firstQueue)); assertEquals(40L, retainedBytesMethod.invoke(secondQueue)); assertFalse(firstQueue.offer(createSizedRequest(2L, 10L, 30))); + assertEquals(2, firstQueueWakeups.get()); assertEquals( SubscriptionQueueRejectionReason.CONSENSUS_REQUEST_MEMORY_LIMIT, ((SubscriptionQueueAdmission) firstQueue).getLastRejectionReason());