Delay implementation in Swing, JavaFx, and scheduled executors is fixed to avoid an extra dispatch

CancellableContinuation.resumeUndispatched is introduced to make this efficient implementation possible
Additional documentation for CoroutineDispatcher implementors

Author: elizarov

kotlinx-coroutines-core/pom.xml

@@ -45,6 +45,17 @@
                     <argLine>-ea -Xmx1g -Xms1g</argLine>
                 </configuration>
             </plugin>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-jar-plugin</artifactId>
+                <executions>
+                    <execution>
+                        <goals>
+                            <goal>test-jar</goal>
+                        </goals>
+                    </execution>
+                </executions>
+            </plugin>
         </plugins>
     </build>
 </project>

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

@@ -49,32 +49,54 @@ public interface CancellableContinuation<in T> : Continuation<T>, Job {
      *
      * It implies that [isActive] is `false` and [isCompleted] is `true`.
      */
-    val isCancelled: Boolean
+    public val isCancelled: Boolean
 
     /**
      * Tries to resume this continuation with a given value and returns non-null object token if it was successful,
      * or `null` otherwise (it was already resumed or cancelled). When non-null object was returned,
      * [completeResume] must be invoked with it.
      */
-    fun tryResume(value: T): Any?
+    public fun tryResume(value: T): Any?
 
     /**
      * Tries to resume this continuation with a given exception and returns non-null object token if it was successful,
      * or `null` otherwise (it was already resumed or cancelled). When non-null object was returned,
      * [completeResume] must be invoked with it.
      */
-    fun tryResumeWithException(exception: Throwable): Any?
+    public fun tryResumeWithException(exception: Throwable): Any?
 
     /**
      * Completes the execution of [tryResume] or [tryResumeWithException] on its non-null result.
      */
-    fun completeResume(token: Any)
+    public fun completeResume(token: Any)
 
     /**
      * Makes this continuation cancellable. Use it with `holdCancellability` optional parameter to
      * [suspendCancellableCoroutine] function. It throws [IllegalStateException] if invoked more than once.
      */
-    fun initCancellability()
+    public fun initCancellability()
+
+    /**
+     * Resumes this continuation with a given [value] in the invoker thread without going though
+     * [dispatch][CoroutineDispatcher.dispatch] function of the [CoroutineDispatcher] in the [context].
+     * This function is designed to be used only by the [CoroutineDispatcher] implementations themselves.
+     * **It should not be used in general code**.
+     *
+     * The receiver [CoroutineDispatcher] of this function be equal to the context dispatcher or
+     * [IllegalArgumentException] if thrown.
+     */
+    public fun CoroutineDispatcher.resumeUndispatched(value: T)
+
+    /**
+     * Resumes this continuation with a given [exception] in the invoker thread without going though
+     * [dispatch][CoroutineDispatcher.dispatch] function of the [CoroutineDispatcher] in the [context].
+     * This function is designed to be used only by the [CoroutineDispatcher] implementations themselves.
+     * **It should not be used in general code**.
+     *
+     * The receiver [CoroutineDispatcher] of this function be equal to the context dispatcher or
+     * [IllegalArgumentException] if thrown.
+     */
+    public fun CoroutineDispatcher.resumeUndispatchedWithException(exception: Throwable)
 }
 
 /**
@@ -123,6 +145,7 @@ internal class SafeCancellableContinuation<in T>(
         const val SUSPENDED = 1
         const val RESUMED = 2
         const val YIELD = 3 // used by cancellable "yield"
+        const val UNDISPATCHED = 4 // used by "undispatchedXXX"
     }
 
     override fun initCancellability() {
@@ -173,12 +196,21 @@ internal class SafeCancellableContinuation<in T>(
         val decision = this.decision // volatile read
         if (decision == UNDECIDED && DECISION.compareAndSet(this, UNDECIDED, RESUMED)) return // will get result in getResult
         // otherwise, getResult has already commenced, i.e. it was resumed later or in other thread
+        when {
+            decision == UNDISPATCHED -> undispatchedCompletion(state)
+            state is CompletedExceptionally -> delegate.resumeWithException(state.exception)
+            decision == YIELD && delegate is DispatchedContinuation -> delegate.resumeYield(parentJob, state as T)
+            else -> delegate.resume(state as T)
+        }
+    }
+
+    @Suppress("UNCHECKED_CAST")
+    private fun undispatchedCompletion(state: Any?) {
+        delegate as DispatchedContinuation // type assertion -- was checked in resumeUndispatched
         if (state is CompletedExceptionally)
-            delegate.resumeWithException(state.exception)
-        else if (decision == YIELD && delegate is DispatchedContinuation)
-            delegate.resumeYield(parentJob, state as T)
+            delegate.resumeUndispatchedWithException(state.exception)
         else
-            delegate.resume(state as T)
+            delegate.resumeUndispatched(state as T)
     }
 
     // can only be invoked in the same thread as getResult (see "yield"), afterCompletion may be concurrent
@@ -187,4 +219,18 @@ internal class SafeCancellableContinuation<in T>(
             DECISION.compareAndSet(this, UNDECIDED, YIELD) // try mark as needing dispatch
         resume(value)
     }
+
+    override fun CoroutineDispatcher.resumeUndispatched(value: T) {
+        val dc = delegate as? DispatchedContinuation ?: throw IllegalArgumentException("Must be used with DispatchedContinuation")
+        check(dc.dispatcher === this) { "Must be invoked from the context CoroutineDispatcher"}
+        DECISION.compareAndSet(this@SafeCancellableContinuation, SUSPENDED, UNDISPATCHED)
+        resume(value)
+    }
+
+    override fun CoroutineDispatcher.resumeUndispatchedWithException(exception: Throwable) {
+        val dc = delegate as? DispatchedContinuation ?: throw IllegalArgumentException("Must be used with DispatchedContinuation")
+        check(dc.dispatcher === this) { "Must be invoked from the context CoroutineDispatcher"}
+        DECISION.compareAndSet(this@SafeCancellableContinuation, SUSPENDED, UNDISPATCHED)
+        resumeWithException(exception)
+    }
 }

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

@@ -41,8 +41,31 @@ import kotlin.coroutines.experimental.CoroutineContext
 public abstract class CoroutineDispatcher :
         AbstractCoroutineContextElement(ContinuationInterceptor), ContinuationInterceptor {
     /**
-     * Return `true` if execution shall be dispatched onto another thread.
+     * Returns `true` if execution shall be dispatched onto another thread.
      * The default behaviour for most dispatchers is to return `true`.
+     *
+     * UI dispatchers _should not_ override `isDispatchNeeded`, but leave a default implementation that
+     * returns `true`. To understand the rationale beyond this recommendation, consider the following code:
+     *
+     * ```kotlin
+     * fun asyncUpdateUI() = async(MainThread) {
+     *     // do something here that updates something in UI
+     * }
+     * ```
+     *
+     * When you invoke `asyncUpdateUI` in some background thread, it immediately continues to the next
+     * line, while UI update happens asynchronously in the UI thread. However, if you invoke
+     * it in the UI thread itself, it updates UI _synchronously_ if your `isDispatchNeeded` is
+     * overridden with a thread check. Checking if we are already in the UI thread seems more
+     * efficient (and it might indeed save a few CPU cycles), but this subtle and context-sensitive
+     * difference in behavior makes the resulting async code harder to debug.
+     *
+     * Basically, the choice here is between "JS-style" asynchronous approach (async actions
+     * are always postponed to be executed later in the even dispatch thread) and "C#-style" approach
+     * (async actions are executed in the invoker thread until the first suspension point).
+     * While, C# approach seems to be more efficient, it ends up with recommendations like
+     * "use `yield` if you need to ....". This is error-prone. JS-style approach is more consistent
+     * and does not require programmers to think about whether they need to yield or not.
      */
     public open fun isDispatchNeeded(context: CoroutineContext): Boolean = true
 
