Configure default Kotlin coroutine context in Spring MVC

Question

I need to configure default coroutine context for all requests in Spring MVC. For example MDCContext (similar question as this but for MVC not WebFlux).

What I have tried

1. Hook into Spring - the coroutine code is here but there is no way to change the default behavior (need to change InvocableHandlerMethod.doInvoke implementation)
2. Use AOP - AOP and coroutines do not play well together

Any ideas?

Answer 1

This seems to work:

@Configuration
class ContextConfig: WebMvcRegistrations {

    override fun getRequestMappingHandlerAdapter(): RequestMappingHandlerAdapter {
        return object: RequestMappingHandlerAdapter() {
            override fun createInvocableHandlerMethod(handlerMethod: HandlerMethod): ServletInvocableHandlerMethod {
                return object : ServletInvocableHandlerMethod(handlerMethod) {
                    override fun doInvoke(vararg args: Any?): Any? {
                        val method = bridgedMethod
                        ReflectionUtils.makeAccessible(method)
                        if (KotlinDetector.isSuspendingFunction(method)) {
                            // Exception handling skipped for brevity, copy it from super.doInvoke()
                            return invokeSuspendingFunctionX(method, bean, *args)
                        }
                        return super.doInvoke(*args)
                    }

                    /**
                     * Copied from CoroutinesUtils in order to be able to set CoroutineContext
                     */
                    @Suppress("UNCHECKED_CAST")
                    private fun invokeSuspendingFunctionX(method: Method, target: Any, vararg args: Any?): Publisher<*> {
                        val function = method.kotlinFunction!!
                        val mono = mono(YOUR_CONTEXT_HERE) {
                            function.callSuspend(target, *args.sliceArray(0..(args.size-2))).let { if (it == Unit) null else it }
                        }.onErrorMap(InvocationTargetException::class.java) { it.targetException }
                        return if (function.returnType.classifier == Flow::class) {
                            mono.flatMapMany { (it as Flow<Any>).asFlux() }
                        }
                        else {
                            mono
                        }
                    }
                }
            }
        }
    }
}