[incubator-celeborn] branch main updated: [CELEBORN-333][FLINK] bypass unexpected backlog message when stream closed. (#1268)

[ethanfeng]

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

ethanfeng pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git


The following commit(s) were added to refs/heads/main by this push:
     new a963aa4b [CELEBORN-333][FLINK] bypass unexpected backlog message when 
stream closed. (#1268)
a963aa4b is described below

commit a963aa4b4f6e4df3dd793e0ac97bcef01cca26cc
Author: Shuang <lvshuang...@gmail.com>
AuthorDate: Thu Feb 23 18:49:14 2023 +0800

    [CELEBORN-333][FLINK] bypass unexpected backlog message when stream closed. 
(#1268)
---
 .../plugin/flink/network/ReadClientHandler.java    |  5 +++--
 .../plugin/flink/RemoteBufferStreamReader.java     | 25 ++++++++++++++++------
 2 files changed, 22 insertions(+), 8 deletions(-)

diff --git 
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java
 
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java
index 0da9f67d..8f768c26 100644
--- 
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java
+++ 
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java
@@ -66,10 +66,11 @@ public class ReadClientHandler extends BaseMessageHandler {
       case BACKLOG_ANNOUNCEMENT:
         BacklogAnnouncement backlogAnnouncement = (BacklogAnnouncement) msg;
         streamId = backlogAnnouncement.getStreamId();
-        if (streamHandlers.containsKey(streamId)) {
+        Consumer<RequestMessage> consumer = streamHandlers.get(streamId);
+        if (consumer != null) {
           logger.debug(
               "received streamId: {}, backlog: {}", streamId, 
backlogAnnouncement.getBacklog());
-          streamHandlers.get(streamId).accept(msg);
+          consumer.accept(msg);
         } else {
           logger.warn("Unexpected streamId received: {}", streamId);
         }
diff --git 
a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReader.java
 
b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReader.java
index d423f4d8..d55441a9 100644
--- 
a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReader.java
+++ 
b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReader.java
@@ -45,7 +45,7 @@ public class RemoteBufferStreamReader extends CreditListener {
   private Consumer<ByteBuf> dataListener;
   private Consumer<Throwable> failureListener;
   private RssBufferStream bufferStream;
-  private boolean closed = false;
+  private volatile boolean closed = false;
   private Consumer<RequestMessage> messageConsumer;
 
   public RemoteBufferStreamReader(
@@ -91,8 +91,17 @@ public class RemoteBufferStreamReader extends CreditListener 
{
 
   public void close() {
     isOpen = false;
-    
client.getReadClientHandler().removeHandler(this.bufferStream.getStreamId());
-    bufferStream.close();
+    // need set closed first before remove Handler
+    closed = true;
+    if (this.bufferStream != null) {
+      
client.getReadClientHandler().removeHandler(this.bufferStream.getStreamId());
+      bufferStream.close();
+    } else {
+      logger.warn(
+          "bufferStream is null when closed, shuffleId: {}, partitionId: {}",
+          shuffleId,
+          partitionId);
+    }
   }
 
   public boolean isOpened() {
@@ -100,8 +109,10 @@ public class RemoteBufferStreamReader extends 
CreditListener {
   }
 
   public void notifyAvailableCredits(int numCredits) {
-    ReadAddCredit addCredit = new 
ReadAddCredit(this.bufferStream.getStreamId(), numCredits);
-    bufferStream.addCredit(addCredit);
+    if (!closed) {
+      ReadAddCredit addCredit = new 
ReadAddCredit(this.bufferStream.getStreamId(), numCredits);
+      bufferStream.addCredit(addCredit);
+    }
   }
 
   public ByteBuf requestBuffer() {
@@ -109,7 +120,9 @@ public class RemoteBufferStreamReader extends 
CreditListener {
   }
 
   public void backlogReceived(int backlog) {
-    bufferPool.reserveBuffers(this, backlog);
+    if (!closed) {
+      bufferPool.reserveBuffers(this, backlog);
+    }
   }
 
   public void dataReceived(ReadData readData) {