Optimize creation of DispatchTask for cancellable continuations

Author: elizarov

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

@@ -28,10 +28,11 @@ private const val RESUMED = 2
 /**
  * @suppress **This is unstable API and it is subject to change.**
  */
+// Note: it also works directly as DispatchTask for this delegate
 internal abstract class AbstractContinuation<in T>(
     @JvmField protected val delegate: Continuation<T>,
     @JvmField protected val resumeMode: Int
-) : JobSupport(true), Continuation<T> {
+) : JobSupport(true), Continuation<T>, Runnable {
     private val _decision = atomic(UNDECIDED)
 
     /* decision state machine
@@ -81,10 +82,23 @@ internal abstract class AbstractContinuation<in 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
+        var useMode = mode
+        if (mode.isDispatchedMode && delegate is DispatchedContinuation<*> && mode.isCancellableMode == resumeMode.isCancellableMode) {
+            // dispatch directly using this instance's Runnable implementation
+            val dispatcher = delegate.dispatcher
+            val context = delegate.context
+            if (dispatcher.isDispatchNeeded(context)) {
+                dispatcher.dispatch(context, this)
+                return // and that's it -- dispatched via fast-path
+            } else {
+                useMode = MODE_UNDISPATCHED
+            }
+        }
+        // slow-path - use delegate
         if (state is CompletedExceptionally) {
-            delegate.resumeWithExceptionMode(state.exception, mode)
+            delegate.resumeWithExceptionMode(state.exception, useMode)
         } else {
-            delegate.resumeMode(getSuccessfulResult(state), mode)
+            delegate.resumeMode(getSuccessfulResult(state), useMode)
         }
     }
 
@@ -118,4 +132,23 @@ internal abstract class AbstractContinuation<in T>(
     override fun handleException(exception: Throwable) {
         handleCoroutineException(context, exception)
     }
+
+    // see all DispatchTask.run with the same logic
+    override fun run() {
+        check(delegate is DispatchedContinuation)
+        try {
+            val context = delegate.context
+            val job = if (resumeMode.isCancellableMode) context[Job] else null
+            val state = this.state
+            withCoroutineContext(context) {
+                when {
+                    job != null && !job.isActive -> delegate.resumeWithException(job.getCancellationException())
+                    state is CompletedExceptionally -> delegate.resumeWithException(state.exception)
+                    else -> delegate.resume(getSuccessfulResult(state))
+                }
+            }
+        } catch (e: Throwable) {
+            throw RuntimeException("Unexpected exception running $this", e)
+        }
+    }
 }

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

@@ -214,7 +214,7 @@ private class RemoveOnCancel(
 internal class CancellableContinuationImpl<in T>(
     delegate: Continuation<T>,
     resumeMode: Int
-) : AbstractContinuation<T>(delegate, resumeMode), CancellableContinuation<T> {
+) : AbstractContinuation<T>(delegate, resumeMode), CancellableContinuation<T>, Runnable {
     @Volatile // just in case -- we don't want an extra data race, even benign one
     private var _context: CoroutineContext? = null // created on first need
 

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

@@ -110,6 +110,7 @@ public actual abstract class CoroutineDispatcher actual constructor() :
 public actual typealias Runnable = java.lang.Runnable
 
 // named class for ease of debugging, better stack-traces and optimize the number of anonymous classes
+// note that CancellableContinuationImpl directly works as DispatchTask
 internal class DispatchTask<in T>(
     private val continuation: Continuation<T>,
     private val value: Any?, // T | Throwable

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

@@ -24,6 +24,9 @@ import kotlin.coroutines.experimental.Continuation
 @PublishedApi internal const val MODE_UNDISPATCHED = 3   // when the thread is right, but need to mark it with current coroutine
 @PublishedApi internal const val MODE_IGNORE = 4         // don't do anything
 
+internal val Int.isCancellableMode get() = this == MODE_CANCELLABLE
+internal val Int.isDispatchedMode get() = this == MODE_ATOMIC_DEFAULT || this == MODE_CANCELLABLE
+
 internal fun <T> Continuation<T>.resumeMode(value: T, mode: Int) {
     when (mode) {
         MODE_ATOMIC_DEFAULT -> resume(value)