From 440d558f40bc30481283a982b16e996a72c4c63a Mon Sep 17 00:00:00 2001 From: Gian Merlino Date: Tue, 6 Oct 2026 22:04:48 -0700 Subject: [PATCH] minor: Complete CollectedAsyncResource immediately on failure. Rather than waiting for all source resources to complete, we can immediately complete the target resource if any source resource fails. This prevents doing needless work gathering up all the other resources, and allows them to be closed (and therefore potentially canceled) more quickly. --- .../common/asyncresource/AsyncResources.java | 3 + .../asyncresource/CollectAsyncResource.java | 25 +++++++- .../asyncresource/AsyncResourcesTest.java | 18 ++++++ .../CollectAsyncResourceTest.java | 57 ++++++++++++++++++- 4 files changed, 99 insertions(+), 4 deletions(-) diff --git a/processing/src/main/java/org/apache/druid/common/asyncresource/AsyncResources.java b/processing/src/main/java/org/apache/druid/common/asyncresource/AsyncResources.java index d4cf84ffa6c1..e4f23726eb9d 100644 --- a/processing/src/main/java/org/apache/druid/common/asyncresource/AsyncResources.java +++ b/processing/src/main/java/org/apache/druid/common/asyncresource/AsyncResources.java @@ -118,6 +118,9 @@ public static AsyncResource transform( /** * Returns an {@link AsyncResource} that collects a list of underlying resources into a single lifecycle. * + *

The returned resource becomes ready once every resource in {@code asyncResources} has succeeded, or as soon as + * any of them fails. + * *

Once this method returns, the returned {@link AsyncResource} is the caller's to close, and closing it also * closes every resource in {@code asyncResources}, so the caller must not close them itself. If this method throws, * nothing has been taken over and the caller still owns all of them. diff --git a/processing/src/main/java/org/apache/druid/common/asyncresource/CollectAsyncResource.java b/processing/src/main/java/org/apache/druid/common/asyncresource/CollectAsyncResource.java index 3955e0b25a48..600d441a1754 100644 --- a/processing/src/main/java/org/apache/druid/common/asyncresource/CollectAsyncResource.java +++ b/processing/src/main/java/org/apache/druid/common/asyncresource/CollectAsyncResource.java @@ -24,6 +24,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; /** @@ -37,6 +38,12 @@ public class CollectAsyncResource implements AsyncResource> private final AtomicInteger readyCount = new AtomicInteger(0); private final SettableAsyncResource> targetResource = new SettableAsyncResource<>(); + /** + * Whether a source has failed. Only the first failure is reported, since {@link #targetResource} can only be + * completed once. + */ + private final AtomicBoolean failed = new AtomicBoolean(false); + /** * Constructor. Can also be created with {@link AsyncResources#collect(List)}. */ @@ -48,7 +55,7 @@ public class CollectAsyncResource implements AsyncResource> targetResource.set(List.of(), null); } else { for (final AsyncResource asyncResource : sourceResources) { - asyncResource.addReadyCallback(this::onOneSourceReady); + asyncResource.addReadyCallback(() -> onOneSourceReady(asyncResource)); } } } @@ -80,9 +87,23 @@ public void close() CloseableUtils.closeAndWrapExceptions(closer); } - private void onOneSourceReady() + private void onOneSourceReady(final AsyncResource sourceResource) { + try { + sourceResource.get(); + } + catch (Throwable e) { + // Fail without waiting for the other sources, but deliberately leave them open: closing them is up to the + // caller, via close(). + if (failed.compareAndSet(false, true)) { + targetResource.setException(e); + } + return; + } + if (readyCount.incrementAndGet() == sourceResources.size()) { + // Will only enter this block if all source resources succeeded, because a failed resource + // does not increment readyCount. try { final List resources = new ArrayList<>(sourceResources.size()); for (final AsyncResource asyncResource : sourceResources) { diff --git a/processing/src/test/java/org/apache/druid/common/asyncresource/AsyncResourcesTest.java b/processing/src/test/java/org/apache/druid/common/asyncresource/AsyncResourcesTest.java index a612cbf3366e..1da563c7289b 100644 --- a/processing/src/test/java/org/apache/druid/common/asyncresource/AsyncResourcesTest.java +++ b/processing/src/test/java/org/apache/druid/common/asyncresource/AsyncResourcesTest.java @@ -114,6 +114,24 @@ public void testCollectOfFromFutureReadyWhenAllComplete() collected.close(); } + @Test + public void testCollectOfFromFutureFailsWithoutWaitingAndCancelsOthersOnClose() + { + final SettableFuture a = SettableFuture.create(); + final SettableFuture b = SettableFuture.create(); + final AsyncResource> collected = + AsyncResources.collect(List.of(AsyncResources.fromFutureUnmanaged(a), AsyncResources.fromFutureUnmanaged(b))); + + final RuntimeException failure = new RuntimeException("boom"); + b.setException(failure); + Assertions.assertTrue(collected.isReady(), "one failure must not wait for the other sources"); + Assertions.assertSame(failure, Assertions.assertThrows(RuntimeException.class, collected::get)); + Assertions.assertFalse(a.isCancelled(), "the other sources are left to close()"); + + collected.close(); + Assertions.assertTrue(a.isCancelled()); + } + @Test public void testFromFutureGetBeforeReadyThrows() { diff --git a/processing/src/test/java/org/apache/druid/common/asyncresource/CollectAsyncResourceTest.java b/processing/src/test/java/org/apache/druid/common/asyncresource/CollectAsyncResourceTest.java index e71e19c8054f..5fd0dd7cb074 100644 --- a/processing/src/test/java/org/apache/druid/common/asyncresource/CollectAsyncResourceTest.java +++ b/processing/src/test/java/org/apache/druid/common/asyncresource/CollectAsyncResourceTest.java @@ -96,6 +96,59 @@ public void testSourceFailurePropagates() collected.close(); } + @Test + public void testSourceFailureFailsWithoutWaitingForOtherSources() + { + final AtomicInteger fired = new AtomicInteger(); + final AtomicInteger aCancel = new AtomicInteger(); + final AtomicInteger bClose = new AtomicInteger(); + final SettableAsyncResource a = new SettableAsyncResource<>(); + final SettableAsyncResource b = new SettableAsyncResource<>(); + final SettableAsyncResource c = new SettableAsyncResource<>(); + a.setCanceler(aCancel::incrementAndGet); + b.set("b", bClose::incrementAndGet); + + final AsyncResource> collected = AsyncResources.collect(List.of(a, b, c)); + collected.addReadyCallback(fired::incrementAndGet); + + // a is still pending, but c's failure must not wait for it. + final RuntimeException failure = new RuntimeException("boom"); + c.setException(failure); + + Assertions.assertTrue(collected.isReady()); + Assertions.assertEquals(1, fired.get()); + Assertions.assertSame(failure, Assertions.assertThrows(RuntimeException.class, collected::get)); + Assertions.assertEquals(0, aCancel.get(), "sources must stay open until the collected resource is closed"); + Assertions.assertEquals(0, bClose.get(), "sources must stay open until the collected resource is closed"); + + collected.close(); + Assertions.assertEquals(1, aCancel.get(), "close must cancel the pending source"); + Assertions.assertEquals(1, bClose.get(), "close must close the ready source"); + } + + @Test + public void testAlreadyFailedSourcesFailOnConstructionWithTheFirstFailure() + { + final AtomicInteger cCancel = new AtomicInteger(); + final SettableAsyncResource a = new SettableAsyncResource<>(); + final SettableAsyncResource b = new SettableAsyncResource<>(); + final SettableAsyncResource c = new SettableAsyncResource<>(); + final RuntimeException failure = new RuntimeException("boom"); + a.setException(failure); + b.setException(new RuntimeException("second")); + c.setCanceler(cCancel::incrementAndGet); + + // Both failures fire inline while registering; failing the collected resource twice would throw out of collect. + final AsyncResource> collected = AsyncResources.collect(List.of(a, b, c)); + + Assertions.assertTrue(collected.isReady()); + Assertions.assertSame(failure, Assertions.assertThrows(RuntimeException.class, collected::get)); + Assertions.assertEquals(0, cCancel.get()); + + collected.close(); + Assertions.assertEquals(1, cCancel.get()); + } + @Test public void testCloseClosesAllReadySources() { @@ -162,8 +215,8 @@ public void testCloseWithAMixOfReadyAndPendingSourcesReportsCancellation() collected.addReadyCallback(fired::incrementAndGet); Assertions.assertFalse(collected.isReady(), "one source is still pending"); - // The already-ready source counted toward readiness, so closing the pending one brings the internal count to the - // source count and runs the collect body against sources this close just tore down. That must stay harmless. + // Closing the pending source fires its ready callback while this close is still tearing the sources down. That + // must stay harmless. collected.close(); Assertions.assertEquals(1, fired.get());