@@ -78,28 +101,34 @@ internal class DispatchedContinuation<in T>(
         val context = continuation.context
         if (dispatcher.isDispatchNeeded(context))
             dispatcher.dispatch(context, Runnable {
-                withCoroutineContext(context) {
-                    continuation.resume(value)
-                }
+                resumeUndispatched(value)
             })
         else
-            withCoroutineContext(context) {
-                continuation.resume(value)
-            }
+            resumeUndispatched(value)
+    }
+
+    @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)
+        }
     }
 
     override fun resumeWithException(exception: Throwable) {
         val context = continuation.context
         if (dispatcher.isDispatchNeeded(context))
             dispatcher.dispatch(context, Runnable {
-                withCoroutineContext(context) {
-                    continuation.resumeWithException(exception)
-                }
+                resumeUndispatchedWithException(exception)
             })
         else
-            withCoroutineContext(context) {
-                continuation.resumeWithException(exception)
-            }
+            resumeUndispatchedWithException(exception)
+    }
+
+    @Suppress("NOTHING_TO_INLINE") // we need it inline to save us an entry on the stack
+    inline fun resumeUndispatchedWithException(exception: Throwable) {
+        withCoroutineContext(context) {
+            continuation.resumeWithException(exception)
+        }
     }
 
     // used by "yield" implementation

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

