fix

apache · Aug 29, 2023 · f8df6c3 · f8df6c3
1 parent f7cf444
commit f8df6c3
Showing 1 changed file with 2 additions and 7 deletions.
diff --git a/client-spark/spark3/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java b/client-spark/spark3/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java
@@ -326,23 +326,18 @@ protected void checkBlockSendResult(Set<Long> blockIds) {
           if (blockIds.isEmpty()) {
             break;
           }
-          LOG.info("Wait " + blockIds.size() + " blocks sent to shuffle server");
           if (finishEventQueue.isEmpty()) {
+            remainingMs = Math.max(end - System.currentTimeMillis(), 0);
             Object event = finishEventQueue.poll(remainingMs, TimeUnit.MILLISECONDS);
             if (event == null) {
               break;
             }
           }
         } catch (InterruptedException e) {
           interrupted = true;
-        } finally {
-          remainingMs = end - System.currentTimeMillis();
-          if (remainingMs < 0) {
-            remainingMs = 0;
-          }
         }
       }
-      if (remainingMs <= 0) {
+      if (end - System.currentTimeMillis() <= 0) {
         String errorMsg =
             "Timeout: Task["
                 + taskId