diff --git a/components/camel-aws/camel-aws2-athena/pom.xml b/components/camel-aws/camel-aws2-athena/pom.xml index b9cf0dd87c1de..1048efd4f173f 100644 --- a/components/camel-aws/camel-aws2-athena/pom.xml +++ b/components/camel-aws/camel-aws2-athena/pom.xml @@ -76,6 +76,11 @@ ${mockito-version} test + + org.assertj + assertj-core + test + diff --git a/components/camel-aws/camel-aws2-athena/src/main/java/org/apache/camel/component/aws2/athena/Athena2QueryHelper.java b/components/camel-aws/camel-aws2-athena/src/main/java/org/apache/camel/component/aws2/athena/Athena2QueryHelper.java index f4cc50e95c062..9bfa3c48f0519 100644 --- a/components/camel-aws/camel-aws2-athena/src/main/java/org/apache/camel/component/aws2/athena/Athena2QueryHelper.java +++ b/components/camel-aws/camel-aws2-athena/src/main/java/org/apache/camel/component/aws2/athena/Athena2QueryHelper.java @@ -169,7 +169,7 @@ boolean shouldWait() { void doWait() { // Use Camel's task API for polling delay instead of Thread.sleep() // We use initialDelay for the actual delay, and maxIterations(1) to run once - Tasks.foregroundTask() + boolean completed = Tasks.foregroundTask() .withBudget(Budgets.iterationBudget() .withMaxIterations(1) .withInitialDelay(Duration.ofMillis(this.currentDelay)) @@ -179,6 +179,14 @@ void doWait() { .build() .run(exchange.getContext(), () -> true); + if (!completed && Thread.currentThread().isInterrupted()) { + // the task API already restored the interrupt status, so only record it here to let the + // attempt and wait loops bail out at the earliest opportunity + this.interrupted = true; + LOG.trace( + "AWS Athena start query execution wait thread was interrupted; will return at earliest opportunity"); + } + this.currentDelay = this.delay; } diff --git a/components/camel-aws/camel-aws2-athena/src/test/java/org/apache/camel/component/aws2/athena/Athena2QueryHelperTest.java b/components/camel-aws/camel-aws2-athena/src/test/java/org/apache/camel/component/aws2/athena/Athena2QueryHelperTest.java index db81f368106c0..2d668168f85d6 100644 --- a/components/camel-aws/camel-aws2-athena/src/test/java/org/apache/camel/component/aws2/athena/Athena2QueryHelperTest.java +++ b/components/camel-aws/camel-aws2-athena/src/test/java/org/apache/camel/component/aws2/athena/Athena2QueryHelperTest.java @@ -28,6 +28,7 @@ import software.amazon.awssdk.services.athena.model.QueryExecutionState; import software.amazon.awssdk.services.athena.model.QueryExecutionStatus; +import static org.assertj.core.api.Assertions.assertThat; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -155,6 +156,36 @@ public void testStartQueryExecutionHappyPath() { assertFalse(helper.shouldAttempt()); } + @Test + void doWaitRecordsThreadInterruptionSoTheLoopsStop() { + Athena2Configuration configuration = new Athena2Configuration(); + configuration.setMaxAttempts(3); + configuration.setWaitTimeout(60_000); + configuration.setInitialDelay(10_000); + configuration.setDelay(10_000); + + Athena2QueryHelper helper = new Athena2QueryHelper( + new DefaultExchange(new DefaultCamelContext()), + configuration); + + helper.markAttempt(); + assertThat(helper.shouldWait()).isTrue(); + + Thread.currentThread().interrupt(); + try { + helper.doWait(); + + // without this, the wait loop keeps polling Athena with no delay at all, because every + // subsequent sleep returns immediately while the interrupt status is still set + assertThat(helper.isInterrupted()).isTrue(); + assertThat(helper.shouldWait()).isFalse(); + assertThat(helper.shouldAttempt()).isFalse(); + } finally { + // clear the interrupt status so it does not leak into the following tests + Thread.interrupted(); + } + } + @Test public void isComplete() { Athena2QueryHelper helper = defaultAthena2QueryHelper();