Fixed CoroutineScope of withContext block for structured concurrency

* Both are optimized and are rewritten via AbstractCoroutine
* Support for cancelling state is dropped from AbstractContinuation
  and it is now faster, too

Fixes #553
Fixes #617

Author: elizarov

common/kotlinx-coroutines-core-common/src/AbstractContinuation.kt

@@ -31,14 +31,6 @@ internal abstract class AbstractContinuation<in T>(
      * 4) Its cancellation listeners cannot be deregistered
      * As a consequence it has much simpler state machine, more lightweight machinery and
      * less dependencies.
-     *
-     * Cancelling state
-     * If useCancellingState is true, then this continuation can have additional cancelling state,
-     * which is transition from Active to Cancelled. This is specific state to support withContext(ctx)
-     * construction: block in withContext can be cancelled from withing or even before stepping into withContext,
-     * but we still want to properly run it (e.g. when it has atomic cancellation mode) and run its completion listener
-     * after.
-     * During cancellation all pending exceptions are aggregated and thrown during transition to final state
      */
 
     /* decision state machine
@@ -63,7 +55,6 @@ internal abstract class AbstractContinuation<in T>(
        ------      ------------         ------------    -----------
        ACTIVE      Active               : Active        active, no listeners
        SINGLE_A    CancelHandler        : Active        active, one cancellation listener
-       CANCELLING  Cancelling           : Active        in the process of cancellation due to cancellation of parent job
        CANCELLED   Cancelled            : Cancelled     cancelled (final state)
        COMPLETED   any                  : Completed     produced some result or threw an exception (final state)
      */
