Kotlin coroutine context inherited by 'async' block

Question

Please consider the following scenarios in which I'm using a ThreadLocal as a Kotlin coroutine context element.

Scenario 1 - use withContext() to set the value of a ThreadLocal on an inner block:

val threadLocal = ThreadLocal<String?>()

fun main() = runBlocking(Dispatchers.IO + threadLocal.asContextElement("main")) {
    println("#1: ${threadLocal.get()}")

    withContext(coroutineContext + threadLocal.asContextElement("inner")) {
        println("#2: ${threadLocal.get()}")

        async {
            println("#3: ${threadLocal.get()}")
        }.await()
        println("#4: ${threadLocal.get()}")
    }
    println("#5: ${threadLocal.get()}")
}

This produces the following output:

#1: main
#2: inner
#3: inner
#4: inner
#5: main

Using the ThreadLocal as a coroutine context element, the value is first set to "main", then changed to "inner" with an inner context block. In particular note that #3 gets the value "inner" as well. In my understanding, this is how coroutine context inheritance is supposed to work. A child coroutine inherits the context from its parent.

Now consider this minor variation on the same setup, in which the call to withContext() is made from another function:

Scenario 2 - withContext() called from a function

fun main() = runBlocking(Dispatchers.IO + threadLocal.asContextElement("main")) {
    println("#1: ${threadLocal.get()}")

    withThreadLocalValue("inner") {          // <-- my function wrapping call to 'withContext()'
        println("#2: ${threadLocal.get()}")

        async {
            println("#3: ${threadLocal.get()}")
        }.await()
        println("#4: ${threadLocal.get()}")
    }
    println("#5: ${threadLocal.get()}")
}

suspend inline fun <T> withThreadLocalValue(value: String, crossinline block: suspend () -> T):T =
    withContext(coroutineContext + threadLocal.asContextElement(value)) {
        block()
    }

In this case, instead of using withContext() directly to set the ThreadLocal "inner" context, the same call to withContext() is made from a separate (inline) function, that then invokes the enclosing block.

This produces the following output:

#1: main
#2: inner
#3: main
#4: inner
#5: main

Note that, as expected, both #2 and #4 get the "inner" ThreadLocal value. #3 though, inside the async block, does not inherit that "inner" value on its context, but instead has the "main" value from the surrounding context.

So my questions are:

1. Is this the expected behavior? I was not expecting any functional difference between scenarios 1 & 2.
2. If this is expected, is there a correct way to get the behavior of scenario 1, while still making the call to withContext() from another function?
3. If this is not expected, should this be considered a bug in the coroutine framework?

Answer 1

You made a small mistake in your implementation of withThreadLocalValue. Please note the original withContext function accepts a block: suspend CoroutineScope.() -> T. Our lambda receives a new coroutine scope - scope of the inner coroutine, so async is related to the inner coroutine. withThreadLocalValue doesn't provide the inner coroutine scope to the lambda, so the lambda still uses the outer scope - async creates a coroutine in runBlocking, not in the inner coroutine.

To fix the problem, we just need to provide the inner coroutine scope to the lambda:

suspend inline fun <T> withThreadLocalValue(value: String, crossinline block: suspend CoroutineScope.() -> T):T =

Then it prints inner for #3.