@@ -38,6 +38,18 @@ public interface Delay {
 
     /**
      * Schedules resume of a specified [continuation] after a specified delay [time].
+     *
+     * Continuation **must be scheduled** to resume even if it is already cancelled, because a cancellation is just
+     * an exception that the coroutine that used `delay` might wanted to catch and process. It might
+     * need to close some resources in its `finally` blocks, for example.
+     *
+     * This implementation is supposed to use dispatcher's native ability for scheduled execution in its thread(s).
+     * In order to avoid an extra delay of execution, the following code shall be used to resume this
+     * [continuation] when the code is already executing in the appropriate thread:
+     *
+     * ```kotlin
+     * with(continuation) { resumeUndispatched(Unit) }
+     * ```
      */
     fun scheduleResumeAfterDelay(time: Long, unit: TimeUnit, continuation: CancellableContinuation<Unit>)
 }
@@ -59,6 +71,7 @@ suspend fun delay(time: Long, unit: TimeUnit = TimeUnit.MILLISECONDS) {
             scheduleResumeAfterDelay(time, unit, cont)
             return@sc
         }
-        scheduledExecutor.scheduleResumeAfterDelay(time, unit, cont)
+        val timeout = scheduledExecutor.schedule(ResumeRunnable(cont), time, unit)
+        cont.cancelFutureOnCompletion(timeout)
     }
 }

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

@@ -17,9 +17,9 @@
 package kotlinx.coroutines.experimental
 
 import java.util.concurrent.Executor
-import java.util.concurrent.ExecutorService
 import java.util.concurrent.ScheduledExecutorService
 import java.util.concurrent.TimeUnit
