diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigNodeRuntimeAgent.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigNodeRuntimeAgent.java index f5a864cd974b3..d8636988f9635 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigNodeRuntimeAgent.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigNodeRuntimeAgent.java @@ -24,6 +24,7 @@ import org.apache.iotdb.commons.log.LoggerPeriodicalLogReducer; import org.apache.iotdb.commons.pipe.agent.runtime.PipePeriodicalJobExecutor; import org.apache.iotdb.commons.pipe.agent.runtime.PipePeriodicalPhantomReferenceCleaner; +import org.apache.iotdb.commons.pipe.agent.task.meta.PipeMeta; import org.apache.iotdb.commons.pipe.agent.task.meta.PipeTaskMeta; import org.apache.iotdb.commons.pipe.config.PipeConfig; import org.apache.iotdb.commons.pipe.event.EnrichedEvent; @@ -126,6 +127,11 @@ public void decreaseListenerReference(final PipeParameters parameters) regionListener.decreaseReference(parameters); } + public void reconcileListenerReferences(final Iterable pipeMetas) + throws IllegalPathException { + regionListener.reconcileReferences(pipeMetas); + } + /** Notify the region listener that the leader is ready to allow pipe operations. */ public void notifyLeaderReady() { regionListener.notifyLeaderReady(); diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigRegionListener.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigRegionListener.java index 278e133494b50..3fb5202f11e05 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigRegionListener.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigRegionListener.java @@ -20,6 +20,7 @@ package org.apache.iotdb.confignode.manager.pipe.agent.runtime; import org.apache.iotdb.commons.exception.IllegalPathException; +import org.apache.iotdb.commons.pipe.agent.task.meta.PipeMeta; import org.apache.iotdb.confignode.manager.pipe.source.ConfigRegionListeningFilter; import org.apache.iotdb.confignode.manager.pipe.source.ConfigRegionListeningQueue; import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameters; @@ -57,6 +58,24 @@ public synchronized void decreaseReference(final PipeParameters parameters) } } + public synchronized void reconcileReferences(final Iterable pipeMetas) + throws IllegalPathException { + int referenceCount = 0; + for (final PipeMeta pipeMeta : pipeMetas) { + if (!ConfigRegionListeningFilter.parseListeningPlanTypeSet( + pipeMeta.getStaticMeta().getExtractorParameters()) + .isEmpty()) { + referenceCount++; + } + } + listeningQueueReferenceCount = referenceCount; + if (referenceCount == 0) { + listeningQueue.close(); + } else { + listeningQueue.open(); + } + } + public synchronized boolean isLeaderReady() { return isLeaderReady.get(); } diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/executor/ConfigPlanExecutor.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/executor/ConfigPlanExecutor.java index 60ec748fb5478..07efa5bc24f42 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/executor/ConfigPlanExecutor.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/executor/ConfigPlanExecutor.java @@ -22,6 +22,7 @@ import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.common.rpc.thrift.TSchemaNode; import org.apache.iotdb.commons.auth.AuthException; +import org.apache.iotdb.commons.exception.IllegalPathException; import org.apache.iotdb.commons.path.PartialPath; import org.apache.iotdb.commons.schema.node.MNodeType; import org.apache.iotdb.commons.schema.ttl.TTLCache; @@ -663,9 +664,16 @@ public void loadSnapshot(File latestSnapshotRootDir) { } }); if (result.get()) { - LOGGER.info( - "[ConfigNodeSnapshot] Load snapshot success, latestSnapshotRootDir: {}", - latestSnapshotRootDir); + try { + PipeConfigNodeAgent.runtime() + .reconcileListenerReferences(pipeInfo.getPipeTaskInfo().getPipeMetaList()); + LOGGER.info( + "[ConfigNodeSnapshot] Load snapshot success, latestSnapshotRootDir: {}", + latestSnapshotRootDir); + } catch (final IllegalPathException e) { + result.set(false); + LOGGER.error(e.getMessage(), e); + } } } diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/pipe/PipeInfo.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/pipe/PipeInfo.java index 032534f9161e6..1da72eae7c80a 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/pipe/PipeInfo.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/pipe/PipeInfo.java @@ -290,11 +290,6 @@ public void processLoadSnapshot(final File snapshotDir) throws IOException { try { pipeTaskInfo.processLoadSnapshot(snapshotDir); - - for (final PipeMeta pipeMeta : pipeTaskInfo.getPipeMetaList()) { - PipeConfigNodeAgent.runtime() - .increaseListenerReference(pipeMeta.getStaticMeta().getExtractorParameters()); - } } catch (final Exception ex) { LOGGER.error("Failed to load pipe task info from snapshot", ex); loadPipeTaskInfoException = ex; diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/quota/QuotaInfo.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/quota/QuotaInfo.java index 93d2b27210910..f4b68c03e8380 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/quota/QuotaInfo.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/quota/QuotaInfo.java @@ -258,6 +258,7 @@ public Map getThrottleQuotaLimit() { public void clear() { spaceQuotaLimit.clear(); + spaceQuotaUsage.clear(); throttleQuotaLimit.clear(); } } diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/schema/TemplatePreSetTable.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/schema/TemplatePreSetTable.java index 251100c00718a..bd9d0bcc67b47 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/schema/TemplatePreSetTable.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/schema/TemplatePreSetTable.java @@ -146,6 +146,8 @@ public void processLoadSnapshot(File snapshotDir) throws IOException { try { File snapshotFile = new File(snapshotDir, SNAPSHOT_FILENAME); if (!snapshotFile.exists()) { + // Empty preset tables are represented by the absence of a snapshot file. + templatePreSetMap.clear(); return; } diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/schema/TemplateTable.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/schema/TemplateTable.java index 23cde4e9447d0..04a09632409c2 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/schema/TemplateTable.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/schema/TemplateTable.java @@ -260,9 +260,8 @@ public void processLoadSnapshot(File snapshotDir) throws IOException { BufferedInputStream bufferedInputStream = new BufferedInputStream(fileInputStream)) { // Load snapshot of template this.templateMap.clear(); + this.templateIdMap.clear(); deserialize(bufferedInputStream); - bufferedInputStream.close(); - fileInputStream.close(); } finally { templateReadWriteLock.writeLock().unlock(); } @@ -270,6 +269,13 @@ public void processLoadSnapshot(File snapshotDir) throws IOException { @TestOnly public void clear() { - this.templateMap.clear(); + templateReadWriteLock.writeLock().lock(); + try { + this.templateMap.clear(); + this.templateIdMap.clear(); + this.templateIdGenerator.set(0); + } finally { + templateReadWriteLock.writeLock().unlock(); + } } } diff --git a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigRegionListenerSnapshotTest.java b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigRegionListenerSnapshotTest.java new file mode 100644 index 0000000000000..3d7b6529a3a09 --- /dev/null +++ b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigRegionListenerSnapshotTest.java @@ -0,0 +1,59 @@ +/* + * 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.runtime; + +import org.apache.iotdb.commons.pipe.agent.task.meta.PipeMeta; +import org.apache.iotdb.commons.pipe.agent.task.meta.PipeStaticMeta; +import org.apache.iotdb.commons.pipe.config.constant.PipeSourceConstant; + +import org.junit.Assert; +import org.junit.Test; +import org.mockito.Mockito; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +public class PipeConfigRegionListenerSnapshotTest { + + @Test + public void testReconcileReferencesReplacesSnapshotState() throws Exception { + final Map sourceAttributes = new HashMap<>(); + sourceAttributes.put(PipeSourceConstant.EXTRACTOR_INCLUSION_KEY, "schema.database.create"); + final PipeStaticMeta staticMeta = + new PipeStaticMeta( + "pipe", 1, sourceAttributes, Collections.emptyMap(), Collections.emptyMap()); + final PipeMeta pipeMeta = Mockito.mock(PipeMeta.class); + Mockito.when(pipeMeta.getStaticMeta()).thenReturn(staticMeta); + + final PipeConfigRegionListener listener = new PipeConfigRegionListener(); + listener.increaseReference(staticMeta.getExtractorParameters()); + listener.listener().close(); + + listener.reconcileReferences(Collections.emptyList()); + listener.increaseReference(staticMeta.getExtractorParameters()); + Assert.assertTrue(listener.listener().isOpened()); + + listener.reconcileReferences(Collections.singletonList(pipeMeta)); + listener.reconcileReferences(Collections.singletonList(pipeMeta)); + listener.decreaseReference(staticMeta.getExtractorParameters()); + Assert.assertFalse(listener.listener().isOpened()); + } +} diff --git a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/QuotaInfoTest.java b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/QuotaInfoTest.java index 134a0a203a9a0..c48b766b04640 100644 --- a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/QuotaInfoTest.java +++ b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/QuotaInfoTest.java @@ -37,6 +37,7 @@ import java.io.File; import java.io.IOException; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -97,6 +98,19 @@ public void testSnapshot() throws TException, IOException { QuotaInfo quotaInfo2 = new QuotaInfo(); quotaInfo2.processLoadSnapshot(snapshotDir); + final TSpaceQuota staleQuota = new TSpaceQuota(); + staleQuota.setDeviceNum(1); + staleQuota.setTimeserieNum(1); + staleQuota.setDiskSize(1); + quotaInfo2.setSpaceQuota( + new SetSpaceQuotaPlan(Collections.singletonList("root.stale"), staleQuota)); + Assert.assertTrue(quotaInfo2.getSpaceQuotaUsage().containsKey("root.stale")); + + quotaInfo2.processLoadSnapshot(snapshotDir); + + Assert.assertFalse(quotaInfo2.getSpaceQuotaLimit().containsKey("root.stale")); + Assert.assertFalse(quotaInfo2.getSpaceQuotaUsage().containsKey("root.stale")); + Assert.assertEquals(quotaInfo.getSpaceQuotaLimit(), quotaInfo2.getSpaceQuotaLimit()); Assert.assertEquals(quotaInfo.getThrottleQuotaLimit(), quotaInfo2.getThrottleQuotaLimit()); } diff --git a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/schema/TemplatePreSetTableTest.java b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/schema/TemplatePreSetTableTest.java index 769359525dfde..9d90a71a97ad9 100644 --- a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/schema/TemplatePreSetTableTest.java +++ b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/schema/TemplatePreSetTableTest.java @@ -104,11 +104,24 @@ public void testSnapshot() throws IllegalPathException { newTemplatePreSetTable = new TemplatePreSetTable(); newTemplatePreSetTable.processLoadSnapshot(snapshotDir); - Assert.assertTrue(templatePreSetTable.isPreSet(templateId1, templateSetPath1)); - Assert.assertTrue(templatePreSetTable.isPreSet(templateId2, templateSetPath1)); - Assert.assertTrue(templatePreSetTable.isPreSet(templateId2, templateSetPath2)); + Assert.assertTrue(newTemplatePreSetTable.isPreSet(templateId1, templateSetPath1)); + Assert.assertTrue(newTemplatePreSetTable.isPreSet(templateId2, templateSetPath1)); + Assert.assertTrue(newTemplatePreSetTable.isPreSet(templateId2, templateSetPath2)); } catch (IOException e) { Assert.fail(); } } + + @Test + public void testEmptySnapshotReplacesExistingState() throws IOException, IllegalPathException { + final int templateId = 5; + final PartialPath templateSetPath = new PartialPath("root.db.t1"); + + Assert.assertTrue(templatePreSetTable.processTakeSnapshot(snapshotDir)); + templatePreSetTable.preSetTemplate(templateId, templateSetPath); + + templatePreSetTable.processLoadSnapshot(snapshotDir); + + Assert.assertFalse(templatePreSetTable.isPreSet(templateId, templateSetPath)); + } } diff --git a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/schema/TemplateTableTest.java b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/schema/TemplateTableTest.java index c4b8aab849abe..2f3c2bfd6d4e7 100644 --- a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/schema/TemplateTableTest.java +++ b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/schema/TemplateTableTest.java @@ -77,7 +77,8 @@ public void testSnapshot() throws IOException, MetadataException { } templateTable.processTakeSnapshot(snapshotDir); - templateTable.clear(); + Template staleTemplate = newSchemaTemplate("stale_template"); + templateTable.createTemplate(staleTemplate); templateTable.processLoadSnapshot(snapshotDir); // show nodes in schemaengine template @@ -86,6 +87,13 @@ public void testSnapshot() throws IOException, MetadataException { Template template = templates.get(i); Assert.assertEquals(template, templateTable.getTemplate(templateNameTmp)); } + + try { + templateTable.getTemplate(staleTemplate.getId()); + Assert.fail("Template created after the snapshot should be removed when loading it"); + } catch (MetadataException expected) { + // expected + } } @Test diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSink.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSink.java index ed6f127213c2b..8226eeac0f73e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSink.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSink.java @@ -256,19 +256,24 @@ private void doTransfer( final List sealedFiles = batchToTransfer.sealTsFiles(); final Map, Double> pipe2WeightMap = batchToTransfer.deepCopyPipe2WeightMap(); - for (final File tsFile : sealedFiles) { - doTransfer(pipe2WeightMap, socket, tsFile, null, null, tsFile.getName()); - try { - RetryUtils.retryOnException( - () -> { - FileUtils.delete(tsFile); - return null; - }); - } catch (final NoSuchFileException e) { - LOGGER.info("The file {} is not found, may already be deleted.", tsFile); - } catch (final Exception e) { - LOGGER.warn( - "Failed to delete batch file {}, this file should be deleted manually later", tsFile); + try { + for (final File tsFile : sealedFiles) { + doTransfer(pipe2WeightMap, socket, tsFile, null, null, tsFile.getName()); + } + } finally { + for (final File tsFile : sealedFiles) { + try { + RetryUtils.retryOnException( + () -> { + FileUtils.delete(tsFile); + return null; + }); + } catch (final NoSuchFileException e) { + LOGGER.info("The file {} is not found, may already be deleted.", tsFile); + } catch (final Exception e) { + LOGGER.warn( + "Failed to delete batch file {}, this file should be deleted manually later", tsFile); + } } } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java index 77d3d02864f62..abc10f4fcc275 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java @@ -64,6 +64,7 @@ import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq; import com.google.common.collect.ImmutableSet; +import org.apache.commons.io.FileUtils; import org.apache.tsfile.exception.write.WriteProcessException; import org.apache.tsfile.utils.Pair; import org.slf4j.Logger; @@ -249,6 +250,7 @@ private void transferInBatchWithoutCheck( final AtomicInteger eventsReferenceCount = new AtomicInteger(sealedFiles.size()); final AtomicBoolean eventsHadBeenAddedToRetryQueue = new AtomicBoolean(false); + int transferredFileCount = 0; try { for (final File sealedFile : sealedFiles) { transfer( @@ -262,8 +264,12 @@ private void transferInBatchWithoutCheck( null, false, null)); + transferredFileCount++; } } catch (final Exception e) { + for (int i = transferredFileCount; i < sealedFiles.size(); i++) { + FileUtils.deleteQuietly(sealedFiles.get(i)); + } PipeLogger.log(LOGGER::warn, e, "Failed to transfer tsfile batch (%s).", sealedFiles); if (eventsHadBeenAddedToRetryQueue.compareAndSet(false, true)) { addFailureEventsToRetryQueue(events, e); @@ -421,22 +427,28 @@ private boolean transferWithoutCheck(final TsFileInsertionEvent tsFileInsertionE private void transfer(final PipeTransferTsFileHandler pipeTransferTsFileHandler) { transferTsFileCounter.incrementAndGet(); - CompletableFuture completableFuture = - CompletableFuture.supplyAsync( - () -> { - AsyncPipeDataTransferServiceClient client = null; - try { - client = transferTsFileClientManager.borrowClient(); - pipeTransferTsFileHandler.transfer(transferTsFileClientManager, client); - } catch (final Exception ex) { - logOnClientException(client, ex); - pipeTransferTsFileHandler.onError(ex); - } finally { - transferTsFileCounter.decrementAndGet(); - } - return null; - }, - transferTsFileClientManager.getExecutor()); + final CompletableFuture completableFuture; + try { + completableFuture = + CompletableFuture.supplyAsync( + () -> { + AsyncPipeDataTransferServiceClient client = null; + try { + client = transferTsFileClientManager.borrowClient(); + pipeTransferTsFileHandler.transfer(transferTsFileClientManager, client); + } catch (final Exception ex) { + logOnClientException(client, ex); + pipeTransferTsFileHandler.onError(ex); + } finally { + transferTsFileCounter.decrementAndGet(); + } + return null; + }, + transferTsFileClientManager.getExecutor()); + } catch (final RuntimeException e) { + transferTsFileCounter.decrementAndGet(); + throw e; + } if (PipeConfig.getInstance().isTransferTsFileSync()) { try { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandler.java index ba9f1e54b1e24..b6936485770f1 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandler.java @@ -38,6 +38,7 @@ import org.apache.iotdb.db.pipe.sink.payload.evolvable.request.PipeTransferTsFileSealReq; import org.apache.iotdb.db.pipe.sink.payload.evolvable.request.PipeTransferTsFileSealWithModReq; import org.apache.iotdb.db.pipe.sink.protocol.thrift.async.IoTDBDataRegionAsyncSink; +import org.apache.iotdb.pipe.api.exception.PipeConnectionException; import org.apache.iotdb.pipe.api.exception.PipeException; import org.apache.iotdb.rpc.TSStatusCode; import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq; @@ -154,6 +155,7 @@ public void transfer( "Client has been returned to the pool. Current handler status is %s. Will not transfer %s.", sink.isClosed() ? "CLOSED" : "NOT CLOSED", tsFile); + onError(new PipeConnectionException("Client has been returned to the pool.")); return; } @@ -469,8 +471,26 @@ public void clearEventsReferenceCount() { @Override public void close() { - super.close(); - releaseReadBufferMemoryBlock(); + try { + if (reader != null) { + reader.close(); + reader = null; + } + + if (currentFile.exists() + && events.stream().anyMatch(event -> !(event instanceof PipeTsFileInsertionEvent))) { + RetryUtils.retryOnException( + () -> { + FileUtils.delete(currentFile); + return null; + }); + } + } catch (final IOException e) { + LOGGER.warn("Failed to close file reader or delete generated batch file.", e); + } finally { + super.close(); + releaseReadBufferMemoryBlock(); + } } private void releaseReadBufferMemoryBlock() { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/sync/IoTDBDataRegionSyncSink.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/sync/IoTDBDataRegionSyncSink.java index 2d3ec89756ef1..7a1fe1463f638 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/sync/IoTDBDataRegionSyncSink.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/sync/IoTDBDataRegionSyncSink.java @@ -279,19 +279,24 @@ private void doTransfer(final PipeTabletEventTsFileBatch batchToTransfer) final List sealedFiles = batchToTransfer.sealTsFiles(); final Map, Double> pipe2WeightMap = batchToTransfer.deepCopyPipe2WeightMap(); - for (final File tsFile : sealedFiles) { - doTransfer(pipe2WeightMap, tsFile, null, null); - try { - RetryUtils.retryOnException( - () -> { - FileUtils.delete(tsFile); - return null; - }); - } catch (final NoSuchFileException e) { - LOGGER.info("The file {} is not found, may already be deleted.", tsFile); - } catch (final Exception e) { - LOGGER.warn( - "Failed to delete batch file {}, this file should be deleted manually later", tsFile); + try { + for (final File tsFile : sealedFiles) { + doTransfer(pipe2WeightMap, tsFile, null, null); + } + } finally { + for (final File tsFile : sealedFiles) { + try { + RetryUtils.retryOnException( + () -> { + FileUtils.delete(tsFile); + return null; + }); + } catch (final NoSuchFileException e) { + LOGGER.info("The file {} is not found, may already be deleted.", tsFile); + } catch (final Exception e) { + LOGGER.warn( + "Failed to delete batch file {}, this file should be deleted manually later", tsFile); + } } } } 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 5568aa4015c6d..47032a273a80f 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 @@ -709,7 +709,7 @@ protected void upgradeAndUpdateDeviceLastFlushTime( // checked above //noinspection OptionalGetWithoutIsPresent long endTime = resource.getEndTime(deviceId).get(); - endTimeMap.put(deviceId, endTime); + endTimeMap.merge(deviceId, endTime, Math::max); } } if (config.isEnableSeparateData()) { @@ -1275,6 +1275,10 @@ private boolean insertTabletToTsFileProcessor( try { tsFileProcessor.insertTablet(insertTabletNode, start, end, results); } catch (WriteProcessRejectException e) { + final TSStatus failureStatus = RpcUtils.getStatus(e.getErrorCode(), e.getMessage()); + for (int i = start; i < end; i++) { + results[i] = failureStatus; + } logger.warn("insert to TsFileProcessor rejected, {}", e.getMessage()); return false; } catch (WriteProcessException e) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/HashLastFlushTimeMap.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/HashLastFlushTimeMap.java index f29c158cc0dfd..2452cfd011125 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/HashLastFlushTimeMap.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/HashLastFlushTimeMap.java @@ -152,6 +152,7 @@ public long getFlushedTime(long timePartitionId, IDeviceID deviceId) { @Override public void clearFlushedTime() { partitionLatestFlushedTime.clear(); + memCostForEachPartition.clear(); } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/TsFileProcessor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/TsFileProcessor.java index 6bbcb427f8798..da62e80252bab 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/TsFileProcessor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/TsFileProcessor.java @@ -376,7 +376,13 @@ public void insert(InsertRowsNode insertRowsNode, long[] costsForMetrics) } } long[] alignedMemIncrements = checkAlignedMemCostAndAddToTspInfoForRows(alignedList); - long[] nonAlignedMemIncrements = checkMemCostAndAddToTspInfoForRows(nonAlignedList); + final long[] nonAlignedMemIncrements; + try { + nonAlignedMemIncrements = checkMemCostAndAddToTspInfoForRows(nonAlignedList); + } catch (final WriteProcessException e) { + rollbackMemoryInfoIfNeeded(alignedMemIncrements); + throw e; + } memIncrements = new long[3]; for (int i = 0; i < 3; i++) { memIncrements[i] = alignedMemIncrements[i] + nonAlignedMemIncrements[i]; @@ -1034,6 +1040,15 @@ private void rollbackMemoryInfo(long[] memIncrements) { workMemTable.releaseTextDataSize(textDataIncrement); } + private void rollbackMemoryInfoIfNeeded(final long[] memIncrements) { + for (final long memIncrement : memIncrements) { + if (memIncrement != 0) { + rollbackMemoryInfo(memIncrements); + return; + } + } + } + /** * Delete data which belongs to the timeseries `deviceId.measurementId` and the timestamp of which * <= 'timestamp' in the deletion.
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTsFileEventBatch.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTsFileEventBatch.java index 8e3eee58ad04f..47ee100f7c9a3 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTsFileEventBatch.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTsFileEventBatch.java @@ -20,6 +20,7 @@ package org.apache.iotdb.db.subscription.event.batch; import org.apache.iotdb.commons.pipe.event.EnrichedEvent; +import org.apache.iotdb.commons.utils.TestOnly; import org.apache.iotdb.db.pipe.sink.payload.evolvable.batch.PipeTabletEventTsFileBatch; import org.apache.iotdb.db.subscription.broker.SubscriptionPrefetchingTsFileQueue; import org.apache.iotdb.db.subscription.event.SubscriptionEvent; @@ -28,6 +29,7 @@ import org.apache.iotdb.pipe.api.event.dml.insertion.TsFileInsertionEvent; import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext; +import org.apache.commons.io.FileUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -42,6 +44,7 @@ public class SubscriptionPipeTsFileEventBatch extends SubscriptionPipeEventBatch LoggerFactory.getLogger(SubscriptionPipeTsFileEventBatch.class); private final PipeTabletEventTsFileBatch batch; + private final List sealedFiles = new ArrayList<>(); public SubscriptionPipeTsFileEventBatch( final int regionId, @@ -52,6 +55,17 @@ public SubscriptionPipeTsFileEventBatch( this.batch = new PipeTabletEventTsFileBatch(maxDelayInMs, maxBatchSizeInBytes); } + @TestOnly + SubscriptionPipeTsFileEventBatch( + final int regionId, + final SubscriptionPrefetchingTsFileQueue prefetchingQueue, + final int maxDelayInMs, + final long maxBatchSizeInBytes, + final PipeTabletEventTsFileBatch batch) { + super(regionId, prefetchingQueue, maxDelayInMs, maxBatchSizeInBytes); + this.batch = batch; + } + @Override public synchronized void ack() { batch.decreaseEventsReferenceCount(this.getClass().getName(), true); @@ -59,9 +73,14 @@ public synchronized void ack() { @Override public synchronized void cleanUp(final boolean force) { - // close batch, it includes clearing the reference count of events - batch.close(); - enrichedEvents.clear(); + try { + // close batch, it includes clearing the reference count of events + batch.close(); + } finally { + sealedFiles.forEach(FileUtils::deleteQuietly); + sealedFiles.clear(); + enrichedEvents.clear(); + } } /////////////////////////////// utility /////////////////////////////// @@ -95,6 +114,7 @@ protected List generateSubscriptionEvents() throws Exception final List events = new ArrayList<>(); final List tsFiles = batch.sealTsFiles(); + sealedFiles.addAll(tsFiles); final AtomicInteger ackReferenceCount = new AtomicInteger(tsFiles.size()); final AtomicInteger cleanReferenceCount = new AtomicInteger(tsFiles.size()); for (final File tsFile : tsFiles) { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandlerCleanupTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandlerCleanupTest.java new file mode 100644 index 0000000000000..aca597f0463ff --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandlerCleanupTest.java @@ -0,0 +1,86 @@ +/* + * 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.pipe.sink.protocol.thrift.async.handler; + +import org.apache.iotdb.commons.pipe.event.EnrichedEvent; +import org.apache.iotdb.db.pipe.event.common.tsfile.PipeTsFileInsertionEvent; +import org.apache.iotdb.db.pipe.sink.protocol.thrift.async.IoTDBDataRegionAsyncSink; + +import org.junit.Assert; +import org.junit.Test; +import org.mockito.Mockito; + +import java.io.File; +import java.nio.file.Files; +import java.util.Collections; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +public class PipeTransferTsFileHandlerCleanupTest { + + @Test + public void testCloseDeletesBatchFile() throws Exception { + final File file = Files.createTempFile("pipe-transfer-batch", ".tsfile").toFile(); + final EnrichedEvent event = Mockito.mock(EnrichedEvent.class); + + createHandler(file, event).close(); + + Assert.assertFalse(file.exists()); + } + + @Test + public void testNullClientDeletesBatchFile() throws Exception { + final File file = Files.createTempFile("pipe-transfer-null-client", ".tsfile").toFile(); + final EnrichedEvent event = Mockito.mock(EnrichedEvent.class); + final PipeTransferTsFileHandler handler = createHandler(file, event); + + handler.transfer(null, null); + + Assert.assertFalse(file.exists()); + } + + @Test + public void testCloseKeepsSourceTsFile() throws Exception { + final File file = Files.createTempFile("pipe-transfer-source", ".tsfile").toFile(); + final PipeTsFileInsertionEvent event = Mockito.mock(PipeTsFileInsertionEvent.class); + try { + createHandler(file, event).close(); + Assert.assertTrue(file.exists()); + } finally { + if (file.exists()) { + Assert.assertTrue(file.delete()); + } + } + } + + private PipeTransferTsFileHandler createHandler(final File file, final EnrichedEvent event) + throws Exception { + return new PipeTransferTsFileHandler( + Mockito.mock(IoTDBDataRegionAsyncSink.class), + Collections.emptyMap(), + Collections.singletonList(event), + new AtomicInteger(1), + new AtomicBoolean(false), + file, + null, + false, + null); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/LastFlushTimeMapTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/LastFlushTimeMapTest.java index b260445a5bd76..877c71b15006c 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/LastFlushTimeMapTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/LastFlushTimeMapTest.java @@ -39,9 +39,13 @@ import org.junit.Assert; import org.junit.Before; import org.junit.Test; +import org.mockito.Mockito; import java.io.File; import java.io.IOException; +import java.util.Arrays; +import java.util.Collections; +import java.util.Optional; public class LastFlushTimeMapTest { @@ -130,6 +134,35 @@ record = new TSRecord(j, "root.vehicle.d0"); dataRegion.getLastFlushTimeMap().getFlushedTime(0, new PlainDeviceID("root.vehicle.d1"))); } + @Test + public void testClearFlushedTimeClearsMemoryCost() { + HashLastFlushTimeMap lastFlushTimeMap = new HashLastFlushTimeMap(); + IDeviceID device = new PlainDeviceID("root.vehicle.d0"); + + lastFlushTimeMap.updateMultiDeviceFlushedTime(0, Collections.singletonMap(device, 100L)); + Assert.assertTrue(lastFlushTimeMap.getMemSize(0) > 0); + lastFlushTimeMap.clearFlushedTime(); + Assert.assertEquals(0, lastFlushTimeMap.getMemSize(0)); + } + + @Test + public void testUpgradeDeviceFlushTimeUsesMaximumAcrossResources() { + final IDeviceID device = new PlainDeviceID("root.vehicle.d0"); + final TsFileResource newerEndTimeResource = Mockito.mock(TsFileResource.class); + final TsFileResource olderEndTimeResource = Mockito.mock(TsFileResource.class); + Mockito.when(newerEndTimeResource.getDevices()).thenReturn(Collections.singleton(device)); + Mockito.when(newerEndTimeResource.getEndTime(device)).thenReturn(Optional.of(100L)); + Mockito.when(olderEndTimeResource.getDevices()).thenReturn(Collections.singleton(device)); + Mockito.when(olderEndTimeResource.getEndTime(device)).thenReturn(Optional.of(50L)); + + dataRegion.getLastFlushTimeMap().clearFlushedTime(); + dataRegion.getLastFlushTimeMap().checkAndCreateFlushedTimePartition(0, true); + dataRegion.upgradeAndUpdateDeviceLastFlushTime( + 0, Arrays.asList(newerEndTimeResource, olderEndTimeResource)); + + Assert.assertEquals(100L, dataRegion.getLastFlushTimeMap().getFlushedTime(0, device)); + } + @Test public void testRecoverLastFlushTimeMap() throws IOException, IllegalPathException, WriteProcessException, DataRegionException { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTsFileEventBatchCleanupTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTsFileEventBatchCleanupTest.java new file mode 100644 index 0000000000000..9496ac6169f91 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTsFileEventBatchCleanupTest.java @@ -0,0 +1,76 @@ +/* + * 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.subscription.event.batch; + +import org.apache.iotdb.db.pipe.sink.payload.evolvable.batch.PipeTabletEventTsFileBatch; +import org.apache.iotdb.db.subscription.broker.SubscriptionPrefetchingTsFileQueue; +import org.apache.iotdb.db.subscription.event.SubscriptionEvent; +import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext; + +import org.apache.commons.io.FileUtils; +import org.junit.Assert; +import org.junit.Test; +import org.mockito.Mockito; + +import java.io.File; +import java.nio.file.Files; +import java.util.Arrays; +import java.util.List; + +public class SubscriptionPipeTsFileEventBatchCleanupTest { + + @Test + public void testDeletesSealedFilesAfterAllSharedEventsAreCleaned() throws Exception { + final File temporaryDirectory = + Files.createTempDirectory("subscription-tsfile-cleanup").toFile(); + final File firstFile = new File(temporaryDirectory, "first.tsfile"); + final File secondFile = new File(temporaryDirectory, "second.tsfile"); + Assert.assertTrue(firstFile.createNewFile()); + Assert.assertTrue(secondFile.createNewFile()); + + final PipeTabletEventTsFileBatch innerBatch = Mockito.mock(PipeTabletEventTsFileBatch.class); + final SubscriptionPrefetchingTsFileQueue queue = + Mockito.mock(SubscriptionPrefetchingTsFileQueue.class); + Mockito.when(innerBatch.isEmpty()).thenReturn(false); + Mockito.when(innerBatch.sealTsFiles()).thenReturn(Arrays.asList(firstFile, secondFile)); + Mockito.when(queue.generateSubscriptionCommitContext()) + .thenReturn( + new SubscriptionCommitContext(1, 1, "topic", "group", 1), + new SubscriptionCommitContext(1, 1, "topic", "group", 2)); + + try { + final SubscriptionPipeTsFileEventBatch batch = + new SubscriptionPipeTsFileEventBatch(1, queue, 1, 1, innerBatch); + final List events = batch.generateSubscriptionEvents(); + + Assert.assertEquals(2, events.size()); + events.get(0).cleanUp(false); + Assert.assertTrue(firstFile.exists()); + Assert.assertTrue(secondFile.exists()); + + events.get(1).cleanUp(false); + Assert.assertFalse(firstFile.exists()); + Assert.assertFalse(secondFile.exists()); + Mockito.verify(innerBatch, Mockito.times(1)).close(); + } finally { + FileUtils.deleteDirectory(temporaryDirectory); + } + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/datastructure/queue/ConcurrentIterableLinkedQueue.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/datastructure/queue/ConcurrentIterableLinkedQueue.java index 306b5491f32ee..95608a38e58d3 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/datastructure/queue/ConcurrentIterableLinkedQueue.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/datastructure/queue/ConcurrentIterableLinkedQueue.java @@ -136,6 +136,7 @@ public long tryRemoveBefore(long newFirstIndex) { if (firstNode == null) { firstNode = pilotNode; lastNode = pilotNode; + pilotNode.next = null; } // Update iterators if necessary @@ -207,9 +208,7 @@ public void setFirstIndex(final long firstIndex) { lock.writeLock().lock(); try { this.firstIndex = firstIndex; - if (tailIndex < firstIndex) { - tailIndex = firstIndex; - } + tailIndex = firstIndex; } finally { lock.writeLock().unlock(); } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/datastructure/queue/listening/AbstractSerializableListeningQueue.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/datastructure/queue/listening/AbstractSerializableListeningQueue.java index f398bf7e0abda..1371d856c5071 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/datastructure/queue/listening/AbstractSerializableListeningQueue.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/datastructure/queue/listening/AbstractSerializableListeningQueue.java @@ -138,7 +138,7 @@ public synchronized void deserializeFromFile(final File snapshotName) throws IOE return; } - queue.clear(); + close(); try (final FileInputStream inputStream = new FileInputStream(snapshotFile)) { isClosed.set(ReadWriteIOUtils.readBool(inputStream)); final QueueSerializerType type = diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/datastructure/ConcurrentIterableLinkedQueueTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/datastructure/ConcurrentIterableLinkedQueueTest.java index a7a7310b113b8..60917a710c0b3 100644 --- a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/datastructure/ConcurrentIterableLinkedQueueTest.java +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/datastructure/ConcurrentIterableLinkedQueueTest.java @@ -480,4 +480,25 @@ public void testIteratorSeek() { Assert.assertTrue(itr.hasNext()); Assert.assertEquals(Integer.valueOf(5), itr.next()); } + + @Test(timeout = 60000) + public void testSetFirstIndexReplacesEmptyQueueIndexes() { + queue.add(1); + queue.add(2); + queue.clear(); + + queue.setFirstIndex(1); + + Assert.assertEquals(1, queue.getFirstIndex()); + Assert.assertEquals(1, queue.getTailIndex()); + queue.add(3); + Assert.assertEquals(2, queue.getTailIndex()); + + try (final ConcurrentIterableLinkedQueue.DynamicIterator itr = + queue.iterateFromEarliest()) { + Assert.assertTrue(itr.hasNext()); + Assert.assertEquals(Integer.valueOf(3), itr.next()); + Assert.assertFalse(itr.hasNext()); + } + } } diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/datastructure/queue/listening/AbstractPipeListeningQueueTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/datastructure/queue/listening/AbstractPipeListeningQueueTest.java new file mode 100644 index 0000000000000..10755b4178731 --- /dev/null +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/datastructure/queue/listening/AbstractPipeListeningQueueTest.java @@ -0,0 +1,91 @@ +/* + * 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.commons.pipe.datastructure.queue.listening; + +import org.apache.iotdb.commons.pipe.event.EnrichedEvent; +import org.apache.iotdb.commons.pipe.event.PipeSnapshotEvent; +import org.apache.iotdb.pipe.api.event.Event; + +import org.junit.Assert; +import org.junit.Test; +import org.mockito.Mockito; + +import java.io.File; +import java.nio.ByteBuffer; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Collections; +import java.util.List; + +public class AbstractPipeListeningQueueTest { + + @Test + public void testDeserializeReleasesExistingQueueAndSnapshotCache() throws Exception { + final Path temporaryDirectory = Files.createTempDirectory("pipe-listening-queue-snapshot"); + final File snapshotFile = temporaryDirectory.resolve("queue.snapshot").toFile(); + final TestPipeListeningQueue emptyQueue = new TestPipeListeningQueue(); + final TestPipeListeningQueue queue = new TestPipeListeningQueue(); + final EnrichedEvent queuedEvent = Mockito.mock(EnrichedEvent.class); + final PipeSnapshotEvent cachedSnapshot = Mockito.mock(PipeSnapshotEvent.class); + try { + Assert.assertTrue(emptyQueue.serializeToFile(snapshotFile)); + + queue.open(); + queue.addEvent(queuedEvent); + queue.addSnapshots(Collections.singletonList(cachedSnapshot)); + Assert.assertEquals(1, queue.getSize()); + Assert.assertEquals(1, queue.findAvailableSnapshots(false).getRight().size()); + + queue.deserializeFromFile(snapshotFile); + + Mockito.verify(queuedEvent) + .decreaseReferenceCount(AbstractPipeListeningQueue.class.getName(), false); + Mockito.verify(cachedSnapshot) + .decreaseReferenceCount(AbstractPipeListeningQueue.class.getName(), false); + Assert.assertEquals(0, queue.getSize()); + Assert.assertTrue(queue.findAvailableSnapshots(false).getRight().isEmpty()); + Assert.assertFalse(queue.isOpened()); + } finally { + Files.deleteIfExists(snapshotFile.toPath()); + Files.deleteIfExists(temporaryDirectory); + } + } + + private static final class TestPipeListeningQueue extends AbstractPipeListeningQueue { + + private void addEvent(final EnrichedEvent event) { + tryListen(event); + } + + private void addSnapshots(final List snapshots) { + tryListen(snapshots); + } + + @Override + protected ByteBuffer serializeToByteBuffer(final Event event) { + return ByteBuffer.allocate(0); + } + + @Override + protected Event deserializeFromByteBuffer(final ByteBuffer byteBuffer) { + return null; + } + } +}