memory for CteScanOperator / CteScanReader

Author: shizy818

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/CteScanOperator.java

@@ -21,10 +21,10 @@
 
 package org.apache.iotdb.db.queryengine.execution.operator.source.relational;
 
-import org.apache.iotdb.commons.utils.TestOnly;
 import org.apache.iotdb.db.queryengine.execution.MemoryEstimationHelper;
 import org.apache.iotdb.db.queryengine.execution.operator.OperatorContext;
 import org.apache.iotdb.db.queryengine.execution.operator.source.SourceOperator;
+import org.apache.iotdb.db.queryengine.plan.planner.memory.MemoryReservationManager;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeId;
 import org.apache.iotdb.db.utils.cte.CteDataReader;
 import org.apache.iotdb.db.utils.cte.CteDataStore;
@@ -36,6 +36,8 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import static java.util.Objects.requireNonNull;
+
 public class CteScanOperator implements SourceOperator {
   private static final Logger LOGGER = LoggerFactory.getLogger(CteScanOperator.class);
   private static final long INSTANCE_SIZE =
@@ -46,42 +48,33 @@ public class CteScanOperator implements SourceOperator {
 
   private final OperatorContext operatorContext;
   private final PlanNodeId sourceId;
-
-  private final CteDataStore dataStore;
   private final CteDataReader dataReader;
-  private final int dataStoreRefCount;
 
   public CteScanOperator(
-      OperatorContext operatorContext, PlanNodeId sourceId, CteDataStore dataStore) {
+      OperatorContext operatorContext,
+      PlanNodeId sourceId,
+      CteDataStore dataStore,
+      MemoryReservationManager memoryReservationManager) {
+    requireNonNull(dataStore, "dataStore is null");
     this.operatorContext = operatorContext;
     this.sourceId = sourceId;
-    this.dataStore = dataStore;
-    this.dataReader = new MemoryReader(dataStore.getCachedData());
-    this.dataStoreRefCount = dataStore.increaseRefCount();
+    this.dataReader = new MemoryReader(dataStore, memoryReservationManager);
   }
 
   @Override
   public TsBlock next() throws Exception {
-    if (dataReader == null) {
-      return null;
-    }
     return dataReader.next();
   }
 
   @Override
   public boolean hasNext() throws Exception {
-    if (dataReader == null) {
-      return false;
-    }
     return dataReader.hasNext();
   }
 
   @Override
   public void close() throws Exception {
     try {
-      if (dataReader != null) {
-        dataReader.close();
-      }
+      dataReader.close();
     } catch (Exception e) {
       LOGGER.error("Fail to close CteDataReader", e);
     }
@@ -109,15 +102,9 @@ public long calculateRetainedSizeAfterCallingNext() {
 
   @Override
   public long ramBytesUsed() {
-    long bytes =
-        INSTANCE_SIZE
-            + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(operatorContext)
-            + dataReader.bytesUsed();
-    if (dataStoreRefCount == 1) {
-      bytes += dataStore.getCachedBytes();
-    }
-
-    return bytes;
+    return INSTANCE_SIZE
+        + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(operatorContext)
+        + dataReader.ramBytesUsed();
   }
 
   @Override
@@ -129,9 +116,4 @@ public OperatorContext getOperatorContext() {
   public PlanNodeId getSourceId() {
     return sourceId;
   }
-
-  @TestOnly
-  public int getDataStoreRefCount() {
-    return dataStoreRefCount;
-  }
 }

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

@@ -1202,7 +1202,11 @@ public Operator visitCteScan(CteScanNode node, LocalExecutionPlanContext context
                 context.getNextOperatorId(),
                 node.getPlanNodeId(),
                 CteScanOperator.class.getSimpleName());
-    return new CteScanOperator(operatorContext, node.getPlanNodeId(), node.getDataStore());
+    return new CteScanOperator(
+        operatorContext,
+        node.getPlanNodeId(),
+        node.getDataStore(),
+        context.getInstanceContext().getMemoryReservationContext());
   }
 
   @Override

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

@@ -177,7 +177,7 @@ public CteDataStore fetchCteQueryResult(
                 table,
                 "!!! Failed to materialize CTE. The main query falls back to INLINE mode !!!");
           }
-          context.releaseMemoryReservedForFrontEnd(cteDataStore.getCachedBytes());
+          context.releaseMemoryReservedForFrontEnd(cteDataStore.ramBytesUsed());
           cteDataStore.clear();
           return null;
         }
@@ -192,7 +192,7 @@ public CteDataStore fetchCteQueryResult(
       return cteDataStore;
     } catch (final Throwable throwable) {
       if (cteDataStore != null) {
-        context.releaseMemoryReservedForFrontEnd(cteDataStore.getCachedBytes());
+        context.releaseMemoryReservedForFrontEnd(cteDataStore.ramBytesUsed());
         cteDataStore.clear();
       }
       t = throwable;

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/cte/CteDataReader.java

@@ -24,8 +24,9 @@
 import org.apache.iotdb.commons.exception.IoTDBException;
 
 import org.apache.tsfile.read.common.block.TsBlock;
+import org.apache.tsfile.utils.Accountable;
 
-public interface CteDataReader {
+public interface CteDataReader extends Accountable {
   /**
    * Check if there is more data in CteDataReader. DiskSpillerReader may run out of current TsBlocks
    * , then it needs to read from file and cache more data. This method should be called before
@@ -48,11 +49,4 @@ public interface CteDataReader {
    * @throws IoTDBException the error occurs when closing fileChannel
    */
   void close() throws IoTDBException;
-
-  /**
-   * Get the bytes used by this CteDataReader.
-   *
-   * @return the bytes used by this CteDataReader
-   */
-  long bytesUsed();
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/cte/CteDataStore.java

@@ -21,25 +21,31 @@
 
 package org.apache.iotdb.db.utils.cte;
 
+import org.apache.iotdb.commons.utils.TestOnly;
 import org.apache.iotdb.db.conf.IoTDBConfig;
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
 import org.apache.iotdb.db.queryengine.plan.relational.metadata.TableSchema;
 
 import org.apache.tsfile.read.common.block.TsBlock;
+import org.apache.tsfile.utils.Accountable;
+import org.apache.tsfile.utils.RamUsageEstimator;
 
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
 
-public class CteDataStore {
+public class CteDataStore implements Accountable {
+  private static final long INSTANCE_SIZE =
+      RamUsageEstimator.shallowSizeOfInstance(CteDataStore.class);
+
   private final TableSchema tableSchema;
   private final List<Integer> columnIndex2TsBlockColumnIndexList;
 
   private final List<TsBlock> cachedData;
   private long cachedBytes;
   private int cachedRows;
 
-  // reference count by CteScanOperator
+  // reference count by CteScanReader
   private final AtomicInteger count;
 
   public CteDataStore(TableSchema tableSchema, List<Integer> columnIndex2TsBlockColumnIndexList) {
@@ -75,10 +81,6 @@ public List<TsBlock> getCachedData() {
     return cachedData;
   }
 
-  public long getCachedBytes() {
-    return cachedBytes;
-  }
-
   public TableSchema getTableSchema() {
     return tableSchema;
   }
@@ -87,7 +89,21 @@ public List<Integer> getColumnIndex2TsBlockColumnIndexList() {
     return columnIndex2TsBlockColumnIndexList;
   }
 
-  public int increaseRefCount() {
+  public int incrementAndGetCount() {
     return count.incrementAndGet();
   }
+
+  public int decrementAndGetCount() {
+    return count.decrementAndGet();
+  }
+
+  @Override
+  public long ramBytesUsed() {
+    return INSTANCE_SIZE + cachedBytes;
+  }
+
+  @TestOnly
+  public int getCount() {
+    return count.get();
+  }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/cte/MemoryReader.java

@@ -22,43 +22,54 @@
 package org.apache.iotdb.db.utils.cte;
 
 import org.apache.iotdb.commons.exception.IoTDBException;
+import org.apache.iotdb.db.queryengine.plan.planner.memory.MemoryReservationManager;
 
 import org.apache.tsfile.read.common.block.TsBlock;
 import org.apache.tsfile.utils.RamUsageEstimator;
 
-import java.util.List;
-
 public class MemoryReader implements CteDataReader {
   private static final long INSTANCE_SIZE =
       RamUsageEstimator.shallowSizeOfInstance(MemoryReader.class);
 
+  // thread-safe memory manager
+  private final MemoryReservationManager memoryReservationManager;
   // all the data in MemoryReader lies in memory
-  private final List<TsBlock> cachedData;
+  private final CteDataStore dataStore;
   private int tsBlockIndex;
 
-  public MemoryReader(List<TsBlock> cachedTsBlock) {
-    this.cachedData = cachedTsBlock;
+  public MemoryReader(CteDataStore dataStore, MemoryReservationManager memoryReservationManager) {
+    this.dataStore = dataStore;
     this.tsBlockIndex = 0;
+    this.memoryReservationManager = memoryReservationManager;
+    if (dataStore.incrementAndGetCount() == 1) {
+      memoryReservationManager.reserveMemoryCumulatively(dataStore.ramBytesUsed());
+    }
   }
 
   @Override
   public boolean hasNext() throws IoTDBException {
-    return cachedData != null && tsBlockIndex < cachedData.size();
+    return dataStore.getCachedData() != null && tsBlockIndex < dataStore.getCachedData().size();
   }
 
   @Override
   public TsBlock next() throws IoTDBException {
-    if (cachedData == null || tsBlockIndex >= cachedData.size()) {
+    if (dataStore.getCachedData() == null || tsBlockIndex >= dataStore.getCachedData().size()) {
       return null;
     }
-    return cachedData.get(tsBlockIndex++);
+    return dataStore.getCachedData().get(tsBlockIndex++);
   }
 
   @Override
-  public void close() throws IoTDBException {}
+  public void close() throws IoTDBException {
+    if (dataStore.decrementAndGetCount() == 0) {
+      memoryReservationManager.releaseMemoryCumulatively(dataStore.ramBytesUsed());
+    }
+  }
 
   @Override
-  public long bytesUsed() {
+  public long ramBytesUsed() {
+    // The calculation excludes the memory occupied by the CteDataStore.
+    // memory allocate/release for CteDataStore is handled during constructor and close
     return INSTANCE_SIZE;
   }
 }