You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iotdb.apache.org by sh...@apache.org on 2026/08/13 09:04:02 UTC
(iotdb) branch loadModificationFileWithMemControl-1.3 created (now 4fe906fd485)
This is an automated email from the ASF dual-hosted git repository.
shuwenwei pushed a change to branch loadModificationFileWithMemControl-1.3
in repository https://gitbox.apache.org/repos/asf/iotdb.git
at 4fe906fd485 Improve query modification loading memory control (#17788)
This branch includes the following new commits:
new 4fe906fd485 Improve query modification loading memory control (#17788)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(iotdb) 01/01: Improve query modification loading memory control (#17788)
Posted by sh...@apache.org.
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 <55...@users.noreply.github.com>
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;
+ }
+ }
+}