[incubator-celeborn] branch main updated: [CELEBORN-321] When register shuffle failed, DataPushQueue should directly take the task queue to avoid NPE (#1258)

[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 251b923b [CELEBORN-321] When register shuffle failed, DataPushQueue 
should directly take the task queue to avoid NPE (#1258)
251b923b is described below

commit 251b923b5beb0bd13cdf8059a39edc7d6e435364
Author: Angerszhuuuu <angers....@gmail.com>
AuthorDate: Tue Feb 21 17:02:37 2023 +0800

    [CELEBORN-321] When register shuffle failed, DataPushQueue should directly 
take the task queue to avoid NPE (#1258)
---
 .../apache/celeborn/client/write/DataPushQueue.java | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git 
a/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java 
b/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java
index aa654269..5d202bf6 100644
--- a/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java
+++ b/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java
@@ -93,16 +93,19 @@ public class DataPushQueue {
         int partitionId = task.getPartitionId();
         Map<Integer, PartitionLocation> partitionLocationMap =
             client.getPartitionLocation(appId, shuffleId, numMappers, 
numPartitions);
-        PartitionLocation loc = partitionLocationMap.get(partitionId);
-
-        if (!reachLimitWorker.contains(loc.hostAndPushPort())) {
-          boolean reachLimit = pushState.reachLimit(loc.hostAndPushPort(), 
maxInFlight);
-          if (!reachLimit) {
-            iterator.remove();
-            return task;
-          } else {
-            reachLimitWorker.add(loc.hostAndPushPort());
+        if (partitionLocationMap != null) {
+          PartitionLocation loc = partitionLocationMap.get(partitionId);
+          if (!reachLimitWorker.contains(loc.hostAndPushPort())) {
+            boolean reachLimit = pushState.reachLimit(loc.hostAndPushPort(), 
maxInFlight);
+            if (!reachLimit) {
+              iterator.remove();
+              return task;
+            } else {
+              reachLimitWorker.add(loc.hostAndPushPort());
+            }
           }
+        } else {
+          return task;
         }
       }
     }