(celeborn) branch branch-0.5 updated: [CELEBORN-1867][0.5] Fix flink client memory leak of TransportResponseHandler#outstandingRpcs for handling addCredit response

[rexxiong]

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

rexxiong pushed a commit to branch branch-0.5
in repository https://gitbox.apache.org/repos/asf/celeborn.git


The following commit(s) were added to refs/heads/branch-0.5 by this push:
     new 2c926746c [CELEBORN-1867][0.5] Fix flink client memory leak of 
TransportResponseHandler#outstandingRpcs for handling addCredit response
2c926746c is described below

commit 2c926746c33e731d5ea314c1e1e19df52cd2ad62
Author: codenohup <huangxu.wal...@gmail.com>
AuthorDate: Fri Feb 28 10:42:54 2025 +0800

    [CELEBORN-1867][0.5] Fix flink client memory leak of 
TransportResponseHandler#outstandingRpcs for handling addCredit response
    
    ### What changes were proposed in this pull request?
    
    Backport https://github.com/apache/celeborn/pull/3103 to branch-0.5.
    
    When I tested the performance of Flink with Celeborn in session mode, using 
both the shuffle-service plugin and hybrid shuffle integration strategies, I 
noticed that the task heap continuously grew even when no jobs were running.
    
    The issue arises because the Celeborn client sends addCredit or 
notifyRequiredSegment requests, expecting a response. This creates a callback 
and maintains a reference to CelebornBufferStream, SingleInputGate, and 
StreamTask.
    These callbacks are stored in TransportResponseHandler#outstandingRpcs and 
are cleared upon receiving a response.
    
    However, the Worker does not send a response for these two RPC requests, 
which leads to a significant memory leak in the client.
    
    ### Why are the changes needed?
    
    It may cause flink client memory leak.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    Test by rerun TPCDS in Flink session mode.
    
    Closes #3123 from SteNicholas/CELEBORN-1867.
    
    Authored-by: codenohup <huangxu.wal...@gmail.com>
    Signed-off-by: Shuang <lvshuang....@alibaba-inc.com>
---
 .../service/deploy/worker/FetchHandler.scala        | 21 +++++++++++++++++----
 1 file changed, 17 insertions(+), 4 deletions(-)

diff --git 
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/FetchHandler.scala
 
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/FetchHandler.scala
index 934a4e967..42009835e 100644
--- 
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/FetchHandler.scala
+++ 
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/FetchHandler.scala
@@ -35,7 +35,7 @@ import 
org.apache.celeborn.common.CelebornConf.MAX_CHUNKS_BEING_TRANSFERRED
 import org.apache.celeborn.common.exception.CelebornIOException
 import org.apache.celeborn.common.internal.Logging
 import org.apache.celeborn.common.meta.{DiskFileInfo, FileInfo, MapFileMeta, 
MemoryFileInfo, ReduceFileMeta}
-import org.apache.celeborn.common.network.buffer.{FileChunkBuffers, 
MemoryChunkBuffers, NioManagedBuffer}
+import org.apache.celeborn.common.network.buffer.{FileChunkBuffers, 
MemoryChunkBuffers, NettyManagedBuffer, NioManagedBuffer}
 import org.apache.celeborn.common.network.client.{RpcResponseCallback, 
TransportClient}
 import org.apache.celeborn.common.network.protocol._
 import org.apache.celeborn.common.network.server.BaseMessageHandler
@@ -106,7 +106,7 @@ class FetchHandler(
       case r: BufferStreamEnd =>
         handleEndStreamFromClient(client, r.getStreamId)
       case r: ReadAddCredit =>
-        handleReadAddCredit(client, r.getCredit, r.getStreamId)
+        handleReadAddCredit(client, r.getCredit, r.getStreamId, -1)
       case r: ChunkFetchRequest =>
         handleChunkFetchRequest(client, r.streamChunkSlice, r)
       case unknown: RequestMessage =>
@@ -173,7 +173,11 @@ class FetchHandler(
           bufferStreamEnd.getStreamId,
           bufferStreamEnd.getStreamType)
       case readAddCredit: PbReadAddCredit =>
-        handleReadAddCredit(client, readAddCredit.getCredit, 
readAddCredit.getStreamId)
+        handleReadAddCredit(
+          client,
+          readAddCredit.getCredit,
+          readAddCredit.getStreamId,
+          rpcRequest.requestId)
       case chunkFetchRequest: PbChunkFetchRequest =>
         handleChunkFetchRequest(
           client,
@@ -471,13 +475,22 @@ class FetchHandler(
     }
   }
 
-  def handleReadAddCredit(client: TransportClient, credit: Int, streamId: 
Long): Unit = {
+  def handleReadAddCredit(
+      client: TransportClient,
+      credit: Int,
+      streamId: Long,
+      requestId: Long): Unit = {
     val shuffleKey = creditStreamManager.getStreamShuffleKey(streamId)
     if (shuffleKey != null) {
       workerSource.recordAppActiveConnection(
         client,
         shuffleKey)
       creditStreamManager.addCredit(credit, streamId)
+      if (requestId != -1) {
+        client.getChannel.writeAndFlush(new RpcResponse(
+          requestId,
+          NettyManagedBuffer.EmptyBuffer))
+      }
     }
   }