Add a scope for launching background work in tests

kotlinx-coroutines-core/common/src/internal/ThreadSafeHeap.kt

@@ -37,6 +37,16 @@ public open class ThreadSafeHeap<T> : SynchronizedObject() where T: ThreadSafeHe
         _size.value = 0
     }
 
+    public fun find(
+        predicate: (value: T) -> Boolean
+    ): T? = synchronized(this) {
+        for (i in 0 until size) {
+            val value = a?.get(i)!!
+            if (predicate(value)) return value
+        }
+        null
+    }
+
     public fun peek(): T? = synchronized(this) { firstImpl() }
 
     public fun removeFirstOrNull(): T? = synchronized(this) {

kotlinx-coroutines-test/api/kotlinx-coroutines-test.api

@@ -105,6 +105,7 @@ public final class kotlinx/coroutines/test/TestDispatchers {
 }
 
 public abstract interface class kotlinx/coroutines/test/TestScope : kotlinx/coroutines/CoroutineScope {
+	public abstract fun getBackgroundWorkScope ()Lkotlinx/coroutines/CoroutineScope;
 	public abstract fun getTestScheduler ()Lkotlinx/coroutines/test/TestCoroutineScheduler;
 }
 

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

@@ -164,7 +164,11 @@ public fun TestScope.runTest(
 ): TestResult = asSpecificImplementation().let {
     it.enter()
     createTestResult {
-        runTestCoroutine(it, dispatchTimeoutMs, TestScopeImpl::tryGetCompletionCause, testBody) { it.leave() }
+        runTestCoroutine(it, dispatchTimeoutMs, TestScopeImpl::tryGetCompletionCause, testBody) {
+            backgroundWorkScope.cancel()
+            testScheduler.advanceUntilIdle(backgroundIsIdle = false)
+            it.leave()
+        }
     }
 }
 

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

@@ -151,8 +151,7 @@ private class StandardTestDispatcherImpl(
 ) : TestDispatcher() {
 
     override fun dispatch(context: CoroutineContext, block: Runnable) {
-        checkSchedulerInContext(scheduler, context)
-        scheduler.registerEvent(this, 0, block) { false }
+        scheduler.registerEvent(this, 0, block, context) { false }
     }
 
     override fun toString(): String = "${name ?: "StandardTestDispatcher"}[scheduler=$scheduler]"

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

@@ -62,13 +62,16 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
         dispatcher: TestDispatcher,
         timeDeltaMillis: Long,
         marker: T,
+        context: CoroutineContext,
         isCancelled: (T) -> Boolean
     ): DisposableHandle {
         require(timeDeltaMillis >= 0) { "Attempted scheduling an event earlier in time (with the time delta $timeDeltaMillis)" }
+        checkSchedulerInContext(this, context)
         val count = count.getAndIncrement()
+        val isForeground = context[BackgroundWork] === null
         return synchronized(lock) {
             val time = addClamping(currentTime, timeDeltaMillis)
-            val event = TestDispatchEvent(dispatcher, count, time, marker as Any) { isCancelled(marker) }
+            val event = TestDispatchEvent(dispatcher, count, time, marker as Any, isForeground) { isCancelled(marker) }
             events.addLast(event)
             /** can't be moved above: otherwise, [onDispatchEvent] could consume the token sent here before there's
              * actually anything in the event queue. */
@@ -82,10 +85,12 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
     }
 
     /**
-     * Runs the next enqueued task, advancing the virtual time to the time of its scheduled awakening.
+     * Runs the next enqueued task, advancing the virtual time to the time of its scheduled awakening,
+     * unless all the remaining tasks are the background ones.
      */
-    private fun tryRunNextTask(): Boolean {
+    private fun tryRunNextTask(backgroundIsIdle: Boolean): Boolean {
         val event = synchronized(lock) {
+            if (backgroundIsIdle && events.none(TestDispatchEvent<*>::isForeground)) return false
             val event = events.removeFirstOrNull() ?: return false
             if (currentTime > event.time)
                 currentTimeAheadOfEvents()
@@ -105,9 +110,16 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
      * functionality, query [currentTime] before and after the execution to achieve the same result.
      */
     @ExperimentalCoroutinesApi
-    public fun advanceUntilIdle() {
-        while (!synchronized(lock) { events.isEmpty }) {
-            tryRunNextTask()
+    public fun advanceUntilIdle(): Unit = advanceUntilIdle(backgroundIsIdle = true)
+
+    /**
+     * [backgroundIsIdle]: `true` if the background tasks should not be considered
+     * when checking if the scheduler is already idle.
+     */
+    internal fun advanceUntilIdle(backgroundIsIdle: Boolean) {
+        while (true) {
+            if (!tryRunNextTask(backgroundIsIdle = backgroundIsIdle))
+                return
         }
     }
 
@@ -169,18 +181,10 @@ public class TestCoroutineScheduler : AbstractCoroutineContextElement(TestCorout
     /**
      * Checks that the only tasks remaining in the scheduler are cancelled.
      */
-    internal fun isIdle(strict: Boolean = true): Boolean {
+    internal fun isIdle(strict: Boolean = true): Boolean =
         synchronized(lock) {
-            if (strict)
-                return events.isEmpty
-            // TODO: also completely empties the queue, as there's no nondestructive way to iterate over [ThreadSafeHeap]
-            val presentEvents = mutableListOf<TestDispatchEvent<*>>()
-            while (true) {
-                presentEvents += events.removeFirstOrNull() ?: break
-            }
-            return presentEvents.all { it.isCancelled() }
+            if (strict) events.isEmpty else events.none { !it.isCancelled() }
         }
-    }
 
     /**
      * Notifies this scheduler about a dispatch event.
@@ -216,6 +220,8 @@ private class TestDispatchEvent<T>(
     private val count: Long,
     @JvmField val time: Long,
     @JvmField val marker: T,
+    @JvmField val isForeground: Boolean,
+    // TODO: remove once the deprecated API is gone
     @JvmField val isCancelled: () -> Boolean
 ) : Comparable<TestDispatchEvent<*>>, ThreadSafeHeapNode {
     override var heap: ThreadSafeHeap<*>? = null
@@ -238,3 +244,17 @@ internal fun checkSchedulerInContext(scheduler: TestCoroutineScheduler, context:
         }
     }
 }
+
+/**
+ * A coroutine context key denoting that the work is to be executed in the background.
+ * @see [TestScope.backgroundWorkScope]
+ */
+internal object BackgroundWork : CoroutineContext.Key<BackgroundWork>, CoroutineContext.Element {
+    override val key: CoroutineContext.Key<*>
+        get() = this
+
+    override fun toString(): String = "BackgroundWork"
+}
+
+private fun<T> ThreadSafeHeap<T>.none(predicate: (T) -> Boolean) where T: ThreadSafeHeapNode, T: Comparable<T> =
+    find(predicate) == null

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

@@ -10,14 +10,14 @@ import kotlin.jvm.*
 
 /**
  * A test dispatcher that can interface with a [TestCoroutineScheduler].
- * 
+ *
  * The available implementations are:
  * * [StandardTestDispatcher] is a dispatcher that places new tasks into a queue.
  * * [UnconfinedTestDispatcher] is a dispatcher that behaves like [Dispatchers.Unconfined] while allowing to control
  *   the virtual time.
  */
 @ExperimentalCoroutinesApi
-public abstract class TestDispatcher internal constructor(): CoroutineDispatcher(), Delay {
+public abstract class TestDispatcher internal constructor() : CoroutineDispatcher(), Delay {
     /** The scheduler that this dispatcher is linked to. */
     @ExperimentalCoroutinesApi
     public abstract val scheduler: TestCoroutineScheduler
@@ -30,16 +30,13 @@ public abstract class TestDispatcher internal constructor(): CoroutineDispatcher
 
     /** @suppress */
     override fun scheduleResumeAfterDelay(timeMillis: Long, continuation: CancellableContinuation<Unit>) {
-        checkSchedulerInContext(scheduler, continuation.context)
         val timedRunnable = CancellableContinuationRunnable(continuation, this)
-        scheduler.registerEvent(this, timeMillis, timedRunnable, ::cancellableRunnableIsCancelled)
+        scheduler.registerEvent(this, timeMillis, timedRunnable, continuation.context, ::cancellableRunnableIsCancelled)
     }
 
     /** @suppress */
-    override fun invokeOnTimeout(timeMillis: Long, block: Runnable, context: CoroutineContext): DisposableHandle {
-        checkSchedulerInContext(scheduler, context)
-        return scheduler.registerEvent(this, timeMillis, block) { false }
-    }
+    override fun invokeOnTimeout(timeMillis: Long, block: Runnable, context: CoroutineContext): DisposableHandle =
+        scheduler.registerEvent(this, timeMillis, block, context) { false }
 }
 
 /**

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

@@ -46,6 +46,23 @@ public sealed interface TestScope : CoroutineScope {
      */
     @ExperimentalCoroutinesApi
     public val testScheduler: TestCoroutineScheduler
+
+    /**
+     * A scope for background work.
+     *
+     * This scope is automatically cancelled when the test finishes.
+     * Additionally, while the coroutines in this scope are run as usual when
+     * using [advanceTimeBy] and [runCurrent], [advanceUntilIdle] will stop advancing the virtual time
+     * once only the coroutines in this scope are left unprocessed.
+     *
+     * Failures in coroutines in this scope do not terminate the test.
+     * Instead, they are reported at the end of the test.
+     *
+     * A typical use case for this scope is to launch tasks that would outlive the tested code in
+     * the production environment.
+     */
+    @ExperimentalCoroutinesApi
+    public val backgroundWorkScope: CoroutineScope
 }
 
 /**
@@ -170,6 +187,9 @@ internal class TestScopeImpl(context: CoroutineContext) :
     private val uncaughtExceptions = mutableListOf<Throwable>()
     private val lock = SynchronizedObject()
 
+    override val backgroundWorkScope: CoroutineScope =
+        CoroutineScope(coroutineContext + SupervisorJob() + BackgroundWork)
+
     /** Called upon entry to [runTest]. Will throw if called more than once. */
     fun enter() {
         val exceptions = synchronized(lock) {
@@ -233,6 +253,7 @@ internal class TestScopeImpl(context: CoroutineContext) :
 }
 
 /** Use the knowledge that any [TestScope] that we receive is necessarily a [TestScopeImpl]. */
+@Suppress("NO_ELSE_IN_WHEN") // TODO: a problem with `sealed` in MPP not allowing total pattern-matching
 internal fun TestScope.asSpecificImplementation(): TestScopeImpl = when (this) {
     is TestScopeImpl -> this
 }