[CELEBORN-656][FOLLOWUP] Fix wrong message call when revive return ST…

…AGE_END ### What changes were proposed in this pull request? As title ### Why are the changes needed? ### Does this PR introduce _any_ user-facing change? ### How was this patch tested? Closes #1755 from AngersZhuuuu/CELEBORN-656-FOLLOWUP. Authored-by: Angerszhuuuu <[email protected]> Signed-off-by: zky.zhoukeyong <[email protected]>
apache · Jul 24, 2023 · 67c18e6 · 67c18e6
1 parent b8cdf36
commit 67c18e6
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -729,7 +729,7 @@ Map<Integer, Integer> reviveBatch(
             pushExcludedWorkers.remove(loc.getPeer().hostAndPushPort());
           }
         } else if (StatusCode.STAGE_ENDED.getValue() == statusCode) {
-          stageEnded(shuffleId);
+          stageEndShuffleSet.add(shuffleId);
           return results;
         } else if (StatusCode.SHUFFLE_NOT_REGISTERED.getValue() == statusCode) {
           logger.error("SHUFFLE_NOT_REGISTERED!");
@@ -1650,10 +1650,10 @@ public void setupLifecycleManagerRef(RpcEndpointRef endpointRef) {
 
   boolean mapperEnded(int shuffleId, int mapId) {
     return (mapperEndMap.containsKey(shuffleId) && mapperEndMap.get(shuffleId).contains(mapId))
-        || stageEnded(shuffleId);
+        || isStageEnded(shuffleId);
   }
 
-  protected boolean stageEnded(int shuffleId) {
+  protected boolean isStageEnded(int shuffleId) {
     return stageEndShuffleSet.contains(shuffleId);
   }