fix: memory to be released is larger than the memory of memory block in TVList owner transfer case

Author: shizy818

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java

@@ -60,6 +60,7 @@
 import org.apache.tsfile.read.common.TimeRange;
 import org.apache.tsfile.read.filter.basic.Filter;
 import org.apache.tsfile.read.filter.factory.FilterFactory;
+import org.apache.tsfile.utils.Pair;
 import org.apache.tsfile.utils.RamUsageEstimator;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -902,19 +903,32 @@ private void releaseTVListOwnedByQuery() {
         queryContextSet.remove(this);
         if (tvList.getOwnerQuery() == this) {
           if (queryContextSet.isEmpty()) {
-            LOGGER.debug(
-                "TVList {} is released by the query, FragmentInstance Id is {}",
-                tvList,
-                this.getId());
+            if (LOGGER.isDebugEnabled()) {
+              LOGGER.debug(
+                  "TVList {} is released by the query, FragmentInstance Id is {}",
+                  tvList,
+                  this.getId());
+            }
             memoryReservationManager.releaseMemoryCumulatively(tvList.calculateRamSize());
             tvList.clear();
           } else {
+            // Transfer memory to next query. It must be exception-safe as this method is called
+            // during FragmentInstanceExecution cleanup. Any exception during this process could
+            // prevent proper resource cleanup and cause memory leaks.
+            Pair<Long, Long> releasedBytes =
+                memoryReservationManager.releaseMemoryVirtually(tvList.calculateRamSize());
             FragmentInstanceContext queryContext =
                 (FragmentInstanceContext) queryContextSet.iterator().next();
-            LOGGER.debug(
-                "TVList {} is now owned by another query, FragmentInstance Id is {}",
-                tvList,
-                queryContext.getId());
+            queryContext
+                .getMemoryReservationContext()
+                .reserveMemoryVirtually(releasedBytes.left, releasedBytes.right);
+
+            if (LOGGER.isDebugEnabled()) {
+              LOGGER.debug(
+                  "TVList {} is now owned by another query, FragmentInstance Id is {}",
+                  tvList,
+                  queryContext.getId());
+            }
             tvList.setOwnerQuery(queryContext);
           }
         }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java

@@ -294,6 +294,10 @@ public void reserveFromFreeMemoryForOperators(
       final long reservedBytes,
       final String queryId,
       final String contextHolder) {
+    if (memoryInBytes <= 0) {
+      throw new IllegalArgumentException(
+          "Bytes to reserve from free memory for operators should be larger than 0");
+    }
     if (OPERATORS_MEMORY_BLOCK.allocate(memoryInBytes)) {
       if (LOGGER.isDebugEnabled()) {
         LOGGER.debug(
@@ -317,6 +321,10 @@ public void reserveFromFreeMemoryForOperators(
   }
 
   public void releaseToFreeMemoryForOperators(final long memoryInBytes) {
+    if (memoryInBytes <= 0) {
+      throw new IllegalArgumentException(
+          "Bytes to release to free memory for operators should be larger than 0");
+    }
     OPERATORS_MEMORY_BLOCK.release(memoryInBytes);
   }
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/FakedMemoryReservationManager.java

@@ -19,6 +19,8 @@
 
 package org.apache.iotdb.db.queryengine.plan.planner.memory;
 
+import org.apache.tsfile.utils.Pair;
+
 public class FakedMemoryReservationManager implements MemoryReservationManager {
 
   @Override
@@ -32,4 +34,13 @@ public void releaseMemoryCumulatively(long size) {}
 
   @Override
   public void releaseAllReservedMemory() {}
+
+  @Override
+  public Pair<Long, Long> releaseMemoryVirtually(final long size) {
+    return null;
+  }
+
+  @Override
+  public void reserveMemoryVirtually(
+      final long releasedBytesInReserved, final long releasedBytesInTotal) {}
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/MemoryReservationManager.java

@@ -19,6 +19,8 @@
 
 package org.apache.iotdb.db.queryengine.plan.planner.memory;
 
+import org.apache.tsfile.utils.Pair;
+
 public interface MemoryReservationManager {
   /**
    * Reserve memory for the given size. The memory reservation request will be accumulated and the
@@ -43,4 +45,23 @@ public interface MemoryReservationManager {
    * this manager ends, Or the memory to be released in the batch may not be released correctly.
    */
   void releaseAllReservedMemory();
+
+  /**
+   * Release memory virtually without actually freeing the memory. This is used for memory
+   * reservation transfer scenarios where memory ownership needs to be transferred between different
+   * FragmentInstances without actual memory deallocation.
+   *
+   * <p>NOTE: When calling this method, it should be guaranteed that bytesToBeReserved +
+   * reservedBytesInTotal >= size to ensure proper memory accounting and prevent negative
+   * reservation values.
+   */
+  Pair<Long, Long> releaseMemoryVirtually(final long size);
+
+  /**
+   * Reserve memory virtually without actually allocating new memory. This is used to transfer
+   * memory ownership from one FragmentInstances to another by reserving the memory that was
+   * previously released virtually. It updates the internal reservation state without changing the
+   * actual memory allocation.
+   */
+  void reserveMemoryVirtually(final long releasedBytesInReserved, final long releasedBytesInTotal);
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/NotThreadSafeMemoryReservationManager.java

@@ -22,6 +22,8 @@
 import org.apache.iotdb.db.queryengine.common.QueryId;
 import org.apache.iotdb.db.queryengine.plan.planner.LocalExecutionPlanner;
 
+import org.apache.tsfile.utils.Pair;
+
 import javax.annotation.concurrent.NotThreadSafe;
 
 @NotThreadSafe
@@ -91,4 +93,25 @@ public void releaseAllReservedMemory() {
       bytesToBeReleased = 0;
     }
   }
+
+  @Override
+  public Pair<Long, Long> releaseMemoryVirtually(final long size) {
+    if (bytesToBeReserved >= size) {
+      bytesToBeReserved -= size;
+      return new Pair<>(size, 0L);
+    } else {
+      long releasedBytesInReserved = bytesToBeReserved;
+      long releasedBytesInTotal = size - bytesToBeReserved;
+      bytesToBeReserved = 0;
+      reservedBytesInTotal -= releasedBytesInTotal;
+      return new Pair<>(releasedBytesInReserved, releasedBytesInTotal);
+    }
+  }
+
+  @Override
+  public void reserveMemoryVirtually(
+      final long releasedBytesInReserved, final long releasedBytesInTotal) {
+    bytesToBeReserved += releasedBytesInReserved;
+    reservedBytesInTotal += releasedBytesInTotal;
+  }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/ThreadSafeMemoryReservationManager.java

@@ -21,6 +21,8 @@
 
 import org.apache.iotdb.db.queryengine.common.QueryId;
 
+import org.apache.tsfile.utils.Pair;
+
 import javax.annotation.concurrent.ThreadSafe;
 
 @ThreadSafe
@@ -48,4 +50,15 @@ public synchronized void releaseMemoryCumulatively(long size) {
   public synchronized void releaseAllReservedMemory() {
     super.releaseAllReservedMemory();
   }
+
+  @Override
+  public Pair<Long, Long> releaseMemoryVirtually(final long size) {
+    return super.releaseMemoryVirtually(size);
+  }
+
+  @Override
+  public void reserveMemoryVirtually(
+      final long releasedBytesInReserved, final long releasedBytesInTotal) {
+    super.reserveMemoryVirtually(releasedBytesInReserved, releasedBytesInTotal);
+  }
 }