Struggling with using MDC inside of flows

[jadlr]

Consider this minimal example:

To run this you need

implementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.7.3")
implementation("org.jetbrains.kotlinx:kotlinx-coroutines-slf4j:1.7.3")
implementation("org.slf4j:slf4j-api:2.0.7")
implementation("ch.qos.logback:logback-classic:1.4.11")

import kotlinx.coroutines.*
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.flow.launchIn
import kotlinx.coroutines.flow.transform
import kotlinx.coroutines.slf4j.MDCContext
import org.slf4j.LoggerFactory
import org.slf4j.MDC

object MDCExample {

    val log = LoggerFactory.getLogger(javaClass)
    suspend fun <T> withMDC(vararg context: Pair<String, String>, block: suspend CoroutineScope.() -> T) =
        withContext(MDCContext(contextMap = (MDC.getCopyOfContextMap() ?: mutableMapOf()) + context), block)

    @JvmStatic
    fun main(args: Array<String>) {

        val flow  = flow {
            var count = 0
            while (currentCoroutineContext().isActive) {
                emit(count++)
                delay(100)
            }
        }.transform {
            withMDC("number" to "$it") {
                if (it % 2 == 0) {
                    emit(it)
                } else {
                    log.warn("detected uneven number")
                }
            }
        }

        runBlocking {
            flow.launchIn(CoroutineScope(Dispatchers.IO)).join()
        }
    }
}

This gives me

Exception in thread "DefaultDispatcher-worker-1" java.lang.IllegalStateException: Flow invariant is violated:
		Flow was collected in [StandaloneCoroutine{Active}@6919444e, Dispatchers.IO],
		but emission happened in [kotlinx.coroutines.slf4j.MDCContext@d433365, kotlinx.coroutines.UndispatchedMarker@56f3a43a, UndispatchedCoroutine{Active}@4c3eefbf, Dispatchers.IO].
		Please refer to 'flow' documentation or use 'flowOn' instead

Is there a way to fix this? I'm working on a complex pipeline with multiple steps and I'd like to use slf4j MDC to add meta to the logs.

I'm struggling to understand what the issue is, since when you exchange the transform with map the issue disappears.

[jadlr]

This feels somehow related:
#1421

[jadlr]

This looks like a bug 🤔 Whenever I change my implementation to map / filter the problems vanish.

[qwwdfsad]

The exception already contains an answer: withContext breaks the context preservation property (you are emitting from a context different that the upstream is) and the solution (also from the exception message) is to use .flowOn(your mdc context).

map and filter are not prone to this error because they are not emitting from the new context, but only transforming an element.
They are roughly equivalent to the following:

transform {
    val elementOrNull = withMDC("number" to "$it") {
        if (it % 2 == 0) {
            it
        } else {
            log.warn("detected uneven number")
            null
        }
    }
    if (elementOrNull !== null) emit(elementOrNull)
}

[jadlr]

Thanks! This makes sense.