Fixed exception unwrapping in fast-path of CompletableFuture.await()

Author: elizarov

kotlinx-coroutines-jdk8/src/main/kotlin/kotlinx/coroutines/experimental/future/Future.kt

@@ -50,9 +50,19 @@ public fun <T> Deferred<T>.toCompletableFuture(): CompletableFuture<T> {
  * If the [Job] of the current coroutine is completed while this suspending function is waiting, this function
  * immediately resumes with [CancellationException] .
  */
-public suspend fun <T> CompletableFuture<T>.await(): T =
-    // quick check if already complete (avoid extra object creation)
-    if (isDone) get() else suspendCancellableCoroutine { cont: CancellableContinuation<T> ->
+public suspend fun <T> CompletableFuture<T>.await(): T {
+    if (isDone) {
+        // then only way to get unwrapped exception from the CompletableFuture...
+        var result: T? = null
+        var exception: Throwable? = null
+        whenComplete { r, e ->
+            result = r
+            exception = e
+        }
+        if (exception != null) throw exception!!
+        return result as T
+    }
+    return suspendCancellableCoroutine { cont: CancellableContinuation<T> ->
         val completionFuture = whenComplete { result, exception ->
             if (exception == null) // the future has been completed normally
                 cont.resume(result)
@@ -62,6 +72,7 @@ public suspend fun <T> CompletableFuture<T>.await(): T =
         cont.cancelFutureOnCompletion(completionFuture)
         Unit
     }
+}
 
 private class CompletableFutureCoroutine<T>(
     override val context: CoroutineContext

kotlinx-coroutines-jdk8/src/test/kotlin/kotlinx/coroutines/experimental/future/FutureTest.kt

@@ -37,6 +37,22 @@ class FutureTest {
         assertEquals("OK", future.get())
     }
 
+    @Test
+    fun testDoneFutureCompletedExceptionally() {
+        val toAwait = CompletableFuture<String>()
+        toAwait.completeExceptionally(RuntimeException("O"))
+        val future = future<String> {
+            try {
+                toAwait.await()
+            } catch (e: RuntimeException) {
+                e.message!!
+            } + "K"
+        }
+
+        assertFalse(future.isDone)
+        assertEquals("OK", future.get())
+    }
+
     @Test
     fun testAwaitedFutureCompletedExceptionally() {
         val toAwait = CompletableFuture<String>()