da-daken commented on code in PR #926:
URL: https://github.com/apache/flink-agents/pull/926#discussion_r3792071970
##########
runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java:
##########
@@ -148,6 +162,142 @@ public <T> T executeAsync(ContinuationContext context,
Supplier<T> supplier) thr
return (T) context.getAsyncResultRef().get();
}
+ /**
+ * Executes all suppliers as one async batch and returns one {@link
Outcome} per supplier.
+ * Supplier failures are captured in their own outcome so one failed
supplier does not abort the
+ * whole batch.
+ *
+ * @param context the continuation context for this action
+ * @param suppliers the suppliers to execute
+ * @param timeout the timeout for the whole batch; null or non-positive
means no timeout
+ * @param <T> the result type
+ * @return outcomes in supplier order
+ */
+ @SuppressWarnings("unchecked")
+ public <T> List<Outcome<T>> executeAllAsync(
+ ContinuationContext context,
+ List<Callable<T>> suppliers,
+ Duration timeout,
+ int maxParallelism)
+ throws Exception {
+ context.clearAsyncState();
+ if (suppliers.isEmpty()) {
+ return List.of();
+ }
+
+ final int batchSize = suppliers.size();
+ CompletableFuture<Outcome<T>>[] slots = new
CompletableFuture[batchSize];
+ boolean[] counted = new boolean[batchSize];
+ int completed = 0;
+ int nextToSubmit = 0;
+ int parallelismLimit = Math.min(Math.max(maxParallelism, 1),
batchSize);
+
+ long deadlineNanos = getDeadlineNanos(timeout);
+ CompletableFuture<Void> batchBarrier = new CompletableFuture<>();
+ context.setPendingBatchFuture(batchBarrier, deadlineNanos);
+
+ while (completed < batchSize) {
+ if (System.nanoTime() >= deadlineNanos) {
+ TimeoutException exception =
+ new TimeoutException(
+ "Async durable batch execution timed out after
" + timeout);
+ batchBarrier.cancel(true);
+ context.setPendingBatchFuture(null);
+ return collectBatchOutcomesOnTimeout(slots, exception);
+ }
+
+ while (nextToSubmit < batchSize && countInFlight(slots,
nextToSubmit) < parallelismLimit) {
+ int index = nextToSubmit++;
+ Callable<T> supplier = suppliers.get(index);
+ slots[index] =
+ CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ return
Outcome.success(supplier.call());
+ } catch (Exception e) {
+ return Outcome.failure(e);
+ }
+ }, asyncExecutor);
+ }
+
+ for (int i = 0; i < nextToSubmit; i++) {
+ if (!counted[i] && slots[i].isDone()) {
+ counted[i] = true;
+ completed++;
+ }
+ }
+
+ if (completed < batchSize) {
+ Continuation.yield(SCOPE);
+ }
+ }
+
+ batchBarrier.complete(null);
+ context.setPendingBatchFuture(null);
+ return collectBatchOutcomes(Arrays.asList(slots));
+ }
+
+ private static <T> int countInFlight(
+ CompletableFuture<Outcome<T>>[] slots, int submittedCount) {
+ int inFlight = 0;
+ for (int i = 0; i < submittedCount; i++) {
+ if (!slots[i].isDone()) {
+ inFlight++;
+ }
+ }
+ return inFlight;
+ }
+
+
+ /**
+ * Collects per-slot outcomes after the batch barrier completes normally.
+ *
+ * <p>Each supplier already wraps success and failure into an {@link
Outcome}, so {@code join()}
+ * returns that outcome rather than throwing for ordinary tool exceptions.
+ */
+ private static <T> List<Outcome<T>> collectBatchOutcomes(
+ List<CompletableFuture<Outcome<T>>> futures) {
+ List<Outcome<T>> results = new ArrayList<>(futures.size());
+ for (CompletableFuture<Outcome<T>> future : futures) {
+ results.add(future.join());
+ }
+ return results;
+ }
+
+ /**
+ * Collects per-slot outcomes when the batch deadline elapses.
+ *
+ * <p>Completed slots keep their success or failure outcome. Only slots
that are still running
+ * (or become cancelled) are finalized as timeout failures. {@code
cancel(true)} is attempted
+ * only for unfinished futures; a future that completes between the check
and cancel stays
+ * non-cancelled and is collected as a normal outcome.
+ */
+ private static <T> List<Outcome<T>> collectBatchOutcomesOnTimeout(
+ CompletableFuture<Outcome<T>>[] futures, TimeoutException
timeoutException) {
+ List<Outcome<T>> results = new ArrayList<>(futures.length);
+ for (CompletableFuture<Outcome<T>> future : futures) {
+ if (future == null) {
+ results.add(Outcome.failure(timeoutException));
Review Comment:
> Would leaving those slots PENDING, so recovery can still run them, fit
better than recording them as failed?
Thanks for catching that – this looks friendlier than just failing
everything. I'll go with your suggestion.
--
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]
