[Fix] Preserve query memory registrations during stale fragment cleanup

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);
         }
       }

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<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() {