Skip to content

Harden query memory reservation deregistration ordering #18272

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Caideyipi wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Harden query memory reservation deregistration ordering #18272

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
44c8493
Fix query memory reservation deregistration race
Caideyipi Jul 22, 2026
7af40af
Handle deregistration winning reservation race
Caideyipi Jul 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2748,8 +2748,8 @@ private DataNodeQueryMessages() {}
public static final String QUERY_EXCEPTION_QUERY_IS_ABORTED_SINCE_IT_REQUESTS_MORE_MEMORY_THAN_CAN_D77C2921 =
"Query is aborted since it requests more memory than can be allocated, bytesToReserve: %sB, "
+ "maxBytesCanReserve: %sB";
public static final String QUERY_EXCEPTION_RELATEDMEMORYRESERVED_CAN_T_BE_NULL_WHEN_FREEING_MEMORY_C80009F2 =
"RelatedMemoryReserved can't be null when freeing memory";
public static final String EXCEPTION_NO_MEMORY_RESERVATION_IS_REGISTERED_FOR_QUERY_ARG_AND_FRAGMENT_INSTANCE_ARG_06016520 =
"No memory reservation is registered for query %s and fragment instance %s.";
public static final String QUERY_EXCEPTION_INTERRUPTED_BY_92FAED2D = "Interrupted By";
public static final String QUERY_EXCEPTION_DRIVER_WAS_INTERRUPTED_737358E4 =
"Driver was interrupted";
Expand Down
4 changes: 2 additions & 2 deletions iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3269,8 +3269,8 @@ private DataNodeQueryMessages() {}


