afterincomparableyum commented on code in PR #3724:
URL: https://github.com/apache/celeborn/pull/3724#discussion_r3408264691


##########
cpp/celeborn/client/ShuffleClient.cpp:
##########
@@ -1096,6 +1176,22 @@ std::optional<std::unordered_map<int, int>> 
ShuffleClientImpl::reviveBatch(
         case protocol::StatusCode::SUCCESS: {
           partitionLocationMap->set(
               partitionInfo.partitionId, partitionInfo.partition);
+          // Revive moved this partition off the failed worker(s); drop the
+          // push exclusion on both the old and new locations.
+          if (pushExcludeWorkerOnFailureEnabled_) {
+            if (auto oldIter = oldLocationMap.find(partitionInfo.partitionId);
+                oldIter != oldLocationMap.end() && oldIter->second) {
+              pushExcludedWorkers_.erase(oldIter->second->hostAndPushPort());

Review Comment:
   So here we only erase on SUCCESS. However, java client the old-location 
removal sits before and outside the SUCCESS check. We should copy this java 
behavior in c++ as well. 



##########
cpp/celeborn/client/ShuffleClient.cpp:
##########
@@ -1145,6 +1241,116 @@ bool ShuffleClientImpl::revive(
   return false;
 }
 
+namespace {
+// Mirrors org.apache.celeborn.common.util.ExceptionUtils#connectFail.
+bool connectFail(const std::string& message) {
+  return (message.find("Connection from ") != std::string::npos &&
+          message.find(" closed") != std::string::npos) ||
+      message.find("Connection reset by peer") != std::string::npos ||
+      message.find("Failed to send request ") != std::string::npos;
+}
+} // namespace
+
+protocol::StatusCode ShuffleClientImpl::getPushDataFailCause(
+    const std::string& message) {
+  VLOG(1) << "Push data failed cause message: " << message;
+  if (message.empty()) {
+    LOG(ERROR) << "Push data throw unexpected exception";
+    return protocol::StatusCode::PUSH_DATA_FAIL_NON_CRITICAL_CAUSE;
+  }
+  // The transport wraps the worker's error, so match the StatusCode name as a
+  // substring (via protocol::toString) rather than a prefix. Order follows the
+  // Java client so more specific causes win.
+  static constexpr std::array<protocol::StatusCode, 13> kCandidateCauses = {
+      protocol::StatusCode::PUSH_DATA_FAIL_NON_CRITICAL_CAUSE_REPLICA,
+      protocol::StatusCode::PUSH_DATA_WRITE_FAIL_REPLICA,
+      protocol::StatusCode::PUSH_DATA_WRITE_FAIL_PRIMARY,
+      protocol::StatusCode::PUSH_DATA_CREATE_CONNECTION_FAIL_PRIMARY,
+      protocol::StatusCode::PUSH_DATA_CREATE_CONNECTION_FAIL_REPLICA,
+      protocol::StatusCode::PUSH_DATA_CONNECTION_EXCEPTION_PRIMARY,
+      protocol::StatusCode::PUSH_DATA_CONNECTION_EXCEPTION_REPLICA,
+      protocol::StatusCode::PUSH_DATA_TIMEOUT_PRIMARY,
+      protocol::StatusCode::PUSH_DATA_TIMEOUT_REPLICA,
+      protocol::StatusCode::REPLICATE_DATA_FAILED,
+      protocol::StatusCode::PUSH_DATA_PRIMARY_WORKER_EXCLUDED,
+      protocol::StatusCode::PUSH_DATA_REPLICA_WORKER_EXCLUDED,
+      protocol::StatusCode::PUSH_DATA_FAIL_PARTITION_NOT_FOUND,
+  };
+  for (auto cause : kCandidateCauses) {
+    if (message.find(protocol::toString(cause)) != std::string::npos) {
+      return cause;
+    }
+  }
+  if (message.find("Timed out") != std::string::npos) {
+    // A client-side push timeout surfaces as a folly FutureTimeout ("Timed
+    // out") with no StatusCode token; classify it as a push timeout so worker
+    // exclusion engages like the Java client.
+    return protocol::StatusCode::PUSH_DATA_TIMEOUT_PRIMARY;
+  }
+  if (connectFail(message)) {
+    // Thrown when push to primary worker connection fails.
+    return protocol::StatusCode::PUSH_DATA_CONNECTION_EXCEPTION_PRIMARY;
+  }
+  return protocol::StatusCode::PUSH_DATA_FAIL_NON_CRITICAL_CAUSE;
+}
+
+std::optional<protocol::StatusCode> ShuffleClientImpl::classifyPushFailure(
+    const std::string& errorMsg,
+    int remainingReviveTimes,
+    PushState& pushState) {
+  const protocol::StatusCode cause = getPushDataFailCause(errorMsg);
+  if (remainingReviveTimes <= 0) {
+    // Out of revive attempts: surface a terminal failure annotated with the
+    // cause, like Java's `new CelebornIOException(cause, e)`.
+    pushState.setException(std::make_unique<std::runtime_error>(
+        protocol::toString(cause) + ": " + errorMsg));

Review Comment:
   NIT: Doesn't errorMsg already begin with the StatusCode name? Is 
toString(cause) needed?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to