(beam) branch master updated: [Java] Simplify MoreFutures usage in case it is cause of issues with ForkJoin pool stuckness (#35673)

[sc...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

scwhittle pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new 4043c901e99 [Java] Simplify MoreFutures usage in case it is cause of issues with ForkJoin pool stuckness (#35673)
4043c901e99 is described below

commit 4043c901e99464396a318911a996a2338418e799
Author: Sam Whittle <sc...@users.noreply.github.com>
AuthorDate: Mon Jul 28 18:07:31 2025 +0200

    [Java] Simplify MoreFutures usage in case it is cause of issues with ForkJoin pool stuckness (#35673)
---
 .../java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java | 6 ++----
 .../core/src/main/java/org/apache/beam/sdk/util/MoreFutures.java   | 7 +++++++
 .../main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtil.java | 2 +-
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java
index c1b56a2b445..8404f15842e 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java
@@ -1250,9 +1250,8 @@ public abstract class WriteFiles<UserT, DestinationT, OutputT>
       // before we return from this processElement call. This allows us to perform the writes/closes
       // in parallel with the prior elements close calls and bounds the amount of data buffered to
       // limit the number of OOMs.
-      CompletionStage<List<Void>> pastCloseFutures = MoreFutures.allAsList(closeFutures);
+      CompletionStage<Void> pastCloseFutures = MoreFutures.allOf(closeFutures);
       closeFutures.clear();
-
       // Close all writers in the background
       for (Map.Entry<DestinationT, Writer<DestinationT, OutputT>> entry : writers.entrySet()) {
         int shard = c.element().getKey().getShardNumber();
@@ -1267,7 +1266,6 @@ public abstract class WriteFiles<UserT, DestinationT, OutputT>
                 new FileResult<>(writer.getOutputFile(), shard, window, c.pane(), entry.getKey())));
         closeWriterInBackground(writer);
       }
-
       // Block on completing the past closes before returning. We do so after starting the current
       // closes in the background so that they can happen in parallel.
       MoreFutures.get(pastCloseFutures);
@@ -1293,7 +1291,7 @@ public abstract class WriteFiles<UserT, DestinationT, OutputT>
     @FinishBundle
     public void finishBundle(FinishBundleContext c) throws Exception {
       try {
-        MoreFutures.get(MoreFutures.allAsList(closeFutures));
+        MoreFutures.get(MoreFutures.allOf(closeFutures));
         // If all writers were closed without exception, output the results to the next stage.
         for (KV<Instant, FileResult<DestinationT>> result : deferredOutput) {
           c.output(result.getValue(), result.getKey(), result.getValue().getWindow());
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/MoreFutures.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/MoreFutures.java
index 0999f2ad077..441e604af3a 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/MoreFutures.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/MoreFutures.java
@@ -159,6 +159,13 @@ public class MoreFutures {
         nothing -> Arrays.stream(f).map(CompletableFuture::join).collect(Collectors.toList()));
   }
 
+  public static <T> CompletionStage<Void> allOf(
+      Collection<? extends CompletionStage<? extends T>> futures) {
+    // CompletableFuture.allOf completes exceptionally if any of the futures do.
+    CompletableFuture<? extends T>[] f = futuresToCompletableFutures(futures);
+    return CompletableFuture.allOf(f);
+  }
+
   /**
    * An object that represents either a result or an exceptional termination.
    *
diff --git a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtil.java b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtil.java
index c56f5be0000..77670eafbb4 100644
--- a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtil.java
+++ b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtil.java
@@ -990,7 +990,7 @@ public class GcsUtil {
 
     try {
       try {
-        MoreFutures.get(MoreFutures.allAsList(futures));
+        MoreFutures.get(MoreFutures.allOf(futures));
       } catch (ExecutionException e) {
         if (e.getCause() instanceof FileNotFoundException) {
           throw (FileNotFoundException) e.getCause();