This is an automated email from the ASF dual-hosted git repository. shuwenwei pushed a commit to branch loadModificationFileWithMemControl-1.3 in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 4fe906fd485809b66c8e2444f6084f752f6a3fc2 Author: shuwenwei <[email protected]> AuthorDate: Fri Jun 5 13:56:15 2026 +0800 Improve query modification loading memory control (#17788) --- .../fragment/FragmentInstanceContext.java | 57 ++++ .../fragment/QueryModificationLoader.java | 302 +++++++++++++++++ .../memory/FakedMemoryReservationManager.java | 3 + .../planner/memory/MemoryReservationManager.java | 8 + .../fragment/QueryModificationLoaderTest.java | 372 +++++++++++++++++++++ 5 files changed, 742 insertions(+) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java index 24b233efcdf..b8c95d8e6ce 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java @@ -20,6 +20,7 @@ package org.apache.iotdb.db.queryengine.execution.fragment; import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.exception.IllegalPathException; import org.apache.iotdb.commons.exception.IoTDBException; import org.apache.iotdb.commons.path.AlignedPath; import org.apache.iotdb.commons.path.PartialPath; @@ -76,6 +77,7 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CountDownLatch; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Predicate; import java.util.stream.Collectors; import static org.apache.iotdb.db.queryengine.metric.DriverSchedulerMetricSet.BLOCK_QUEUED_TIME; @@ -91,6 +93,7 @@ public class FragmentInstanceContext extends QueryContext { private static final Logger LOGGER = LoggerFactory.getLogger(FragmentInstanceContext.class); private static final IoTDBConfig CONFIG = IoTDBDescriptor.getInstance().getConfig(); private static final long END_TIME_INITIAL_VALUE = -1L; + private static final int MODS_MEMORY_ESTIMATE_READ_INTERVAL = 10_000; // wait over 5s for driver to close is abnormal private static final long LONG_WAIT_DURATION = 5_000_000_000L; private final FragmentInstanceId id; @@ -399,6 +402,60 @@ public class FragmentInstanceContext extends QueryContext { lastExecutionStartTime.set(now); } + @Override + public List<Modification> getPathModifications( + TsFileResource tsFileResource, IDeviceID deviceID, String measurement) + throws IllegalPathException { + if (!checkIfModificationExists(tsFileResource)) { + return Collections.emptyList(); + } + if (memoryReservationManager == null) { + return super.getPathModifications(tsFileResource, deviceID, measurement); + } + PartialPath path = new PartialPath(deviceID, measurement); + try (QueryModificationLoader modificationLoader = + getQueryModificationLoader( + tsFileResource, + modification -> modification.getPath().overlapWith(path), + mods -> getPathModifications(mods, path))) { + return modificationLoader.getPathModifications(); + } + } + + @Override + public List<Modification> getPathModifications(TsFileResource tsFileResource, PartialPath path) { + if (!checkIfModificationExists(tsFileResource)) { + return Collections.emptyList(); + } + if (memoryReservationManager == null) { + return super.getPathModifications(tsFileResource, path); + } + try (QueryModificationLoader modificationLoader = + getQueryModificationLoader( + tsFileResource, + modification -> modification.getPath().overlapWith(path), + mods -> getPathModifications(mods, path))) { + return modificationLoader.getPathModifications(); + } catch (IllegalPathException e) { + throw new IllegalStateException(e); + } + } + + private QueryModificationLoader getQueryModificationLoader( + TsFileResource tsFileResource, + Predicate<Modification> fallbackModificationMatcher, + QueryModificationLoader.ModsTreeMatcher modsTreeMatcher) { + return new QueryModificationLoader( + tsFileResource, + memoryReservationManager, + CONFIG.getModsCacheSizeLimitPerFI(), + MODS_MEMORY_ESTIMATE_READ_INTERVAL, + fileModCache, + cachedModEntriesSize, + fallbackModificationMatcher, + modsTreeMatcher); + } + @Override protected boolean checkIfModificationExists(TsFileResource tsFileResource) { if (isSingleSourcePath()) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryModificationLoader.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryModificationLoader.java new file mode 100644 index 00000000000..3172bb94be9 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryModificationLoader.java @@ -0,0 +1,302 @@ +/* + * 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.execution.fragment; + +import org.apache.iotdb.commons.exception.IllegalPathException; +import org.apache.iotdb.commons.path.PatternTreeMap; +import org.apache.iotdb.db.queryengine.exception.MemoryNotEnoughException; +import org.apache.iotdb.db.queryengine.plan.planner.memory.MemoryReservationManager; +import org.apache.iotdb.db.storageengine.dataregion.modification.Modification; +import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileID; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.utils.datastructure.PatternTreeMapFactory; + +import org.apache.tsfile.utils.RamUsageEstimator; + +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Predicate; + +class QueryModificationLoader implements AutoCloseable { + + private final TsFileResource resource; + private final MemoryReservationManager memoryReservationManager; + private final long modsCacheSizeLimitPerFI; + private final int modsMemoryEstimateReadInterval; + private final Map<TsFileID, PatternTreeMap<Modification, PatternTreeMapFactory.ModsSerializer>> + fileModCache; + private final AtomicLong cachedModEntriesSize; + private final Predicate<Modification> modificationMatcher; + private final ModsTreeMatcher modsTreeMatcher; + + private Iterator<Modification> currentIterator; + + QueryModificationLoader( + TsFileResource resource, + MemoryReservationManager memoryReservationManager, + long modsCacheSizeLimitPerFI, + int modsMemoryEstimateReadInterval, + Map<TsFileID, PatternTreeMap<Modification, PatternTreeMapFactory.ModsSerializer>> + fileModCache, + AtomicLong cachedModEntriesSize, + Predicate<Modification> modificationMatcher, + ModsTreeMatcher modsTreeMatcher) { + this.resource = resource; + this.memoryReservationManager = memoryReservationManager; + this.modsCacheSizeLimitPerFI = modsCacheSizeLimitPerFI; + this.modsMemoryEstimateReadInterval = modsMemoryEstimateReadInterval; + this.fileModCache = fileModCache; + this.cachedModEntriesSize = cachedModEntriesSize; + this.modificationMatcher = modificationMatcher; + this.modsTreeMatcher = modsTreeMatcher; + } + + List<Modification> getPathModifications() throws IllegalPathException { + AtomicReference<LoadModsResult> loadedResult = new AtomicReference<>(); + PatternTreeMap<Modification, PatternTreeMapFactory.ModsSerializer> cachedMods = + fileModCache.computeIfAbsent( + resource.getTsFileID(), ignored -> loadAllModificationsForCache(loadedResult)); + if (cachedMods != null) { + return modsTreeMatcher.match(cachedMods); + } + + LoadModsResult result = loadedResult.get(); + try { + if (result.loadedAllModEntries) { + return fallbackByMatchLoadedPatternTree(result); + } else { + return fallbackByMatchedScan(result); + } + } finally { + close(); + } + } + + private PatternTreeMap<Modification, PatternTreeMapFactory.ModsSerializer> + loadAllModificationsForCache(AtomicReference<LoadModsResult> loadedResult) { + LoadModsResult result = loadAllModificationsWithQuotaControl(); + loadedResult.set(result); + if (!result.cacheable) { + return null; + } + + closeCurrentIterator(); + return result.mods; + } + + private LoadModsResult loadAllModificationsWithQuotaControl() { + PatternTreeMap<Modification, PatternTreeMapFactory.ModsSerializer> modifications = + PatternTreeMapFactory.getModsPatternTreeMap(); + LoadModsResult result = new LoadModsResult(modifications); + if (resource.getModFile().getSize() > getRemainingCacheQuota()) { + currentIterator = resource.getModFile().getModificationsIter().iterator(); + result.loadedAllModEntries = false; + result.cacheable = false; + return result; + } + + currentIterator = resource.getModFile().getModificationsIter().iterator(); + + int appendedModCount = 0; + boolean estimatedAfterLastAppend = false; + + while (currentIterator.hasNext()) { + Modification modification = currentIterator.next(); + modifications.append(modification.getPath(), modification); + appendedModCount++; + estimatedAfterLastAppend = false; + + if (appendedModCount % modsMemoryEstimateReadInterval == 0) { + if (!tryEstimateAndReserveTreeMemory(result)) { + result.loadedAllModEntries = false; + result.cacheable = false; + return result; + } + estimatedAfterLastAppend = true; + } + } + + if (!estimatedAfterLastAppend) { + result.cacheable = tryEstimateAndReserveTreeMemory(result); + } else { + result.cacheable = true; + } + + result.loadedAllModEntries = true; + return result; + } + + private boolean tryEstimateAndReserveTreeMemory(LoadModsResult result) { + long currentEstimatedSize = estimateModsTreeMemory(result.mods); + long delta = currentEstimatedSize - result.reservedTreeMemoryBytes; + if (delta < 0) { + throw new IllegalStateException( + String.format( + "Estimated mods tree size decreased from %d to %d for TsFile %s.", + result.reservedTreeMemoryBytes, currentEstimatedSize, resource)); + } + if (delta == 0) { + return true; + } + + if (!tryClaimCacheQuota(delta)) { + return false; + } + result.cacheQuotaBytes += delta; + + try { + memoryReservationManager.reserveMemoryImmediately(delta); + } catch (MemoryNotEnoughException e) { + return false; + } + + result.reservedTreeMemoryBytes = currentEstimatedSize; + return true; + } + + private boolean tryClaimCacheQuota(long delta) { + if (delta <= 0) { + return true; + } + + long alreadyUsedMemoryForCachedModEntries = cachedModEntriesSize.get(); + while (alreadyUsedMemoryForCachedModEntries + delta <= modsCacheSizeLimitPerFI) { + if (cachedModEntriesSize.compareAndSet( + alreadyUsedMemoryForCachedModEntries, alreadyUsedMemoryForCachedModEntries + delta)) { + return true; + } + alreadyUsedMemoryForCachedModEntries = cachedModEntriesSize.get(); + } + return false; + } + + private long getRemainingCacheQuota() { + return modsCacheSizeLimitPerFI - cachedModEntriesSize.get(); + } + + private List<Modification> fallbackByMatchedScan(LoadModsResult partialTree) + throws IllegalPathException { + List<Modification> matchedMods = matchLoadedTreeAndRelease(partialTree); + long reservedMatchedModsMemoryBytes = reserveMatchedModsMemory(matchedMods); + int matchedModCount = matchedMods.size(); + + while (currentIterator.hasNext()) { + Modification modification = currentIterator.next(); + if (modificationMatcher.test(modification)) { + matchedMods.add(modification); + matchedModCount++; + if (matchedModCount % modsMemoryEstimateReadInterval == 0) { + reservedMatchedModsMemoryBytes = + reserveMatchedModsMemoryIncrementally(matchedMods, reservedMatchedModsMemoryBytes); + } + } + } + + List<Modification> sortedAndMergedMods = ModificationFile.sortAndMerge(matchedMods); + adjustMatchedModsMemoryReservation(sortedAndMergedMods, reservedMatchedModsMemoryBytes); + return sortedAndMergedMods; + } + + private List<Modification> fallbackByMatchLoadedPatternTree(LoadModsResult loadedTree) + throws IllegalPathException { + List<Modification> matchedMods = matchLoadedTreeAndRelease(loadedTree); + reserveMatchedModsMemory(matchedMods); + return matchedMods; + } + + private List<Modification> matchLoadedTreeAndRelease(LoadModsResult loadedTree) + throws IllegalPathException { + try { + return new ArrayList<>(modsTreeMatcher.match(loadedTree.mods)); + } finally { + loadedTree.mods = null; + cachedModEntriesSize.addAndGet(-loadedTree.cacheQuotaBytes); + loadedTree.cacheQuotaBytes = 0; + memoryReservationManager.releaseMemoryCumulatively(loadedTree.reservedTreeMemoryBytes); + loadedTree.reservedTreeMemoryBytes = 0; + } + } + + private long reserveMatchedModsMemory(List<Modification> matchedMods) { + long estimatedSize = RamUsageEstimator.sizeOfArrayList(matchedMods); + memoryReservationManager.reserveMemoryCumulatively(estimatedSize); + return estimatedSize; + } + + private long reserveMatchedModsMemoryIncrementally( + List<Modification> matchedMods, long reservedMatchedModsMemoryBytes) { + long currentEstimatedSize = RamUsageEstimator.sizeOfArrayList(matchedMods); + long delta = currentEstimatedSize - reservedMatchedModsMemoryBytes; + memoryReservationManager.reserveMemoryCumulatively(delta); + return currentEstimatedSize; + } + + private void adjustMatchedModsMemoryReservation( + List<Modification> matchedMods, long reservedMatchedModsMemoryBytes) { + long currentEstimatedSize = RamUsageEstimator.sizeOfArrayList(matchedMods); + long delta = currentEstimatedSize - reservedMatchedModsMemoryBytes; + if (delta >= 0) { + memoryReservationManager.reserveMemoryCumulatively(delta); + } else { + memoryReservationManager.releaseMemoryCumulatively(-delta); + } + } + + private long estimateModsTreeMemory( + PatternTreeMap<Modification, PatternTreeMapFactory.ModsSerializer> modifications) { + return RamUsageEstimator.sizeOfObject(modifications); + } + + @Override + public void close() { + closeCurrentIterator(); + } + + private void closeCurrentIterator() { + currentIterator = null; + } + + private static class LoadModsResult { + + private PatternTreeMap<Modification, PatternTreeMapFactory.ModsSerializer> mods; + private long cacheQuotaBytes; + private long reservedTreeMemoryBytes; + private boolean loadedAllModEntries; + private boolean cacheable; + + private LoadModsResult( + PatternTreeMap<Modification, PatternTreeMapFactory.ModsSerializer> mods) { + this.mods = mods; + } + } + + @FunctionalInterface + interface ModsTreeMatcher { + + List<Modification> match( + PatternTreeMap<Modification, PatternTreeMapFactory.ModsSerializer> modsTree) + throws IllegalPathException; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/FakedMemoryReservationManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/FakedMemoryReservationManager.java index 1742a8070b3..2ea39471f40 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/FakedMemoryReservationManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/FakedMemoryReservationManager.java @@ -29,6 +29,9 @@ public class FakedMemoryReservationManager implements MemoryReservationManager { @Override public void reserveMemoryImmediately() {} + @Override + public void reserveMemoryImmediately(final long size) {} + @Override public void releaseMemoryCumulatively(long size) {} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/MemoryReservationManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/MemoryReservationManager.java index 9a9036b1d97..1013167ac68 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/MemoryReservationManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/MemoryReservationManager.java @@ -33,6 +33,14 @@ public interface MemoryReservationManager { /** Reserve memory for the accumulated memory size immediately. */ void reserveMemoryImmediately(); + /** + * Reserve memory for the given size immediately without changing the accumulated pending + * reservation size maintained by this manager. + * + * @param size the size of memory to reserve immediately + */ + void reserveMemoryImmediately(final long size); + /** * Release memory for the given size. * diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryModificationLoaderTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryModificationLoaderTest.java new file mode 100644 index 00000000000..dc76e4123c9 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryModificationLoaderTest.java @@ -0,0 +1,372 @@ +/* + * 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.execution.fragment; + +import org.apache.iotdb.commons.exception.IllegalPathException; +import org.apache.iotdb.commons.path.MeasurementPath; +import org.apache.iotdb.commons.path.PartialPath; +import org.apache.iotdb.commons.path.PatternTreeMap; +import org.apache.iotdb.db.queryengine.exception.MemoryNotEnoughException; +import org.apache.iotdb.db.queryengine.plan.planner.memory.MemoryReservationManager; +import org.apache.iotdb.db.storageengine.dataregion.modification.Deletion; +import org.apache.iotdb.db.storageengine.dataregion.modification.Modification; +import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileID; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.generator.TsFileNameGenerator; +import org.apache.iotdb.db.utils.constant.TestConstant; +import org.apache.iotdb.db.utils.datastructure.PatternTreeMapFactory; + +import org.apache.commons.io.FileUtils; +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.file.metadata.PlainDeviceID; +import org.apache.tsfile.utils.Pair; +import org.junit.After; +import org.junit.Test; + +import java.io.File; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicLong; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +public class QueryModificationLoaderTest { + + private static final IDeviceID DEVICE_ID = new PlainDeviceID("root.sg.d1"); + + private File testDir; + + @After + public void tearDown() throws Exception { + if (testDir != null && testDir.exists()) { + FileUtils.deleteDirectory(testDir); + } + } + + @Test + public void testCacheLoadedModsTreeWhenQuotaEnough() throws Exception { + TsFileResource resource = prepareResource("cache"); + writeMods(resource, deletion("root.sg.d1.s1", 0, 10), deletion("root.sg.d2.s1", 20, 30)); + + Map<TsFileID, PatternTreeMap<Modification, PatternTreeMapFactory.ModsSerializer>> fileModCache = + new ConcurrentHashMap<>(); + AtomicLong cachedModEntriesSize = new AtomicLong(); + CountingMemoryReservationManager memoryReservationManager = + new CountingMemoryReservationManager(); + + try (QueryModificationLoader loader = + newLoader( + resource, + Long.MAX_VALUE, + fileModCache, + cachedModEntriesSize, + memoryReservationManager, + 1)) { + List<Modification> result = loader.getPathModifications(); + + assertEquals(1, result.size()); + assertTrue(fileModCache.containsKey(resource.getTsFileID())); + assertTrue(cachedModEntriesSize.get() > 0); + assertTrue(memoryReservationManager.getReservedBytes() >= cachedModEntriesSize.get()); + assertTrue(memoryReservationManager.getImmediateReservationCount() > 0); + } + } + + @Test + public void testFallbackScansModsWhenFileSizeExceedsRemainingQuotaBeforeLoad() throws Exception { + TsFileResource resource = prepareResource("file-size-precheck-fallback"); + writeMods( + resource, + deletion("root.sg.d1.s1", 0, 10), + deletion("root.sg.d2.s1", 20, 30), + deletion("root.sg.d1.s1", 40, 50)); + + Map<TsFileID, PatternTreeMap<Modification, PatternTreeMapFactory.ModsSerializer>> fileModCache = + new ConcurrentHashMap<>(); + AtomicLong cachedModEntriesSize = new AtomicLong(); + CountingMemoryReservationManager memoryReservationManager = + new CountingMemoryReservationManager(); + + try (QueryModificationLoader loader = + newLoader(resource, 1, fileModCache, cachedModEntriesSize, memoryReservationManager, 1)) { + List<Modification> result = loader.getPathModifications(); + + assertEquals(2, result.size()); + assertFalse(fileModCache.containsKey(resource.getTsFileID())); + assertEquals(0, cachedModEntriesSize.get()); + assertTrue(memoryReservationManager.getReservedBytes() > 0); + assertEquals(0, memoryReservationManager.getRemainingImmediateFailures()); + assertEquals(0, memoryReservationManager.getImmediateReservationCount()); + } + } + + @Test + public void testFallbackScansRemainingModsWhenEstimatedTreeExceedsQuota() throws Exception { + TsFileResource resource = prepareResource("estimated-tree-quota-fallback"); + writeMods( + resource, + deletion("root.sg.d1.s1", 0, 10), + deletion("root.sg.d2.s1", 20, 30), + deletion("root.sg.d1.s1", 40, 50)); + + Map<TsFileID, PatternTreeMap<Modification, PatternTreeMapFactory.ModsSerializer>> fileModCache = + new ConcurrentHashMap<>(); + AtomicLong cachedModEntriesSize = new AtomicLong(); + CountingMemoryReservationManager memoryReservationManager = + new CountingMemoryReservationManager(); + + try (QueryModificationLoader loader = + newLoader( + resource, + resource.getModFile().getSize() + 1, + fileModCache, + cachedModEntriesSize, + memoryReservationManager, + 1)) { + List<Modification> result = loader.getPathModifications(); + + assertEquals(2, result.size()); + assertFalse(fileModCache.containsKey(resource.getTsFileID())); + assertEquals(0, cachedModEntriesSize.get()); + assertTrue(memoryReservationManager.getReservedBytes() > 0); + assertTrue(memoryReservationManager.getCumulativeReleaseCount() > 0); + } + } + + @Test + public void testFallbackMatchesLoadedTreeWhenFinalReservationFailed() throws Exception { + TsFileResource resource = prepareResource("final-reserve-failed"); + writeMods( + resource, + deletion("root.sg.d1.s1", 0, 10), + deletion("root.sg.d2.s1", 20, 30), + deletion("root.sg.d1.s1", 40, 50)); + + Map<TsFileID, PatternTreeMap<Modification, PatternTreeMapFactory.ModsSerializer>> fileModCache = + new ConcurrentHashMap<>(); + AtomicLong cachedModEntriesSize = new AtomicLong(); + CountingMemoryReservationManager memoryReservationManager = + new CountingMemoryReservationManager(1); + + try (QueryModificationLoader loader = + newLoader( + resource, + Long.MAX_VALUE, + fileModCache, + cachedModEntriesSize, + memoryReservationManager, + 100)) { + List<Modification> result = loader.getPathModifications(); + + assertEquals(2, result.size()); + assertFalse(fileModCache.containsKey(resource.getTsFileID())); + assertEquals(0, cachedModEntriesSize.get()); + assertTrue(memoryReservationManager.getReservedBytes() > 0); + assertEquals(0, memoryReservationManager.getRemainingImmediateFailures()); + assertEquals(1, memoryReservationManager.getImmediateReservationCount()); + } + } + + @Test + public void testFallbackReservesMatchedModsCumulativelyWhenQuotaExceeded() throws Exception { + TsFileResource resource = prepareResource("fallback-cumulative-reserve"); + writeMods( + resource, + deletion("root.sg.d1.s1", 0, 10), + deletion("root.sg.d2.s1", 20, 30), + deletion("root.sg.d1.s1", 40, 50)); + + Map<TsFileID, PatternTreeMap<Modification, PatternTreeMapFactory.ModsSerializer>> fileModCache = + new ConcurrentHashMap<>(); + AtomicLong cachedModEntriesSize = new AtomicLong(); + CountingMemoryReservationManager memoryReservationManager = + new CountingMemoryReservationManager(1); + + try (QueryModificationLoader loader = + newLoader(resource, 1, fileModCache, cachedModEntriesSize, memoryReservationManager, 1)) { + List<Modification> result = loader.getPathModifications(); + + assertEquals(2, result.size()); + assertFalse(fileModCache.containsKey(resource.getTsFileID())); + assertEquals(0, cachedModEntriesSize.get()); + assertTrue(memoryReservationManager.getReservedBytes() > 0); + assertEquals(1, memoryReservationManager.getRemainingImmediateFailures()); + assertEquals(0, memoryReservationManager.getImmediateReservationCount()); + } + } + + @Test + public void testFallbackAdjustsReservedMemoryAfterSortAndMerge() throws Exception { + TsFileResource resource = prepareResource("fallback-sort-merge-adjust"); + writeMods( + resource, + deletion("root.sg.d1.s1", 0, 10), + deletion("root.sg.d1.s1", 5, 15), + deletion("root.sg.d2.s1", 20, 30)); + + Map<TsFileID, PatternTreeMap<Modification, PatternTreeMapFactory.ModsSerializer>> fileModCache = + new ConcurrentHashMap<>(); + AtomicLong cachedModEntriesSize = new AtomicLong(); + CountingMemoryReservationManager memoryReservationManager = + new CountingMemoryReservationManager(); + + try (QueryModificationLoader loader = + newLoader(resource, 1, fileModCache, cachedModEntriesSize, memoryReservationManager, 1)) { + List<Modification> result = loader.getPathModifications(); + + assertEquals(1, result.size()); + assertFalse(fileModCache.containsKey(resource.getTsFileID())); + assertEquals(0, cachedModEntriesSize.get()); + assertTrue(memoryReservationManager.getReservedBytes() > 0); + assertTrue(memoryReservationManager.getCumulativeReleaseCount() > 0); + } + } + + private QueryModificationLoader newLoader( + TsFileResource resource, + long modsCacheSizeLimitPerFI, + Map<TsFileID, PatternTreeMap<Modification, PatternTreeMapFactory.ModsSerializer>> + fileModCache, + AtomicLong cachedModEntriesSize, + MemoryReservationManager memoryReservationManager, + int modsMemoryEstimateReadInterval) + throws IllegalPathException { + QueryContext queryContext = new QueryContext(false); + PartialPath queryPath = new PartialPath(DEVICE_ID, "s1"); + return new QueryModificationLoader( + resource, + memoryReservationManager, + modsCacheSizeLimitPerFI, + modsMemoryEstimateReadInterval, + fileModCache, + cachedModEntriesSize, + modification -> modification.getPath().overlapWith(queryPath), + modsTree -> queryContext.getPathModifications(modsTree, queryPath)); + } + + private TsFileResource prepareResource(String name) { + testDir = new File(TestConstant.BASE_OUTPUT_PATH, "QueryModificationLoaderTest-" + name); + testDir.mkdirs(); + File tsFile = + new File(TsFileNameGenerator.generateNewTsFilePath(testDir.getAbsolutePath(), 1, 1, 0, 0)); + return new TsFileResource(tsFile); + } + + private static Deletion deletion(String path, long startTime, long endTime) + throws IllegalPathException { + return new Deletion(new MeasurementPath(path), 0, startTime, endTime); + } + + private void writeMods(TsFileResource resource, Deletion... modifications) throws Exception { + try (ModificationFile modificationFile = resource.getModFile()) { + for (Deletion modification : modifications) { + modificationFile.write(modification); + } + } + } + + private static class CountingMemoryReservationManager implements MemoryReservationManager { + + private long reservedBytes; + private int remainingImmediateFailures; + private int immediateReservationCount; + private int cumulativeReleaseCount; + + private CountingMemoryReservationManager() {} + + private CountingMemoryReservationManager(int remainingImmediateFailures) { + this.remainingImmediateFailures = remainingImmediateFailures; + } + + @Override + public void reserveMemoryCumulatively(long size) { + reservedBytes += size; + } + + @Override + public void reserveMemoryImmediately() { + immediateReservationCount++; + if (remainingImmediateFailures > 0) { + remainingImmediateFailures--; + throw new MemoryNotEnoughException("Mock memory reservation failure."); + } + } + + @Override + public void reserveMemoryImmediately(long size) { + immediateReservationCount++; + if (remainingImmediateFailures > 0) { + remainingImmediateFailures--; + throw new MemoryNotEnoughException("Mock memory reservation failure."); + } + reservedBytes += size; + } + + @Override + public void releaseMemoryCumulatively(long size) { + cumulativeReleaseCount++; + reservedBytes -= size; + } + + @Override + public void releaseMemoryImmediately(long size) { + reservedBytes -= size; + } + + @Override + public void releaseAllReservedMemory() { + reservedBytes = 0; + } + + @Override + public Pair<Long, Long> releaseMemoryVirtually(long size) { + reservedBytes -= size; + return new Pair<>(size, 0L); + } + + @Override + public void reserveMemoryVirtually(long bytesToBeReserved, long bytesAlreadyReserved) { + reservedBytes += bytesToBeReserved + bytesAlreadyReserved; + } + + @Override + public void setHighestPriority(boolean isHighestPriority) {} + + private long getReservedBytes() { + return reservedBytes; + } + + private int getRemainingImmediateFailures() { + return remainingImmediateFailures; + } + + private int getImmediateReservationCount() { + return immediateReservationCount; + } + + private int getCumulativeReleaseCount() { + return cumulativeReleaseCount; + } + } +}
