Skip to content

Commit 2fe40b5

Browse files
shizy818shizy818
authored
fix: reserve memory for sorting indices during query execution (#16959)
1 parent 345efb7 commit 2fe40b5

File tree

8 files changed

+195
-10
lines changed

8 files changed

+195
-10
lines changed

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -869,26 +869,33 @@ public void releaseResourceWhenAllDriversAreClosed() {
869869
*/
870870
private void releaseTVListOwnedByQuery() {
871871
for (TVList tvList : tvListSet) {
872+
long tvListRamSize = tvList.calculateRamSize();
872873
tvList.lockQueryList();
873874
Set<QueryContext> queryContextSet = tvList.getQueryContextSet();
874875
try {
875876
queryContextSet.remove(this);
876877
if (tvList.getOwnerQuery() == this) {
878+
if (tvList.getReservedMemoryBytes() != tvListRamSize) {
879+
LOGGER.warn(
880+
"Release TVList owned by query: allocate size {}, release size {}",
881+
tvList.getReservedMemoryBytes(),
882+
tvListRamSize);
883+
}
877884
if (queryContextSet.isEmpty()) {
878885
if (LOGGER.isDebugEnabled()) {
879886
LOGGER.debug(
880887
"TVList {} is released by the query, FragmentInstance Id is {}",
881888
tvList,
882889
this.getId());
883890
}
884-
memoryReservationManager.releaseMemoryCumulatively(tvList.calculateRamSize());
891+
memoryReservationManager.releaseMemoryCumulatively(tvList.getReservedMemoryBytes());
885892
tvList.clear();
886893
} else {
887894
// Transfer memory to next query. It must be exception-safe as this method is called
888895
// during FragmentInstanceExecution cleanup. Any exception during this process could
889896
// prevent proper resource cleanup and cause memory leaks.
890897
Pair<Long, Long> releasedBytes =
891-
memoryReservationManager.releaseMemoryVirtually(tvList.calculateRamSize());
898+
memoryReservationManager.releaseMemoryVirtually(tvList.getReservedMemoryBytes());
892899
FragmentInstanceContext queryContext =
893900
(FragmentInstanceContext) queryContextSet.iterator().next();
894901
queryContext

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,7 @@ protected Map<TVList, Integer> prepareTvListMapForQuery(
148148
// mutable tvlist
149149
TVList list = memChunk.getWorkingTVList();
150150
TVList cloneList = null;
151+
long tvListRamSize = list.calculateRamSize();
151152
list.lockQueryList();
152153
try {
153154
if (copyTimeFilter != null
@@ -188,7 +189,8 @@ protected Map<TVList, Integer> prepareTvListMapForQuery(
188189
if (firstQuery instanceof FragmentInstanceContext) {
189190
MemoryReservationManager memoryReservationManager =
190191
((FragmentInstanceContext) firstQuery).getMemoryReservationContext();
191-
memoryReservationManager.reserveMemoryCumulatively(list.calculateRamSize());
192+
memoryReservationManager.reserveMemoryCumulatively(tvListRamSize);
193+
list.setReservedMemoryBytes(tvListRamSize);
192194
}
193195
list.setOwnerQuery(firstQuery);
194196

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractWritableMemChunk.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@ protected void maybeReleaseTvList(TVList tvList) {
9999
}
100100

101101
private void tryReleaseTvList(TVList tvList) {
102+
long tvListRamSize = tvList.calculateRamSize();
102103
tvList.lockQueryList();
103104
try {
104105
if (tvList.getQueryContextSet().isEmpty()) {
@@ -110,7 +111,8 @@ private void tryReleaseTvList(TVList tvList) {
110111
if (firstQuery instanceof FragmentInstanceContext) {
111112
MemoryReservationManager memoryReservationManager =
112113
((FragmentInstanceContext) firstQuery).getMemoryReservationContext();
113-
memoryReservationManager.reserveMemoryCumulatively(tvList.calculateRamSize());
114+
memoryReservationManager.reserveMemoryCumulatively(tvListRamSize);
115+
tvList.setReservedMemoryBytes(tvListRamSize);
114116
}
115117
// update current TVList owner to first query in the list
116118
tvList.setOwnerQuery(firstQuery);

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AlignedReadOnlyMemChunk.java‎

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
package org.apache.iotdb.db.storageengine.dataregion.memtable;
2121

22+
import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext;
2223
import org.apache.iotdb.db.queryengine.execution.fragment.QueryContext;
2324
import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering;
2425
import org.apache.iotdb.db.storageengine.dataregion.read.reader.chunk.MemAlignedChunkLoader;
@@ -113,6 +114,21 @@ public void sortTvLists() {
113114
int queryRowCount = entry.getValue();
114115
if (!alignedTvList.isSorted() && queryRowCount > alignedTvList.seqRowCount()) {
115116
alignedTvList.sort();
117+
long alignedTvListRamSize = alignedTvList.calculateRamSize();
118+
alignedTvList.lockQueryList();
119+
try {
120+
FragmentInstanceContext ownerQuery =
121+
(FragmentInstanceContext) alignedTvList.getOwnerQuery();
122+
if (ownerQuery != null) {
123+
long deltaBytes = alignedTvListRamSize - alignedTvList.getReservedMemoryBytes();
124+
if (deltaBytes > 0) {
125+
ownerQuery.getMemoryReservationContext().reserveMemoryCumulatively(deltaBytes);
126+
alignedTvList.addReservedMemoryBytes(deltaBytes);
127+
}
128+
}
129+
} finally {
130+
alignedTvList.unlockQueryList();
131+
}
116132
}
117133
}
118134
}
@@ -339,10 +355,25 @@ public boolean isEmpty() {
339355
@Override
340356
public IPointReader getPointReader() {
341357
for (Map.Entry<TVList, Integer> entry : alignedTvListQueryMap.entrySet()) {
342-
AlignedTVList tvList = (AlignedTVList) entry.getKey();
358+
AlignedTVList alignedTvList = (AlignedTVList) entry.getKey();
343359
int queryLength = entry.getValue();
344-
if (!tvList.isSorted() && queryLength > tvList.seqRowCount()) {
345-
tvList.sort();
360+
if (!alignedTvList.isSorted() && queryLength > alignedTvList.seqRowCount()) {
361+
alignedTvList.sort();
362+
long alignedTvListRamSize = alignedTvList.calculateRamSize();
363+
alignedTvList.lockQueryList();
364+
try {
365+
FragmentInstanceContext ownerQuery =
366+
(FragmentInstanceContext) alignedTvList.getOwnerQuery();
367+
if (ownerQuery != null) {
368+
long deltaBytes = alignedTvListRamSize - alignedTvList.getReservedMemoryBytes();
369+
if (deltaBytes > 0) {
370+
ownerQuery.getMemoryReservationContext().reserveMemoryCumulatively(deltaBytes);
371+
alignedTvList.addReservedMemoryBytes(deltaBytes);
372+
}
373+
}
374+
} finally {
375+
alignedTvList.unlockQueryList();
376+
}
346377
}
347378
}
348379
TsBlock tsBlock = buildTsBlock();

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/ReadOnlyMemChunk.java‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
import org.apache.iotdb.commons.utils.TestOnly;
2323
import org.apache.iotdb.db.exception.query.QueryProcessException;
24+
import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext;
2425
import org.apache.iotdb.db.queryengine.execution.fragment.QueryContext;
2526
import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering;
2627
import org.apache.iotdb.db.storageengine.dataregion.read.reader.chunk.MemChunkLoader;
@@ -135,6 +136,20 @@ public void sortTvLists() {
135136
int queryRowCount = entry.getValue();
136137
if (!tvList.isSorted() && queryRowCount > tvList.seqRowCount()) {
137138
tvList.sort();
139+
long tvListRamSize = tvList.calculateRamSize();
140+
tvList.lockQueryList();
141+
try {
142+
FragmentInstanceContext ownerQuery = (FragmentInstanceContext) tvList.getOwnerQuery();
143+
if (ownerQuery != null) {
144+
long deltaBytes = tvListRamSize - tvList.getReservedMemoryBytes();
145+
if (deltaBytes > 0) {
146+
ownerQuery.getMemoryReservationContext().reserveMemoryCumulatively(deltaBytes);
147+
tvList.addReservedMemoryBytes(deltaBytes);
148+
}
149+
}
150+
} finally {
151+
tvList.unlockQueryList();
152+
}
138153
}
139154
}
140155
}
@@ -273,6 +288,20 @@ public IPointReader getPointReader() {
273288
int queryLength = entry.getValue();
274289
if (!tvList.isSorted() && queryLength > tvList.seqRowCount()) {
275290
tvList.sort();
291+
long tvListRamSize = tvList.calculateRamSize();
292+
tvList.lockQueryList();
293+
try {
294+
FragmentInstanceContext ownerQuery = (FragmentInstanceContext) tvList.getOwnerQuery();
295+
if (ownerQuery != null) {
296+
long deltaBytes = tvListRamSize - tvList.getReservedMemoryBytes();
297+
if (deltaBytes > 0) {
298+
ownerQuery.getMemoryReservationContext().reserveMemoryCumulatively(deltaBytes);
299+
tvList.addReservedMemoryBytes(deltaBytes);
300+
}
301+
}
302+
} finally {
303+
tvList.unlockQueryList();
304+
}
276305
}
277306
}
278307
TsBlock tsBlock = buildTsBlock();

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/AlignedTVList.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -836,7 +836,7 @@ public TSDataType getDataType() {
836836
}
837837

838838
@Override
839-
public long calculateRamSize() {
839+
public synchronized long calculateRamSize() {
840840
return timestamps.size() * alignedTvListArrayMemCost();
841841
}
842842

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/TVList.java‎

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,15 +75,20 @@ public abstract class TVList implements WALEntryValue {
7575
// Index relation: arrayIndex -> elementIndex
7676
protected List<BitMap> bitMap;
7777

78-
// lock to provide synchronization for query list
78+
// Guards queryContextSet, ownerQuery, and reservedMemoryBytes.
79+
// Always acquire this lock before accessing/modifying these fields.
7980
private final ReentrantLock queryListLock = new ReentrantLock();
81+
8082
// set of query that this TVList is used
8183
protected final Set<QueryContext> queryContextSet;
8284

8385
// the owner query which is obligated to release the TVList.
8486
// When it is null, the TVList is owned by insert thread and released after flush.
8587
protected QueryContext ownerQuery;
8688

89+
// Reserved memory by the query. Ensure to acquire queryListLock before update.
90+
protected long reservedMemoryBytes = 0L;
91+
8792
protected boolean sorted = true;
8893
protected long maxTime;
8994
protected long minTime;
@@ -151,14 +156,26 @@ public static long tvListArrayMemCost(TSDataType type) {
151156
return size;
152157
}
153158

154-
public long calculateRamSize() {
159+
public synchronized long calculateRamSize() {
155160
return timestamps.size() * tvListArrayMemCost();
156161
}
157162

158163
public synchronized boolean isSorted() {
159164
return sorted;
160165
}
161166

167+
public void setReservedMemoryBytes(long bytes) {
168+
this.reservedMemoryBytes = bytes;
169+
}
170+
171+
public void addReservedMemoryBytes(long bytes) {
172+
this.reservedMemoryBytes += bytes;
173+
}
174+
175+
public long getReservedMemoryBytes() {
176+
return reservedMemoryBytes;
177+
}
178+
162179
public abstract void sort();
163180

164181
public void increaseReferenceCount() {

‎iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,12 @@
2020
package org.apache.iotdb.db.queryengine.execution.fragment;
2121

2222
import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory;
23+
import org.apache.iotdb.commons.exception.IllegalPathException;
24+
import org.apache.iotdb.commons.exception.MetadataException;
25+
import org.apache.iotdb.commons.path.MeasurementPath;
26+
import org.apache.iotdb.commons.path.PartialPath;
2327
import org.apache.iotdb.db.conf.IoTDBDescriptor;
28+
import org.apache.iotdb.db.exception.query.QueryProcessException;
2429
import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
2530
import org.apache.iotdb.db.queryengine.common.PlanFragmentId;
2631
import org.apache.iotdb.db.queryengine.exception.CpuNotEnoughException;
@@ -30,15 +35,26 @@
3035
import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISink;
3136
import org.apache.iotdb.db.queryengine.execution.schedule.IDriverScheduler;
3237
import org.apache.iotdb.db.storageengine.dataregion.DataRegion;
38+
import org.apache.iotdb.db.storageengine.dataregion.memtable.DeviceIDFactory;
39+
import org.apache.iotdb.db.storageengine.dataregion.memtable.IMemTable;
40+
import org.apache.iotdb.db.storageengine.dataregion.memtable.IWritableMemChunk;
41+
import org.apache.iotdb.db.storageengine.dataregion.memtable.IWritableMemChunkGroup;
42+
import org.apache.iotdb.db.storageengine.dataregion.memtable.PrimitiveMemTable;
43+
import org.apache.iotdb.db.storageengine.dataregion.memtable.ReadOnlyMemChunk;
3344
import org.apache.iotdb.db.utils.datastructure.AlignedTVList;
3445
import org.apache.iotdb.db.utils.datastructure.TVList;
3546

3647
import com.google.common.collect.ImmutableMap;
3748
import org.apache.tsfile.enums.TSDataType;
49+
import org.apache.tsfile.file.metadata.enums.CompressionType;
50+
import org.apache.tsfile.file.metadata.enums.TSEncoding;
51+
import org.apache.tsfile.read.reader.IPointReader;
52+
import org.apache.tsfile.write.schema.MeasurementSchema;
3853
import org.junit.Test;
3954
import org.mockito.Mockito;
4055

4156
import java.io.ByteArrayOutputStream;
57+
import java.io.IOException;
4258
import java.io.PrintStream;
4359
import java.util.ArrayList;
4460
import java.util.Collections;
@@ -49,6 +65,7 @@
4965
import static org.apache.iotdb.db.queryengine.common.QueryId.MOCK_QUERY_ID;
5066
import static org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext.createFragmentInstanceContext;
5167
import static org.junit.Assert.assertEquals;
68+
import static org.junit.Assert.assertFalse;
5269
import static org.junit.Assert.assertTrue;
5370
import static org.junit.Assert.fail;
5471

@@ -157,6 +174,70 @@ public void testTVListOwnerTransfer() throws InterruptedException {
157174
}
158175
}
159176

177+
@Test
178+
public void testTVListCloneForQuery() {
179+
IoTDBDescriptor.getInstance().getConfig().setDataNodeId(1);
180+
181+
ExecutorService instanceNotificationExecutor =
182+
IoTDBThreadPoolFactory.newFixedThreadPool(1, "test-instance-notification");
183+
184+
try {
185+
String deviceId = "d1";
186+
String measurementId = "s1";
187+
IMemTable memTable = createMemTable(deviceId, measurementId);
188+
assertEquals(1, memTable.getMemTableMap().size());
189+
IWritableMemChunkGroup memChunkGroup = memTable.getMemTableMap().values().iterator().next();
190+
assertEquals(1, memChunkGroup.getMemChunkMap().size());
191+
IWritableMemChunk memChunk = memChunkGroup.getMemChunkMap().values().iterator().next();
192+
TVList tvList = memChunk.getWorkingTVList();
193+
assertFalse(tvList.isSorted());
194+
195+
// FragmentInstance Context
196+
FragmentInstanceId id1 = new FragmentInstanceId(new PlanFragmentId(MOCK_QUERY_ID, 1), "1");
197+
FragmentInstanceStateMachine stateMachine1 =
198+
new FragmentInstanceStateMachine(id1, instanceNotificationExecutor);
199+
FragmentInstanceContext fragmentInstanceContext1 =
200+
createFragmentInstanceContext(id1, stateMachine1);
201+
202+
FragmentInstanceId id2 = new FragmentInstanceId(new PlanFragmentId(MOCK_QUERY_ID, 2), "2");
203+
FragmentInstanceStateMachine stateMachine2 =
204+
new FragmentInstanceStateMachine(id2, instanceNotificationExecutor);
205+
FragmentInstanceContext fragmentInstanceContext2 =
206+
createFragmentInstanceContext(id2, stateMachine2);
207+
208+
// query on memtable
209+
MeasurementPath fullPath =
210+
new MeasurementPath(
211+
deviceId,
212+
measurementId,
213+
new MeasurementSchema(
214+
measurementId,
215+
TSDataType.INT32,
216+
TSEncoding.RLE,
217+
CompressionType.UNCOMPRESSED,
218+
Collections.emptyMap()));
219+
ReadOnlyMemChunk readOnlyMemChunk1 =
220+
memTable.query(fragmentInstanceContext1, fullPath, Long.MIN_VALUE, null, null);
221+
ReadOnlyMemChunk readOnlyMemChunk2 =
222+
memTable.query(fragmentInstanceContext2, fullPath, Long.MIN_VALUE, null, null);
223+
224+
IPointReader pointReader = readOnlyMemChunk1.getPointReader();
225+
while (pointReader.hasNextTimeValuePair()) {
226+
pointReader.nextTimeValuePair();
227+
}
228+
assertTrue(tvList.isSorted());
229+
assertEquals(tvList.calculateRamSize(), tvList.getReservedMemoryBytes());
230+
} catch (QueryProcessException
231+
| IOException
232+
| MetadataException
233+
| MemoryNotEnoughException
234+
| IllegalArgumentException e) {
235+
fail(e.getMessage());
236+
} finally {
237+
instanceNotificationExecutor.shutdown();
238+
}
239+
}
240+
160241
private FragmentInstanceExecution createFragmentInstanceExecution(int id, Executor executor)
161242
throws CpuNotEnoughException {
162243
IDriverScheduler scheduler = Mockito.mock(IDriverScheduler.class);
@@ -201,4 +282,20 @@ private TVList buildTVList() {
201282
}
202283
return tvList;
203284
}
285+
286+
private IMemTable createMemTable(String deviceId, String measurementId)
287+
throws IllegalPathException {
288+
IMemTable memTable = new PrimitiveMemTable("root.test", "1");
289+
290+
int rows = 100;
291+
for (int i = 0; i < 100; i++) {
292+
memTable.write(
293+
DeviceIDFactory.getInstance().getDeviceID(new PartialPath(deviceId)),
294+
Collections.singletonList(
295+
new MeasurementSchema(measurementId, TSDataType.INT32, TSEncoding.PLAIN)),
296+
rows - i - 1,
297+
new Object[] {i + 10});
298+
}
299+
return memTable;
300+
}
204301
}

0 commit comments

Comments
 (0)