withContext should copy CopyableThreadContextElement

[dovchinnikov]

val measurer = if (measure) MyCoroutineMeasurer() else null // CopyableThreadContextElement
runBlocking(Dispatchers.Default + CoroutineName("root") + measurer ?: EmptyCoroutineContext) {
  
  launch(CoroutineName("task 1") { ... } // ok
  
  async(CoroutineName("task 2") { ... } // ok
  
  // not copied but #updateThreadContext is still invoked
  withContext(CoroutineName("task 3")) { ... } 
}
measurer?.report()

I know that withContext {} is not equivalent to async {}.await(), but it does create a child coroutine (although UndispatchedCoroutine, but it still a separate instance) and introduces a CoroutineScope (it can also have children).

async(CoroutineName("task 3"), start = UNDISPATCHED) {}.await() kind of does what I want, but I also want the production code to look idiomatic since the measurer might or might not be present in the context.

[qwwdfsad]

The rationale behind the current behaviour is the following:

• We do not consider withContext semantically as creating a new coroutine. We do consider it a natural continuation (sorry for the terminology clash here) of the current computation and, e.g., do not render it in the debugger as well. It also does not change the control flow, similarly to how coroutineScope or supervisorScope don't. Yet it creates its own child scope, and it may be an additional source of confusion
• It has serious performance implications: the righthand side of withContext is always small -- a few elements at most while the lefthand side is arbitrary big and typically contains dozen of elements. Linearly scanning them for a copyable element will eventually become a performance killer for all users, including ones that do not use this API.

Also, if we are going to change the behaviour of withContext, we will have to change it in *Scope {} builders as well, as they can be used almost interchangeably and it may impose additional challenges for CopyableThreadContextElement users.

It seems like the best solution is to provide your own clearly-named alias to async{}.await(), e.g. stage or something like that (depending on how widespread the API is) and use it instead.

I would also like to hear feedback regarding this API and this particular behaviour from @yorickhenning if he is ready to elaborate on that matter

[yorickhenning]

I think I share @qwwdfsad's understanding of what a coroutine "is".

I'd expect three copied MyCoroutineMeasurers to each measure one of:

1. The runBlocking {} block
2. The launch {} block
3. The async {} block

Those being the three pseudo-lightweight-thread blocks that execute parallelisable statement lists.

I double checked a prototype and added these:

    @Test
    fun testWithContextCopies() = runTest {
        assertNull(coroutineContext[IntegerContextElement.Key])

        launch(AddElementInterceptor()) {
            withContext(Dispatchers.IO) {
                assertEquals(
                    1,
                    coroutineContext[IntegerContextElement.Key]!!.number
                )
            }
        }
    }

    @Test
    fun testCoroutineScopeCopies() = runTest {
        assertNull(coroutineContext[IntegerContextElement.Key])

        launch(AddElementInterceptor()) {
            coroutineScope {
                assertEquals(
                    1,
                    coroutineContext[IntegerContextElement.Key]!!.number
                )
            }
        }
    }

They fail with NPE at []!!. So I don't think I ever considered propagation at withContext {} and coroutineScope {} to be necessary.

IIRC the reason ContinuationInterceptor didn't work for trace propagation was because runBlocking {} and UNDISPATCHED launches construct and resume coroutines in a thread without calling outside of the Coroutines library's globals.

Given how things are, if measuring an undispatched launch were interesting, I'd add a span. I think that's what @qwwdfsad suggested as stage. stage or span, the function looks like:

suspend fun businessLogic() = coroutineScope {
  span("task 1") { // starts timer
    launch (coroutineStart = UNDISPATCHED) {
      suspendyBusinessLogicOptimizedToStartWithoutYielding()
    }
  } // stops timer and records duration
  
  moreBusinessLogic()
}

We do not consider withContext semantically as creating a new coroutine. We do consider it a natural continuation (sorry for the terminology clash here) of the current computation and, e.g., do not render it in the debugger as well. It also does not change the control flow, similarly to how coroutineScope or supervisorScope don't. Yet it creates its own child scope, and it may be an additional source of confusion

Intended semantics aside, I think calling withContext {} calls AbstractCoroutine's constructor via a super call from one of various subclasses? That could be understood to be "creating a coroutine".

An API that intercepts all AbstractCoroutine constructor calls and from there every resumeWith {} would probably enable a (previously-discussed?) superset of use cases for CopyableThreadContextElement. I expect it'd be rather simpler to implement against. I don't know for sure, but I suspect such an API would work for moving mutable ThreadLocals around and so work for thread safe trace propagation.

At that level though, I think withContext {} looks exactly like launch {} or async {}. So systems that'd like to treat those differently and automatically so would find that awkward. From the sound of it, that'd include the Coroutines debugger.

• It has serious performance implications: the righthand side of withContext is always small -- a few elements at most while the lefthand side is arbitrary big and typically contains dozen of elements. Linearly scanning them for a copyable element will eventually become a performance killer for all users, including ones that do not use this API.

In the systems I work on, both sides are consistently about four elements long. They both have a constant upper bound on length. I don't know whether those systems are typical.

I believe we're paying a linear search cost in every context in our production coroutine systems by including a CopyableThreadContextElement. The overhead seems acceptable, so far. So, I don't think small n O(n) there is a performance killer for all users. Some, possibly.

Benchmarks would certainly show a difference.

[dovchinnikov]

I'd be happy with stage or span or whatever name it might have if it will be a part of the coroutines library. This way it can be used to write code which is not aware of any kind of telemetry, in other words the code will be stdlib-idiomatic.
It can be optimised to work exactly like withContext(CoroutineName("xx")) {} in the best-case scenario.