[iotdb] branch rel/1.1 updated: [To rel/1.1] Fix potential memory leak in MemoryPool

[ja...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

jackietien pushed a commit to branch rel/1.1
in repository https://gitbox.apache.org/repos/asf/iotdb.git


The following commit(s) were added to refs/heads/rel/1.1 by this push:
     new af8f02ef829 [To rel/1.1] Fix potential memory leak in MemoryPool
af8f02ef829 is described below

commit af8f02ef8298e2c778c9e233f2bb71360ff06a6a
Author: Liao Lanyu <14...@qq.com>
AuthorDate: Wed Jul 26 15:56:22 2023 +0800

    [To rel/1.1] Fix potential memory leak in MemoryPool
---
 .../mpp/execution/exchange/sink/SinkChannel.java   |  2 +-
 .../iotdb/db/mpp/execution/memory/MemoryPool.java  | 35 ++++++++++------------
 2 files changed, 17 insertions(+), 20 deletions(-)

diff --git a/server/src/main/java/org/apache/iotdb/db/mpp/execution/exchange/sink/SinkChannel.java b/server/src/main/java/org/apache/iotdb/db/mpp/execution/exchange/sink/SinkChannel.java
index 94522a55648..630dd5b244c 100644
--- a/server/src/main/java/org/apache/iotdb/db/mpp/execution/exchange/sink/SinkChannel.java
+++ b/server/src/main/java/org/apache/iotdb/db/mpp/execution/exchange/sink/SinkChannel.java
@@ -251,7 +251,7 @@ public class SinkChannel implements ISinkChannel {
     }
     sequenceIdToTsBlock.clear();
     if (blocked != null) {
-      bufferRetainedSizeInBytes -= localMemoryManager.getQueryPool().tryComplete(blocked);
+      bufferRetainedSizeInBytes -= localMemoryManager.getQueryPool().tryCancel(blocked);
     }
     if (bufferRetainedSizeInBytes > 0) {
       localMemoryManager
diff --git a/server/src/main/java/org/apache/iotdb/db/mpp/execution/memory/MemoryPool.java b/server/src/main/java/org/apache/iotdb/db/mpp/execution/memory/MemoryPool.java
index 3c9dbb10e3a..0ff44714082 100644
--- a/server/src/main/java/org/apache/iotdb/db/mpp/execution/memory/MemoryPool.java
+++ b/server/src/main/java/org/apache/iotdb/db/mpp/execution/memory/MemoryPool.java
@@ -57,6 +57,9 @@ public class MemoryPool {
      */
     private final long maxBytesCanReserve;
 
+    /** Being marked means this future will be completed shortly. */
+    private boolean isMarked = false;
+
     private MemoryReservationFuture(
         String queryId,
         String fragmentInstanceId,
@@ -93,6 +96,14 @@ public class MemoryPool {
       return maxBytesCanReserve;
     }
 
+    public boolean isMarked() {
+      return isMarked;
+    }
+
+    public void setMarked(boolean marked) {
+      isMarked = marked;
+    }
+
     public static <V> MemoryReservationFuture<V> create(
         String queryId,
         String fragmentInstanceId,
@@ -257,27 +268,12 @@ public class MemoryPool {
     Validate.isTrue(
         future instanceof MemoryReservationFuture,
         "invalid future type " + future.getClass().getSimpleName());
-    future.cancel(true);
-    return ((MemoryReservationFuture<Void>) future).getBytesToReserve();
-  }
-
-  /**
-   * Complete the specified memory reservation. If the reservation has finished, do nothing.
-   *
-   * @param future The future returned from {@link #reserve(String, String, String, long, long)}
-   * @return If the future has not complete, return the number of bytes being reserved. Otherwise,
-   *     return 0.
-   */
-  public synchronized long tryComplete(ListenableFuture<Void> future) {
-    Validate.notNull(future);
-    // If the future is not a MemoryReservationFuture, it must have been completed.
-    if (future.isDone()) {
+    // The MemoryReservationFuture has been marked, which means that the related Memory Reservation
+    // has been recorded in MemoryPool.
+    if (((MemoryReservationFuture<Void>) future).isMarked()) {
       return 0L;
     }
-    Validate.isTrue(
-        future instanceof MemoryReservationFuture,
-        "invalid future type " + future.getClass().getSimpleName());
-    ((MemoryReservationFuture<Void>) future).set(null);
+    future.cancel(true);
     return ((MemoryReservationFuture<Void>) future).getBytesToReserve();
   }
 
@@ -332,6 +328,7 @@ public class MemoryPool {
               .computeIfAbsent(curQueryId, x -> new HashMap<>())
               .computeIfAbsent(curFragmentInstanceId, x -> new HashMap<>())
               .merge(curPlanNodeId, bytesToReserve, Long::sum);
+          future.setMarked(true);
           futureList.add(future);
           iterator.remove();
         }