refactor memory management

Author: Young-Leo

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java

@@ -194,7 +194,6 @@ public void setDatabaseName(@Nullable String databaseName) {
    * @param statementName the name of the prepared statement
    * @return the removed prepared statement info, or null if not found
    */
-  @Nullable
   public abstract PreparedStatementInfo removePreparedStatement(String statementName);
 
   /**
@@ -203,7 +202,6 @@ public void setDatabaseName(@Nullable String databaseName) {
    * @param statementName the name of the prepared statement
    * @return the prepared statement info, or null if not found
    */
-  @Nullable
   public abstract PreparedStatementInfo getPreparedStatement(String statementName);
 
   /**

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/PreparedStatementInfo.java

@@ -19,7 +19,6 @@
 
 package org.apache.iotdb.db.protocol.session;
 
-import org.apache.iotdb.commons.memory.IMemoryBlock;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Statement;
 
 import java.util.Objects;
@@ -28,28 +27,28 @@
 
 /**
  * Information about a prepared statement stored in a session. The AST is cached here to avoid
- * re-parsing on EXECUTE.
+ * reparsing on EXECUTE.
  */
 public class PreparedStatementInfo {
 
   private final String statementName;
   private final Statement sql; // Cached AST (contains Parameter nodes)
   private final long createTime;
-  private final IMemoryBlock memoryBlock; // Memory block allocated for this PreparedStatement
+  private final long memorySizeInBytes; // Memory size allocated for this PreparedStatement
 
-  public PreparedStatementInfo(String statementName, Statement sql, IMemoryBlock memoryBlock) {
+  public PreparedStatementInfo(String statementName, Statement sql, long memorySizeInBytes) {
     this.statementName = requireNonNull(statementName, "statementName is null");
     this.sql = requireNonNull(sql, "sql is null");
     this.createTime = System.currentTimeMillis();
-    this.memoryBlock = memoryBlock;
+    this.memorySizeInBytes = memorySizeInBytes;
   }
 
   public PreparedStatementInfo(
-      String statementName, Statement sql, long createTime, IMemoryBlock memoryBlock) {
+      String statementName, Statement sql, long createTime, long memorySizeInBytes) {
     this.statementName = requireNonNull(statementName, "statementName is null");
     this.sql = requireNonNull(sql, "sql is null");
     this.createTime = createTime;
-    this.memoryBlock = memoryBlock;
+    this.memorySizeInBytes = memorySizeInBytes;
   }
 
   public String getStatementName() {
@@ -64,8 +63,8 @@ public long getCreateTime() {
     return createTime;
   }
 
-  public IMemoryBlock getMemoryBlock() {
-    return memoryBlock;
+  public long getMemorySizeInBytes() {
+    return memorySizeInBytes;
   }
 
   @Override

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

@@ -29,6 +29,8 @@
 import org.apache.iotdb.commons.conf.CommonConfig;
 import org.apache.iotdb.commons.conf.CommonDescriptor;
 import org.apache.iotdb.commons.conf.IoTDBConstant;
+import org.apache.iotdb.commons.memory.IMemoryBlock;
+import org.apache.iotdb.commons.memory.MemoryBlockType;
 import org.apache.iotdb.db.auth.AuthorityChecker;
 import org.apache.iotdb.db.conf.IoTDBConfig;
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
@@ -195,13 +197,26 @@ public class Coordinator {
 
   private static final Coordinator INSTANCE = new Coordinator();
 
+  private static final IMemoryBlock coordinatorMemoryBlock;
+
   private final ConcurrentHashMap<Long, IQueryExecution> queryExecutionMap;
 
   private final StatementRewrite statementRewrite;
   private final List<PlanOptimizer> logicalPlanOptimizers;
   private final List<PlanOptimizer> distributionPlanOptimizers;
   private final DataNodeLocationSupplierFactory.DataNodeLocationSupplier dataNodeLocationSupplier;
 
+  static {
+    coordinatorMemoryBlock =
+        IoTDBDescriptor.getInstance()
+            .getMemoryConfig()
+            .getCoordinatorMemoryManager()
+            .exactAllocate("Coordinator", MemoryBlockType.DYNAMIC);
+    LOGGER.debug(
+        "Initialized shared MemoryBlock 'Coordinator' with all available memory: {} bytes",
+        coordinatorMemoryBlock.getTotalMemorySizeInBytes());
+  }
+
   private Coordinator() {
     this.queryExecutionMap = new ConcurrentHashMap<>();
     this.executor = getQueryExecutor();
@@ -665,6 +680,10 @@ public static Coordinator getInstance() {
     return INSTANCE;
   }
 
+  public static IMemoryBlock getCoordinatorMemoryBlock() {
+    return coordinatorMemoryBlock;
+  }
+
   public void recordExecutionTime(long queryId, long executionTime) {
     IQueryExecution queryExecution = getQueryExecution(queryId);
     if (queryExecution != null) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/session/DeallocateTask.java

@@ -47,28 +47,26 @@ public DeallocateTask(String statementName) {
   public ListenableFuture<ConfigTaskResult> execute(IConfigTaskExecutor configTaskExecutor)
       throws InterruptedException {
     SettableFuture<ConfigTaskResult> future = SettableFuture.create();
-    try {
-      IClientSession session = SessionManager.getInstance().getCurrSession();
-      if (session == null) {
-        future.setException(
-            new IllegalStateException("No current session available for DEALLOCATE statement"));
-        return future;
-      }
-
-      // Remove the prepared statement
-      PreparedStatementInfo removedInfo = session.removePreparedStatement(statementName);
-      if (removedInfo == null) {
-        throw new SemanticException(
-            String.format("Prepared statement '%s' does not exist", statementName));
-      } else {
-        // Release the memory allocated for this PreparedStatement
-        PreparedStatementMemoryManager.getInstance().release(removedInfo.getMemoryBlock());
-      }
+    IClientSession session = SessionManager.getInstance().getCurrSession();
+    if (session == null) {
+      future.setException(
+          new IllegalStateException("No current session available for DEALLOCATE statement"));
+      return future;
+    }
 
-      future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS));
-    } catch (Exception e) {
-      future.setException(e);
+    // Remove the prepared statement
+    PreparedStatementInfo removedInfo = session.removePreparedStatement(statementName);
+    if (removedInfo == null) {
+      future.setException(
+          new SemanticException(
+              String.format("Prepared statement '%s' does not exist", statementName)));
+      return future;
     }
+
+    // Release the memory allocated for this PreparedStatement from the shared MemoryBlock
+    PreparedStatementMemoryManager.getInstance().release(removedInfo.getMemorySizeInBytes());
+
+    future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS));
     return future;
   }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/session/PrepareTask.java

@@ -19,7 +19,6 @@
 
 package org.apache.iotdb.db.queryengine.plan.execution.config.session;
 
-import org.apache.iotdb.commons.memory.IMemoryBlock;
 import org.apache.iotdb.db.exception.sql.SemanticException;
 import org.apache.iotdb.db.protocol.session.IClientSession;
 import org.apache.iotdb.db.protocol.session.PreparedStatementInfo;
@@ -36,7 +35,7 @@
 
 /**
  * Task for executing PREPARE statement. Stores the prepared statement AST in the session. The AST
- * is cached to avoid re-parsing on EXECUTE (skipping Parser phase). Memory is allocated from
+ * is cached to avoid reparsing on EXECUTE (skipping Parser phase). Memory is allocated from
  * CoordinatorMemoryManager and shared across all sessions.
  */
 public class PrepareTask implements IConfigTask {
@@ -53,42 +52,34 @@ public PrepareTask(String statementName, Statement sql) {
   public ListenableFuture<ConfigTaskResult> execute(IConfigTaskExecutor configTaskExecutor)
       throws InterruptedException {
     SettableFuture<ConfigTaskResult> future = SettableFuture.create();
-    IMemoryBlock memoryBlock = null;
-    try {
-      IClientSession session = SessionManager.getInstance().getCurrSession();
-      if (session == null) {
-        future.setException(
-            new IllegalStateException("No current session available for PREPARE statement"));
-        return future;
-      }
+    IClientSession session = SessionManager.getInstance().getCurrSession();
+    if (session == null) {
+      future.setException(
+          new IllegalStateException("No current session available for PREPARE statement"));
+      return future;
+    }
 
-      // Check if prepared statement with the same name already exists
-      PreparedStatementInfo existingInfo = session.getPreparedStatement(statementName);
-      if (existingInfo != null) {
-        throw new SemanticException(
-            String.format("Prepared statement '%s' already exists.", statementName));
-      }
+    // Check if prepared statement with the same name already exists
+    PreparedStatementInfo existingInfo = session.getPreparedStatement(statementName);
+    if (existingInfo != null) {
+      future.setException(
+          new SemanticException(
+              String.format("Prepared statement '%s' already exists.", statementName)));
+      return future;
+    }
 
-      // Estimate memory size of the AST
-      long memorySizeInBytes = AstMemoryEstimator.estimateMemorySize(sql);
+    // Estimate memory size of the AST
+    long memorySizeInBytes = AstMemoryEstimator.estimateMemorySize(sql);
 
-      // Allocate memory from CoordinatorMemoryManager
-      // This memory is shared across all sessions
-      memoryBlock =
-          PreparedStatementMemoryManager.getInstance().allocate(statementName, memorySizeInBytes);
+    // Allocate memory from CoordinatorMemoryManager
+    // This memory is shared across all sessions using a single MemoryBlock
+    PreparedStatementMemoryManager.getInstance().allocate(statementName, memorySizeInBytes);
 
-      // Create and store the prepared statement info (AST is cached)
-      PreparedStatementInfo info = new PreparedStatementInfo(statementName, sql, memoryBlock);
-      session.addPreparedStatement(statementName, info);
+    // Create and store the prepared statement info (AST is cached)
+    PreparedStatementInfo info = new PreparedStatementInfo(statementName, sql, memorySizeInBytes);
+    session.addPreparedStatement(statementName, info);
 
-      future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS));
-    } catch (Exception e) {
-      // If memory allocation succeeded but something else failed, release the memory
-      if (memoryBlock != null) {
-        PreparedStatementMemoryManager.getInstance().release(memoryBlock);
-      }
-      future.setException(e);
-    }
+    future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS));
     return future;
   }
 }