+import kotlin.coroutines.experimental.Continuation
 import kotlin.coroutines.experimental.CoroutineContext
 
 /**
@@ -32,16 +32,26 @@ internal open class ExecutorCoroutineDispatcher(val executor: Executor) : Corout
     override fun dispatch(context: CoroutineContext, block: Runnable) = executor.execute(block)
 
     override fun scheduleResumeAfterDelay(time: Long, unit: TimeUnit, continuation: CancellableContinuation<Unit>) {
-        (executor as? ScheduledExecutorService ?: scheduledExecutor).scheduleResumeAfterDelay(time, unit, continuation)
+        val timeout = if (executor is ScheduledExecutorService)
+            executor.schedule(ResumeUndispatchedRunnable(this, continuation), time, unit) else
+            scheduledExecutor.schedule(ResumeRunnable(continuation), time, unit)
+        continuation.cancelFutureOnCompletion(timeout)
     }
 }
 
-internal fun ExecutorService.scheduleResume(cont: CancellableContinuation<Unit>) {
-    val future = submit { cont.resume(Unit) }
-    cont.cancelFutureOnCompletion(future)
+// --- reusing these classes in other places ---
+
+internal class ResumeUndispatchedRunnable(
+    val dispatcher: CoroutineDispatcher,
+    val continuation: CancellableContinuation<Unit>
+) : Runnable {
+    override fun run() {
+        with(continuation) { dispatcher.resumeUndispatched(Unit) }
+    }
 }
 
-internal fun ScheduledExecutorService.scheduleResumeAfterDelay(time: Long, unit: TimeUnit, cont: CancellableContinuation<Unit>) {
-    val timeout = schedule({ cont.resume(Unit) }, time, unit)
-    cont.cancelFutureOnCompletion(timeout)
+internal class ResumeRunnable(val continuation: Continuation<Unit>) : Runnable {
+    override fun run() {
+        continuation.resume(Unit)
+    }
 }

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

@@ -62,6 +62,7 @@ private class ThreadPoolDispatcher(
     override fun dispatch(context: CoroutineContext, block: Runnable) = executor.execute(block)
 
     override fun scheduleResumeAfterDelay(time: Long, unit: TimeUnit, continuation: CancellableContinuation<Unit>) {
-        executor.scheduleResumeAfterDelay(time, unit, continuation)
+        val timeout = executor.schedule(ResumeUndispatchedRunnable(this, continuation), time, unit)
+        continuation.cancelFutureOnCompletion(timeout)
     }
 }

kotlinx-coroutines-javafx/src/main/kotlin/kotlinx/coroutines/experimental/javafx/JavaFx.kt

@@ -50,24 +50,28 @@ object JavaFx : CoroutineDispatcher(), Delay {
     }
 
     override fun scheduleResumeAfterDelay(time: Long, unit: TimeUnit, continuation: CancellableContinuation<Unit>) {
-        val timeline = Timeline(KeyFrame(Duration.millis(unit.toMillis(time).toDouble()),
-                EventHandler<ActionEvent> { continuation.resume(Unit) }))
+        val handler = EventHandler<ActionEvent> {
+            with(continuation) { resumeUndispatched(Unit) }
+        }
+        val timeline = Timeline(KeyFrame(Duration.millis(unit.toMillis(time).toDouble()), handler))
         timeline.play()
         continuation.onCompletion { timeline.stop() }
     }
-}
 
-private class PulseTimer : AnimationTimer() {
-    val next = CopyOnWriteArrayList<CancellableContinuation<Long>>()
+    private class PulseTimer : AnimationTimer() {
+        val next = CopyOnWriteArrayList<CancellableContinuation<Long>>()
 
-    override fun handle(now: Long) {
-        val cur = next.toTypedArray()
-        next.clear()
-        for (cont in cur)
-            cont.resume(now)
-    }
+        override fun handle(now: Long) {
+            val cur = next.toTypedArray()
+            next.clear()
+            for (cont in cur)
+                with (cont) { resumeUndispatched(now) }
+        }
 
-    fun onNext(cont: CancellableContinuation<Long>) {
-        next += cont
+        fun onNext(cont: CancellableContinuation<Long>) {
+            next += cont
+        }
     }
+
+    override fun toString() = "JavaFx"
 }

kotlinx-coroutines-swing/pom.xml

@@ -40,6 +40,13 @@
             <version>${project.version}</version>
             <scope>compile</scope>
         </dependency>
+        <dependency>
+            <groupId>org.jetbrains.kotlinx</groupId>
+            <artifactId>kotlinx-coroutines-core</artifactId>
+            <version>${project.version}</version>
+            <classifier>tests</classifier>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>org.jetbrains.kotlinx</groupId>
             <artifactId>kotlinx-coroutines-jdk8</artifactId>

kotlinx-coroutines-swing/src/main/kotlin/kotlinx/coroutines/experimental/swing/Swing.kt

@@ -30,15 +30,19 @@ import kotlin.coroutines.experimental.CoroutineContext
  * Dispatches execution onto Swing event dispatching thread and provides native [delay] support.
  */
 object Swing : CoroutineDispatcher(), Delay {
-    override fun isDispatchNeeded(context: CoroutineContext): Boolean = !SwingUtilities.isEventDispatchThread()
     override fun dispatch(context: CoroutineContext, block: Runnable) = SwingUtilities.invokeLater(block)
 
     override fun scheduleResumeAfterDelay(time: Long, unit: TimeUnit, continuation: CancellableContinuation<Unit>) {
         val timerTime = unit.toMillis(time).coerceAtMost(Int.MAX_VALUE.toLong()).toInt()
-        val timer = Timer(timerTime, ActionListener { continuation.resume(Unit) }).apply {
+        val action = ActionListener {
+            with(continuation) { resumeUndispatched(Unit) }
+        }
+        val timer = Timer(timerTime, action).apply {
             isRepeats = false
             start()
         }
         continuation.onCompletion { timer.stop() }
     }
+
+    override fun toString() = "Swing"
 }