Eagerly enter launch and async blocks with unconfined dispatcher (#3011)

Also, fix `Dispatchers.Main` not delegating `Delay` methods and
discover that, on JS, `Dispatchers.Main` gets reset during the test
if it is reset in `AfterTest`.

Author: dkhalanskyjb

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

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

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

@@ -219,7 +219,7 @@ public fun runTest(
     return createTestResult {
         /** TODO: moving this [AbstractCoroutine.start] call outside [createTestResult] fails on Native with
          * [TestCoroutineDispatcher], because the event loop is not started. */
-        testScope.start(CoroutineStart.DEFAULT, testScope) {
+        testScope.start(CoroutineStart.UNDISPATCHED, testScope) {
             testBody()
         }
         var completed = false

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

@@ -7,6 +7,8 @@ package kotlinx.coroutines.test
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
 import kotlinx.coroutines.flow.*
+import kotlinx.coroutines.test.internal.*
+import kotlinx.coroutines.test.internal.TestMainDispatcher
 import kotlin.coroutines.*
 
 /**
@@ -15,10 +17,32 @@ import kotlin.coroutines.*
  * This dispatcher is similar to [Dispatchers.Unconfined]: the tasks that it executes are not confined to any particular
  * thread and form an event loop; it's different in that it skips delays, as all [TestDispatcher]s do.
  *
+ * Like [Dispatchers.Unconfined], this one does not provide guarantees about the execution order when several coroutines
+ * are queued in this dispatcher. However, we ensure that the [launch] and [async] blocks at the top level of [runTest]
+ * are entered eagerly. This allows launching child coroutines and not calling [runCurrent] for them to start executing.
+ *
+ * ```
+ * @Test
+ * fun testEagerlyEnteringChildCoroutines() = runTest(UnconfinedTestDispatcher()) {
+ *   var entered = false
+ *   val deferred = CompletableDeferred<Unit>()
+ *   var completed = false
+ *   launch {
+ *     entered = true
+ *     deferred.await()
+ *     completed = true
+ *   }
+ *   assertTrue(entered) // `entered = true` already executed.
+ *   assertFalse(completed) // however, the child coroutine then suspended, so it is enqueued.
+ *   deferred.complete(Unit) // resume the coroutine.
+ *   assertTrue(completed) // now the child coroutine is immediately completed.
+ * }
+ * ```
+ *
  * Using this [TestDispatcher] can greatly simplify writing tests where it's not important which thread is used when and
  * in which order the queued coroutines are executed.
- * The typical use case for this is launching child coroutines that are resumed immediately, without going through a
- * dispatch; this can be helpful for testing [Channel] and [StateFlow] usages.
+ * Another typical use case for this dispatcher is launching child coroutines that are resumed immediately, without
+ * going through a dispatch; this can be helpful for testing [Channel] and [StateFlow] usages.
  *
  * ```
  * @Test
@@ -40,14 +64,16 @@ import kotlin.coroutines.*
  * }
  * ```
  *
- * However, please be aware that, like [Dispatchers.Unconfined], this is a specific dispatcher with execution order
+ * Please be aware that, like [Dispatchers.Unconfined], this is a specific dispatcher with execution order
  * guarantees that are unusual and not shared by most other dispatchers, so it can only be used reliably for testing
  * functionality, not the specific order of actions.
  * See [Dispatchers.Unconfined] for a discussion of the execution order guarantees.
  *
  * In order to support delay skipping, this dispatcher is linked to a [TestCoroutineScheduler], which is used to control
- * the virtual time and can be shared among many test dispatchers. If no [scheduler] is passed as an argument, a new one
- * is created.
+ * the virtual time and can be shared among many test dispatchers.
+ * If no [scheduler] is passed as an argument, [Dispatchers.Main] is checked, and if it was mocked with a
+ * [TestDispatcher] via [Dispatchers.setMain], the [TestDispatcher.scheduler] of the mock dispatcher is used; if
+ * [Dispatchers.Main] is not mocked with a [TestDispatcher], a new [TestCoroutineScheduler] is created.
  *
  * Additionally, [name] can be set to distinguish each dispatcher instance when debugging.
  *
@@ -56,14 +82,14 @@ import kotlin.coroutines.*
 @ExperimentalCoroutinesApi
 @Suppress("FunctionName")
 public fun UnconfinedTestDispatcher(
-    scheduler: TestCoroutineScheduler = TestCoroutineScheduler(),
+    scheduler: TestCoroutineScheduler? = null,
     name: String? = null
-): TestDispatcher = UnconfinedTestDispatcherImpl(scheduler, name)
+): TestDispatcher = UnconfinedTestDispatcherImpl(scheduler ?: mainTestScheduler ?: TestCoroutineScheduler(), name)
 
 private class UnconfinedTestDispatcherImpl(
     override val scheduler: TestCoroutineScheduler,
     private val name: String? = null
-): TestDispatcher() {
+) : TestDispatcher() {
 
     override fun isDispatchNeeded(context: CoroutineContext): Boolean = false
 
@@ -103,22 +129,24 @@ private class UnconfinedTestDispatcherImpl(
  * run these pending tasks, which will block until there are no more tasks scheduled at this point in time, or, when
  * inside [runTest], call [yield] to yield the (only) thread used by [runTest] to the newly-launched coroutines.
  *
- * If a [scheduler] is not passed as an argument, a new one is created.
+ * If no [scheduler] is passed as an argument, [Dispatchers.Main] is checked, and if it was mocked with a
+ * [TestDispatcher] via [Dispatchers.setMain], the [TestDispatcher.scheduler] of the mock dispatcher is used; if
+ * [Dispatchers.Main] is not mocked with a [TestDispatcher], a new [TestCoroutineScheduler] is created.
  *
  * One can additionally pass a [name] in order to more easily distinguish this dispatcher during debugging.
  *
  * @see UnconfinedTestDispatcher for a dispatcher that is not confined to any particular thread.
  */
 @Suppress("FunctionName")
 public fun StandardTestDispatcher(
-    scheduler: TestCoroutineScheduler = TestCoroutineScheduler(),
+    scheduler: TestCoroutineScheduler? = null,
     name: String? = null
-): TestDispatcher = StandardTestDispatcherImpl(scheduler, name)
+): TestDispatcher = StandardTestDispatcherImpl(scheduler ?: mainTestScheduler ?: TestCoroutineScheduler(), name)
 
 private class StandardTestDispatcherImpl(
     override val scheduler: TestCoroutineScheduler = TestCoroutineScheduler(),
     private val name: String? = null
-): TestDispatcher() {
+) : TestDispatcher() {
 
     override fun dispatch(context: CoroutineContext, block: Runnable) {
         checkSchedulerInContext(scheduler, context)
@@ -127,3 +155,6 @@ private class StandardTestDispatcherImpl(
 
     override fun toString(): String = "${name ?: "StandardTestDispatcher"}[scheduler=$scheduler]"
 }
+
+private val mainTestScheduler
+    get() = ((Dispatchers.Main as? TestMainDispatcher)?.delegate as? TestDispatcher)?.scheduler

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

@@ -6,8 +6,6 @@ package kotlinx.coroutines.test
 
 import kotlinx.coroutines.*
 import kotlinx.coroutines.internal.*
-import kotlinx.coroutines.test.internal.*
-import kotlinx.coroutines.test.internal.TestMainDispatcher
 import kotlin.coroutines.*
 
 /**
@@ -172,12 +170,7 @@ public fun createTestCoroutineScope(context: CoroutineContext = EmptyCoroutineCo
             }
             dispatcher
         }
-        null -> {
-            val mainDispatcherScheduler =
-                ((Dispatchers.Main as? TestMainDispatcher)?.delegate as? TestDispatcher)?.scheduler
-            scheduler = context[TestCoroutineScheduler] ?: mainDispatcherScheduler ?: TestCoroutineScheduler()
-            StandardTestDispatcher(scheduler)
-        }
+        null -> StandardTestDispatcher(context[TestCoroutineScheduler]).also { scheduler = it.scheduler }
         else -> throw IllegalArgumentException("Dispatcher must implement TestDispatcher: $dispatcher")
     }
     var scope: TestCoroutineScopeImpl? = null

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

@@ -12,10 +12,13 @@ import kotlin.coroutines.*
  */
 internal class TestMainDispatcher(var delegate: CoroutineDispatcher):
     MainCoroutineDispatcher(),
-    Delay by (delegate as? Delay ?: defaultDelay)
+    Delay
 {
     private val mainDispatcher = delegate // the initial value passed to the constructor
 
+    private val delay
+        get() = delegate as? Delay ?: defaultDelay
+
     override val immediate: MainCoroutineDispatcher
         get() = (delegate as? MainCoroutineDispatcher)?.immediate ?: this
 
@@ -28,6 +31,12 @@ internal class TestMainDispatcher(var delegate: CoroutineDispatcher):
     fun resetDispatcher() {
         delegate = mainDispatcher
     }
+
+    override fun scheduleResumeAfterDelay(timeMillis: Long, continuation: CancellableContinuation<Unit>) =
+        delay.scheduleResumeAfterDelay(timeMillis, continuation)
+
+    override fun invokeOnTimeout(timeMillis: Long, block: Runnable, context: CoroutineContext): DisposableHandle =
+        delay.invokeOnTimeout(timeMillis, block, context)
 }
 
 @Suppress("INVISIBLE_MEMBER")

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

@@ -67,5 +67,8 @@ open class OrderedExecutionTestBase {
 
 internal fun <T> T.void() { }
 
+@OptionalExpectation
+expect annotation class NoJs()
+
 @OptionalExpectation
 expect annotation class NoNative()

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

@@ -84,7 +84,7 @@ class RunTestTest {
 
     /** Tests that too low of a dispatch timeout causes crashes. */
     @Test
-    @Ignore // TODO: timeout leads to `Cannot execute task because event loop was shut down` on Native
+    @NoNative // TODO: timeout leads to `Cannot execute task because event loop was shut down` on Native
     fun testRunTestWithSmallTimeout() = testResultMap({ fn ->
         assertFailsWith<UncompletedCoroutinesError> { fn() }
     }) {
@@ -107,7 +107,7 @@ class RunTestTest {
 
     /** Tests uncaught exceptions taking priority over dispatch timeout in error reports. */
     @Test
-    @Ignore // TODO: timeout leads to `Cannot execute task because event loop was shut down` on Native
+    @NoNative // TODO: timeout leads to `Cannot execute task because event loop was shut down` on Native
     fun testRunTestTimingOutAndThrowing() = testResultMap({ fn ->
         assertFailsWith<IllegalArgumentException> { fn() }
     }) {
@@ -174,12 +174,12 @@ class RunTestTest {
 
     /** Tests that, once the test body has thrown, the child coroutines are cancelled. */
     @Test
-    fun testChildrenCancellationOnTestBodyFailure() {
+    fun testChildrenCancellationOnTestBodyFailure(): TestResult {
         var job: Job? = null
-        testResultMap({
+        return testResultMap({
             assertFailsWith<AssertionError> { it() }
             assertTrue(job!!.isCancelled)
-        }, {
+        }) {
             runTest {
                 job = launch {
                     while (true) {
@@ -188,34 +188,34 @@ class RunTestTest {
                 }
                 throw AssertionError()
             }
-        })
+        }
     }
 
     /** Tests that [runTest] reports [TimeoutCancellationException]. */
     @Test
     fun testTimeout() = testResultMap({
         assertFailsWith<TimeoutCancellationException> { it() }
-    }, {
+    }) {
         runTest {
             withTimeout(50) {
                 launch {
                     delay(1000)
                 }
             }
         }
-    })
+    }
 
     /** Checks that [runTest] throws the root cause and not [JobCancellationException] when a child coroutine throws. */
     @Test
     fun testRunTestThrowsRootCause() = testResultMap({
         assertFailsWith<TestException> { it() }
-    }, {
+    }) {
         runTest {
             launch {
                 throw TestException()
             }
         }
-    })
+    }
 
     /** Tests that [runTest] completes its job. */
     @Test
@@ -224,13 +224,13 @@ class RunTestTest {
         return testResultMap({
             it()
             assertTrue(handlerCalled)
-        }, {
+        }) {
             runTest {
                 coroutineContext.job.invokeOnCompletion {
                     handlerCalled = true
                 }
             }
-        })
+        }
     }
 
     /** Tests that [runTest] doesn't complete the job that was passed to it as an argument. */
@@ -245,11 +245,11 @@ class RunTestTest {
             it()
             assertFalse(handlerCalled)
             assertEquals(0, job.children.filter { it.isActive }.count())
-        }, {
+        }) {
             runTest(job) {
                 assertTrue(coroutineContext.job in job.children)
             }
-        })
+        }
     }
 
     /** Tests that, when the test body fails, the reported exceptions are suppressed. */
@@ -267,14 +267,14 @@ class RunTestTest {
             assertEquals("y", suppressed[1].message)
             assertEquals("z", suppressed[2].message)
         }
-    }, {
+    }) {
         runTest {
             launch(SupervisorJob()) { throw TestException("x") }
             launch(SupervisorJob()) { throw TestException("y") }
             launch(SupervisorJob()) { throw TestException("z") }
             throw TestException("w")
         }
-    })
+    }
 
     /** Tests that [TestCoroutineScope.runTest] does not inherit the exception handler and works. */
     @Test
@@ -287,10 +287,10 @@ class RunTestTest {
             } catch (e: TestException) {
                 scope.cleanupTestCoroutines() // should not fail
             }
-        }, {
+        }) {
             scope.runTest {
                 launch(SupervisorJob()) { throw TestException("x") }
             }
-        })
+        }
     }
 }

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

@@ -54,4 +54,17 @@ class StandardTestDispatcherTest: OrderedExecutionTestBase() {
         expect(5)
     }.void()
 
+    /** Tests that the [TestCoroutineScheduler] used for [Dispatchers.Main] gets used by default. */
+    @Test
+    fun testSchedulerReuse() {
+        val dispatcher1 = StandardTestDispatcher()
+        Dispatchers.setMain(dispatcher1)
+        try {
+            val dispatcher2 = StandardTestDispatcher()
+            assertSame(dispatcher1.scheduler, dispatcher2.scheduler)
+        } finally {
+            Dispatchers.resetMain()
+        }
+    }
+
 }