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());