run uses cancelling state & propagates exceptions when cancelled

* When coroutine that was switched into a different dispatcher using
`run` is cancelled, the run invocation does not complete immediately,
but waits until the body completes.
* If the body completes with exception, then this exception is
propagated.

Fixes #147

Author: elizarov

core/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/AbstractContinuation.kt

@@ -19,7 +19,7 @@ package kotlinx.coroutines.experimental
 import kotlinx.atomicfu.atomic
 import kotlinx.atomicfu.loop
 import kotlin.coroutines.experimental.Continuation
-
+import kotlin.coroutines.experimental.intrinsics.COROUTINE_SUSPENDED
 
 private const val UNDECIDED = 0
 private const val SUSPENDED = 1
@@ -29,9 +29,9 @@ private const val RESUMED = 2
  * @suppress **This is unstable API and it is subject to change.**
  */
 internal abstract class AbstractContinuation<in T>(
-    active: Boolean,
+    @JvmField protected val delegate: Continuation<T>,
     @JvmField protected val resumeMode: Int
-) : JobSupport(active), Continuation<T> {
+) : JobSupport(true), Continuation<T> {
     private val _decision = atomic(UNDECIDED)
 
     /* decision state machine
@@ -49,7 +49,7 @@ internal abstract class AbstractContinuation<in T>(
         Note: both tryResume and trySuspend can be invoked at most once, first invocation wins
      */
 
-    protected fun trySuspend(): Boolean {
+    private fun trySuspend(): Boolean {
         _decision.loop { decision ->
             when (decision) {
                 UNDECIDED -> if (this._decision.compareAndSet(UNDECIDED, SUSPENDED)) return true
@@ -59,7 +59,7 @@ internal abstract class AbstractContinuation<in T>(
         }
     }
 
-    protected fun tryResume(): Boolean {
+    private fun tryResume(): Boolean {
         _decision.loop { decision ->
             when (decision) {
                 UNDECIDED -> if (this._decision.compareAndSet(UNDECIDED, RESUMED)) return true
@@ -69,32 +69,48 @@ internal abstract class AbstractContinuation<in T>(
         }
     }
 
-    override fun resume(value: T) = resumeImpl(value, resumeMode)
+    @PublishedApi
+    internal fun getResult(): Any? {
+        if (trySuspend()) return COROUTINE_SUSPENDED
+        // otherwise, afterCompletion was already invoked & invoked tryResume, and the result is in the state
+        val state = this.state
+        if (state is CompletedExceptionally) throw state.exception
+        return getSuccessfulResult(state)
+    }
 
-    protected fun resumeImpl(value: T, resumeMode: Int) {
-        loopOnState { state ->
-            when (state) {
-                is Incomplete -> if (updateState(state, value, resumeMode)) return
-                is Cancelled -> return // ignore resumes on cancelled continuation
-                else -> error("Already resumed, but got value $value")
-            }
+    override fun afterCompletion(state: Any?, mode: Int) {
+        if (tryResume()) return // completed before getResult invocation -- bail out
+        // otherwise, getResult has already commenced, i.e. completed later or in other thread
+        if (state is CompletedExceptionally) {
+            delegate.resumeWithExceptionMode(state.exception, mode)
+        } else {
+            delegate.resumeMode(getSuccessfulResult(state), mode)
         }
     }
 
-    override fun resumeWithException(exception: Throwable) = resumeWithExceptionImpl(exception, resumeMode)
+    @Suppress("UNCHECKED_CAST")
+    protected open fun <T> getSuccessfulResult(state: Any?): T =
+        state as T
+
+    override fun resume(value: T) =
+        resumeImpl(value, resumeMode)
+
+    override fun resumeWithException(exception: Throwable) =
+        resumeImpl(CompletedExceptionally(exception), resumeMode)
 
-    protected fun resumeWithExceptionImpl(exception: Throwable, resumeMode: Int) {
+    protected fun resumeImpl(proposedUpdate: Any?, resumeMode: Int) {
         loopOnState { state ->
             when (state) {
                 is Incomplete -> {
-                    if (updateState(state, CompletedExceptionally(exception), resumeMode)) return
+                    if (updateState(state, proposedUpdate, resumeMode)) return
                 }
                 is Cancelled -> {
-                    // ignore resumes on cancelled continuation, but handle exception if a different one is here
-                    if (exception != state.exception) handleCoroutineException(context, exception)
+                    // Ignore resumes in cancelled coroutines, but handle exception if a different one here
+                    if (proposedUpdate is CompletedExceptionally && proposedUpdate.exception != state.exception)
+                        handleException(proposedUpdate.exception)
                     return
                 }
-                else -> throw IllegalStateException("Already resumed, but got exception $exception", exception)
+                else -> error("Already resumed, but got $proposedUpdate")
             }
         }
     }

core/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/Builders.kt

@@ -18,7 +18,6 @@ package kotlinx.coroutines.experimental
 
 import java.util.concurrent.locks.LockSupport
 import kotlin.coroutines.experimental.*
-import kotlin.coroutines.experimental.intrinsics.COROUTINE_SUSPENDED
 import kotlin.coroutines.experimental.intrinsics.startCoroutineUninterceptedOrReturn
 import kotlin.coroutines.experimental.intrinsics.suspendCoroutineOrReturn
 
@@ -179,29 +178,15 @@ private class RunContinuationDirect<in T>(
     continuation: Continuation<T>
 ) : Continuation<T> by continuation
 
+
 @Suppress("UNCHECKED_CAST")
 private class RunCompletion<in T>(
     override val context: CoroutineContext,
-    private val delegate: Continuation<T>,
+    delegate: Continuation<T>,
     resumeMode: Int
-) : AbstractContinuation<T>(true, resumeMode) {
-    @PublishedApi
-    internal fun getResult(): Any? {
-        if (trySuspend()) return COROUTINE_SUSPENDED
-        // otherwise, afterCompletion was already invoked & invoked tryResume, and the result is in the state
-        val state = this.state
-        if (state is CompletedExceptionally) throw state.exception
-        return state as T
-    }
-
-    override fun afterCompletion(state: Any?, mode: Int) {
-        if (tryResume()) return // completed before getResult invocation -- bail out
-        // otherwise, getResult has already commenced, i.e. completed later or in other thread
-        if (state is CompletedExceptionally)
-            delegate.resumeWithExceptionMode(state.exception, mode)
-        else
-            delegate.resumeMode(state as T, mode)
-    }
+) : AbstractContinuation<T>(delegate, resumeMode) {
+    override val hasCancellingState: Boolean
+        get() = true
 }
 
 private class BlockingCoroutine<T>(

core/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/CancellableContinuation.kt

@@ -19,7 +19,6 @@ package kotlinx.coroutines.experimental
 import kotlinx.coroutines.experimental.internal.LockFreeLinkedListNode
 import kotlin.coroutines.experimental.Continuation
 import kotlin.coroutines.experimental.CoroutineContext
-import kotlin.coroutines.experimental.intrinsics.COROUTINE_SUSPENDED
 import kotlin.coroutines.experimental.intrinsics.suspendCoroutineOrReturn
 import kotlin.coroutines.experimental.suspendCoroutine
 
@@ -179,9 +178,9 @@ private class RemoveOnCancel(
 
 @PublishedApi
 internal class CancellableContinuationImpl<in T>(
-    private val delegate: Continuation<T>,
+    delegate: Continuation<T>,
     resumeMode: Int
-) : AbstractContinuation<T>(true, resumeMode), CancellableContinuation<T> {
+) : AbstractContinuation<T>(delegate, resumeMode), CancellableContinuation<T> {
     @Volatile // just in case -- we don't want an extra data race, even benign one
     private var _context: CoroutineContext? = null // created on first need
 
@@ -192,25 +191,6 @@ internal class CancellableContinuationImpl<in T>(
         initParentJob(delegate.context[Job])
     }
 
-    @PublishedApi
-    internal fun getResult(): Any? {
-        if (trySuspend()) return COROUTINE_SUSPENDED
-        // otherwise, afterCompletion was already invoked & invoked tryResume, and the result is in the state
-        val state = this.state
-        if (state is CompletedExceptionally) throw state.exception
-        return getSuccessfulResult(state)
-    }
-
-    override fun afterCompletion(state: Any?, mode: Int) {
-        if (tryResume()) return // completed before getResult invocation -- bail out
-        // otherwise, getResult has already commenced, i.e. completed later or in other thread
-        if (state is CompletedExceptionally) {
-            delegate.resumeWithExceptionMode(state.exception, mode)
-        } else {
-            delegate.resumeMode(getSuccessfulResult<T>(state), mode)
-        }
-    }
-
     override fun tryResume(value: T, idempotent: Any?): Any? {
         while (true) { // lock-free loop on state
             val state = this.state // atomic read
@@ -255,9 +235,13 @@ internal class CancellableContinuationImpl<in T>(
 
     override fun CoroutineDispatcher.resumeUndispatchedWithException(exception: Throwable) {
         val dc = delegate as? DispatchedContinuation
-        resumeWithExceptionImpl(exception, if (dc?.dispatcher === this) MODE_UNDISPATCHED else resumeMode)
+        resumeImpl(CompletedExceptionally(exception), if (dc?.dispatcher === this) MODE_UNDISPATCHED else resumeMode)
     }
 
+    @Suppress("UNCHECKED_CAST")
+    override fun <T> getSuccessfulResult(state: Any?): T =
+        if (state is CompletedIdempotentResult) state.result as T else state as T
+
     override fun nameString(): String =
         "CancellableContinuation(${delegate.toDebugString()})"
 }
@@ -270,6 +254,3 @@ private class CompletedIdempotentResult(
     override fun toString(): String = "CompletedIdempotentResult[$result]"
 }
 
-@Suppress("UNCHECKED_CAST")
-private fun <T> getSuccessfulResult(state: Any?): T =
-    if (state is CompletedIdempotentResult) state.result as T else state as T

core/kotlinx-coroutines-core/src/test/kotlin/kotlinx/coroutines/experimental/RunTest.kt

@@ -19,6 +19,7 @@ package kotlinx.coroutines.experimental
 import org.hamcrest.MatcherAssert.assertThat
 import org.hamcrest.core.IsEqual
 import org.junit.Test
+import java.io.IOException
 import kotlin.coroutines.experimental.ContinuationInterceptor
 import kotlin.coroutines.experimental.CoroutineContext
 
@@ -168,6 +169,33 @@ class RunTest : TestBase() {
         }
     }
 
+    @Test
+    fun testRunWithException() = runTest {
+        expect(1)
+        var job: Job? = null
+        job = launch(coroutineContext) {
+            try {
+                expect(3)
+                run(wrapperDispatcher(coroutineContext)) {
+                    expect(5)
+                    job!!.cancel() // cancel itself
+                    throw IOException() // but throw a different exception
+                }
+            } catch (e: Throwable) {
+                expect(7)
+                // make sure IOException, not CancellationException is thrown!
+                check(e is IOException)
+            }
+        }
+        expect(2)
+        yield() // to the launched job
+        expect(4)
+        yield() // again to the job
+        expect(6)
+        yield() // again to exception handler
+        finish(8)
+    }
+
     private fun wrapperDispatcher(context: CoroutineContext): CoroutineContext {
         val dispatcher = context[ContinuationInterceptor] as CoroutineDispatcher
         return object : CoroutineDispatcher() {