[incubator-celeborn] branch main updated: [CELEBORN-331] submitRetryPushData should throw PUSH_DATA_CREATE_CONNECTION_FAIL_MASTER too (#1266)

[angerszhuuuu]

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

angerszhuuuu 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 3067efcf [CELEBORN-331] submitRetryPushData should throw 
PUSH_DATA_CREATE_CONNECTION_FAIL_MASTER too (#1266)
3067efcf is described below

commit 3067efcfd31475fc44c38fad0813fd3350a55f57
Author: Angerszhuuuu <angers....@gmail.com>
AuthorDate: Thu Feb 23 14:57:11 2023 +0800

    [CELEBORN-331] submitRetryPushData should throw 
PUSH_DATA_CREATE_CONNECTION_FAIL_MASTER too (#1266)
    
    * [CELEBORN-331] submitRetryPushData should throw 
PUSH_DATA_CREATE_CONNECTION_FAIL_MASTER too
---
 .../org/apache/celeborn/client/ShuffleClientImpl.java  | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 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
index ad95aab6..5b71e2f3 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -176,14 +176,14 @@ public class ShuffleClientImpl extends ShuffleClient {
       byte[] body,
       int batchId,
       PartitionLocation loc,
-      RpcResponseCallback callback,
+      RpcResponseCallback wrappedCallback,
       PushState pushState,
       StatusCode cause,
       int remainReviveTimes) {
     int partitionId = loc.getId();
     if (!revive(
         applicationId, shuffleId, mapId, attemptId, partitionId, 
loc.getEpoch(), loc, cause)) {
-      callback.onFailure(
+      wrappedCallback.onFailure(
           new CelebornIOException("Revive Failed, remain revive times " + 
remainReviveTimes));
     } else if (mapperEnded(shuffleId, mapId, attemptId)) {
       logger.debug(
@@ -201,20 +201,26 @@ public class ShuffleClientImpl extends ShuffleClient {
 
           PushData newPushData =
               new PushData(MASTER_MODE, shuffleKey, newLoc.getUniqueId(), 
newBuffer);
-          client.pushData(newPushData, pushDataTimeout, callback);
+          client.pushData(newPushData, pushDataTimeout, wrappedCallback);
         } else {
           throw new RuntimeException(
               "Mock push data submit retry failed. remainReviveTimes = " + 
remainReviveTimes + ".");
         }
-      } catch (Exception ex) {
+      } catch (Exception e) {
         logger.warn(
             "Exception raised while pushing data for shuffle {} map {} attempt 
{}" + " batch {}.",
             shuffleId,
             mapId,
             attemptId,
             batchId,
-            ex);
-        callback.onFailure(ex);
+            e);
+        wrappedCallback.onFailure(
+            new Exception(
+                StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_MASTER.getMessage()
+                    + "! "
+                    + e.getMessage()
+                    + ". "
+                    + newLoc));
       }
     }
   }