From b0b96f61f624cb246750925832ac3b5029c473d8 Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Mon, 12 Oct 2026 16:34:27 +0800 Subject: [PATCH] Fix stale fragment cleanup removing query memory registrations --- .../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 8e1c8d1319cfa..b2d35568fa31c 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 @@ -243,7 +243,10 @@ public void deRegisterFragmentInstanceFromQueryMemoryMap( } 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 8614d39f20a54..5f2285e4ca910 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 @@ -20,6 +20,7 @@ package org.apache.iotdb.db.queryengine.execution.memory; import org.apache.iotdb.commons.memory.MemoryManager; +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; @@ -28,6 +29,14 @@ import org.junit.Test; import org.mockito.Mockito; +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; @@ -44,6 +53,125 @@ public void before() { 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 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() {