Missing updateThreadContext call when emit resumes

[cdli-corp]

Hi Kotlin team,

IIUC, each coroutine should have a dedicated CopyableThreadContextElement, and use this element to call updateThreadContext and restoreThreadContext at the beginning and end of each sub-routine, respectively. However, we found that is not always the case, and here is an example: https://pl.kotl.in/USpmRjCNL.

The following details my observation of the code's execution:

val flow = 
    flow {
        // 1. Here updateThreadContext is called on 
        //    TestElement("flowOn")'s child (the distinction is "flowOn" v.s. "collect", where child is immaterial, 
        //    so will skip mentioning child in the following explanation)
        log("before emit")
        // 2. We call the block in collect below, and suspend
        emit(1)
        // 4. As we reach here, TestElement("flowOn").updateThreadContext is not called
        //    although the current context is (has, but use "is" for the simplicity of explanation) TestElement("flowOn")
        log("after emit")
    }
    .flowOn(TestElement("flowOn", null))

launch(TestElement("collect", null)) {
    flow.collect {
        reached.set(true)
        log("collect about to lock")
        mutex.lock()
        // 3. When resuming, updateThreadContext is called on TestElement("collect")
        log("collect resumed")
    }
}

Here is the highlighted output to show so:

[13] updateThreadContext: TestElement(name=child[0], parent=TestElement(name=collect, parent=null))
[13] collect resumed
// <- no updateThreadContext for TestElement("flowOn")
[13] after emit
[13] restoreThreadContext: TestElement(name=child[0], parent=TestElement(name=collect, parent=null))

For context we have a ThreadLocal-based tracing system, which supports coroutine tracing with the help of CopyableThreadContextElement, and we have clients broken by this because the correct trace is not restored (for the missing updateThreadContext call). So post the finding here and see if we can get this fixed, thanks!

[Esteth]

I think something like what I describe below is happening, but I'm not overly familiar with the mechanics of coroutines at this level, so please forgive me if I have a misunderstanding.

In the flowOn/ChannelFlow fast path we avoid dispatching to the collector when the emitter emits if the collector and emitter use the same dispatcher. the fast path knows that the dispatcher would have still called the ThreadContextElement hooks in this case, and so the fast path does so here.

If we suspend inside the collector (say when we try to lock the mutex in the repro) then we rightfully return out of the withContextUndispatched fast path with COROUTINE_SUSPENDED and tear down the current stack of ThreadContextElement update/restore calls. So far, so good.

Now at some later time, the continuation for the collector becomes wakeable (in our case, the other thread unlocks the mutex) and the dispatcher is given the continuation to resume. It calls the update method on the ThreadContextElement in that continuation, then the second half of the collector runs. We're still fine.

This is where my understanding gets a little shaky...

When the collector block finishes, we're still using the "fast path" ChannelFlow, so the stack for the coroutine has underneath it the emitter. We pop the current stack frame and head back to after the emit call of the emitter. Because there was no dispatcher involved and we're no longer in the withContextUndispatched block we were last time, nothing knows to call the ThreadContextElement update/restore hooks, and we end up with the global state corruption we observe.

If my understanding is correct, then I'm not sure exactly what we can do. Could we save the "stack" of ThreadContextElements in the context before we pop it somehow like this?

1. Check the return value of startCoroutineUninterceptedOrReturn here.
2. If it's COROUTINE_SUSPENDED then attach the stack of "parent" updateThreadContext contexts to the current ThreadContextElement
3. When we would resume a coroutine with ThreadContextElements, install all of the ThreadContextElements in the stored stack.

I'm not confident that's even possible, but I'm far from an expert here.

[Esteth]

I've cleaned up cdli's repro into an actual pass/fail test case for FlowOnTest in d1c4296