Make Dispatchers.setMain affect all delays in tests.

kotlinx-coroutines-core/api/kotlinx-coroutines-core.api

@@ -249,6 +249,18 @@ public final class kotlinx/coroutines/DebugKt {
 	public static final field DEBUG_PROPERTY_VALUE_ON Ljava/lang/String;
 }
 
+public final class kotlinx/coroutines/DefaultExecutor : java/lang/Runnable {
+	public static final field INSTANCE Lkotlinx/coroutines/DefaultExecutor;
+	public static final field THREAD_NAME Ljava/lang/String;
+	public static synthetic fun decrementUseCount$default (Lkotlinx/coroutines/EventLoop;ZILjava/lang/Object;)V
+	public fun enqueue (Ljava/lang/Runnable;)V
+	public static synthetic fun incrementUseCount$default (Lkotlinx/coroutines/EventLoop;ZILjava/lang/Object;)V
+	public fun invokeOnTimeout (JLjava/lang/Runnable;Lkotlin/coroutines/CoroutineContext;)Lkotlinx/coroutines/DisposableHandle;
+	public fun run ()V
+	public fun shutdown ()V
+	public final fun shutdownForTests (J)V
+}
+
 public abstract interface class kotlinx/coroutines/Deferred : kotlinx/coroutines/Job {
 	public abstract fun await (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public abstract fun getCompleted ()Ljava/lang/Object;

kotlinx-coroutines-core/common/src/CoroutineContext.common.kt

@@ -12,7 +12,6 @@ import kotlin.coroutines.*
  */
 public expect fun CoroutineScope.newCoroutineContext(context: CoroutineContext): CoroutineContext
 
-@PublishedApi
 @Suppress("PropertyName")
 internal expect val DefaultDelay: Delay
 

kotlinx-coroutines-core/js/src/CoroutineContext.kt

@@ -34,7 +34,7 @@ private fun isJsdom() = jsTypeOf(navigator) != UNDEFINED &&
     navigator.userAgent.match("\\bjsdom\\b")
 
 internal actual val DefaultDelay: Delay
-    get() = Dispatchers.Default as Delay
+    get() = Dispatchers.Main as? Delay ?: Dispatchers.Default as Delay
 
 public actual fun CoroutineScope.newCoroutineContext(context: CoroutineContext): CoroutineContext {
     val combined = coroutineContext + context

kotlinx-coroutines-core/js/src/Dispatchers.kt

@@ -19,11 +19,6 @@ public actual object Dispatchers {
     internal fun injectMain(dispatcher: MainCoroutineDispatcher) {
         injectedMainDispatcher = dispatcher
     }
-
-    @PublishedApi
-    internal fun resetInjectedMain() {
-        injectedMainDispatcher = null
-    }
 }
 
 private class JsMainDispatcher(

kotlinx-coroutines-core/jvm/src/DefaultExecutor.kt

@@ -8,10 +8,10 @@ import kotlinx.coroutines.internal.*
 import java.util.concurrent.*
 import kotlin.coroutines.*
 
-internal actual val DefaultDelay: Delay = initializeDefaultDelay()
-
 private val defaultMainDelayOptIn = systemProp("kotlinx.coroutines.main.delay", true)
 
+internal actual val DefaultDelay: Delay = initializeDefaultDelay()
+
 private fun initializeDefaultDelay(): Delay {
     // Opt-out flag
     if (!defaultMainDelayOptIn) return DefaultExecutor
@@ -21,10 +21,11 @@ private fun initializeDefaultDelay(): Delay {
      * no sense to create a separate thread with timer that cannot be controller
      * by the UI runtime.
      */
-    return if (main.isMissing() || main !is Delay) DefaultExecutor else main
+    return if (main !is Delay) DefaultExecutor else main
 }
 
 @Suppress("PLATFORM_CLASS_MAPPED_TO_KOTLIN")
+@PublishedApi
 internal actual object DefaultExecutor : EventLoopImplBase(), Runnable {
     const val THREAD_NAME = "kotlinx.coroutines.DefaultExecutor"
 

kotlinx-coroutines-core/jvm/src/internal/MainDispatchers.kt

@@ -61,7 +61,12 @@ public fun MainDispatcherFactory.tryCreateDispatcher(factories: List<MainDispatc
 
 /** @suppress */
 @InternalCoroutinesApi
-public fun MainCoroutineDispatcher.isMissing(): Boolean = this is MissingMainCoroutineDispatcher
+public fun MainCoroutineDispatcher.exceptionIfMissing(): Throwable? =
+    try {
+        (this as? MissingMainCoroutineDispatcher)?.isDispatchNeeded(EmptyCoroutineContext)
+    } catch (e: Throwable) {
+        e
+    }
 
 // R8 optimization hook, not const on purpose to enable R8 optimizations via "assumenosideeffects"
 @Suppress("MayBeConstant")
@@ -86,26 +91,17 @@ internal fun throwMissingMainDispatcherException(): Nothing {
 private class MissingMainCoroutineDispatcher(
     private val cause: Throwable?,
     private val errorHint: String? = null
-) : MainCoroutineDispatcher(), Delay {
+) : MainCoroutineDispatcher() {
 
     override val immediate: MainCoroutineDispatcher get() = this
 
-    override fun isDispatchNeeded(context: CoroutineContext): Boolean =
-        missing()
-
-    override fun limitedParallelism(parallelism: Int): CoroutineDispatcher =
-        missing()
-
-    override suspend fun delay(time: Long) =
-        missing()
-
-    override fun invokeOnTimeout(timeMillis: Long, block: Runnable, context: CoroutineContext): DisposableHandle =
+    override fun isDispatchNeeded(context: CoroutineContext): Nothing =
         missing()
 
-    override fun dispatch(context: CoroutineContext, block: Runnable) =
+    override fun limitedParallelism(parallelism: Int): Nothing =
         missing()
 
-    override fun scheduleResumeAfterDelay(timeMillis: Long, continuation: CancellableContinuation<Unit>) =
+    override fun dispatch(context: CoroutineContext, block: Runnable): Nothing =
         missing()
 
     private fun missing(): Nothing {
@@ -124,6 +120,7 @@ private class MissingMainCoroutineDispatcher(
  * @suppress
  */
 @InternalCoroutinesApi
+@Deprecated("No longer used", level = DeprecationLevel.WARNING) // WARNING in 1.6, ERROR in 1.7, removed in 1.8
 public object MissingMainCoroutineDispatcherFactory : MainDispatcherFactory {
     override val loadPriority: Int
         get() = -1

kotlinx-coroutines-core/native/src/CoroutineContext.kt

@@ -41,7 +41,11 @@ internal actual object DefaultExecutor : CoroutineDispatcher(), Delay {
 internal expect fun createDefaultDispatcher(): CoroutineDispatcher
 
 @SharedImmutable
-internal actual val DefaultDelay: Delay = if (multithreadingSupported) DefaultExecutor else OldDefaultExecutor
+@PublishedApi
+internal val NonMockedDefaultDelay: Delay = if (multithreadingSupported) DefaultExecutor else OldDefaultExecutor
+
+internal actual val DefaultDelay: Delay
+    get() = Dispatchers.Main as? Delay ?: NonMockedDefaultDelay
 
 public actual fun CoroutineScope.newCoroutineContext(context: CoroutineContext): CoroutineContext {
     val combined = coroutineContext + context

kotlinx-coroutines-core/native/src/Dispatchers.kt

@@ -6,6 +6,7 @@ package kotlinx.coroutines
 
 import kotlinx.coroutines.internal.multithreadingSupported
 import kotlin.coroutines.*
+import kotlin.native.concurrent.*
 
 public actual object Dispatchers {
     public actual val Default: CoroutineDispatcher = createDefaultDispatcherBasedOnMm()
@@ -24,13 +25,15 @@ public actual object Dispatchers {
         }
         injectedMainDispatcher = dispatcher
     }
-
-    @PublishedApi
-    internal fun resetInjectedMain() {
-        injectedMainDispatcher = null
-    }
 }
 
+/**
+ * Does this platform define a Main dispatcher?
+ */
+@PublishedApi
+@SharedImmutable
+internal expect val mainDispatcherIsPresent: Boolean
+
 internal expect fun createMainDispatcher(default: CoroutineDispatcher): MainCoroutineDispatcher
 
 private fun createDefaultDispatcherBasedOnMm(): CoroutineDispatcher {

kotlinx-coroutines-core/nativeDarwin/src/Dispatchers.kt

@@ -17,6 +17,9 @@ internal fun isMainThread(): Boolean = CFRunLoopGetCurrent() == CFRunLoopGetMain
 internal actual fun createMainDispatcher(default: CoroutineDispatcher): MainCoroutineDispatcher =
     if (multithreadingSupported) DarwinMainDispatcher(false) else OldMainDispatcher(Dispatchers.Default)
 
+@SharedImmutable
+internal actual val mainDispatcherIsPresent: Boolean = true
+
 internal actual fun createDefaultDispatcher(): CoroutineDispatcher = DarwinGlobalQueueDispatcher
 
 private object DarwinGlobalQueueDispatcher : CoroutineDispatcher() {

kotlinx-coroutines-core/nativeOther/src/Dispatchers.kt

@@ -5,10 +5,14 @@
 package kotlinx.coroutines
 
 import kotlin.coroutines.*
+import kotlin.native.concurrent.*
 
 internal actual fun createMainDispatcher(default: CoroutineDispatcher): MainCoroutineDispatcher =
     MissingMainDispatcher
 
+@SharedImmutable
+internal actual val mainDispatcherIsPresent: Boolean = false
+
 internal actual fun createDefaultDispatcher(): CoroutineDispatcher = DefaultDispatcher
 
 private object DefaultDispatcher : CoroutineDispatcher() {

kotlinx-coroutines-test/README.md

@@ -74,8 +74,12 @@ class SomeTest {
 
 Calling `setMain` or `resetMain` immediately changes the `Main` dispatcher globally.
 
-If `Main` is overridden with a [TestDispatcher], then its [TestCoroutineScheduler] is used when new [TestDispatcher] or
-[TestScope] instances are created without [TestCoroutineScheduler] being passed as an argument.
+If `Main` is overridden with a [TestDispatcher], then
+* its [TestCoroutineScheduler] is used when new [TestDispatcher] or [TestScope] instances are created without
+  [TestCoroutineScheduler] being passed as an argument.
+* even the code running on other dispatchers, like `Dispatchers.IO` or `Dispatchers.Default`,
+  will run its delays using its [TestCoroutineScheduler].
+  The only exception is the code running on dispatchers that implement `Delay`.
 
 ## runTest
 
@@ -336,8 +340,9 @@ fun testFooWithTimeout() = runTest {
 ## Virtual time support with other dispatchers
 
 Calls to `withContext(Dispatchers.IO)`, `withContext(Dispatchers.Default)` ,and `withContext(Dispatchers.Main)` are
-common in coroutines-based code bases. Unfortunately, just executing code in a test will not lead to these dispatchers
-using the virtual time source, so delays will not be skipped in them.
+common in coroutines-based code bases.
+Unfortunately, just executing code in a test will not lead to these dispatchers using the virtual time source, so delays
+will not be skipped in them.
 
 ```kotlin
 suspend fun veryExpensiveFunction() = withContext(Dispatchers.Default) {
@@ -351,9 +356,13 @@ fun testExpensiveFunction() = runTest {
 }
 ```
 
-Tests should, when possible, replace these dispatchers with a [TestDispatcher] if the `withContext` calls `delay` in the
-function under test. For example, `veryExpensiveFunction` above should allow mocking with a [TestDispatcher] using
-either dependency injection, a service locator, or a default parameter, if it is to be used with virtual time.
+To ensure that these delays are skipped, do one of the following:
+* Call [Dispatchers.setMain] with a [TestDispatcher].
+  Then, all delays that happen outside of dispatchers implementing `Delay` will use its [TestCoroutineScheduler].
+* Replace these dispatchers with a [TestDispatcher] if the `withContext` calls `delay` in the
+  function under test.
+  For example, `veryExpensiveFunction` above should allow mocking with a [TestDispatcher] using either dependency
+  injection, a service locator, or a default parameter, if it is to be used with virtual time.
 
 ### Status of the API
 

kotlinx-coroutines-test/common/src/TestBuilders.kt

@@ -8,6 +8,7 @@ package kotlinx.coroutines.test
 
 import kotlinx.coroutines.*
 import kotlinx.coroutines.selects.*
+import kotlinx.coroutines.test.internal.*
 import kotlin.coroutines.*
 import kotlin.jvm.*
 
@@ -214,21 +215,24 @@ internal suspend fun <T: AbstractCoroutine<Unit>> runTestCoroutine(
             completed = true
             continue
         }
-        select<Unit> {
-            coroutine.onJoin {
-                completed = true
-            }
-            scheduler.onDispatchEvent {
-                // we received knowledge that `scheduler` observed a dispatch event, so we reset the timeout
-            }
-            onTimeout(dispatchTimeoutMs) {
-                try {
-                    cleanup()
-                } catch (e: UncompletedCoroutinesError) {
-                    // we expect these and will instead throw a more informative exception just below.
-                    emptyList()
-                }.throwAll()
-                throw UncompletedCoroutinesError("The test coroutine was not completed after waiting for $dispatchTimeoutMs ms")
+        withContext(nonMockedDelay as ContinuationInterceptor) {
+            select<Unit> {
+                coroutine.onJoin {
+                    completed = true
+                }
+                scheduler.onDispatchEvent {
+                    // we received knowledge that `scheduler` observed a dispatch event, so we reset the timeout
+                }
+                onTimeout(dispatchTimeoutMs) {
+                    try {
+                        cleanup()
+                    } catch (e: UncompletedCoroutinesError) {
+                        // we expect these and will instead throw a more informative exception just below.
+                        emptyList()
+                    }.throwAll()
+                    throw UncompletedCoroutinesError(
+                        "The test coroutine was not completed after waiting for $dispatchTimeoutMs ms")
+                }
             }
         }
     }

kotlinx-coroutines-test/common/src/internal/TestMainDispatcher.kt

@@ -13,15 +13,15 @@ import kotlin.coroutines.*
  * The testable main dispatcher used by kotlinx-coroutines-test.
  * It is a [MainCoroutineDispatcher] that delegates all actions to a settable delegate.
  */
-internal class TestMainDispatcher(delegate: CoroutineDispatcher):
+internal class TestMainDispatcher(delegate: CoroutineDispatcher?, mainInitException: Throwable? = null) :
     MainCoroutineDispatcher(),
-    Delay
-{
-    private val mainDispatcher = delegate
+    Delay {
+    private val mainDispatcher = delegate ?: UnsetMainDispatcher(mainInitException)
+
     private var delegate = NonConcurrentlyModifiable(mainDispatcher, "Dispatchers.Main")
 
     private val delay
-        get() = delegate.value as? Delay ?: defaultDelay
+        get() = delegate.value as? Delay ?: nonMockedDelay
 
     override val immediate: MainCoroutineDispatcher
         get() = (delegate.value as? MainCoroutineDispatcher)?.immediate ?: this
@@ -30,7 +30,8 @@ internal class TestMainDispatcher(delegate: CoroutineDispatcher):
 
     override fun isDispatchNeeded(context: CoroutineContext): Boolean = delegate.value.isDispatchNeeded(context)
 
-    override fun dispatchYield(context: CoroutineContext, block: Runnable) = delegate.value.dispatchYield(context, block)
+    override fun dispatchYield(context: CoroutineContext, block: Runnable) =
+        delegate.value.dispatchYield(context, block)
 
     fun setDispatcher(dispatcher: CoroutineDispatcher) {
         delegate.value = dispatcher
@@ -46,6 +47,8 @@ internal class TestMainDispatcher(delegate: CoroutineDispatcher):
     override fun invokeOnTimeout(timeMillis: Long, block: Runnable, context: CoroutineContext): DisposableHandle =
         delay.invokeOnTimeout(timeMillis, block, context)
 
+    override fun toString(): String = "TestMainDispatcher[delegate=${delegate.value}]"
+
     companion object {
         internal val currentTestDispatcher
             get() = (Dispatchers.Main as? TestMainDispatcher)?.delegate?.value as? TestDispatcher
@@ -86,11 +89,25 @@ internal class TestMainDispatcher(delegate: CoroutineDispatcher):
                 if (readers.value != 0) throw concurrentRW()
             }
     }
+
+    private class UnsetMainDispatcher(private val mainInitException: Throwable?) : MainCoroutineDispatcher() {
+        override val immediate: MainCoroutineDispatcher get() = this
+        override fun isDispatchNeeded(context: CoroutineContext): Boolean = missing()
+        override fun limitedParallelism(parallelism: Int): CoroutineDispatcher = missing()
+        override fun dispatch(context: CoroutineContext, block: Runnable) = missing()
+
+        private fun missing(): Nothing {
+            val message = if (mainInitException == null)
+                "Dispatchers.Main is not available and was not provided for tests via Dispatchers.setMain."
+            else
+                "Dispatchers.Main failed to initialize and was not replaced for tests via Dispatchers.setMain."
+            throw IllegalStateException(message, mainInitException)
+        }
+
+        override fun toString(): String = "missing"
+    }
 }
 
-@Suppress("INVISIBLE_MEMBER")
-private val defaultDelay
-    inline get() = DefaultDelay
+internal expect val nonMockedDelay: Delay
 
-@Suppress("INVISIBLE_MEMBER")
-internal expect fun Dispatchers.getTestMainDispatcher(): TestMainDispatcher
+internal expect fun Dispatchers.getTestMainDispatcher(): TestMainDispatcher

kotlinx-coroutines-test/common/test/RunTestTest.kt

@@ -6,6 +6,7 @@ package kotlinx.coroutines.test
 
 import kotlinx.coroutines.*
 import kotlinx.coroutines.flow.*
+import kotlinx.coroutines.selects.*
 import kotlin.coroutines.*
 import kotlin.test.*
 
@@ -70,6 +71,7 @@ class RunTestTest {
 
     /** Tests that a dispatch timeout of `0` will fail the test if there are some dispatches outside the scheduler. */
     @Test
+    @NoNative // the event loop was shut down
     fun testRunTestWithZeroTimeoutWithUncontrolledDispatches() = testResultMap({ fn ->
         assertFailsWith<UncompletedCoroutinesError> { fn() }
     }) {
@@ -105,6 +107,23 @@ class RunTestTest {
         }
     }
 
+    /** Tests that [onTimeout] executes quickly. */
+    @Test
+    fun testOnTimeout() = runTest {
+        assertRunsFast {
+            val deferred = CompletableDeferred<Unit>()
+            val result = select<Boolean> {
+                onTimeout(SLOW) {
+                    true
+                }
+                deferred.onJoin {
+                    fail("unreached")
+                }
+            }
+            assertTrue(result)
+        }
+    }
+
     /** Tests uncaught exceptions taking priority over dispatch timeout in error reports. */
     @Test
     @NoNative // TODO: timeout leads to `Cannot execute task because event loop was shut down` on Native