From c097c9e76cbe9ffeac814570a995c2ccd85b570c Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Sat, 10 Oct 2026 19:00:39 +0800 Subject: [PATCH] [Subscription] Release exhausted TsFile parsers before batch ACK --- .../SubscriptionPipeTabletEventBatch.java | 21 +- .../SubscriptionPipeTabletEventBatchTest.java | 256 ++++++++++++++++++ 2 files changed, 271 insertions(+), 6 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTabletEventBatchTest.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTabletEventBatch.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTabletEventBatch.java index 67321b0465168..1d3987a6dd277 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTabletEventBatch.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTabletEventBatch.java @@ -405,9 +405,7 @@ public synchronized boolean hasNext() { if (currentTabletInsertionEventsIterator.hasNext()) { return true; } else { - // reset - currentTabletInsertionEventsIterator = null; - currentTsFileInsertionEvent = null; + finishCurrentTsFileInsertionEvent(); return hasNext(); } } @@ -467,12 +465,11 @@ private Pair> nextInternal() { } } if (!currentTabletInsertionEventsIterator.hasNext()) { - iterationSnapshot.addIteratedEnrichedEvent((EnrichedEvent) currentTsFileInsertionEvent); + finishCurrentTsFileInsertionEvent(); } return convertToTablets(tabletInsertionEvent); } else { - currentTabletInsertionEventsIterator = null; - currentTsFileInsertionEvent = null; + finishCurrentTsFileInsertionEvent(); } } @@ -517,4 +514,16 @@ private Pair> nextInternal() { return null; } } + + private void finishCurrentTsFileInsertionEvent() { + final PipeTsFileInsertionEvent event = (PipeTsFileInsertionEvent) currentTsFileInsertionEvent; + iterationSnapshot.addIteratedEnrichedEvent(event); + + // Release the parser reservation before opening the next file in the same region. Waiting for + // the whole batch to be ACKed can block that file behind this exhausted parser. The parsed + // tablet events remain referenced by the iteration snapshot until ACK or cleanup. + event.close(); + currentTabletInsertionEventsIterator = null; + currentTsFileInsertionEvent = null; + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTabletEventBatchTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTabletEventBatchTest.java new file mode 100644 index 0000000000000..d0d67f4e9ac76 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/event/batch/SubscriptionPipeTabletEventBatchTest.java @@ -0,0 +1,256 @@ +/* + * 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.commons.conf.CommonConfig; +import org.apache.iotdb.commons.conf.CommonDescriptor; +import org.apache.iotdb.commons.pipe.datastructure.pattern.TablePattern; +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.memory.PipeMemoryManager; +import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryManager.TsFileParserMemoryReservation; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResourceStatus; +import org.apache.iotdb.db.subscription.broker.SubscriptionPrefetchingTabletQueue; +import org.apache.iotdb.db.subscription.event.pipe.SubscriptionPipeTabletBatchEvents; + +import org.apache.tsfile.enums.ColumnCategory; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.external.commons.io.FileUtils; +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.file.metadata.TableSchema; +import org.apache.tsfile.utils.Pair; +import org.apache.tsfile.write.TsFileWriter; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mockito; + +import java.io.File; +import java.nio.file.Files; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +public class SubscriptionPipeTabletEventBatchTest { + + private static final String REGION_ID = "7"; + + private final CommonConfig config = CommonDescriptor.getInstance().getConfig(); + private final PipeMemoryManager memoryManager = PipeDataNodeResourceManager.memory(); + private final List tsFileEvents = new ArrayList<>(); + + private long originalParserMemory; + private int originalGlobalLimit; + private int originalRegionLimit; + private File temporaryDirectory; + private SubscriptionPipeTabletEventBatch batch; + + @Before + public void setUp() throws Exception { + originalParserMemory = config.getPipeTsFileParserMemory(); + originalGlobalLimit = config.getPipeTsFileParserInFlightMaxNum(); + originalRegionLimit = config.getPipeTsFileParserInFlightMaxNumPerPipeRegion(); + config.setPipeTsFileParserMemory(1); + config.setPipeTsFileParserInFlightMaxNum(4); + config.setPipeTsFileParserInFlightMaxNumPerPipeRegion(1); + + temporaryDirectory = Files.createTempDirectory("subscription-tablet-batch").toFile(); + final SubscriptionPrefetchingTabletQueue queue = + Mockito.mock(SubscriptionPrefetchingTabletQueue.class); + Mockito.when(queue.getTopicName()).thenReturn("topic"); + Mockito.when(queue.getConsumerGroupId()).thenReturn("group"); + batch = + new SubscriptionPipeTabletEventBatch(Integer.parseInt(REGION_ID), queue, 20, 1024 * 1024); + } + + @After + public void tearDown() throws Exception { + if (batch != null) { + batch.cleanUp(true); + } + for (final PipeTsFileInsertionEvent event : tsFileEvents) { + event.close(); + } + config.setPipeTsFileParserMemory(originalParserMemory); + config.setPipeTsFileParserInFlightMaxNum(originalGlobalLimit); + config.setPipeTsFileParserInFlightMaxNumPerPipeRegion(originalRegionLimit); + if (temporaryDirectory != null) { + FileUtils.deleteDirectory(temporaryDirectory); + } + } + + @Test(timeout = 10000) + public void testSameRegionFilesReleaseParserBeforeBatchAck() throws Exception { + final long baselineTabletMemory = memoryManager.getUsedMemorySizeInBytesOfTablets(); + batch.enrichedEvents.add(createTsFileEvent("first.tsfile", 1, false, false)); + batch.enrichedEvents.add(createTsFileEvent("second.tsfile", 2, false, false)); + batch.resetForIteration(); + + Assert.assertTrue(batch.hasNext()); + final Pair> first = batch.next(); + assertTablet(first, "table_0", 1); + Assert.assertTrue("The exhausted file must release its parser before ACK", canReserveParser()); + + Assert.assertTrue(batch.hasNext()); + assertTablet(batch.next(), "table_0", 2); + Assert.assertFalse(batch.hasNext()); + Assert.assertTrue(canReserveParser()); + + // Closing a parser must leave its materialized tablets available until the response is ACKed. + assertTablet(first, "table_0", 1); + final SubscriptionPipeTabletIterationSnapshot snapshot = batch.sendIterationSnapshot(); + Assert.assertEquals(2, snapshot.getIteratedEnrichedEvents().size()); + snapshot.ack(); + Assert.assertEquals(baselineTabletMemory, memoryManager.getUsedMemorySizeInBytesOfTablets()); + } + + @Test(timeout = 10000) + public void testFilteredFileReleasesParserWithoutProducingTablet() throws Exception { + final long baselineTabletMemory = memoryManager.getUsedMemorySizeInBytesOfTablets(); + batch.enrichedEvents.add(createTsFileEvent("filtered.tsfile", 1, true, false)); + batch.resetForIteration(); + + Assert.assertTrue(batch.hasNext()); + Assert.assertNull(batch.next()); + Assert.assertTrue("An empty iteration must also release its parser", canReserveParser()); + Assert.assertFalse(batch.hasNext()); + + final SubscriptionPipeTabletIterationSnapshot snapshot = batch.sendIterationSnapshot(); + Assert.assertEquals(1, snapshot.getIteratedEnrichedEvents().size()); + snapshot.ack(); + Assert.assertEquals(baselineTabletMemory, memoryManager.getUsedMemorySizeInBytesOfTablets()); + } + + @Test(timeout = 10000) + public void testDetachedSnapshotCleanupClosesPartiallyParsedFile() throws Exception { + final long baselineTabletMemory = memoryManager.getUsedMemorySizeInBytesOfTablets(); + batch.enrichedEvents.add(createTsFileEvent("partial.tsfile", 1, false, true)); + batch.resetForIteration(); + + Assert.assertNotNull(batch.next()); + Assert.assertTrue(batch.hasNext()); + Assert.assertFalse("A file still being parsed must retain its parser", canReserveParser()); + + final SubscriptionPipeTabletBatchEvents events = new SubscriptionPipeTabletBatchEvents(batch); + events.receiveIterationSnapshot(batch.sendIterationSnapshot()); + events.cleanUp(true); + Assert.assertTrue(canReserveParser()); + Assert.assertFalse(batch.hasNext()); + Assert.assertEquals(baselineTabletMemory, memoryManager.getUsedMemorySizeInBytesOfTablets()); + } + + private boolean canReserveParser() { + final TsFileParserMemoryReservation reservation = new TsFileParserMemoryReservation(); + final boolean reserved = + memoryManager.tryReserveTsFileParserMemory(null, 0, REGION_ID, reservation); + try { + return reserved; + } finally { + if (reserved) { + memoryManager.releaseTsFileParserMemory(null, 0, REGION_ID); + } else { + memoryManager.cancelTsFileParserMemoryReservation(null, 0, REGION_ID, reservation); + } + } + } + + private PipeTsFileInsertionEvent createTsFileEvent( + final String fileName, final long timestamp, final boolean filtered, final boolean twoTables) + throws Exception { + final File partitionDirectory = new File(temporaryDirectory, REGION_ID + File.separator + "0"); + Assert.assertTrue(partitionDirectory.isDirectory() || partitionDirectory.mkdirs()); + final File tsFile = new File(partitionDirectory, fileName); + try (final TsFileWriter writer = new TsFileWriter(tsFile)) { + writeTable(writer, "table_0", timestamp); + if (twoTables) { + writeTable(writer, "table_1", timestamp + 1); + } + } + final TsFileResource resource = new TsFileResource(tsFile); + resource.setStatusForTest(TsFileResourceStatus.NORMAL); + final IDeviceID device = + IDeviceID.Factory.DEFAULT_FACTORY.create(new String[] {"table_0", "device_0"}); + resource.updateStartTime(device, timestamp); + resource.updateEndTime(device, timestamp + 1); + final PipeTsFileInsertionEvent event = + new PipeTsFileInsertionEvent( + true, + "test_sg_0", + resource, + null, + false, + false, + false, + Collections.singleton("table_0"), + null, + 0, + null, + null, + new TablePattern(true, null, filtered ? "other_table" : null), + null, + null, + null, + true, + Long.MIN_VALUE, + Long.MAX_VALUE); + tsFileEvents.add(event); + return event; + } + + private static void writeTable(final TsFileWriter writer, final String tableName, final long time) + throws Exception { + final List categories = Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD); + writer.registerTableSchema( + new TableSchema( + tableName, + Arrays.asList( + new MeasurementSchema("device_id", TSDataType.STRING), + new MeasurementSchema("sensor", TSDataType.INT64)), + categories)); + final Tablet tablet = + new Tablet( + tableName, + Arrays.asList("device_id", "sensor"), + Arrays.asList(TSDataType.STRING, TSDataType.INT64), + categories, + 1); + tablet.addTimestamp(0, time); + tablet.addValue(0, 0, "device_0"); + tablet.addValue(0, 1, time); + writer.writeTable(tablet); + } + + private static void assertTablet( + final Pair> tablets, final String tableName, final long time) { + Assert.assertNotNull(tablets); + Assert.assertEquals("test_sg_0", tablets.left); + Assert.assertEquals(1, tablets.right.size()); + final Tablet tablet = tablets.right.get(0); + Assert.assertEquals(tableName, tablet.getTableName()); + Assert.assertEquals(1, tablet.getRowSize()); + Assert.assertEquals(time, tablet.getTimestamps()[0]); + Assert.assertEquals(time, ((long[]) tablet.getValues()[1])[0]); + } +}