This is an automated email from the ASF dual-hosted git repository. JackieTien97 pushed a commit to branch codex/cp-pr-18834-dev-1.3 in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 4f2243a0c174308ee291dccb44d5c6ed8464c4cb Author: JackieTien97 <[email protected]> AuthorDate: Mon Oct 12 16:34:27 2026 +0800 Fix stale fragment cleanup removing query memory registrations (cherry picked from commit b0b96f61f624cb246750925832ac3b5029c473d8) --- .../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() {
