fix: reserve memory for sorting indices during query execution

Author: shizy818

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

@@ -43,4 +43,9 @@ 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,6 +19,8 @@
 
 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 {
@@ -72,4 +74,13 @@ 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,4 +114,9 @@ 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,4 +61,9 @@ 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/utils/datastructure/TVList.java

@@ -21,7 +21,9 @@
 
 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;
@@ -325,6 +327,13 @@ 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;
   }

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java

@@ -20,7 +20,12 @@
 package org.apache.iotdb.db.queryengine.execution.fragment;
 
 import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory;
+import org.apache.iotdb.commons.exception.IllegalPathException;
+import org.apache.iotdb.commons.exception.MetadataException;
+import org.apache.iotdb.commons.path.NonAlignedFullPath;
+import org.apache.iotdb.commons.path.PartialPath;
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
+import org.apache.iotdb.db.exception.query.QueryProcessException;
 import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
 import org.apache.iotdb.db.queryengine.common.PlanFragmentId;
 import org.apache.iotdb.db.queryengine.exception.CpuNotEnoughException;
@@ -30,15 +35,27 @@
 import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISink;
 import org.apache.iotdb.db.queryengine.execution.schedule.IDriverScheduler;
 import org.apache.iotdb.db.storageengine.dataregion.DataRegion;
+import org.apache.iotdb.db.storageengine.dataregion.memtable.DeviceIDFactory;
+import org.apache.iotdb.db.storageengine.dataregion.memtable.IMemTable;
+import org.apache.iotdb.db.storageengine.dataregion.memtable.IWritableMemChunk;
+import org.apache.iotdb.db.storageengine.dataregion.memtable.IWritableMemChunkGroup;
+import org.apache.iotdb.db.storageengine.dataregion.memtable.PrimitiveMemTable;
+import org.apache.iotdb.db.storageengine.dataregion.memtable.ReadOnlyMemChunk;
 import org.apache.iotdb.db.utils.datastructure.AlignedTVList;
 import org.apache.iotdb.db.utils.datastructure.TVList;
 
 import com.google.common.collect.ImmutableMap;
 import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.file.metadata.IDeviceID;
+import org.apache.tsfile.file.metadata.enums.CompressionType;
+import org.apache.tsfile.file.metadata.enums.TSEncoding;
+import org.apache.tsfile.read.reader.IPointReader;
+import org.apache.tsfile.write.schema.MeasurementSchema;
 import org.junit.Test;
 import org.mockito.Mockito;
 
 import java.io.ByteArrayOutputStream;
+import java.io.IOException;
 import java.io.PrintStream;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -49,6 +66,7 @@
 import static org.apache.iotdb.db.queryengine.common.QueryId.MOCK_QUERY_ID;
 import static org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext.createFragmentInstanceContext;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -157,6 +175,71 @@ public void testTVListOwnerTransfer() throws InterruptedException {
     }
   }
 
+  @Test
+  public void testTVListCloneForQuery() {
+    IoTDBDescriptor.getInstance().getConfig().setDataNodeId(1);
+
+    ExecutorService instanceNotificationExecutor =
+        IoTDBThreadPoolFactory.newFixedThreadPool(1, "test-instance-notification");
+
+    try {
+      String deviceId = "d1";
+      String measurementId = "s1";
+      IMemTable memTable = createMemTable(deviceId, measurementId);
+      assertEquals(1, memTable.getMemTableMap().size());
+      IWritableMemChunkGroup memChunkGroup = memTable.getMemTableMap().values().iterator().next();
+      assertEquals(1, memChunkGroup.getMemChunkMap().size());
+      IWritableMemChunk memChunk = memChunkGroup.getMemChunkMap().values().iterator().next();
+      TVList tvList = memChunk.getWorkingTVList();
+      assertFalse(tvList.isSorted());
+
+      // FragmentInstance Context
+      FragmentInstanceId id1 = new FragmentInstanceId(new PlanFragmentId(MOCK_QUERY_ID, 1), "1");
+      FragmentInstanceStateMachine stateMachine1 =
+          new FragmentInstanceStateMachine(id1, instanceNotificationExecutor);
+      FragmentInstanceContext fragmentInstanceContext1 =
+          createFragmentInstanceContext(id1, stateMachine1);
+
+      FragmentInstanceId id2 = new FragmentInstanceId(new PlanFragmentId(MOCK_QUERY_ID, 2), "2");
+      FragmentInstanceStateMachine stateMachine2 =
+          new FragmentInstanceStateMachine(id2, instanceNotificationExecutor);
+      FragmentInstanceContext fragmentInstanceContext2 =
+          createFragmentInstanceContext(id2, stateMachine2);
+
+      // query on memtable
+      NonAlignedFullPath fullPath =
+          new NonAlignedFullPath(
+              IDeviceID.Factory.DEFAULT_FACTORY.create(deviceId),
+              new MeasurementSchema(
+                  measurementId,
+                  TSDataType.INT32,
+                  TSEncoding.RLE,
+                  CompressionType.UNCOMPRESSED,
+                  Collections.emptyMap()));
+      ReadOnlyMemChunk readOnlyMemChunk1 =
+          memTable.query(fragmentInstanceContext1, fullPath, Long.MIN_VALUE, null, null);
+      ReadOnlyMemChunk readOnlyMemChunk2 =
+          memTable.query(fragmentInstanceContext2, fullPath, Long.MIN_VALUE, null, null);
+
+      IPointReader pointReader = readOnlyMemChunk1.getPointReader();
+      while (pointReader.hasNextTimeValuePair()) {
+        pointReader.nextTimeValuePair();
+      }
+      assertTrue(tvList.isSorted());
+      assertEquals(
+          tvList.calculateRamSize(),
+          fragmentInstanceContext1.getMemoryReservationContext().getReservedMemory());
+    } catch (QueryProcessException
+        | IOException
+        | MetadataException
+        | MemoryNotEnoughException
+        | IllegalArgumentException e) {
+      fail(e.getMessage());
+    } finally {
+      instanceNotificationExecutor.shutdown();
+    }
+  }
+
   private FragmentInstanceExecution createFragmentInstanceExecution(int id, Executor executor)
       throws CpuNotEnoughException {
     IDriverScheduler scheduler = Mockito.mock(IDriverScheduler.class);
@@ -201,4 +284,20 @@ private TVList buildTVList() {
     }
     return tvList;
   }
+
+  private IMemTable createMemTable(String deviceId, String measurementId)
+      throws IllegalPathException {
+    IMemTable memTable = new PrimitiveMemTable("root.test", "1");
+
+    int rows = 100;
+    for (int i = 0; i < 100; i++) {
+      memTable.write(
+          DeviceIDFactory.getInstance().getDeviceID(new PartialPath(deviceId)),
+          Collections.singletonList(
+              new MeasurementSchema(measurementId, TSDataType.INT32, TSEncoding.PLAIN)),
+          rows - i - 1,
+          new Object[] {i + 10});
+    }
+    return memTable;
+  }
 }