Job.cancel and CompletableDeferred.complete support cancelling/completing

states and properly wait for their children to complete on join/await.

Fixes #199

Author: elizarov

common/kotlinx-coroutines-core-common/src/test/kotlin/kotlinx/coroutines/experimental/CommonCompletableDeferredTest.kt

@@ -174,6 +174,47 @@ class CommonCompletableDeferredTest : TestBase() {
         finish(7)
     }
 
+    @Test
+    fun testCancelAndAwaitParentWaitChildren() = runTest {
+        expect(1)
+        val parent = CompletableDeferred<String>()
+        launch(coroutineContext, start = CoroutineStart.UNDISPATCHED, parent = parent) {
+            expect(2)
+            try {
+                yield() // will get cancelled
+            } finally {
+                expect(5)
+            }
+        }
+        expect(3)
+        parent.cancel()
+        expect(4)
+        try {
+            parent.await()
+        } catch (e: CancellationException) {
+            finish(6)
+        }
+    }
+
+    @Test
+    fun testCompleteAndAwaitParentWaitChildren() = runTest {
+        expect(1)
+        val parent = CompletableDeferred<String>()
+        launch(coroutineContext, start = CoroutineStart.UNDISPATCHED, parent = parent) {
+            expect(2)
+            try {
+                yield() // will get cancelled
+            } finally {
+                expect(5)
+            }
+        }
+        expect(3)
+        parent.complete("OK")
+        expect(4)
+        assertEquals("OK", parent.await())
+        finish(6)
+    }
+
     private inline fun <reified T: Throwable> assertThrows(block: () -> Unit) {
         try {
             block()

common/kotlinx-coroutines-core-common/src/test/kotlin/kotlinx/coroutines/experimental/CommonJobTest.kt

@@ -149,4 +149,23 @@ class CommonJobTest : TestBase() {
         job.cancel()
         assertEquals(0, fireCount)
     }
+
+    @Test
+    fun testCancelAndJoinParentWaitChildren() = runTest {
+        expect(1)
+        val parent = Job()
+        launch(coroutineContext, start = CoroutineStart.UNDISPATCHED, parent = parent) {
+            expect(2)
+            try {
+                yield() // will get cancelled
+            } finally {
+                expect(5)
+            }
+        }
+        expect(3)
+        parent.cancel()
+        expect(4)
+        parent.join()
+        finish(6)
+    }
 }

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

@@ -38,17 +38,14 @@ public abstract class AbstractCoroutine<in T>(
     public final override val context: CoroutineContext = parentContext + this
     public final override val coroutineContext: CoroutineContext get() = context
 
-    // all coroutines are cancelled through an intermediate cancelling state
-    final override val hasCancellingState: Boolean get() = true
-
     protected open val defaultResumeMode: Int get() = MODE_ATOMIC_DEFAULT
 
     final override fun resume(value: T) {
-        makeCompleting(value, defaultResumeMode)
+        makeCompletingOnce(value, defaultResumeMode)
     }
 
     final override fun resumeWithException(exception: Throwable) {
-        makeCompleting(CompletedExceptionally(exception), defaultResumeMode)
+        makeCompletingOnce(CompletedExceptionally(exception), defaultResumeMode)
     }
 
     final override fun handleException(exception: Throwable) {

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

@@ -206,10 +206,7 @@ private class RunCompletion<in T>(
     override val context: CoroutineContext,
     delegate: Continuation<T>,
     resumeMode: Int
-) : AbstractContinuation<T>(delegate, resumeMode) {
-    override val hasCancellingState: Boolean
-        get() = true
-}
+) : AbstractContinuation<T>(delegate, resumeMode)
 
 private class BlockingCoroutine<T>(
     parentContext: CoroutineContext,

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

@@ -225,6 +225,8 @@ internal class CancellableContinuationImpl<in T>(
         initParentJob(delegate.context[Job])
     }
 
+    override val onCancelMode: Int get() = ON_CANCEL_MAKE_CANCELLED
+
     override fun tryResume(value: T, idempotent: Any?): Any? {
         while (true) { // lock-free loop on state
             val state = this.state // atomic read

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

@@ -22,31 +22,9 @@ import kotlinx.coroutines.experimental.selects.SelectClause1
  * A [Deferred] that can be completed via public functions
  * [complete], [completeExceptionally], and [cancel].
  *
- * Completion functions return `false` when this deferred value is already complete.
+ * Completion functions return `false` when this deferred value is already complete or completing.
  *
- * A completable deferred value has the following states:
- *
- * | **State**                 | [isActive] | [isCompleted] | [isCompletedExceptionally] | [isCancelled] |
- * | ------------------------- | ---------- | ------------- | -------------------------- | ------------- |
- * | _Active_ (initial state)  | `true`     | `false`       | `false`                    | `false`       |
- * | _Cancelled_ (final state) | `false`    | `true`        | `true`                     | `true`        |
- * | _Resolved_  (final state) | `false`    | `true`        | `false`                    | `false`       |
- * | _Failed_    (final state) | `false`    | `true`        | `true`                     | `false`       |
- *
- * A an instance of completable deferred can be created by `CompletableDeferred()` function in _active_ state.
- *
- *  ```
- *  +--------+   complete   +-----------+
- *  | Active | ---------+-> | Resolved  |
- *  +--------+          |   |(completed)|
- *       |              |   +-----------+
- *       | cancel       |
- *       V              |   +-----------+
- *  +-----------+       +-> |  Failed   |
- *  | Cancelled |           |(completed)|
- *  |(completed)|           +-----------+
- *  +-----------+
- * ```
+ * An instance of completable deferred can be created by `CompletableDeferred()` function in _active_ state.
  *
  * All functions on this interface and on all interfaces derived from it are **thread-safe** and can
  * be safely invoked from concurrent coroutines without external synchronization.
@@ -95,34 +73,16 @@ private class CompletableDeferredImpl<T>(
     parent: Job?
 ) : JobSupport(true), CompletableDeferred<T> {
     init { initParentJob(parent) }
+    override val onCancelMode: Int get() = ON_CANCEL_MAKE_COMPLETING
+
     override fun getCompleted(): T = getCompletedInternal() as T
     override suspend fun await(): T = awaitInternal() as T
     override val onAwait: SelectClause1<T>
         get() = this as SelectClause1<T>
 
-    override fun complete(value: T): Boolean {
-        loopOnState { state ->
-            when (state) {
-                is Incomplete -> {
-                    // actually, we don't care about the mode here at all, so just use a default
-                    if (updateState(state, value, mode = MODE_ATOMIC_DEFAULT))
-                        return true
-                }
-                else -> return false
-            }
-        }
-    }
+    override fun complete(value: T): Boolean =
+        makeCompleting(value)
 
-    override fun completeExceptionally(exception: Throwable): Boolean {
-        loopOnState { state ->
-            when (state) {
-                is Incomplete -> {
-                    // actually, we don't care about the mode here at all, so just use a default
-                    if (updateState(state, CompletedExceptionally(exception), mode = MODE_ATOMIC_DEFAULT))
-                        return true
-                }
-                else -> return false
-            }
-        }
-    }
+    override fun completeExceptionally(exception: Throwable): Boolean =
+        makeCompleting(CompletedExceptionally(exception))
 }

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

@@ -45,10 +45,10 @@ import kotlin.coroutines.experimental.CoroutineContext
  * Such a deferred can be be made _active_ by invoking [start], [join], or [await].
  *
  * A deferred can be _cancelled_ at any time with [cancel] function that forces it to transition to
- * _cancelling_ state immediately. A simple implementation of deferred -- [CompletableDeferred],
- * that is not backed by a coroutine, does not have a _cancelling_ state, but becomes _cancelled_
- * on [cancel] immediately. Coroutines, on the other hand, become _cancelled_ only when they finish
- * executing their code and after all their [children] complete.
+ * _cancelling_ state immediately. Deferred that is not backed by a coroutine (see [CompletableDeferred]) and does not have
+ * [children] becomes _cancelled_ on [cancel] immediately.
+ * Otherwise, deferred becomes _cancelled_  when it finishes executing its code and
+ * when all its children [complete][isCompleted].
  *
  * ```
  *                                                     wait children

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

@@ -60,7 +60,7 @@ import kotlin.coroutines.experimental.intrinsics.suspendCoroutineOrReturn
  * [CoroutineStart.LAZY]. Such a job can be made _active_ by invoking [start] or [join].
  *
  * A job can be _cancelled_ at any time with [cancel] function that forces it to transition to
- * _cancelling_ state immediately. Job that is not backed by a coroutine and does not have
+ * _cancelling_ state immediately. Job that is not backed by a coroutine (see `Job()` function) and does not have
  * [children] becomes _cancelled_ on [cancel] immediately.
  * Otherwise, job becomes _cancelled_  when it finishes executing its code and
  * when all its children [complete][isCompleted].
@@ -841,7 +841,7 @@ public open class JobSupport(active: Boolean) : Job, SelectClause0, SelectClause
                         promoteSingleToNodeList(state as JobNode<*>)
                     } else {
                         if (state is Finishing && state.cancelled != null && onCancelling) {
-                            check(hasCancellingState) // cannot be in this state unless were support cancelling state
+                            check(onCancelMode != ON_CANCEL_MAKE_CANCELLED) // cannot be in this state unless were support cancelling state
                             // installing cancellation handler on job that is being cancelled
                             if (invokeImmediately) handler(state.cancelled.cause)
                             return NonDisposableHandle
@@ -858,14 +858,15 @@ public open class JobSupport(active: Boolean) : Job, SelectClause0, SelectClause
         }
     }
 
-    private fun makeNode(handler: CompletionHandler, onCancelling: Boolean): JobNode<*> =
-        if (onCancelling && hasCancellingState)
+    private fun makeNode(handler: CompletionHandler, onCancelling: Boolean): JobNode<*> {
+        val hasCancellingState = onCancelMode != ON_CANCEL_MAKE_CANCELLED
+        return if (onCancelling && hasCancellingState)
             (handler as? JobCancellationNode<*>)?.also { require(it.job === this) }
                 ?: InvokeOnCancellation(this, handler)
         else
             (handler as? JobNode<*>)?.also { require(it.job === this && (!hasCancellingState || it !is JobCancellationNode)) }
                 ?: InvokeOnCompletion(this, handler)
-
+    }
 
     private fun addLastAtomic(expect: Any, list: NodeList, node: JobNode<*>) =
         list.addLastIf(node) { this.state === expect }
@@ -948,12 +949,14 @@ public open class JobSupport(active: Boolean) : Job, SelectClause0, SelectClause
         }
     }
 
-    protected open val hasCancellingState: Boolean get() = false
+    protected open val onCancelMode: Int get() = ON_CANCEL_MAKE_CANCELLING
 
-    public override fun cancel(cause: Throwable?): Boolean =
-        if (hasCancellingState)
-            makeCancelling(cause) else
-            makeCancelled(cause)
+    public override fun cancel(cause: Throwable?): Boolean = when (onCancelMode) {
+        ON_CANCEL_MAKE_CANCELLED -> makeCancelled(cause)
+        ON_CANCEL_MAKE_CANCELLING -> makeCancelling(cause)
+        ON_CANCEL_MAKE_COMPLETING -> makeCompletingOnCancel(cause)
+        else -> error("Invalid onCancelMode $onCancelMode")
+    }
 
     // we will be dispatching coroutine to process its cancellation exception, so there is no need for
     // an extra check for Job status in MODE_CANCELLABLE
@@ -1013,6 +1016,15 @@ public open class JobSupport(active: Boolean) : Job, SelectClause0, SelectClause
         return true
     }
 
+    private fun makeCompletingOnCancel(cause: Throwable?): Boolean =
+        makeCompleting(Cancelled(this, cause))
+    
+    internal fun makeCompleting(proposedUpdate: Any?): Boolean =
+        when (makeCompletingInternal(proposedUpdate, mode = MODE_ATOMIC_DEFAULT)) {
+            COMPLETING_ALREADY_COMPLETING -> false
+            else -> true
+        }
+
     /**
      * Returns:
      * * `true` if state was updated to completed/cancelled;
@@ -1021,30 +1033,38 @@ public open class JobSupport(active: Boolean) : Job, SelectClause0, SelectClause
      * @throws IllegalStateException if job is already complete or completing
      * @suppress **This is unstable API and it is subject to change.**
      */
-    internal fun makeCompleting(proposedUpdate: Any?, mode: Int): Boolean {
+    internal fun makeCompletingOnce(proposedUpdate: Any?, mode: Int): Boolean =
+        when (makeCompletingInternal(proposedUpdate, mode)) {
+            COMPLETING_COMPLETED -> true
+            COMPLETING_WAITING_CHILDREN -> false
+            else -> throw IllegalStateException("Job $this is already complete or completing, " +
+                "but is being completed with $proposedUpdate", proposedUpdate.exceptionOrNull)
+        }
+
+    private fun makeCompletingInternal(proposedUpdate: Any?, mode: Int): Int {
         loopOnState { state ->
             if (state !is Incomplete)
-                throw IllegalStateException("Job $this is already complete, but is being completed with $proposedUpdate", proposedUpdate.exceptionOrNull)
+                return COMPLETING_ALREADY_COMPLETING
             if (state is Finishing && state.completing)
-                throw IllegalStateException("Job $this is already completing, but is being completed with $proposedUpdate", proposedUpdate.exceptionOrNull)
+                return COMPLETING_ALREADY_COMPLETING
             val child: Child = firstChild(state) ?: // or else complete immediately w/o children
-                if (updateState(state, proposedUpdate, mode)) return true else return@loopOnState
+                if (updateState(state, proposedUpdate, mode)) return COMPLETING_COMPLETED else return@loopOnState
             // must promote to list to correct operate on child lists
             if (state is JobNode<*>) {
                 promoteSingleToNodeList(state)
                 return@loopOnState // retry
             }
             // cancel all children in list on exceptional completion
-            if (proposedUpdate is CompletedExceptionally
-            )
+            if (proposedUpdate is CompletedExceptionally)
                 child.cancelChildrenInternal(proposedUpdate.exception)
             // switch to completing state
-            val completing = Finishing(state.list!!, (state as? Finishing)?.cancelled, true)
+            val cancelled = (state as? Finishing)?.cancelled ?: (proposedUpdate as? Cancelled)
+            val completing = Finishing(state.list!!, cancelled, true)
             if (_state.compareAndSet(state, completing)) {
                 if (tryWaitForChild(child, proposedUpdate))
-                    return false
-                if (updateState(completing, proposedUpdate, MODE_ATOMIC_DEFAULT))
-                    return true
+                    return COMPLETING_WAITING_CHILDREN
+                if (updateState(completing, proposedUpdate, mode = MODE_ATOMIC_DEFAULT))
+                    return COMPLETING_COMPLETED
             }
         }
     }
@@ -1251,8 +1271,7 @@ public open class JobSupport(active: Boolean) : Job, SelectClause0, SelectClause
         cont.disposeOnCompletion(invokeOnCompletion {
             val state = this.state
             check(state !is Incomplete)
-            if (state is CompletedExceptionally
-            )
+            if (state is CompletedExceptionally)
                 cont.resumeWithException(state.exception)
             else
                 cont.resume(state)
@@ -1267,8 +1286,7 @@ public open class JobSupport(active: Boolean) : Job, SelectClause0, SelectClause
             if (state !is Incomplete) {
                 // already complete -- select result
                 if (select.trySelect(null)) {
-                    if (state is CompletedExceptionally
-                    )
+                    if (state is CompletedExceptionally)
                         select.resumeSelectCancellableWithException(state.exception)
                     else
                         block.startCoroutineUndispatched(state, select.completion)
@@ -1286,14 +1304,21 @@ public open class JobSupport(active: Boolean) : Job, SelectClause0, SelectClause
     internal fun <R> selectAwaitCompletion(select: SelectInstance<R>, block: suspend (Any?) -> R) {
         val state = this.state
         // Note: await is non-atomic (can be cancelled while dispatched)
-        if (state is CompletedExceptionally
-        )
+        if (state is CompletedExceptionally)
             select.resumeSelectCancellableWithException(state.exception)
         else
             block.startCoroutineCancellable(state, select.completion)
     }
 }
 
+internal const val ON_CANCEL_MAKE_CANCELLED = 0
+internal const val ON_CANCEL_MAKE_CANCELLING = 1
+internal const val ON_CANCEL_MAKE_COMPLETING = 2
+
+private const val COMPLETING_ALREADY_COMPLETING = 0
+private const val COMPLETING_COMPLETED = 1
+private const val COMPLETING_WAITING_CHILDREN = 2
+
 private const val RETRY = -1
 private const val FALSE = 0
 private const val TRUE = 1
@@ -1310,6 +1335,7 @@ private class Empty(override val isActive: Boolean) : JobSupport.Incomplete {
 
 private class JobImpl(parent: Job? = null) : JobSupport(true) {
     init { initParentJob(parent) }
+    override val onCancelMode: Int get() = ON_CANCEL_MAKE_COMPLETING
 }
 
 // -------- invokeOnCompletion nodes

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

@@ -88,7 +88,7 @@ private fun <U, T: U> setupTimeout(
     }
     return when {
         result == COROUTINE_SUSPENDED -> COROUTINE_SUSPENDED
-        coroutine.makeCompleting(result, MODE_IGNORE) -> {
+        coroutine.makeCompletingOnce(result, MODE_IGNORE) -> {
             if (result is CompletedExceptionally) throw result.exception else result
         }
         else -> COROUTINE_SUSPENDED

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

@@ -90,9 +90,5 @@ class JobDisposeTest: TestBase() {
         exception?.let { throw it }
     }
 
-    class TestJob : JobSupport(active = true) {
-        // The bug was triggering only with cancelling state
-        override val hasCancellingState: Boolean
-            get() = true
-    }
+    class TestJob : JobSupport(active = true)
 }