When using withContext in Kotlin, how to keep the new context after calling another suspend function?

For example: https://pl.kotl.in/fkrnEkYZP

fun main() = runBlocking {
    val anotherBlock = suspend {
        assertSame(Dispatchers.IO, coroutineContext[CoroutineDispatcher], "current CoroutineDispatcher should be still Dispatchers.IO")
        // Exception in thread "main" java.lang.AssertionError: current CoroutineDispatcher should be still Dispatchers.IO. Expected <Dispatchers.IO>, actual <BlockingEventLoop@6863cde3> is not same.
    }
    withContext(Dispatchers.IO) {
        assertSame(Dispatchers.IO, coroutineContext[CoroutineDispatcher], "current CoroutineDispatcher should be Dispatchers.IO")
        anotherBlock()
    }
}

Why is the coroutineContext in anotherBlock different with the parent one?

CodePudding user response：

coroutineContext is a property of CoroutineScope. Since your anotherBlock suspend function does not have a CoroutineScope receiver, you are calling coroutineContext on the scope of the runBlocking lambda, so it's getting the runBlocking's dispatcher.

Change your declaration of the function to use extension function syntax instead of a lambda so it can have its own CoroutineScope.

suspend fun CoroutineScope.anotherBlock() {
    assertSame(Dispatchers.IO, coroutineContext[CoroutineDispatcher], "current CoroutineDispatcher should be still Dispatchers.IO")
    // Exception in thread "main" java.lang.AssertionError: current CoroutineDispatcher should be still Dispatchers.IO. Expected <Dispatchers.IO>, actual <BlockingEventLoop@6863cde3> is not same.
}

CodePudding user response：

The problem is that you're not looking at the actual current coroutine context when doing the assert in the suspend lambda, you're in fact capturing the coroutineContext from runBlocking's CoroutineScope.

The reason why it works in withContext directly is because withContext provides a CoroutineScope as this - which shadows the one given by runBlocking, so luckily you get the correct one.

To avoid this pitfall (which occurs regularly when using both suspend functions and CoroutineScope), you should use currentCoroutineContext() when you want to know the coroutine context in which the current code is executing, which is your case here.

If you do this, you'll see the dispatcher is properly inherited:

fun main() = runBlocking {
    val anotherBlock = suspend {
        assertSame(Dispatchers.IO, currentCoroutineContext()[CoroutineDispatcher], "current CoroutineDispatcher should be still Dispatchers.IO in suspend lambda")
    }
    withContext(Dispatchers.IO) {
        assertSame(Dispatchers.IO, coroutineContext[CoroutineDispatcher], "current CoroutineDispatcher should be Dispatchers.IO")
        anotherBlock()
    }
}