@@ -80,8 +71,6 @@ internal abstract class AbstractContinuation<in T>(
 
     public val isCancelled: Boolean get() = state is CancelledContinuation
 
-    protected open val useCancellingState: Boolean get() = false
-
     internal fun initParentJobInternal(parent: Job?) {
         check(parentHandle == null)
         if (parent == null) {
@@ -105,7 +94,6 @@ internal abstract class AbstractContinuation<in T>(
     public fun cancel(cause: Throwable?): Boolean {
         loopOnState { state ->
             if (state !is NotCompleted) return false // quit if already complete
-            if (state is Cancelling) return false // someone else succeeded
             if (tryCancel(state, cause)) return true
         }
     }
@@ -170,7 +158,6 @@ internal abstract class AbstractContinuation<in T>(
                     handler.invokeIt((state as? CompletedExceptionally)?.cause)
                     return
                 }
-                is Cancelling -> error("Cancellation handlers for continuations with 'Cancelling' state are not supported")
                 else -> return
             }
         }
@@ -179,14 +166,8 @@ internal abstract class AbstractContinuation<in T>(
     private fun makeHandler(handler: CompletionHandler): CancelHandler =
         if (handler is CancelHandler) handler else InvokeOnCancel(handler)
 
-    private fun tryCancel(state: NotCompleted, cause: Throwable?): Boolean {
-        if (useCancellingState) {
-            require(state !is CancelHandler) { "Invariant: 'Cancelling' state and cancellation handlers cannot be used together" }
-            return _state.compareAndSet(state, Cancelling(CancelledContinuation(this, cause)))
-        }
-
-        return updateStateToFinal(state, CancelledContinuation(this, cause), mode = MODE_ATOMIC_DEFAULT)
-    }
+    private fun tryCancel(state: NotCompleted, cause: Throwable?): Boolean =
+        updateStateToFinal(state, CancelledContinuation(this, cause), mode = MODE_ATOMIC_DEFAULT)
 
     private fun dispatchResume(mode: Int) {
         if (tryResume()) return // completed before getResult invocation -- bail out
@@ -203,75 +184,6 @@ internal abstract class AbstractContinuation<in T>(
     protected fun resumeImpl(proposedUpdate: Any?, resumeMode: Int) {
         loopOnState { state ->
             when (state) {
-                is Cancelling -> { // withContext() support
-                    /*
-                     * If already cancelled block is resumed with non-exception,
-                     * resume it with cancellation exception.
-                     * E.g.
-                     * ```
-                     * val value = withContext(ctx) {
-                     *   outerJob.cancel() // -> cancelling
-                     *   42 // -> cancelled
-                     * }
-                     * ```
-                     * should throw cancellation exception instead of returning 42
-                     */
-                    if (proposedUpdate !is CompletedExceptionally) {
-                        val update = state.cancel
-                        if (updateStateToFinal(state, update, resumeMode)) return
-                    } else {
-                        /*
-                         * If already cancelled block is resumed with an exception,
-                         * then we should properly merge them to avoid information loss.
-                         *
-                         * General rule:
-                         * Thrown exception always becomes a result and cancellation reason
-                         * is added to suppressed exceptions if necessary.
-                         * Basic duplicate/cycles check is performed
-                         */
-                        val update: CompletedExceptionally
-
-                        /*
-                         * Proposed update is another CancellationException.
-                         * e.g.
-                         * ```
-                         * T1: ctxJob.cancel(e1) // -> cancelling
-                         * T2:
-                         * withContext(ctx, Mode.ATOMIC) {
-                         *   // -> resumed with cancellation exception
-                         * }
-                         * ```
-                         */
-                        if (proposedUpdate.cause is CancellationException) {
-                            // Keep original cancellation cause and try add to suppressed exception from proposed cancel
-                            update = proposedUpdate
-                            coerceWithException(state, update)
-                        } else {
-                            /*
-                             * Proposed update is exception => transition to terminal state
-                             * E.g.
-                             * ```
-                             * withContext(ctx) {
-                             *   outerJob.cancel() // -> cancelling
-                             *   throw Exception() // -> completed exceptionally
-                             * }
-                             * ```
-                             */
-                            val exception = proposedUpdate.cause
-                            val currentException = state.cancel.cause
-                            // Add to suppressed if original cancellation differs from proposed exception
-                            if (currentException !is CancellationException || currentException.cause !== exception) {
-                                exception.addSuppressedThrowable(currentException)
-                            }
-
-                            update = CompletedExceptionally(exception)
-                        }
-
-                        if (updateStateToFinal(state, update, resumeMode)) {
-                            return
-                        }
-                    }
-                }
                 is NotCompleted -> {
                     if (updateStateToFinal(state, proposedUpdate, resumeMode)) return
                 }
@@ -286,35 +198,20 @@ internal abstract class AbstractContinuation<in T>(
                     if (proposedUpdate is CompletedExceptionally) {
                         handleException(proposedUpdate.cause)
                     }
-
                     return
                 }
                 else -> error("Already resumed, but proposed with update $proposedUpdate")
             }
         }
     }
 
-    // Coerce current cancelling state with proposed exception
-    private fun coerceWithException(state: Cancelling, proposedUpdate: CompletedExceptionally) {
-        val originalCancellation = state.cancel
-        val originalException = originalCancellation.cause
-        val updateCause = proposedUpdate.cause
-        // Cause of proposed update is present and differs from one in current state
-        val isSameCancellation = originalCancellation.cause is CancellationException
-                && originalException.cause === updateCause.cause
-        if (!isSameCancellation && (originalException.cause !== updateCause)) {
-            proposedUpdate.cause.addSuppressedThrowable(originalException)
-        }
-    }
-
     /**
      * Tries to make transition from active to cancelled or completed state and invokes cancellation handler if necessary
      */
     private fun updateStateToFinal(expect: NotCompleted, proposedUpdate: Any?, mode: Int): Boolean {
         if (!tryUpdateStateToFinal(expect, proposedUpdate)) {
             return false
         }
-
         completeStateUpdate(expect, proposedUpdate, mode)
         return true
     }
@@ -374,9 +271,6 @@ internal interface NotCompleted
 private class Active : NotCompleted
 private val ACTIVE: Active = Active()
 
-// In progress of cancellation
-internal class Cancelling(@JvmField val cancel: CancelledContinuation) : NotCompleted
-
 internal abstract class CancelHandler : CancelHandlerBase(), NotCompleted
 
 // Wrapper for lambdas, for the performance sake CancelHandler can be subclassed directly

common/kotlinx-coroutines-core-common/src/Builders.common.kt

@@ -7,6 +7,7 @@
 
 package kotlinx.coroutines.experimental
 
+import kotlinx.atomicfu.*
 import kotlinx.coroutines.experimental.internal.*
 import kotlinx.coroutines.experimental.intrinsics.*
 import kotlinx.coroutines.experimental.selects.*
@@ -193,13 +194,30 @@ public fun launch(context: CoroutineContext, start: Boolean, block: suspend Coro
 public suspend fun <T> withContext(
     context: CoroutineContext,
     block: suspend CoroutineScope.() -> T
-): T =
-    // todo: optimize fast-path to work without allocation (when there is a already a coroutine implementing scope)
-    withContextImpl(context, start = CoroutineStart.DEFAULT) {
-        currentScope {
-            block()
+): T = suspendCoroutineUninterceptedOrReturn sc@ { uCont ->
+    // compute new context
+    val oldContext = uCont.context
+    val newContext = oldContext + context
+    // FAST PATH #1 -- new context is the same as the old one
+    if (newContext === oldContext) {
+        val coroutine = ScopeCoroutine(newContext, uCont) // MODE_DIRECT
+        return@sc coroutine.startUndispatchedOrReturn(coroutine, block)
+    }
+    // FAST PATH #2 -- the new dispatcher is the same as the old one (something else changed)
+    // `equals` is used by design (see equals implementation is wrapper context like ExecutorCoroutineDispatcher)
+    if (newContext[ContinuationInterceptor] == oldContext[ContinuationInterceptor]) {
+        val coroutine = UndispatchedCoroutine(newContext, uCont) // MODE_UNDISPATCHED
+        // There are changes in the context, so this thread needs to be updated
+        withCoroutineContext(newContext) {
+            return@sc coroutine.startUndispatchedOrReturn(coroutine, block)
         }
     }
+    // SLOW PATH -- use new dispatcher
+    val coroutine = DispatchedCoroutine(newContext, uCont) // MODE_DISPATCHED
+    coroutine.initParentJob()
+    block.startCoroutineCancellable(coroutine, coroutine)
+    coroutine.getResult()
+}
 
 /**
  * @suppress **Deprecated**: start parameter is deprecated, no replacement.
@@ -210,48 +228,7 @@ public suspend fun <T> withContext(
     start: CoroutineStart = CoroutineStart.DEFAULT,
     block: suspend CoroutineScope.() -> T
 ): T =
-    // todo: optimize fast-path to work without allocation (when there is a already a coroutine implementing scope)
-    withContextImpl(context, start) {
-        currentScope {
-            block()
-        }
-    }
-
-// todo: optimize it to reduce allocations
-private suspend fun <T> withContextImpl(
-    context: CoroutineContext,
-    start: CoroutineStart = CoroutineStart.DEFAULT,
-    block: suspend () -> T
-): T = suspendCoroutineUninterceptedOrReturn sc@ { uCont ->
-    val oldContext = uCont.context
-    // fast path #1 if there is no change in the actual context:
-    if (context === oldContext || context is CoroutineContext.Element && oldContext[context.key] === context)
-        return@sc block.startCoroutineUninterceptedOrReturn(uCont)
-    // compute new context
-    val newContext = oldContext + context
-    // fast path #2 if the result is actually the same
-    if (newContext === oldContext)
-        return@sc block.startCoroutineUninterceptedOrReturn(uCont)
-    // fast path #3 if the new dispatcher is the same as the old one.
-    // `equals` is used by design (see equals implementation is wrapper context like ExecutorCoroutineDispatcher)
-    if (newContext[ContinuationInterceptor] == oldContext[ContinuationInterceptor]) {
-        val newContinuation = RunContinuationUnintercepted(newContext, uCont)
-        // There are some other changes in the context, so this thread needs to be updated
-        withCoroutineContext(newContext) {
-            return@sc block.startCoroutineUninterceptedOrReturn(newContinuation)
-        }
-    }
-    // slowest path otherwise -- use new interceptor, sync to its result via a full-blown instance of RunCompletion
-    require(!start.isLazy) { "$start start is not supported" }
-    val completion = RunCompletion(
-        context = newContext,
-        delegate = uCont.intercepted(), // delegate to continuation intercepted with old dispatcher on completion
-        resumeMode = if (start == CoroutineStart.ATOMIC) MODE_ATOMIC_DEFAULT else MODE_CANCELLABLE
-    )
-    completion.initParentJobInternal(newContext[Job]) // attach to job
-    start(block, completion)
-    completion.getResult()
-}
+    withContext(context, block)
 
 /** @suppress **Deprecated**: Binary compatibility */
 @Deprecated(level = DeprecationLevel.HIDDEN, message = "Binary compatibility")
@@ -261,7 +238,7 @@ public suspend fun <T> withContext0(
     start: CoroutineStart = CoroutineStart.DEFAULT,
     block: suspend () -> T
 ): T =
-    withContextImpl(context, start, block)
+    withContext(context) { block() }
 
 /** @suppress **Deprecated**: Renamed to [withContext]. */
 @Deprecated(message = "Renamed to `withContext`", level=DeprecationLevel.WARNING,
@@ -271,12 +248,12 @@ public suspend fun <T> run(
     start: CoroutineStart = CoroutineStart.DEFAULT,
     block: suspend () -> T
 ): T =
-    withContextImpl(context, start, block)
+    withContext(context) { block() }
 
 /** @suppress **Deprecated** */
 @Deprecated(message = "It is here for binary compatibility only", level=DeprecationLevel.HIDDEN)
 public suspend fun <T> run(context: CoroutineContext, block: suspend () -> T): T =
-    withContextImpl(context, start = CoroutineStart.ATOMIC, block = block)
+    withContext(context) { block() }
 
 // --------------- implementation ---------------
 
@@ -297,29 +274,61 @@ private class LazyStandaloneCoroutine(
     }
 }
 
-private class RunContinuationUnintercepted<in T>(
-    override val context: CoroutineContext,
-    private val continuation: Continuation<T>
-): Continuation<T> {
-    override fun resume(value: T) {
-        withCoroutineContext(continuation.context) {
-            continuation.resume(value)
+// Used by withContext when context changes, but dispatcher stays the same
+private class UndispatchedCoroutine<in T>(
+    context: CoroutineContext,
+    uCont: Continuation<T>
+) : ScopeCoroutine<T>(context, uCont) {
+    override val defaultResumeMode: Int get() = MODE_UNDISPATCHED
+}
+
+private const val UNDECIDED = 0
+private const val SUSPENDED = 1
+private const val RESUMED = 2
+
+// Used by withContext when context dispatcher changes
+private class DispatchedCoroutine<in T>(
+    context: CoroutineContext,
+    uCont: Continuation<T>
+) : ScopeCoroutine<T>(context, uCont) {
+    override val defaultResumeMode: Int get() = MODE_CANCELLABLE
+
+    // this is copy-and-paste of a decision state machine inside AbstractionContinuation
+    // todo: we may some-how abstract it via inline class
+    private val _decision = atomic(UNDECIDED)
+
+    private fun trySuspend(): Boolean {
+        _decision.loop { decision ->
+            when (decision) {
+                UNDECIDED -> if (this._decision.compareAndSet(UNDECIDED, SUSPENDED)) return true
+                RESUMED -> return false
+                else -> error("Already suspended")
+            }
         }
     }
 
-    override fun resumeWithException(exception: Throwable) {
-        withCoroutineContext(continuation.context) {
-            continuation.resumeWithException(exception)
+    private fun tryResume(): Boolean {
+        _decision.loop { decision ->
+            when (decision) {
+                UNDECIDED -> if (this._decision.compareAndSet(UNDECIDED, RESUMED)) return true
+                SUSPENDED -> return false
+                else -> error("Already resumed")
+            }
         }
     }
-}
 
-@Suppress("UNCHECKED_CAST")
-private class RunCompletion<in T>(
-    override val context: CoroutineContext,
-    delegate: Continuation<T>,
-    resumeMode: Int
-) : AbstractContinuation<T>(delegate, resumeMode) {
+    override fun onCompletionInternal(state: Any?, mode: Int, suppressed: Boolean) {
+        if (tryResume()) return // completed before getResult invocation -- bail out
+        // otherwise, getResult has already commenced, i.e. completed later or in other thread
+        super.onCompletionInternal(state, mode, suppressed)
+    }
 
-    override val useCancellingState: Boolean get() = true
+    fun getResult(): Any? {
+        if (trySuspend()) return COROUTINE_SUSPENDED
+        // otherwise, onCompletionInternal was already invoked & invoked tryResume, and the result is in the state
+        val state = this.state
+        if (state is CompletedExceptionally) throw state.cause
+        @Suppress("UNCHECKED_CAST")
+        return state as T
+    }
 }