This is an automated email from the ASF dual-hosted git repository.
JackieTien97 pushed a commit to branch dev/1.3
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/dev/1.3 by this push:
new c056ce6f861 [To dev/1.3] Preserve query memory registrations during
stale fragment cleanup- #18840
c056ce6f861 is described below
commit c056ce6f8615905e3c54938e969a8074d4f1681f
Author: Jackie Tien <[email protected]>
AuthorDate: Tue Oct 13 11:38:56 2026 +0800
[To dev/1.3] Preserve query memory registrations during stale fragment
cleanup- #18840
---
.../queryengine/execution/memory/MemoryPool.java | 5 +-
.../execution/memory/MemoryPoolTest.java | 128 +++++++++++++++++++++
2 files changed, 132 insertions(+), 1 deletion(-)
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java
index 9a00d5a87de..8682aa904e3 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java
@@ -226,7 +226,10 @@ public class MemoryPool {
}
synchronized (queryMemoryReservations) {
queryRelatedMemory.remove(fragmentInstanceId);
- if (queryRelatedMemory.isEmpty()) {
+ // Another fragment may have recreated the query entry while this
cleanup waited for the
+ // lock.
+ if (queryRelatedMemory.isEmpty()
+ && queryMemoryReservations.get(queryId) == queryRelatedMemory) {
queryMemoryReservations.remove(queryId);
}
}
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java
index cf7e15dc71f..283823c10e3 100644
---
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java
@@ -19,6 +19,7 @@
package org.apache.iotdb.db.queryengine.execution.memory;
+import org.apache.iotdb.db.exception.runtime.MemoryLeakException;
import
org.apache.iotdb.db.queryengine.execution.memory.MemoryPool.MemoryReservationResult;
import com.google.common.util.concurrent.ListenableFuture;
@@ -26,6 +27,14 @@ import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
+import java.lang.management.ManagementFactory;
+import java.lang.management.ThreadInfo;
+import java.lang.reflect.Field;
+import java.util.concurrent.FutureTask;
+import java.util.concurrent.TimeUnit;
+
+import static org.awaitility.Awaitility.await;
+
public class MemoryPoolTest {
MemoryPool pool;
@@ -41,6 +50,125 @@ public class MemoryPoolTest {
pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID,
PLAN_NODE_ID);
}
+ @Test
+ public void testDeregisterDoesNotRemoveReregisteredQueryMemory() throws
Exception {
+ Field reservationsField =
MemoryPool.class.getDeclaredField("queryMemoryReservations");
+ reservationsField.setAccessible(true);
+ Object reservationsLock = reservationsField.get(pool);
+ long lockOwnerId = Thread.currentThread().getId();
+ FutureTask<Void> deregistration =
+ new FutureTask<>(
+ () -> {
+ // An empty fragment may finish without registering any memory.
+ pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID,
"empty", true);
+ return null;
+ });
+ Thread cleanupThread = new Thread(deregistration,
"memory-pool-stale-deregister");
+ try {
+ synchronized (reservationsLock) {
+ cleanupThread.start();
+ // Hold the registration lock until the empty fragment cleanup reaches
it.
+ await()
+ .atMost(10, TimeUnit.SECONDS)
+ .until(
+ () -> {
+ ThreadInfo info =
+
ManagementFactory.getThreadMXBean().getThreadInfo(cleanupThread.getId());
+ return info != null
+ && info.getThreadState() == Thread.State.BLOCKED
+ && info.getLockOwnerId() == lockOwnerId
+ && info.getLockInfo().getIdentityHashCode()
+ == System.identityHashCode(reservationsLock);
+ });
+
+ // Replace the query map while the delayed cleanup still refers to the
old one.
+ pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID,
FRAGMENT_INSTANCE_ID, true);
+ Assert.assertEquals(0, pool.getQueryMemoryReservationSize());
+ pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "result",
PLAN_NODE_ID);
+ Assert.assertTrue(
+ pool.tryReserveForTest(QUERY_ID, "result", PLAN_NODE_ID, 256L,
Long.MAX_VALUE));
+ }
+ deregistration.get(10, TimeUnit.SECONDS);
+
+ Assert.assertEquals(1, pool.getQueryMemoryReservationSize());
+ Assert.assertEquals(256L, pool.getQueryMemoryReservedBytes(QUERY_ID));
+ Assert.assertEquals(256L, pool.getReservedBytes());
+ pool.free(QUERY_ID, "result", PLAN_NODE_ID, 256L);
+ pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, "result",
true);
+ Assert.assertEquals(0, pool.getQueryMemoryReservationSize());
+ Assert.assertEquals(0L, pool.getReservedBytes());
+ } finally {
+ cleanupThread.join(TimeUnit.SECONDS.toMillis(10));
+ Assert.assertFalse(cleanupThread.isAlive());
+ }
+ }
+
+ @Test
+ public void testDeregisterMissingFragmentPreservesReservedMemory() {
+ Assert.assertTrue(
+ pool.tryReserveForTest(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID,
256L, Long.MAX_VALUE));
+
+ pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, "empty", true);
+
+ Assert.assertEquals(1, pool.getQueryMemoryReservationSize());
+ Assert.assertEquals(256L, pool.getQueryMemoryReservedBytes(QUERY_ID));
+ pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256L);
+ pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID,
FRAGMENT_INSTANCE_ID, true);
+ Assert.assertEquals(0, pool.getQueryMemoryReservationSize());
+ Assert.assertEquals(0L, pool.getReservedBytes());
+ }
+
+ @Test
+ public void testDeregisterPreservesOtherFragments() {
+ pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "other", PLAN_NODE_ID);
+ Assert.assertTrue(
+ pool.tryReserveForTest(QUERY_ID, "other", PLAN_NODE_ID, 256L,
Long.MAX_VALUE));
+
+ pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID,
FRAGMENT_INSTANCE_ID, true);
+
+ Assert.assertEquals(1, pool.getQueryMemoryReservationSize());
+ Assert.assertEquals(256L, pool.getQueryMemoryReservedBytes(QUERY_ID));
+ pool.free(QUERY_ID, "other", PLAN_NODE_ID, 256L);
+ pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, "other", true);
+ Assert.assertEquals(0, pool.getQueryMemoryReservationSize());
+ Assert.assertEquals(0L, pool.getReservedBytes());
+ }
+
+ @Test
+ public void testDeregisterWaitsForMemoryRelease() {
+ Assert.assertTrue(
+ pool.tryReserveForTest(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID,
256L, Long.MAX_VALUE));
+
+ pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID,
FRAGMENT_INSTANCE_ID, false);
+
+ Assert.assertEquals(1, pool.getQueryMemoryReservationSize());
+ Assert.assertEquals(256L, pool.getQueryMemoryReservedBytes(QUERY_ID));
+ pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256L);
+ pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID,
FRAGMENT_INSTANCE_ID, false);
+ Assert.assertEquals(0, pool.getQueryMemoryReservationSize());
+ Assert.assertEquals(0L, pool.getReservedBytes());
+
+ // Repeated cleanup after the query has been removed remains a no-op.
+ pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID,
FRAGMENT_INSTANCE_ID, true);
+ Assert.assertEquals(0, pool.getQueryMemoryReservationSize());
+ Assert.assertEquals(0L, pool.getReservedBytes());
+ }
+
+ @Test
+ public void testForceDeregisterReportsMemoryLeak() {
+ Assert.assertTrue(
+ pool.tryReserveForTest(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID,
256L, Long.MAX_VALUE));
+
+ Assert.assertThrows(
+ MemoryLeakException.class,
+ () ->
+ pool.deRegisterFragmentInstanceFromQueryMemoryMap(
+ QUERY_ID, FRAGMENT_INSTANCE_ID, true));
+
+ Assert.assertEquals(0, pool.getQueryMemoryReservationSize());
+ Assert.assertEquals(256L, pool.getReservedBytes());
+ }
+
@Test
public void testTryReserve() {
