Don't use Dispatchers.DEFAULT for the test runner

Author: dkhalanskyjb

kotlinx-coroutines-test/build.gradle.kts

@@ -1,9 +1,9 @@
-import org.jetbrains.kotlin.gradle.plugin.mpp.*
-
 /*
  * Copyright 2016-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
 
+import org.jetbrains.kotlin.gradle.plugin.mpp.*
+
 val experimentalAnnotations = listOf(
     "kotlin.Experimental",
     "kotlinx.coroutines.ExperimentalCoroutinesApi",
@@ -19,4 +19,12 @@ kotlin {
             binaryOptions["memoryModel"] = "experimental"
         }
     }
+
+    sourceSets {
+        jvmTest {
+            dependencies {
+                implementation(project(":kotlinx-coroutines-debug"))
+            }
+        }
+    }
 }

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

@@ -303,67 +303,70 @@ public fun runTest(
 public fun TestScope.runTest(
     timeout: Duration = DEFAULT_TIMEOUT,
     testBody: suspend TestScope.() -> Unit
-): TestResult = asSpecificImplementation().let {
-    it.enter()
+): TestResult = asSpecificImplementation().let { scope ->
+    scope.enter()
     createTestResult {
         /** TODO: moving this [AbstractCoroutine.start] call outside [createTestResult] fails on JS. */
-        it.start(CoroutineStart.UNDISPATCHED, it) {
+        scope.start(CoroutineStart.UNDISPATCHED, scope) {
+            yield()
             testBody()
         }
-        /**
-         * We run the tasks in the test coroutine using [Dispatchers.Default]. On JS, this does nothing particularly,
-         * but on the JVM and Native, this means that the timeout can be processed even while the test runner is busy
-         * doing some synchronous work.
-         */
-        val workRunner = launch(Dispatchers.Default + CoroutineName("kotlinx.coroutines.test runner")) {
-            while (true) {
-                val executedSomething = testScheduler.tryRunNextTaskUnless { !isActive }
-                if (executedSomething) {
-                    /** yield to check for cancellation. On JS, we can't use [ensureActive] here, as the cancellation
-                     * procedure needs a chance to run concurrently. */
-                    yield()
-                } else {
-                    // no more tasks, we should suspend until there are some more
-                    testScheduler.receiveDispatchEvent()
-                }
-            }
-        }
         var timeoutError: Throwable? = null
+        var cancellationException: CancellationException? = null
         try {
             withTimeout(timeout) {
-                it.join()
-                workRunner.cancelAndJoin()
+                coroutineContext.job.invokeOnCompletion(onCancelling = true) { exception ->
+                    if (exception is TimeoutCancellationException) {
+                        dumpCoroutines()
+                        val activeChildren = scope.children.filter(Job::isActive).toList()
+                        val completionCause = if (scope.isCancelled) scope.tryGetCompletionCause() else null
+                        var message = "After waiting for $timeout"
+                        if (completionCause == null)
+                            message += ", the test coroutine is not completing"
+                        if (activeChildren.isNotEmpty())
+                            message += ", there were active child jobs: $activeChildren"
+                        if (completionCause != null && activeChildren.isEmpty()) {
+                            message += if (scope.isCompleted)
+                                ", the test coroutine completed"
+                            else
+                                ", the test coroutine was not completed"
+                        }
+                        timeoutError = UncompletedCoroutinesError(message)
+                        cancellationException = CancellationException("The test timed out")
+                        (scope as Job).cancel(cancellationException!!)
+                    }
+                }
+                while (scope.isActive) {
+                    val executedSomething = testScheduler.tryRunNextTaskUnless { !scope.isActive }
+                    if (executedSomething) {
+                        /** yield to check for cancellation. On JS, we can't use [ensureActive] here, as the cancellation
+                         * procedure needs a chance to run concurrently. */
+                        yield()
+                    } else {
+                        // no tasks, we should suspend until there are some more or the test coroutine is completed
+                        select {
+                            scope.onJoin {
+                                // do nothing, just exit the select and the `while` loop.
+                            }
+                            testScheduler.onDispatchEvent {
+                                // do nothing, just exit the select and run the next iteration
+                            }
+                        }
+                    }
+                }
             }
         } catch (_: TimeoutCancellationException) {
-            val activeChildren = it.children.filter(Job::isActive).toList()
-            val completionCause = if (it.isCancelled) it.tryGetCompletionCause() else null
-            var message = "After waiting for $timeout"
-            if (completionCause == null)
-                message += ", the test coroutine is not completing"
-            if (activeChildren.isNotEmpty())
-                message += ", there were active child jobs: $activeChildren"
-            if (completionCause != null && activeChildren.isEmpty()) {
-                message += if (it.isCompleted)
-                    ", the test coroutine completed"
-                else
-                    ", the test coroutine was not completed"
-            }
-            timeoutError = UncompletedCoroutinesError(message)
-            dumpCoroutines()
-            val cancellationException = CancellationException("The test timed out")
-            (it as Job).cancel(cancellationException)
             // we can't abandon the work we're doing, so if it hanged, we'll still hang, despite the timeout.
-            it.join()
-            val completion = it.getCompletionExceptionOrNull()
+            scope.join()
+            val completion = scope.getCompletionExceptionOrNull()
             if (completion != null && completion !== cancellationException) {
-                timeoutError.addSuppressed(completion)
+                timeoutError!!.addSuppressed(completion)
             }
-            workRunner.cancelAndJoin()
         } finally {
             backgroundScope.cancel()
             testScheduler.advanceUntilIdleOr { false }
-            val uncaughtExceptions = it.leave()
-            throwAll(timeoutError ?: it.getCompletionExceptionOrNull(), uncaughtExceptions)
+            val uncaughtExceptions = scope.leave()
+            throwAll(timeoutError ?: scope.getCompletionExceptionOrNull(), uncaughtExceptions)
         }
     }
 }

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

