Quarkus - Vert.x Context Propagation on Kotlin Coroutines

[pkozakCC]

I have some quarkus-based services written in Kotlin with coroutine support, in which I need to propagate some custom context data.

For that reason, I've been using vert.x duplicated context, however it turns out that it's not being properly propagated on coroutines (code sample below).

I've also tried to enhance it using implementation of org.eclipse.microprofile.context.spi.ThreadContextProvider, that would pass the context value between threads, but it doesn't affect vert.x duplicated context either (code sample below)

I could of course use InheritableThreadLocal.asContextElement(), and sync that with vert.x context, however this would require manually wrapping all my service entrypoints with withContext block, like so:

withContext(CONTEXT_HOLDER.asContextElement(coroutineName)) {
   // ...
}

My question is - what is the correct approach here ? Is there a way to merry vert.x duplicated context with kotlin coroutines, without the use of custom coroutine scope ?

Vert.x based context access:

    var contextVal: String?
        get() {
            val local = vertxContext?.getLocal<String?>("CONTEXT_VAL_KEY")
            return local
        }
        set(newContextVal) {
            if (newContextVal != null)
                vertxContext?.putLocal("CONTEXT_VAL_KEY", newContextVal)
            else
                vertxContext?.removeLocal("CONTEXT_VAL_KEY")
        }

    private val vertxContext: Context?
        get() {
            Vertx.currentContext()?.let {
                val dc = VertxContext.getOrCreateDuplicatedContext(it)
                VertxContextSafetyToggle.setContextSafe(dc, true)
                return dc
            }
            return null
        }

ThreadContextProvider implementation:

class ContextValThreadContextProvider : ThreadContextProvider {


    override fun currentContext(props: Map<String, String>): ThreadContextSnapshot {
        val parentCtxVal = contextVal

        return ThreadContextSnapshot {
            val previousCtxVal = contextVal
            contextVal = parentCtxVal

            ThreadContextController {
                contextVal = previousCtxVal
            }
        }
    }

    override fun clearedContext(props: Map<String, String>): ThreadContextSnapshot {
        return EMPTY_CONTEXT_SNAPSHOT
    }

    override fun getThreadContextType(): String {
        return javaClass.name
    }

    companion object {
        val EMPTY_CONTEXT_SNAPSHOT = ThreadContextSnapshot { ThreadContextController {} }

    }

[quarkus-bot[bot]]

/cc @FroMage (context-propagation), @geoand (kotlin), @manovotn (context-propagation)

[mschorsch]

I am not sure if I have understood correctly what you want to achieve. Why do you need the Vert.x context, why not use the Coroutine Context element? Where is the Vertx Context not propagated correctly?

Do you have a small reproducer?

[pkozakCC]

First of all - thank you for replying.

It doesn't necessarily have to be vert.x context, I just need to propagate context information so that it's available everywhere. More specifically, this information is tenant ID, which will determine the DB that will be used.

I could of course use custom coroutine scope, like so:

        public suspend fun <R> contextCoroutineScope(block: suspend CoroutineScope.() -> R): R {
            return withContext(
                ContextHolder.asContextElement()
            ) {
                block.invoke(this)
            }
        }
        
        class ContextHolder {
            companion object {
                private val CONTEXT_HOLDER: InheritableThreadLocal<Context> = InheritableThreadLocal()
                
                var context: Context?
                    get() = CONTEXT_HOLDER.get()
                    set(ctxVal) = CONTEXT_HOLDER.set(ctxVal)

                fun asContextElement(): ThreadContextElement<Context> {
                    return CONTEXT_HOLDER.asContextElement()
                }
            }
        }
    

and access it using ContextHolder.context, however this will required wrapping all of my code with contextCouroutineScope:

    @GET
    @Path("/test")
    suspend fun test(@QueryParam("jobs") jobs: Int) = contextCoroutineScope {
        // ...
    }

I want to avoid using contextCoroutineScope everywhere, as it's prone to forgetting.
There's no way of wrapping this in an filter, for example like it's done in spring.

Is there any other, better way ?

[mschorsch]

Thank you for the clarification.

The best I can think of at the moment is to use a separate CoroutineContext.Element. The CoroutineContext.Element would then have to be set once for each REST endpoint.

Simple example:

import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.withContext
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.coroutineContext

class TenantIdContextElement(val tenantId: String) : CoroutineContext.Element {
    companion object Key : CoroutineContext.Key<TenantIdContextElement>

    override val key: CoroutineContext.Key<*> get() = Key
}

suspend fun tenantId(): String {
    val element = requireNotNull(coroutineContext[TenantIdContextElement.Key]) { "'TenantIdContextElement' not found" }
    return element.tenantId
}

// Example usage

fun main() = runBlocking {
    withContext(TenantIdContextElement("tenantId 1")) {
        doSomething1()
    }
}

suspend fun doSomething1() {
    println(tenantId())
    doSomething2()
}

suspend fun doSomething2() {
    println(tenantId())
    doSomething3()
}

suspend fun doSomething3() {
    withContext(Dispatchers.IO) {
        println(tenantId())

        withContext(TenantIdContextElement("tenantId 2")) {
            doSomething4()
        }
    }
}

suspend fun doSomething4() {
    println(tenantId())
}

REST endpoint example:

@GET
@Path("/test")
suspend fun test(@QueryParam("jobs") jobs: Int) = withContext(TenantIdContextElement("tenantId 1")) {
  doSomething1()
}

private suspend fun doSomething1() {
    println(tenantId())
}