"查询已中止，因为其请求的内存超过可分配上限，bytesToReserve：%sB，maxBytesCanReserve：%sB";
public static final String QUERY_EXCEPTION_RELATEDMEMORYRESERVED_CAN_T_BE_NULL_WHEN_FREEING_MEMORY_C80009F2 =
"释放内存时 RelatedMemoryReserved 不能为 null";
public static final String EXCEPTION_NO_MEMORY_RESERVATION_IS_REGISTERED_FOR_QUERY_ARG_AND_FRAGMENT_INSTANCE_ARG_06016520 =
"查询 %s、FragmentInstance %s 未注册内存预留。";
public static final String QUERY_EXCEPTION_INTERRUPTED_BY_92FAED2D =
"被中断，原因：";
public static final String QUERY_EXCEPTION_DRIVER_WAS_INTERRUPTED_737358E4 =
Expand Down
187 changes: 120 additions & 67 deletions ...e/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -206,10 +206,13 @@ public int getMemoryReservationSize() {
public void registerPlanNodeIdToQueryMemoryMap(
String queryId, String fragmentInstanceId, String planNodeId) {
synchronized (queryMemoryReservations) {
queryMemoryReservations
.computeIfAbsent(queryId, x -> new ConcurrentHashMap<>())
.computeIfAbsent(fragmentInstanceId, x -> new ConcurrentHashMap<>())
.putIfAbsent(planNodeId, 0L);
final Map<String, Map<String, Long>> queryRelatedMemory =
queryMemoryReservations.computeIfAbsent(queryId, x -> new ConcurrentHashMap<>());
final Map<String, Long> fragmentRelatedMemory =
queryRelatedMemory.computeIfAbsent(fragmentInstanceId, x -> new ConcurrentHashMap<>());
synchronized (fragmentRelatedMemory) {
fragmentRelatedMemory.putIfAbsent(planNodeId, 0L);
}
}
}

Expand All @@ -224,44 +227,45 @@ public void registerPlanNodeIdToQueryMemoryMap(
*/
public void deRegisterFragmentInstanceFromQueryMemoryMap(
String queryId, String fragmentInstanceId, boolean forceDeregister) {
Map<String, Map<String, Long>> queryRelatedMemory = queryMemoryReservations.get(queryId);
if (queryRelatedMemory != null) {
Map<String, Long> fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId);
boolean hasPotentialMemoryLeak = false;
// fragmentRelatedMemory could be null if the FI has not reserved any memory(For example,
// next() of root operator returns no data)
if (fragmentRelatedMemory != null) {
hasPotentialMemoryLeak =
fragmentRelatedMemory.values().stream().anyMatch(value -> value != 0L);
}
if (!forceDeregister && hasPotentialMemoryLeak) {
// If hasPotentialMemoryLeak is true, it means that LocalSourceChannel/LocalSourceHandles
// have not been closed.
// We should wait for them to be closed. Make sure this method is called again with
// forceDeregister == true, after all LocalSourceChannel/LocalSourceHandles are closed.
List<Map.Entry<String, Long>> invalidEntryList = null;
synchronized (queryMemoryReservations) {
final Map<String, Map<String, Long>> queryRelatedMemory =
queryMemoryReservations.get(queryId);
if (queryRelatedMemory == null) {
return;
}
synchronized (queryMemoryReservations) {
queryRelatedMemory.remove(fragmentInstanceId);
if (queryRelatedMemory.isEmpty()) {
queryMemoryReservations.remove(queryId);
final Map<String, Long> fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId);
if (fragmentRelatedMemory != null) {
synchronized (fragmentRelatedMemory) {
invalidEntryList =
fragmentRelatedMemory.entrySet().stream()
.filter(entry -> entry.getValue() != 0L)
.collect(Collectors.toList());
if (!forceDeregister && !invalidEntryList.isEmpty()) {
// The LocalSourceChannels/LocalSourceHandles have not been closed yet. The caller will
// invoke this method again with forceDeregister == true after closing them.
return;
}
queryRelatedMemory.remove(fragmentInstanceId, fragmentRelatedMemory);
}
} else {
// fragmentRelatedMemory could be null if the FI has not reserved any memory (for example,
// next() of root operator returns no data).
queryRelatedMemory.remove(fragmentInstanceId);
}
if (hasPotentialMemoryLeak) {
// hasPotentialMemoryLeak means that fragmentRelatedMemory is not null
List<Map.Entry<String, Long>> invalidEntryList =
fragmentRelatedMemory.entrySet().stream()
.filter(entry -> entry.getValue() != 0L)
.collect(Collectors.toList());
throw new MemoryLeakException(
String.format(
DataNodeQueryMessages
.QUERY_EXCEPTION_PLANNODE_RELATED_MEMORY_IS_NOT_ZERO_WHEN_TRYING_TO_DEREGISTER_E01109C5,
queryId,
fragmentInstanceId,
invalidEntryList));
if (queryRelatedMemory.isEmpty()) {
queryMemoryReservations.remove(queryId);
}
}
if (invalidEntryList != null && !invalidEntryList.isEmpty()) {
throw new MemoryLeakException(
String.format(
DataNodeQueryMessages
.QUERY_EXCEPTION_PLANNODE_RELATED_MEMORY_IS_NOT_ZERO_WHEN_TRYING_TO_DEREGISTER_E01109C5,
queryId,
fragmentInstanceId,
invalidEntryList));
}
}

/**
Expand Down Expand Up @@ -387,24 +391,33 @@ public void free(String queryId, String fragmentInstanceId, String planNodeId, l
Validate.notNull(queryId, DataNodeQueryMessages.EXCEPTION_QUERYID_CAN_NOT_BE_NULL_DOT_16639DBE);
Validate.isTrue(bytes > 0L);

try {
queryMemoryReservations
.get(queryId)
.get(fragmentInstanceId)
.computeIfPresent(
planNodeId,
(k, reservedMemory) -> {
if (reservedMemory < bytes) {
throw new IllegalArgumentException(
DataNodeQueryMessages.FREE_MORE_MEMORY_THAN_HAS_BEEN_RESERVED);
}
return reservedMemory - bytes;
});
} catch (NullPointerException e) {
throw new IllegalArgumentException(
DataNodeQueryMessages
.QUERY_EXCEPTION_RELATEDMEMORYRESERVED_CAN_T_BE_NULL_WHEN_FREEING_MEMORY_C80009F2,
e);
while (true) {
final Map<String, Map<String, Long>> queryRelatedMemory =
queryMemoryReservations.get(queryId);
if (queryRelatedMemory == null) {
throw memoryReservationNotRegistered(queryId, fragmentInstanceId);
}
final Map<String, Long> fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId);
if (fragmentRelatedMemory == null) {
throw memoryReservationNotRegistered(queryId, fragmentInstanceId);
}
synchronized (fragmentRelatedMemory) {
// The fragment map may have been removed while this thread was waiting for its lock.
if (queryMemoryReservations.get(queryId) != queryRelatedMemory
|| queryRelatedMemory.get(fragmentInstanceId) != fragmentRelatedMemory) {
continue;
}
fragmentRelatedMemory.computeIfPresent(
planNodeId,
(k, reservedMemory) -> {
if (reservedMemory < bytes) {
throw new IllegalArgumentException(
DataNodeQueryMessages.FREE_MORE_MEMORY_THAN_HAS_BEEN_RESERVED);
}
return reservedMemory - bytes;
});
break;
}
}

memoryBlock.release(bytes);
Expand Down Expand Up @@ -457,22 +470,62 @@ public boolean tryReserve(
String planNodeId,
long bytesToReserve,
long maxBytesCanReserve) {
long tryUsedBytes = memoryBlock.forceAllocateWithoutLimitation(bytesToReserve);
long queryRemainingBytes =
maxBytesCanReserve
- queryMemoryReservations
.get(queryId)
.get(fragmentInstanceId)
.merge(planNodeId, bytesToReserve, Long::sum);
return tryUsedBytes <= memoryBlock.getTotalMemorySizeInBytes() && queryRemainingBytes >= 0;
while (true) {
final Map<String, Map<String, Long>> queryRelatedMemory =
queryMemoryReservations.get(queryId);
if (queryRelatedMemory == null) {
throw memoryReservationNotRegistered(queryId, fragmentInstanceId);
}
final Map<String, Long> fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId);
if (fragmentRelatedMemory == null) {
throw memoryReservationNotRegistered(queryId, fragmentInstanceId);
}
synchronized (fragmentRelatedMemory) {
// Serialize reservation changes with fragment deregistration while allowing reservations
// for other fragments to proceed concurrently.
if (queryMemoryReservations.get(queryId) != queryRelatedMemory
|| queryRelatedMemory.get(fragmentInstanceId) != fragmentRelatedMemory) {
continue;
}
final long tryUsedBytes = memoryBlock.forceAllocateWithoutLimitation(bytesToReserve);
final long queryRemainingBytes =
maxBytesCanReserve - fragmentRelatedMemory.merge(planNodeId, bytesToReserve, Long::sum);
return tryUsedBytes <= memoryBlock.getTotalMemorySizeInBytes() && queryRemainingBytes >= 0;
}
}
}

private void rollbackReserve(
String queryId, String fragmentInstanceId, String planNodeId, long bytesToReserve) {
queryMemoryReservations
.get(queryId)
.get(fragmentInstanceId)
.merge(planNodeId, -bytesToReserve, Long::sum);
memoryBlock.release(bytesToReserve);
while (true) {
final Map<String, Map<String, Long>> queryRelatedMemory =
queryMemoryReservations.get(queryId);
if (queryRelatedMemory == null) {
throw memoryReservationNotRegistered(queryId, fragmentInstanceId);
}
final Map<String, Long> fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId);
if (fragmentRelatedMemory == null) {
throw memoryReservationNotRegistered(queryId, fragmentInstanceId);
}
synchronized (fragmentRelatedMemory) {
if (queryMemoryReservations.get(queryId) != queryRelatedMemory
|| queryRelatedMemory.get(fragmentInstanceId) != fragmentRelatedMemory) {
continue;
}
fragmentRelatedMemory.merge(planNodeId, -bytesToReserve, Long::sum);
memoryBlock.release(bytesToReserve);
return;
}
}
}

private IllegalArgumentException memoryReservationNotRegistered(
String queryId, String fragmentInstanceId) {
return new IllegalArgumentException(
String.format(
DataNodeQueryMessages
.EXCEPTION_NO_MEMORY_RESERVATION_IS_REGISTERED_FOR_QUERY_ARG_AND_FRAGMENT_INSTANCE_ARG_06016520,
queryId,
fragmentInstanceId));
}
}
Loading
Loading