fix(common): FutureUtils:allOf should always throw root cause exception (#18456)

FutureUtils.allOf() has a race condition that causes the original root-cause exception to be silently replaced by a CancellationException, making it impossible to diagnose failures in any code path that uses it — most notably MultipleSparkJobExecutionStrategy.performClustering(), which executes clustering groups in parallel using FutureUtils.allOf(). Fixing the same in this patch




---------

Co-authored-by: Krishen Bhan <“bkrishen@uber.com”>

Author: kbuci

hudi-common/src/main/java/org/apache/hudi/common/util/FutureUtils.java

@@ -20,6 +20,8 @@
 
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 
 /**
@@ -42,24 +44,37 @@ public class FutureUtils {
    */
   public static <T> CompletableFuture<List<T>> allOf(List<CompletableFuture<T>> futures) {
     CompletableFuture<Void> union = CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]));
+    AtomicReference<Throwable> firstFailure = new AtomicReference<>();
 
     futures.forEach(future -> {
       // NOTE: We add a callback to every future, to cancel all the other not yet completed futures,
       //       which will be providing for an early termination semantic: whenever any of the futures
       //       fail other futures will be cancelled and the exception will be returned as a result
       future.whenComplete((ignored, throwable) -> {
         if (throwable != null) {
+          firstFailure.compareAndSet(null, throwable);
+          // Note that {@link CompletableFuture#cancel} does not interrupt the other underlying tasks;
+          // it only marks their futures as cancelled. The tasks will still run to completion.
           futures.forEach(f -> f.cancel(true));
-          union.completeExceptionally(throwable);
         }
       });
     });
 
-    return union.thenApply(aVoid ->
-        futures.stream()
-            // NOTE: This join wouldn't block, since all the
-            //       futures are completed at this point.
-            .map(CompletableFuture::join)
-            .collect(Collectors.toList()));
+    return union.handle((aVoid, throwable) -> {
+      Throwable realCause = firstFailure.get();
+      // Prefer the first real failure captured by the whenComplete callbacks over whatever
+      // exception the {@link CompletableFuture#allOf} BiRelay propagated, since the BiRelay may
+      // receive either the original exception or a CancellationException from cancel depending on
+      // timing. In the unexpected case that realCause is null but throwable is not null, the
+      // subsequent join call will throw it.
+      if (realCause != null) {
+        throw new CompletionException(realCause);
+      }
+      // NOTE: This join wouldn't block, since all the
+      //       futures are completed at this point (allOf guarantees this).
+      return futures.stream()
+          .map(CompletableFuture::join)
+          .collect(Collectors.toList());
+    });
   }
 }

hudi-common/src/test/java/org/apache/hudi/common/util/TestFutureUtils.java

@@ -0,0 +1,128 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hudi.common.util;
+
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * Tests for {@link FutureUtils}.
+ */
+public class TestFutureUtils {
+
+  @Test
+  public void testAllOfSuccess() {
+    CompletableFuture<String> f1 = CompletableFuture.completedFuture("a");
+    CompletableFuture<String> f2 = CompletableFuture.completedFuture("b");
+    CompletableFuture<String> f3 = CompletableFuture.completedFuture("c");
+
+    List<String> results = FutureUtils.allOf(Arrays.asList(f1, f2, f3)).join();
+    assertEquals(Arrays.asList("a", "b", "c"), results);
+  }
+
+  @Test
+  public void testAllOfSingleFailurePreservesOriginalExceptionAndCancelsOthers() {
+    IOException originalCause = new IOException("disk failed");
+
+    CompletableFuture<String> f1 = new CompletableFuture<>();
+    CompletableFuture<String> f2 = new CompletableFuture<>();
+    CompletableFuture<String> f3 = new CompletableFuture<>();
+
+    CompletableFuture<List<String>> result = FutureUtils.allOf(Arrays.asList(f1, f2, f3));
+
+    f1.completeExceptionally(originalCause);
+
+    CompletionException thrown = assertThrows(CompletionException.class, result::join);
+    assertEquals(originalCause, thrown.getCause(),
+        "The original IOException should be preserved, not masked by CancellationException");
+    assertTrue(f2.isCancelled(), "f2 should be cancelled after f1 fails");
+    assertTrue(f3.isCancelled(), "f3 should be cancelled after f1 fails");
+  }
+
+  /**
+   * Ensure that the original exception is preserved even when the futures are completed in a concurrent manner.
+   */
+  @Test
+  public void testAllOfPreservesOriginalExceptionUnderConcurrency() throws Exception {
+    IOException originalCause = new IOException("disk failed");
+    ExecutorService executor = Executors.newFixedThreadPool(4);
+    try {
+      CountDownLatch allStarted = new CountDownLatch(3);
+
+      CompletableFuture<String> f1 = new CompletableFuture<>();
+      CompletableFuture<String> f2 = new CompletableFuture<>();
+      CompletableFuture<String> f3 = new CompletableFuture<>();
+
+      executor.submit(() -> {
+        allStarted.countDown();
+        try {
+          allStarted.await(5, TimeUnit.SECONDS);
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+        }
+        f1.completeExceptionally(originalCause);
+        return null;
+      });
+      executor.submit(() -> {
+        allStarted.countDown();
+        try {
+          allStarted.await(5, TimeUnit.SECONDS);
+          Thread.sleep(50);
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+        }
+        f2.complete("b");
+        return null;
+      });
+      executor.submit(() -> {
+        allStarted.countDown();
+        try {
+          allStarted.await(5, TimeUnit.SECONDS);
+          Thread.sleep(50);
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+        }
+        f3.complete("c");
+        return null;
+      });
+
+      CompletableFuture<List<String>> result = FutureUtils.allOf(Arrays.asList(f1, f2, f3));
+
+      CompletionException thrown = assertThrows(CompletionException.class, result::join);
+      assertInstanceOf(IOException.class, thrown.getCause(),
+          "original IOException must be preserved, got: " + thrown.getCause());
+    } finally {
+      executor.shutdownNow();
+    }
+  }
+}