From a690625e156001654241e92c45ba8d6fada0cae6 Mon Sep 17 00:00:00 2001
From: Zikun Ma <55695098+DanielWang2035@users.noreply.github.com>
Date: Mon, 19 May 2025 14:26:00 +0800
Subject: [PATCH] Load: fix memory leak when failed to send chunk data in first
 phase (#15518)

This commit addresses a memory leak issue in the load process by ensuring that memory usage is reduced even when the dispatch of a piece node fails. The key changes include:

- Storing the result of the dispatch call in a variable (isDispatchSuccess) before reducing memory usage.
- Deducting the memory usage prior to checking the dispatch result to avoid leaks.
- Returning false immediately after the reduction when the dispatch fails.

(cherry picked from commit 417ddd0becd8f85103d0776d47450f9936241e86)
---
 .../plan/scheduler/load/LoadTsFileScheduler.java   | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java
index f45e98052006e..49a9fc9e3be0b 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java
@@ -663,12 +663,8 @@ private boolean addOrSendChunkData(ChunkData chunkData) throws LoadFileException
           if (pieceNode.getDataSize() == 0) { // total data size has been reduced to 0
             break;
           }
-          if (!scheduler.dispatchOnePieceNode(pieceNode, replicaSet)) {
-            return false;
-          }
+          final boolean isDispatchSuccess = scheduler.dispatchOnePieceNode(pieceNode, replicaSet);
 
-          dataSize -= pieceNode.getDataSize();
-          block.reduceMemoryUsage(pieceNode.getDataSize());
           regionId2ReplicaSetAndNode.replace(
               sortedRegionId,
               new Pair<>(
@@ -678,6 +674,14 @@ private boolean addOrSendChunkData(ChunkData chunkData) throws LoadFileException
                       singleTsFileNode
                           .getTsFileResource()
                           .getTsFile()))); // can not just remove, because of deletion
+          dataSize -= pieceNode.getDataSize();
+          block.reduceMemoryUsage(pieceNode.getDataSize());
+
+          if (!isDispatchSuccess) {
+            // Currently there is no retry, so return directly
+            return false;
+          }
+
           if (isMemoryEnough()) {
             break;
           }