Sort vs setOwner

Author: shizy818

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

@@ -902,21 +902,28 @@ private void releaseTVListOwnedByQuery() {
       try {
         queryContextSet.remove(this);
         if (tvList.getOwnerQuery() == this) {
+          long tvListRamSize = tvList.calculateRamSize();
+          if (tvList.getReservedMemoryBytes() != tvListRamSize) {
+            LOGGER.warn(
+                "Release TVList owned by query: allocate size {}, release size {}",
+                tvList.getReservedMemoryBytes(),
+                tvListRamSize);
+          }
           if (queryContextSet.isEmpty()) {
             if (LOGGER.isDebugEnabled()) {
               LOGGER.debug(
                   "TVList {} is released by the query, FragmentInstance Id is {}",
                   tvList,
                   this.getId());
             }
-            memoryReservationManager.releaseMemoryCumulatively(tvList.calculateRamSize());
+            memoryReservationManager.releaseMemoryCumulatively(tvList.getReservedMemoryBytes());
             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());
+                memoryReservationManager.releaseMemoryVirtually(tvList.getReservedMemoryBytes());
             FragmentInstanceContext queryContext =
                 (FragmentInstanceContext) queryContextSet.iterator().next();
             queryContext

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

@@ -43,9 +43,4 @@ public Pair<Long, Long> releaseMemoryVirtually(final long size) {
   @Override
   public void reserveMemoryVirtually(
       final long bytesToBeReserved, final long bytesAlreadyReserved) {}
-
-  @Override
-  public long getReservedMemory() {
-    return 0;
-  }
 }

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

@@ -19,8 +19,6 @@
 
 package org.apache.iotdb.db.queryengine.plan.planner.memory;
 
-import org.apache.iotdb.commons.utils.TestOnly;
-
 import org.apache.tsfile.utils.Pair;
 
 public interface MemoryReservationManager {
@@ -74,13 +72,4 @@ public interface MemoryReservationManager {
    * @param bytesAlreadyReserved the amount of memory that has already been reserved
    */
   void reserveMemoryVirtually(final long bytesToBeReserved, final long bytesAlreadyReserved);
-
-  /**
-   * Get the total amount of memory currently reserved. This includes memory that has been reserved,
-   * plus memory pending to be reserved, minus memory pending to be released.
-   *
-   * @return the total amount of memory in bytes that is currently reserved
-   */
-  @TestOnly
-  long getReservedMemory();
 }

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

@@ -114,9 +114,4 @@ public void reserveMemoryVirtually(
     reservedBytesInTotal += bytesAlreadyReserved;
     reserveMemoryCumulatively(bytesToBeReserved);
   }
-
-  @Override
-  public long getReservedMemory() {
-    return bytesToBeReserved - bytesToBeReleased + reservedBytesInTotal;
-  }
 }

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

@@ -61,9 +61,4 @@ public synchronized void reserveMemoryVirtually(
       final long bytesToBeReserved, final long bytesAlreadyReserved) {
     super.reserveMemoryVirtually(bytesToBeReserved, bytesAlreadyReserved);
   }
-
-  @Override
-  public synchronized long getReservedMemory() {
-    return super.getReservedMemory();
-  }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/schemaengine/schemaregion/utils/ResourceByPathUtils.java

@@ -193,7 +193,9 @@ protected Map<TVList, Integer> prepareTvListMapForQuery(
           if (firstQuery instanceof FragmentInstanceContext) {
             MemoryReservationManager memoryReservationManager =
                 ((FragmentInstanceContext) firstQuery).getMemoryReservationContext();
-            memoryReservationManager.reserveMemoryCumulatively(list.calculateRamSize());
+            long tvListRamSize = list.calculateRamSize();
+            memoryReservationManager.reserveMemoryCumulatively(tvListRamSize);
+            list.setReservedMemoryBytes(tvListRamSize);
           }
           list.setOwnerQuery(firstQuery);
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractWritableMemChunk.java

@@ -112,7 +112,9 @@ private void tryReleaseTvList(TVList tvList) {
         if (firstQuery instanceof FragmentInstanceContext) {
           MemoryReservationManager memoryReservationManager =
               ((FragmentInstanceContext) firstQuery).getMemoryReservationContext();
-          memoryReservationManager.reserveMemoryCumulatively(tvList.calculateRamSize());
+          long tvListRamSize = tvList.calculateRamSize();
+          memoryReservationManager.reserveMemoryCumulatively(tvListRamSize);
+          tvList.setReservedMemoryBytes(tvListRamSize);
         }
         // update current TVList owner to first query in the list
         tvList.setOwnerQuery(firstQuery);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AlignedReadOnlyMemChunk.java

@@ -19,6 +19,7 @@
 
 package org.apache.iotdb.db.storageengine.dataregion.memtable;
 
+import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext;
 import org.apache.iotdb.db.queryengine.execution.fragment.QueryContext;
 import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering;
 import org.apache.iotdb.db.storageengine.dataregion.read.reader.chunk.MemAlignedChunkLoader;
@@ -118,6 +119,21 @@ public void sortTvLists() {
       int queryRowCount = entry.getValue();
       if (!alignedTvList.isSorted() && queryRowCount > alignedTvList.seqRowCount()) {
         alignedTvList.sort();
+        alignedTvList.lockQueryList();
+        try {
+          FragmentInstanceContext ownerQuery =
+              (FragmentInstanceContext) alignedTvList.getOwnerQuery();
+          if (ownerQuery != null) {
+            long deltaBytes =
+                alignedTvList.calculateRamSize() - alignedTvList.getReservedMemoryBytes();
+            if (deltaBytes > 0) {
+              ownerQuery.getMemoryReservationContext().reserveMemoryCumulatively(deltaBytes);
+              alignedTvList.addReservedMemoryBytes(deltaBytes);
+            }
+          }
+        } finally {
+          alignedTvList.unlockQueryList();
+        }
       }
     }
   }
@@ -356,10 +372,25 @@ public boolean isEmpty() {
   @Override
   public IPointReader getPointReader() {
     for (Map.Entry<TVList, Integer> entry : alignedTvListQueryMap.entrySet()) {
-      AlignedTVList tvList = (AlignedTVList) entry.getKey();
+      AlignedTVList alignedTvList = (AlignedTVList) entry.getKey();
       int queryLength = entry.getValue();
-      if (!tvList.isSorted() && queryLength > tvList.seqRowCount()) {
-        tvList.sort();
+      if (!alignedTvList.isSorted() && queryLength > alignedTvList.seqRowCount()) {
+        alignedTvList.sort();
+        alignedTvList.lockQueryList();
+        try {
+          FragmentInstanceContext ownerQuery =
+              (FragmentInstanceContext) alignedTvList.getOwnerQuery();
+          if (ownerQuery != null) {
+            long deltaBytes =
+                alignedTvList.calculateRamSize() - alignedTvList.getReservedMemoryBytes();
+            if (deltaBytes > 0) {
+              ownerQuery.getMemoryReservationContext().reserveMemoryCumulatively(deltaBytes);
+              alignedTvList.addReservedMemoryBytes(deltaBytes);
+            }
+          }
+        } finally {
+          alignedTvList.unlockQueryList();
+        }
       }
     }
     TsBlock tsBlock = buildTsBlock();

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/ReadOnlyMemChunk.java

@@ -21,6 +21,7 @@
 
 import org.apache.iotdb.commons.utils.TestOnly;
 import org.apache.iotdb.db.exception.query.QueryProcessException;
+import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext;
 import org.apache.iotdb.db.queryengine.execution.fragment.QueryContext;
 import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering;
 import org.apache.iotdb.db.storageengine.dataregion.read.reader.chunk.MemChunkLoader;
@@ -135,6 +136,19 @@ public void sortTvLists() {
       int queryRowCount = entry.getValue();
       if (!tvList.isSorted() && queryRowCount > tvList.seqRowCount()) {
         tvList.sort();
+        tvList.lockQueryList();
+        try {
+          FragmentInstanceContext ownerQuery = (FragmentInstanceContext) tvList.getOwnerQuery();
+          if (ownerQuery != null) {
+            long deltaBytes = tvList.calculateRamSize() - tvList.getReservedMemoryBytes();
+            if (deltaBytes > 0) {
+              ownerQuery.getMemoryReservationContext().reserveMemoryCumulatively(deltaBytes);
+              tvList.addReservedMemoryBytes(deltaBytes);
+            }
+          }
+        } finally {
+          tvList.unlockQueryList();
+        }
       }
     }
   }
@@ -274,6 +288,19 @@ public IPointReader getPointReader() {
       int queryLength = entry.getValue();
       if (!tvList.isSorted() && queryLength > tvList.seqRowCount()) {
         tvList.sort();
+        tvList.lockQueryList();
+        try {
+          FragmentInstanceContext ownerQuery = (FragmentInstanceContext) tvList.getOwnerQuery();
+          if (ownerQuery != null) {
+            long deltaBytes = tvList.calculateRamSize() - tvList.getReservedMemoryBytes();
+            if (deltaBytes > 0) {
+              ownerQuery.getMemoryReservationContext().reserveMemoryCumulatively(deltaBytes);
+              tvList.addReservedMemoryBytes(deltaBytes);
+            }
+          }
+        } finally {
+          tvList.unlockQueryList();
+        }
       }
     }
     TsBlock tsBlock = buildTsBlock();

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/TVList.java

@@ -21,9 +21,7 @@
 
 import org.apache.iotdb.common.rpc.thrift.TSStatus;
 import org.apache.iotdb.commons.utils.TestOnly;
-import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext;
 import org.apache.iotdb.db.queryengine.execution.fragment.QueryContext;
-import org.apache.iotdb.db.queryengine.plan.planner.memory.MemoryReservationManager;
 import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering;
 import org.apache.iotdb.db.service.metrics.WritingMetrics;
 import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALEntryValue;
@@ -88,6 +86,9 @@ public abstract class TVList implements WALEntryValue {
   // When it is null, the TVList is owned by insert thread and released after flush.
   protected QueryContext ownerQuery;
 
+  // Reserved memory by the query. Ensure to acquire queryListLock before update.
+  protected long reservedMemoryBytes = 0L;
+
   protected boolean sorted = true;
   protected long maxTime;
   protected long minTime;
@@ -167,6 +168,18 @@ public synchronized boolean isSorted() {
     return sorted;
   }
 
+  public void setReservedMemoryBytes(long bytes) {
+    this.reservedMemoryBytes = bytes;
+  }
+
+  public void addReservedMemoryBytes(long bytes) {
+    this.reservedMemoryBytes += bytes;
+  }
+
+  public long getReservedMemoryBytes() {
+    return reservedMemoryBytes;
+  }
+
   public abstract void sort();
 
   public void increaseReferenceCount() {
@@ -327,13 +340,6 @@ protected void set(int index, long timestamp, int valueIndex) {
         int offset = i * ARRAY_SIZE;
         Arrays.setAll(indices.get(i), j -> offset + j);
       }
-      // Reserve memory for indices if the TVList is owned by a query
-      if (ownerQuery != null) {
-        long indicesBytes = indices.size() * PrimitiveArrayManager.ARRAY_SIZE * 4L;
-        MemoryReservationManager memoryReservationManager =
-            ((FragmentInstanceContext) ownerQuery).getMemoryReservationContext();
-        memoryReservationManager.reserveMemoryCumulatively(indicesBytes);
-      }
     }
     indices.get(arrayIndex)[elementIndex] = valueIndex;
   }