Load: fix memory leak when failed to send chunk data in first phase (apache#15518) (apache#15542)

DanielWang2035 · SyncGithubBot · commit 8ae8f283264e · 2025-05-20T17:17:17.000+08:00
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 417ddd0)
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
@@ -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;
           }