@@ -77,7 +77,7 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
             val time = addClamping(currentTime, timeDeltaMillis)
             val event = TestDispatchEvent(dispatcher, count, time, marker as Any, isForeground) { isCancelled(marker) }
             events.addLast(event)
-            /** can't be moved above: otherwise, [onDispatchEventForeground] or [receiveDispatchEvent] could consume the
+            /** can't be moved above: otherwise, [onDispatchEventForeground] or [onDispatchEvent] could consume the
              * token sent here before there's actually anything in the event queue. */
             sendDispatchEvent(context)
             DisposableHandle {
@@ -214,6 +214,11 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
      */
     internal suspend fun receiveDispatchEvent() = dispatchEvents.receive()
 
+    /**
+     * Consumes the knowledge that a dispatch event happened recently.
+     */
+    internal val onDispatchEvent: SelectClause1<Unit> get() = dispatchEvents.onReceive
+
     /**
      * Consumes the knowledge that a foreground work dispatch event happened recently.
      */

kotlinx-coroutines-test/jvm/test/DumpOnTimeoutTest.kt

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2016-2023 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.test
+
+import kotlinx.coroutines.*
+import kotlinx.coroutines.debug.*
+import org.junit.Test
+import java.io.*
+import kotlin.test.*
+import kotlin.time.Duration.Companion.milliseconds
+
+class DumpOnTimeoutTest {
+    /**
+     * Tests that the dump on timeout contains the correct stacktrace.
+     */
+    @Test
+    fun testDumpOnTimeout() {
+        val oldErr = System.err
+        val baos = ByteArrayOutputStream()
+        try {
+            System.setErr(PrintStream(baos, true))
+            DebugProbes.withDebugProbes {
+                try {
+                    runTest(timeout = 100.milliseconds) {
+                        uniquelyNamedFunction()
+                    }
+                    throw IllegalStateException("unreachable")
+                } catch (e: UncompletedCoroutinesError) {
+                    // do nothing
+                }
+            }
+            baos.toString().let {
+                assertTrue(it.contains("uniquelyNamedFunction"), "Actual trace:\n$it")
+            }
+        } finally {
+            System.setErr(oldErr)
+        }
+    }
+
+    fun CoroutineScope.uniquelyNamedFunction() {
+        while (true) {
+            ensureActive()
+            Thread.sleep(10)
+        }
+    }
+}