LuciferYang commented on code in PR #1265:
URL:
https://github.com/apache/incubator-uniffle/pull/1265#discussion_r1386138470
##########
client-spark/spark3/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java:
##########
@@ -89,6 +89,7 @@ public class RssShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
private final boolean isMemoryShuffleEnabled;
private final Function<String, Boolean> taskFailureCallback;
private final Set<Long> blockIds = Sets.newConcurrentHashSet();
+ private final Set<CompletableFuture> sendingSet =
Sets.newConcurrentHashSet();
Review Comment:
`Set<CompletableFuture<Long>>`?
##########
client-spark/spark3/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java:
##########
@@ -442,6 +442,8 @@ public Option<MapStatus> stop(boolean success) {
return Option.empty();
}
} finally {
+ // cancel all async thread task related
+ sendingSet.stream().forEach(eventTask -> eventTask.cancel(true));
Review Comment:
Is `cancel` certain to succeed? If `cancel` returns false, what negative
impact will it have?
##########
client-spark/spark3/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java:
##########
@@ -304,56 +306,54 @@ protected List<CompletableFuture<Long>> postBlockEvent(
LOG.error("Add event " + event + " to finishEventQueue fail");
}
});
- futures.add(shuffleManager.sendData(event));
+ CompletableFuture<Long> longCompletableFuture =
shuffleManager.sendData(event);
+ sendingSet.add(longCompletableFuture);
+ longCompletableFuture.thenApply(f ->
sendingSet.remove(longCompletableFuture));
+ futures.add(longCompletableFuture);
}
return futures;
}
+ /** @return false if interrupted, true if all blocks are sent, throw
exception if timeout */
@VisibleForTesting
- protected void checkBlockSendResult(Set<Long> blockIds) {
- boolean interrupted = false;
-
- try {
- long remainingMs = sendCheckTimeout;
- long end = System.currentTimeMillis() + remainingMs;
+ protected boolean checkBlockSendResult(Set<Long> blockIds) {
+ long remainingMs = sendCheckTimeout;
+ long end = System.currentTimeMillis() + remainingMs;
- while (true) {
+ while (true) {
+ finishEventQueue.clear();
+ checkIfBlocksFailed();
+ Set<Long> successBlockIds = shuffleManager.getSuccessBlockIds(taskId);
+ blockIds.removeAll(successBlockIds);
+ if (blockIds.isEmpty()) {
+ break;
+ }
+ if (finishEventQueue.isEmpty()) {
+ remainingMs = Math.max(end - System.currentTimeMillis(), 0);
+ Object event = null;
Review Comment:
```suggestion
Object event;
```
##########
client/src/main/java/org/apache/uniffle/client/util/ClientUtils.java:
##########
@@ -120,8 +126,17 @@ public static boolean waitUntilDoneOrFail(
try {
allFutures.get(10, TimeUnit.MILLISECONDS);
+ } catch (TimeoutException e) {
Review Comment:
We can merge these two Exception
##########
client-spark/spark3/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java:
##########
@@ -442,6 +442,8 @@ public Option<MapStatus> stop(boolean success) {
return Option.empty();
}
} finally {
+ // cancel all async thread task related
+ sendingSet.stream().forEach(eventTask -> eventTask.cancel(true));
Review Comment:
```suggestion
sendingSet.forEach(eventTask -> eventTask.cancel(true));
```
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
