CAMEL-24242: camel-aws2-athena - stop ignoring thread interruption while polling a query

components/camel-aws/camel-aws2-athena/pom.xml

@@ -76,6 +76,11 @@
             <version>${mockito-version}</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.assertj</groupId>
+            <artifactId>assertj-core</artifactId>
+            <scope>test</scope>
+        </dependency>
 
         <!-- test infra -->
         <dependency>

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;
     }
 

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