launched jobs and await can be cancelled while waiting in dispatch queue

* suspendAtomicCancellableCoroutine function is introduced for funs like
  send/receive/receiveOrNull that require atomic cancellation
  (they cannot be cancelled after decision was made)
* Coroutines started with default mode (CoroutineStart.ATOMIC) using
  async/launch/actor builders can be cancelled before execution starts.
* CoroutineStart.ATOMIC is introduced as a start mode to specify that
  coroutine cannot be cancelled before its execution start.

Author: elizarov

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

@@ -16,7 +16,6 @@
 
 package kotlinx.coroutines.experimental
 
-import kotlinx.coroutines.experimental.intrinsics.startCoroutineUndispatched
 import java.util.concurrent.locks.LockSupport
 import kotlin.coroutines.experimental.*
 import kotlin.coroutines.experimental.intrinsics.startCoroutineUninterceptedOrReturn
@@ -146,7 +145,7 @@ private class LazyStandaloneCoroutine(
     private val block: suspend CoroutineScope.() -> Unit
 ) : StandaloneCoroutine(parentContext, active = false) {
     override fun onStart() {
-        block.startCoroutine(this, this)
+        block.startCoroutineCancellable(this, this)
     }
 }
 
@@ -158,7 +157,7 @@ private class RunContinuationDirect<in T>(
 private class RunContinuationCoroutine<in T>(
     override val parentContext: CoroutineContext,
     continuation: Continuation<T>
-) : CancellableContinuationImpl<T>(continuation, active = true)
+) : CancellableContinuationImpl<T>(continuation, defaultResumeMode = MODE_CANCELLABLE, active = true)
 
 private class BlockingCoroutine<T>(
     override val parentContext: CoroutineContext,

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

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2016-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kotlinx.coroutines.experimental
+
+import kotlin.coroutines.experimental.Continuation
+import kotlin.coroutines.experimental.intrinsics.createCoroutineUnchecked
+
+/**
+ * Use this function to start coroutine in a cancellable way, so that it can be cancelled
+ * while waiting to be dispatched.
+ *
+ * @suppress **This is unstable API and it is subject to change.**
+ */
+internal fun <T> (suspend () -> T).startCoroutineCancellable(completion: Continuation<T>) =
+    createCoroutineUnchecked(completion).resumeCancellable(Unit)
+
+/**
+ * Use this function to start coroutine in a cancellable way, so that it can be cancelled
+ * while waiting to be dispatched.
+ *
+ * @suppress **This is unstable API and it is subject to change.**
+ */
+internal fun <R, T> (suspend (R) -> T).startCoroutineCancellable(receiver: R, completion: Continuation<T>) =
+    createCoroutineUnchecked(receiver, completion).resumeCancellable(Unit)
+

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

@@ -20,8 +20,10 @@ import kotlinx.coroutines.experimental.internal.LockFreeLinkedListNode
 import java.util.concurrent.atomic.AtomicIntegerFieldUpdater
 import kotlin.coroutines.experimental.Continuation
 import kotlin.coroutines.experimental.CoroutineContext
+import kotlin.coroutines.experimental.createCoroutine
 import kotlin.coroutines.experimental.intrinsics.COROUTINE_SUSPENDED
 import kotlin.coroutines.experimental.intrinsics.suspendCoroutineOrReturn
+import kotlin.coroutines.experimental.intrinsics.createCoroutineUnchecked
 import kotlin.coroutines.experimental.suspendCoroutine
 
 // --------------- cancellable continuations ---------------
@@ -111,23 +113,43 @@ public interface CancellableContinuation<in T> : Continuation<T>, Job {
 }
 
 /**
- * Suspend coroutine similar to [suspendCoroutine], but provide an implementation of [CancellableContinuation] to
+ * Suspends coroutine similar to [suspendCoroutine], but provide an implementation of [CancellableContinuation] to
  * the [block]. This function throws [CancellationException] if the coroutine is cancelled while suspended.
  *
  * If [holdCancellability] optional parameter is `true`, then the coroutine is suspended, but it is not
  * cancellable until [CancellableContinuation.initCancellability] is invoked.
+ *
+ * See [suspendAtomicCancellableCoroutine] for suspending functions that need *atomic cancellation*.
  */
 public inline suspend fun <T> suspendCancellableCoroutine(
     holdCancellability: Boolean = false,
     crossinline block: (CancellableContinuation<T>) -> Unit
 ): T =
     suspendCoroutineOrReturn { cont ->
-        val cancellable = CancellableContinuationImpl(cont, active = true)
+        val cancellable = CancellableContinuationImpl(cont, defaultResumeMode = MODE_CANCELLABLE, active = true)
         if (!holdCancellability) cancellable.initCancellability()
         block(cancellable)
         cancellable.getResult()
     }
 
+/**
+ * Suspends coroutine similar to [suspendCancellableCoroutine], but with *atomic cancellation*.
+ *
+ * When suspended function throws [CancellationException] it means that the continuation was not resumed.
+ * As a side-effect of atomic cancellation, a thread-bound coroutine (to some UI thread, for example) may
+ * continue to execute even after it was cancelled from the same thread in the case when the continuation
+ * was already resumed and was posted for execution to the thread's queue.
+ */
+public inline suspend fun <T> suspendAtomicCancellableCoroutine(
+    holdCancellability: Boolean = false,
+    crossinline block: (CancellableContinuation<T>) -> Unit
+): T =
+    suspendCoroutineOrReturn { cont ->
+        val cancellable = CancellableContinuationImpl(cont, defaultResumeMode = MODE_ATOMIC_DEFAULT, active = true)
+        if (!holdCancellability) cancellable.initCancellability()
+        block(cancellable)
+        cancellable.getResult()
+    }
 
 /**
  * Removes a given node on cancellation.
@@ -149,14 +171,16 @@ private class RemoveOnCancel(
     override fun toString() = "RemoveOnCancel[$node]"
 }
 
-internal const val MODE_DISPATCHED = 0
-internal const val MODE_UNDISPATCHED = 1
-internal const val MODE_DIRECT = 2
+@PublishedApi internal const val MODE_ATOMIC_DEFAULT = 0 // schedule non-cancellable dispatch for suspendCoroutine
+@PublishedApi internal const val MODE_CANCELLABLE = 1    // schedule cancellable dispatch for suspendCancellableCoroutine
+@PublishedApi internal const val MODE_DIRECT = 2         // when the context is right just invoke the delegate continuation direct
+@PublishedApi internal const val MODE_UNDISPATCHED = 3   // when the thread is right, but need to mark it with current coroutine
 
 @PublishedApi
 internal open class CancellableContinuationImpl<in T>(
     @JvmField
     protected val delegate: Continuation<T>,
+    override val defaultResumeMode: Int,
     active: Boolean
 ) : AbstractCoroutine<T>(active), CancellableContinuation<T> {
     @Volatile
@@ -239,17 +263,19 @@ internal open class CancellableContinuationImpl<in T>(
         if (state is CompletedExceptionally) {
             val exception = state.exception
             when (mode) {
-                MODE_DISPATCHED -> delegate.resumeWithException(exception)
-                MODE_UNDISPATCHED -> (delegate as DispatchedContinuation).resumeUndispatchedWithException(exception)
+                MODE_ATOMIC_DEFAULT -> delegate.resumeWithException(exception)
+                MODE_CANCELLABLE -> delegate.resumeCancellableWithException(exception)
                 MODE_DIRECT -> delegate.resumeDirectWithException(exception)
+                MODE_UNDISPATCHED -> (delegate as DispatchedContinuation).resumeUndispatchedWithException(exception)
                 else -> error("Invalid mode $mode")
             }
         } else {
             val value = getSuccessfulResult<T>(state)
             when (mode) {
-                MODE_DISPATCHED -> delegate.resume(value)
-                MODE_UNDISPATCHED -> (delegate as DispatchedContinuation).resumeUndispatched(value)
+                MODE_ATOMIC_DEFAULT -> delegate.resume(value)
+                MODE_CANCELLABLE -> delegate.resumeCancellable(value)
                 MODE_DIRECT -> delegate.resumeDirect(value)
+                MODE_UNDISPATCHED -> (delegate as DispatchedContinuation).resumeUndispatched(value)
                 else -> error("Invalid mode $mode")
             }
         }

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

@@ -102,37 +102,73 @@ public abstract class CoroutineDispatcher :
 
 }
 
+// named class for ease of debugging, better stack-traces and optimize the number of anonymous classes
+internal class DispatchTask<in T>(
+    private val dispatched: DispatchedContinuation<T>,
+    private val value: Any?, // T | Throwable
+    private val exception: Boolean,
+    private val cancellable: Boolean
+) : Runnable {
+    @Suppress("UNCHECKED_CAST")
+    override fun run() {
+        val job = if (cancellable) dispatched.context[Job] else null
+        when {
+            job != null && job.isCompleted ->
+                dispatched.resumeUndispatchedWithException(job.getCompletionException())
+            exception -> dispatched.resumeUndispatchedWithException(value as Throwable)
+            else -> dispatched.resumeUndispatched(value as T)
+        }
+    }
+
+    override fun toString(): String =
+        "DispatchTask[$value, cancellable=$cancellable, $dispatched]"
+}
+
 internal class DispatchedContinuation<in T>(
     @JvmField val dispatcher: CoroutineDispatcher,
     @JvmField val continuation: Continuation<T>
 ): Continuation<T> by continuation {
     override fun resume(value: T) {
         val context = continuation.context
         if (dispatcher.isDispatchNeeded(context))
-            dispatcher.dispatch(context, Runnable {
-                resumeUndispatched(value)
-            })
+            dispatcher.dispatch(context, DispatchTask(this, value, exception = false, cancellable = false))
         else
             resumeUndispatched(value)
     }
 
+    override fun resumeWithException(exception: Throwable) {
+        val context = continuation.context
+        if (dispatcher.isDispatchNeeded(context))
+            dispatcher.dispatch(context, DispatchTask(this, exception, exception = true, cancellable = false))
+        else
+            resumeUndispatchedWithException(exception)
+    }
+
     @Suppress("NOTHING_TO_INLINE") // we need it inline to save us an entry on the stack
-    inline fun resumeUndispatched(value: T) {
-        withCoroutineContext(context) {
-            continuation.resume(value)
-        }
+    inline fun resumeCancellable(value: T) {
+        val context = continuation.context
+        if (dispatcher.isDispatchNeeded(context))
+            dispatcher.dispatch(context, DispatchTask(this, value, exception = false, cancellable = true))
+        else
+            resumeUndispatched(value)
     }
 
-    override fun resumeWithException(exception: Throwable) {
+    @Suppress("NOTHING_TO_INLINE") // we need it inline to save us an entry on the stack
+    inline fun resumeCancellableWithException(exception: Throwable) {
         val context = continuation.context
         if (dispatcher.isDispatchNeeded(context))
-            dispatcher.dispatch(context, Runnable {
-                resumeUndispatchedWithException(exception)
-            })
+            dispatcher.dispatch(context, DispatchTask(this, exception, exception = true, cancellable = true))
         else
             resumeUndispatchedWithException(exception)
     }
 
+    @Suppress("NOTHING_TO_INLINE") // we need it inline to save us an entry on the stack
+    inline fun resumeUndispatched(value: T) {
+        withCoroutineContext(context) {
+            continuation.resume(value)
+        }
+    }
+
     @Suppress("NOTHING_TO_INLINE") // we need it inline to save us an entry on the stack
     inline fun resumeUndispatchedWithException(exception: Throwable) {
         withCoroutineContext(context) {
@@ -152,6 +188,18 @@ internal class DispatchedContinuation<in T>(
             }
         })
     }
+
+    override fun toString(): String = "DispatchedContinuation[$dispatcher, $continuation]"
+}
+
+internal fun <T> Continuation<T>.resumeCancellable(value: T) = when (this) {
+    is DispatchedContinuation -> resumeCancellable(value)
+    else -> resume(value)
+}
+
+internal fun <T> Continuation<T>.resumeCancellableWithException(exception: Throwable) = when (this) {
+    is DispatchedContinuation -> resumeCancellableWithException(exception)
+    else -> resumeWithException(exception)
 }
 
 internal fun <T> Continuation<T>.resumeDirect(value: T) = when (this) {

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

@@ -67,7 +67,7 @@ public abstract class AbstractCoroutine<in T>(
 
     protected open fun createContext() = parentContext + this
 
-    protected open val defaultResumeMode: Int get() = MODE_DISPATCHED
+    protected open val defaultResumeMode: Int get() = MODE_ATOMIC_DEFAULT
 
     protected open val ignoreRepeatedResume: Boolean get() = false
 

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

@@ -35,6 +35,12 @@ public enum class CoroutineStart {
      *
      * Note, that [Unconfined] dispatcher always returns `false` from its [CoroutineDispatcher.isDispatchNeeded]
      * function, so starting coroutine with [Unconfined] dispatcher by [DEFAULT] is the same as using [UNDISPATCHED].
+     *
+     * If coroutine [Job] is cancelled before it even had a chance to start executing, then it will not start its
+     * execution at all, but complete with an exception.
+     *
+     * Cancellability of coroutine at suspension points depends on the particular implementation details of
+     * suspending functions. Use [suspendCancellableCoroutine] to implement cancellable suspending functions.
      */
     DEFAULT,
 
@@ -43,9 +49,21 @@ public enum class CoroutineStart {
      *
      * See the documentation for the corresponding coroutine builders for details:
      * [launch], [async], and [actor][kotlinx.coroutines.experimental.channels.actor].
+     *
+     * If coroutine [Job] is cancelled before it even had a chance to start executing, then it will not start its
+     * execution at all, but complete with an exception.
      */
     LAZY,
 
+    /**
+     * Atomically schedules coroutines for execution according to its context. This is similar to [DEFAULT],
+     * but the coroutine cannot be cancelled before it starts executing.
+     *
+     * Cancellability of coroutine at suspension points depends on the particular implementation details of
+     * suspending functions as in [DEFAULT].
+     */
+    ATOMIC,
+
     /**
      * Immediately executes coroutine until its first suspension point _in the current thread_ as if it the
      * coroutine was started using [Unconfined] dispatcher. However, when coroutine is resumed from suspension
@@ -56,13 +74,15 @@ public enum class CoroutineStart {
     /**
      * Starts the corresponding block as a coroutine with this coroutine start strategy.
      *
-     * * [DEFAULT] uses [startCoroutine].
+     * * [DEFAULT] uses [startCoroutineCancellable].
+     * * [ATOMIC] uses [startCoroutine].
      * * [UNDISPATCHED] uses [startCoroutineUndispatched].
      * * [LAZY] does nothing.
      */
     public operator fun <R, T> invoke(block: suspend R.() -> T, receiver: R, completion: Continuation<T>) =
         when (this) {
-            CoroutineStart.DEFAULT -> block.startCoroutine(receiver, completion)
+            CoroutineStart.DEFAULT -> block.startCoroutineCancellable(receiver, completion)
+            CoroutineStart.ATOMIC -> block.startCoroutine(receiver, completion)
             CoroutineStart.UNDISPATCHED -> block.startCoroutineUndispatched(receiver, completion)
             CoroutineStart.LAZY -> Unit // will start lazily
         }

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

@@ -206,10 +206,11 @@ private open class DeferredCoroutine<T>(
     @Suppress("UNCHECKED_CAST")
     internal fun <R> selectAwaitCompletion(select: SelectInstance<R>, block: suspend (T) -> R, state: Any? = this.state) {
         if (select.trySelect(idempotent = null)) {
+            // Note: await is non-atomic (can be cancelled while dispatched)
             if (state is CompletedExceptionally)
-                select.resumeSelectWithException(state.exception, MODE_DISPATCHED)
+                select.resumeSelectWithException(state.exception, MODE_CANCELLABLE)
             else
-                block.startCoroutine(state as T, select.completion)
+                block.startCoroutineCancellable(state as T, select.completion)
         }
     }
 
@@ -236,6 +237,6 @@ private class LazyDeferredCoroutine<T>(
     private val block: suspend CoroutineScope.() -> T
 ) : DeferredCoroutine<T>(parentContext, active = false) {
     override fun onStart() {
-        block.startCoroutine(this, this)
+        block.startCoroutineCancellable(this, this)
     }
 }