Integrate TestDispatcher with the unconfined event loop

[zach-klippenstein]

Description

It seems that when using runTest with UnconfinedTestDispatcher, when code inside the test wraps the dispatcher with a ContinuationInterceptor that delegates to the dispatcher, advanceUntilIdle will not wait for or cause to execute any continuations. However, yield calls will cause scheduled continuations to resume. This seems like a bug to me since the test dispatcher is still being used, and the continuations are still being scheduled, it's just that advanceUntilIdle doesn't "know" about them.

Kotlin version: 1.7.20
Coroutines version: 1.6.4

Repro

This test code will fail:

@Test
fun minimalRepro() {
    runTest(UnconfinedTestDispatcher()) {
        val parentInterceptor = coroutineContext[ContinuationInterceptor]!!
        val wrappedInterceptor =
            object : AbstractCoroutineContextElement(ContinuationInterceptor),
                ContinuationInterceptor {
                override fun <T> interceptContinuation(continuation: Continuation<T>) =
                    // Return continuation directly and the test passes
                    // (DispatchedContinuation vs Continuation).
                    parentInterceptor.interceptContinuation(continuation)
            }
        var entered = false

        withContext(wrappedInterceptor) {
            launch {
                entered = true
            }

            // OR replace this with yield() and the test passes.
            // This function returns immediately before the launched coroutine gets
            // dispatched.
            advanceUntilIdle()

            // This will fail, entered is still false.
            assertTrue(entered)
        }
    }
}

Investigation

I tried debugging into the launch and continuation code to figure out where the disconnect is happening, but I couldn't pinpoint the issue. It looks like UnconfinedTestDispatcher doesn't actually do any dispatching, always returns true from isDispatchNeeded, and thus relies on the EventLoop installed by runTest to actually dispatch. advanceUntilIdle only looks at the event list in the scheduler. It seems clear enough that they're looking at completely different data structures, but I can't figure out how they should be communicating.

Use case

The reason I'm running into this is because I'm writing a test for a library that wraps an UnconfinedTestDispatcher in a ContinuationInterceptor to defer dispatching continuations in specific circumstances. Moving away from UnconfinedTestDispatcher to a more standard, non-immediate dispatcher would be ideal, but is unfeasible at this time given that we have thousands of tests written against this unconfined dispatcher, and

[zach-klippenstein]

I think this might be related to issue #3253, since this bug doesn't happen if the wrapping interceptor returns the incoming Continuation (not a DispatchedContinuation) directly instead of delegating to the dispatcher's interceptContinuation (returns a DispatchedContinuation).

[dkhalanskyjb]

Wow. Thanks, I hate it! This is the kind of bug that makes you instantly lose all of your strength.

A simple reproducer without any fancy stuff happening:

    @Test
    fun testASoulCrushingRealization() = runTest(UnconfinedTestDispatcher()) {
        launch { // top-level launches are entered eagerly
            var entered = false
            launch {
                entered = true
            }
            advanceUntilIdle()
            assertTrue(entered)
        }
    }