From b0fc859dafb780411a8d8dec418e9cef85f48be2 Mon Sep 17 00:00:00 2001
From: Konstantin Pavlov <{ID}+{username}@users.noreply.github.com>
Date: Wed, 17 Jul 2024 21:56:46 +0300
Subject: [PATCH] Handling InterruptedException should throw
 CompletionException

---
 pom.xml                                       |  3 ++
 src/main/java/me/kpavlov/await4j/Async.java   | 44 +++++++++----------
 .../await4j/AsyncCompletableFutureTest.java   |  2 -
 3 files changed, 24 insertions(+), 25 deletions(-)

diff --git a/pom.xml b/pom.xml
index d04d6e6..c59af73 100644
--- a/pom.xml
+++ b/pom.xml
@@ -72,6 +72,7 @@
             <plugin>
                 <groupId>org.apache.maven.plugins</groupId>
                 <artifactId>maven-source-plugin</artifactId>
+                <version>3.3.1</version>
                 <executions>
                     <execution>
                         <id>attach-sources</id>
@@ -84,6 +85,7 @@
             <plugin>
                 <groupId>org.apache.maven.plugins</groupId>
                 <artifactId>maven-javadoc-plugin</artifactId>
+                <version>3.7.0</version>
                 <executions>
                     <execution>
                         <id>attach-javadocs</id>
@@ -96,6 +98,7 @@
             <plugin>
                 <groupId>org.apache.maven.plugins</groupId>
                 <artifactId>maven-release-plugin</artifactId>
+                <version>3.1.1</version>
                 <configuration>
                     <autoVersionSubmodules>true</autoVersionSubmodules>
                     <useReleaseProfile>true</useReleaseProfile>
diff --git a/src/main/java/me/kpavlov/await4j/Async.java b/src/main/java/me/kpavlov/await4j/Async.java
index 1bbeec0..2b250f1 100644
--- a/src/main/java/me/kpavlov/await4j/Async.java
+++ b/src/main/java/me/kpavlov/await4j/Async.java
@@ -114,7 +114,7 @@ public static <T> T await(Callable<T> block, long millis) {
             }
         } catch (InterruptedException e) {
             Thread.currentThread().interrupt();
-            throw new RuntimeException("Interrupted virtual thread", e);
+            throw new CompletionException("Interrupted virtual thread", e);
         }
     }
 
@@ -145,7 +145,6 @@ public static <T> T await(Future<T> future) {
         return await(() -> future.get());
     }
 
-
     /**
      * Waits for the completion of a Future and returns its result.
      *
@@ -212,34 +211,33 @@ private static <T> Result<T> callWithErrorHandling(Callable<T> block) {
         }
     }
 
+    @SuppressWarnings("java:S1181")
     private static <T> boolean shortCircuitDoneFuture(Future<T> future) {
-        switch (future.state()) {
-            case SUCCESS -> {
-                return true;
-            }
-            case FAILED -> {
-                Throwable throwable = future.exceptionNow();
-                if (throwable instanceof Error e) {
-                    throw e;
+        try {
+            if (future.isDone()) {
+                if (future.isCancelled()) {
+                    throw new CancellationException("Execution is cancelled");
                 }
-                throw toRuntimeException(throwable);
-            }
-            case CANCELLED ->
-                throw new CancellationException("Execution is cancelled");
-            default ->
-            {
-                return false;
+                future.get();
+                return true;
             }
+            return false;
+        } catch (Error e) {
+            throw e;
+        } catch (ExecutionException e) {
+            throw toRuntimeException(e.getCause());
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            throw new CompletionException("Interrupted while waiting for future", e);
         }
     }
 
     private static RuntimeException toRuntimeException(Throwable cause) {
-        if (cause instanceof RuntimeException re) {
-            // re-throw RuntimeException as it is
-            return re;
-        } else {
-            return new CompletionException("Can't execute async task: exception", cause);
-        }
+        return switch (cause) {
+            case RuntimeException re -> re;
+            case Error e -> throw e;
+            default -> new CompletionException("Can't execute async task: exception", cause);
+        };
     }
 
 }
diff --git a/src/test/java/me/kpavlov/await4j/AsyncCompletableFutureTest.java b/src/test/java/me/kpavlov/await4j/AsyncCompletableFutureTest.java
index 51cc2db..cd2f18f 100644
--- a/src/test/java/me/kpavlov/await4j/AsyncCompletableFutureTest.java
+++ b/src/test/java/me/kpavlov/await4j/AsyncCompletableFutureTest.java
@@ -5,11 +5,9 @@
 import org.junit.jupiter.params.provider.MethodSource;
 
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.CompletionException;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import static me.kpavlov.await4j.TestUtils.sleepMillis;
-import static me.kpavlov.await4j.TestUtils.sleepOneSecond;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;