fix bufferstream read stream

apache · Aug 28, 2023 · 57df489 · 57df489
1 parent c2fded8
commit 57df489
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 24 deletions.
diff --git a/...ommon/src/main/java/org/apache/celeborn/plugin/flink/readclient/CelebornBufferStream.java b/...ommon/src/main/java/org/apache/celeborn/plugin/flink/readclient/CelebornBufferStream.java
@@ -149,8 +149,16 @@ public void close() {
   public void moveToNextPartitionIfPossible(long endedStreamId) {
     if (currentLocationIndex < locations.length) {
       if (currentLocationIndex > 0) {
-        logger.debug("Get end streamId {}", endedStreamId);
-        cleanStream(endedStreamId);
+        if (endedStreamId == streamId) {
+          logger.debug("Get end streamId {}", endedStreamId);
+          cleanStream(endedStreamId);
+        } else {
+          logger.warn(
+                  "Received unexpected stream end, current stream id {} received ended stream id {}",
+                  this.streamId,
+                  endedStreamId);
+          return;
+        }
       }
       try {
         openStreamInternal();

diff --git a/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/WordCountTest.scala b/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/WordCountTest.scala
@@ -84,20 +84,29 @@ class WordCountTest extends AnyFunSuite with Logging with MiniClusterFeature
     graph.setJobType(JobType.BATCH)
     if (flinkCluster != null) {
       flinkCluster.close()
-      flinkCluster = null
-      System.gc()
     }
+    flinkCluster = null
+    System.gc()
     val miniClusterConfiguration =
       (new MiniClusterConfiguration.Builder).setConfiguration(configuration).build()
-    flinkCluster = new MiniCluster(miniClusterConfiguration)
-    flinkCluster.start()
+    try {
+      flinkCluster = new MiniCluster(miniClusterConfiguration)
+      flinkCluster.start()
+
+      val jobGraph: JobGraph = StreamingJobGraphGenerator.createJobGraph(graph)
+      val jobID = flinkCluster.submitJob(jobGraph).get.getJobID
+      val jobResult = flinkCluster.requestJobResult(jobID).get
+      if (jobResult.getSerializedThrowable.isPresent)
+        throw new AssertionError(jobResult.getSerializedThrowable.get)
+      checkFlushingFileLength()
+    } finally {
+      if (flinkCluster != null) {
+        flinkCluster.close()
+        flinkCluster = null
+        System.gc()
+      }
+    }
 
-    val jobGraph: JobGraph = StreamingJobGraphGenerator.createJobGraph(graph)
-    val jobID = flinkCluster.submitJob(jobGraph).get.getJobID
-    val jobResult = flinkCluster.requestJobResult(jobID).get
-    if (jobResult.getSerializedThrowable.isPresent)
-      throw new AssertionError(jobResult.getSerializedThrowable.get)
-    checkFlushingFileLength()
   }
 
   private def checkFlushingFileLength(): Unit = {
@@ -135,20 +144,26 @@ class WordCountTest extends AnyFunSuite with Logging with MiniClusterFeature
     SplitHelper.runSplitRead(env)
     if (flinkCluster != null) {
       flinkCluster.close()
-      flinkCluster = null
-      System.gc()
     }
+    flinkCluster = null
+    System.gc()
     val miniClusterConfiguration =
       (new MiniClusterConfiguration.Builder).setConfiguration(configuration).build()
-    flinkCluster = new MiniCluster(miniClusterConfiguration)
-    flinkCluster.start()
-    val graph = env.getStreamGraph
-    graph.setGlobalStreamExchangeMode(GlobalStreamExchangeMode.ALL_EDGES_BLOCKING)
-    graph.setJobType(JobType.BATCH)
-    val jobGraph: JobGraph = StreamingJobGraphGenerator.createJobGraph(graph)
-    val jobID = flinkCluster.submitJob(jobGraph).get.getJobID
-    val jobResult = flinkCluster.requestJobResult(jobID).get
-    if (jobResult.getSerializedThrowable.isPresent)
-      throw new AssertionError(jobResult.getSerializedThrowable.get)
+    try {
+      flinkCluster = new MiniCluster(miniClusterConfiguration)
+      flinkCluster.start()
+      val graph = env.getStreamGraph
+      graph.setGlobalStreamExchangeMode(GlobalStreamExchangeMode.ALL_EDGES_BLOCKING)
+      graph.setJobType(JobType.BATCH)
+      val jobGraph: JobGraph = StreamingJobGraphGenerator.createJobGraph(graph)
+      val jobID = flinkCluster.submitJob(jobGraph).get.getJobID
+      val jobResult = flinkCluster.requestJobResult(jobID).get
+      if (jobResult.getSerializedThrowable.isPresent)
+        throw new AssertionError(jobResult.getSerializedThrowable.get)
+    } finally {
+      flinkCluster.close()
+      flinkCluster = null
+      System.gc()
+    }
   }
 }