Review and optimize usage of CancellableContinuation.invokeOnCancellation

Author: elizarov

common/kotlinx-coroutines-core-common/src/main/kotlin/kotlinx/coroutines/experimental/Await.kt

@@ -74,12 +74,17 @@ private class AwaitAll<T>(private val deferreds: Array<out Deferred<T>>) {
             it.start() // To properly await lazily started deferreds
             it.invokeOnCompletion(AwaitAllNode(cont, it).asHandler)
         }
-        cont.invokeOnCancellation {
+        cont.invokeOnCancellation(handler = DisposeHandlersOnCancel(handlers).asHandler)
+    }
+
+    private class DisposeHandlersOnCancel(private val handlers: List<DisposableHandle>) : CancelHandler() {
+        override fun invoke(cause: Throwable?) {
             handlers.forEach { it.dispose() }
         }
+        override fun toString(): String = "DisposeHandlersOnCancel[$handlers]"
     }
 
-    inner class AwaitAllNode(private val continuation: CancellableContinuation<List<T>>, job: Job) : JobNode<Job>(job) {
+    private inner class AwaitAllNode(private val continuation: CancellableContinuation<List<T>>, job: Job) : JobNode<Job>(job) {
         override fun invoke(cause: Throwable?) {
             if (cause != null) {
                 val token = continuation.tryResumeWithException(cause)

common/kotlinx-coroutines-core-common/src/main/kotlin/kotlinx/coroutines/experimental/channels/AbstractChannel.kt

@@ -769,11 +769,15 @@ public abstract class AbstractChannel<E> : AbstractSendChannel<E>(), Channel<E>
 
     // ------ private ------
 
-    private fun removeReceiveOnCancel(cont: CancellableContinuation<*>, receive: Receive<*>) {
-        cont.invokeOnCancellation {
+    private fun removeReceiveOnCancel(cont: CancellableContinuation<*>, receive: Receive<*>) =
+        cont.invokeOnCancellation(handler = RemoveReceiveOnCancel(receive).asHandler)
+
+    private inner class RemoveReceiveOnCancel(private val receive: Receive<*>) : CancelHandler() {
+        override fun invoke(cause: Throwable?) {
             if (receive.remove())
                 onReceiveDequeued()
         }
+        override fun toString(): String = "RemoveReceiveOnCancel[$receive]"
     }
 
     private class Itr<E>(val channel: AbstractChannel<E>) : ChannelIterator<E> {

integration/kotlinx-coroutines-nio/src/main/kotlin/kotlinx/coroutines/experimental/nio/Nio.kt

@@ -137,12 +137,12 @@ suspend fun AsynchronousSocketChannel.aWrite(
 
 private fun Channel.closeOnCancel(cont: CancellableContinuation<*>) {
     cont.invokeOnCancellation {
-            try {
-                close()
-            } catch (ex: Throwable) {
-                // Specification says that it is Ok to call it any time, but reality is different,
-                // so we have just to ignore exception
-            }
+        try {
+            close()
+        } catch (ex: Throwable) {
+            // Specification says that it is Ok to call it any time, but reality is different,
+            // so we have just to ignore exception
+        }
     }
 }
 

js/kotlinx-coroutines-core-js/src/main/kotlin/kotlinx/coroutines/experimental/JSDispatcher.kt

@@ -28,16 +28,18 @@ internal class NodeDispatcher : CoroutineDispatcher(), Delay {
     override fun scheduleResumeAfterDelay(time: Long, unit: TimeUnit, continuation: CancellableContinuation<Unit>) {
         val handle = setTimeout({ with(continuation) { resumeUndispatched(Unit) } }, time.toIntMillis(unit))
         // Actually on cancellation, but clearTimeout is idempotent
-        continuation.invokeOnCancellation { clearTimeout(handle) }
+        continuation.invokeOnCancellation(handler = ClearTimeout(handle).asHandler)
+    }
+
+    private class ClearTimeout(private val handle: Int) : CancelHandler(), DisposableHandle {
+        override fun dispose() { clearTimeout(handle) }
+        override fun invoke(cause: Throwable?) { dispose() }
+        override fun toString(): String = "ClearTimeout[$handle]"
     }
 
     override fun invokeOnTimeout(time: Long, unit: TimeUnit, block: Runnable): DisposableHandle {
         val handle = setTimeout({ block.run() }, time.toIntMillis(unit))
-        return object : DisposableHandle {
-            override fun dispose() {
-                clearTimeout(handle)
-            }
-        }
+        return ClearTimeout(handle)
     }
 }
 

reactive/kotlinx-coroutines-reactive/src/main/kotlin/kotlinx/coroutines/experimental/reactive/Await.kt

@@ -120,15 +120,15 @@ private suspend fun <T> Publisher<T>.awaitOne(
                 Mode.FIRST, Mode.FIRST_OR_DEFAULT -> {
                     if (!seenValue) {
                         seenValue = true
-                        cont.resume(t)
                         subscription.cancel()
+                        cont.resume(t)
                     }
                 }
                 Mode.LAST, Mode.SINGLE -> {
                     if (mode == Mode.SINGLE && seenValue) {
+                        subscription.cancel()
                         if (cont.isActive)
                             cont.resumeWithException(IllegalArgumentException("More that one onNext value for $mode"))
-                        subscription.cancel()
                     } else {
                         value = t
                         seenValue = true

reactive/kotlinx-coroutines-rx1/src/main/kotlin/kotlinx/coroutines/experimental/rx1/RxAwait.kt

@@ -33,7 +33,7 @@ import rx.*
  */
 public suspend fun Completable.awaitCompleted(): Unit = suspendCancellableCoroutine { cont ->
     subscribe(object : CompletableSubscriber {
-        override fun onSubscribe(s: Subscription) { cont.unsubscribeOnCompletion(s) }
+        override fun onSubscribe(s: Subscription) { cont.unsubscribeOnCancellation(s) }
         override fun onCompleted() { cont.resume(Unit) }
         override fun onError(e: Throwable) { cont.resumeWithException(e) }
     })
@@ -50,7 +50,7 @@ public suspend fun Completable.awaitCompleted(): Unit = suspendCancellableCorout
  * immediately resumes with [CancellationException].
  */
 public suspend fun <T> Single<T>.await(): T = suspendCancellableCoroutine { cont ->
-    cont.unsubscribeOnCompletion(subscribe(object : SingleSubscriber<T>() {
+    cont.unsubscribeOnCancellation(subscribe(object : SingleSubscriber<T>() {
         override fun onSuccess(t: T) { cont.resume(t) }
         override fun onError(error: Throwable) { cont.resumeWithException(error) }
     }))
@@ -128,14 +128,13 @@ public suspend fun <T> Observable<T>.awaitSingle(): T = single().awaitOne()
 // ------------------------ private ------------------------
 
 private suspend fun <T> Observable<T>.awaitOne(): T = suspendCancellableCoroutine { cont ->
-    cont.unsubscribeOnCompletion(subscribe(object : Subscriber<T>() {
+    cont.unsubscribeOnCancellation(subscribe(object : Subscriber<T>() {
         override fun onStart() { request(1) }
         override fun onNext(t: T) { cont.resume(t) }
         override fun onCompleted() { if (cont.isActive) cont.resumeWithException(IllegalStateException("Should have invoked onNext")) }
         override fun onError(e: Throwable) { cont.resumeWithException(e) }
     }))
 }
 
-internal fun <T> CancellableContinuation<T>.unsubscribeOnCompletion(sub: Subscription) {
+internal fun <T> CancellableContinuation<T>.unsubscribeOnCancellation(sub: Subscription) =
     invokeOnCancellation { sub.unsubscribe() }
-}

reactive/kotlinx-coroutines-rx1/src/main/kotlin/kotlinx/coroutines/experimental/rx1/RxScheduler.kt

@@ -46,7 +46,7 @@ public class SchedulerCoroutineDispatcher(private val scheduler: Scheduler) : Co
                     with(continuation) { resumeUndispatched(Unit) }
                  }, time, unit)
                 .let { subscription ->
-                    continuation.unsubscribeOnCompletion(subscription)
+                    continuation.unsubscribeOnCancellation(subscription)
                 }
 
     override fun invokeOnTimeout(time: Long, unit: TimeUnit, block: Runnable): DisposableHandle =

reactive/kotlinx-coroutines-rx2/src/main/kotlin/kotlinx/coroutines/experimental/rx2/RxAwait.kt

@@ -34,7 +34,7 @@ import kotlinx.coroutines.experimental.suspendCancellableCoroutine
  */
 public suspend fun CompletableSource.await(): Unit = suspendCancellableCoroutine { cont ->
     subscribe(object : CompletableObserver {
-        override fun onSubscribe(d: Disposable) { cont.disposeOnCompletion(d) }
+        override fun onSubscribe(d: Disposable) { cont.disposeOnCancellation(d) }
         override fun onComplete() { cont.resume(Unit) }
         override fun onError(e: Throwable) { cont.resumeWithException(e) }
     })
@@ -65,7 +65,7 @@ public suspend fun <T> MaybeSource<T>.await(): T? = (this as MaybeSource<T?>).aw
  */
 public suspend fun <T> MaybeSource<T>.awaitOrDefault(default: T): T = suspendCancellableCoroutine { cont ->
     subscribe(object : MaybeObserver<T> {
-        override fun onSubscribe(d: Disposable) { cont.disposeOnCompletion(d) }
+        override fun onSubscribe(d: Disposable) { cont.disposeOnCancellation(d) }
         override fun onComplete() { cont.resume(default) }
         override fun onSuccess(t: T) { cont.resume(t) }
         override fun onError(error: Throwable) { cont.resumeWithException(error) }
@@ -84,7 +84,7 @@ public suspend fun <T> MaybeSource<T>.awaitOrDefault(default: T): T = suspendCan
  */
 public suspend fun <T> SingleSource<T>.await(): T = suspendCancellableCoroutine { cont ->
     subscribe(object : SingleObserver<T> {
-        override fun onSubscribe(d: Disposable) { cont.disposeOnCompletion(d) }
+        override fun onSubscribe(d: Disposable) { cont.disposeOnCancellation(d) }
         override fun onSuccess(t: T) { cont.resume(t) }
         override fun onError(error: Throwable) { cont.resumeWithException(error) }
     })
@@ -161,7 +161,7 @@ public suspend fun <T> ObservableSource<T>.awaitSingle(): T = awaitOne(Mode.SING
 
 // ------------------------ private ------------------------
 
-internal fun CancellableContinuation<*>.disposeOnCompletion(d: Disposable) =
+internal fun CancellableContinuation<*>.disposeOnCancellation(d: Disposable) =
     invokeOnCancellation { d.dispose() }
 
 private enum class Mode(val s: String) {

reactive/kotlinx-coroutines-rx2/src/main/kotlin/kotlinx/coroutines/experimental/rx2/RxScheduler.kt

@@ -43,7 +43,7 @@ public class SchedulerCoroutineDispatcher(private val scheduler: Scheduler) : Co
         val disposable = scheduler.scheduleDirect({
             with(continuation) { resumeUndispatched(Unit) }
         }, time, unit)
-        continuation.disposeOnCompletion(disposable)
+        continuation.disposeOnCancellation(disposable)
     }
 
     override fun invokeOnTimeout(time: Long, unit: TimeUnit, block: Runnable): DisposableHandle {