Coroutines in Spring API : give me !

Coroutines are a sharply increasing practice among Android developers, but it is seems quite limited to the Android ecosystem, even though the original intent of them is to create an agnostic framework to implement asynchronous and cheap concurrency programs.

There already is a framework that can help combine together coroutines and spring mvc : https://github.com/konrad-kaminski/spring-kotlin-coroutine

This framework is working nicely for hello world or CRUD Apis, but hardly adapts when you have spring security configured, with a LocaleContextHolder, a MDC map, or maybe also a RequestContextHolder attributes object.

This very rough implementation below will help you declare a CoroutineContext bean that you can bind directly in your RestController if you want to use coroutines seamlessly.
It will make a coroutineContext parameter available, and will also make the controller wait until the result has been resolved.

package com.company.myapi

import kotlinx.coroutines.ExecutorCoroutineDispatcher
import org.slf4j.MDC
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.context.annotation.Bean
import org.springframework.context.annotation.Configuration
import org.springframework.context.annotation.Primary
import org.springframework.context.i18n.LocaleContextHolder
import org.springframework.core.MethodParameter
import org.springframework.security.core.context.SecurityContextHolder
import org.springframework.security.core.parameters.DefaultSecurityParameterNameDiscoverer
import org.springframework.web.bind.support.WebDataBinderFactory
import org.springframework.web.context.request.NativeWebRequest
import org.springframework.web.context.request.RequestContextHolder
import org.springframework.web.context.request.async.DeferredResult
import org.springframework.web.method.support.AsyncHandlerMethodReturnValueHandler
import org.springframework.web.method.support.HandlerMethodArgumentResolver
import org.springframework.web.method.support.HandlerMethodReturnValueHandler
import org.springframework.web.method.support.ModelAndViewContainer
import org.springframework.web.servlet.config.annotation.WebMvcConfigurer
import org.springframework.web.servlet.mvc.method.annotation.DeferredResultMethodReturnValueHandler
import java.lang.reflect.Method
import java.util.concurrent.ExecutorService
import java.util.concurrent.Executors
import kotlin.coroutines.Continuation
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.intrinsics.COROUTINE_SUSPENDED
import kotlin.reflect.jvm.kotlinFunction

object CoroutinesInRestController {

@Configuration
internal class CoroutinesInjection {
@Bean
@Primary
fun parameterNameDiscovererWithCoroutines() =
object: DefaultSecurityParameterNameDiscoverer(listOf()) {
override fun getParameterNames(method: Method) =
((super.getParameterNames(method)?.toList() ?: listOf()) +
if (method.isSuspend) listOf("__continuation__") else listOf()).toTypedArray()
}
@Bean
fun coroutineContext() = object : ExecutorCoroutineDispatcher() {
override val executor = Executors.newFixedThreadPool(128)

override fun dispatch(context: CoroutineContext, block: Runnable) {
val securityContext = SecurityContextHolder.getContext()
val requestAttributes = RequestContextHolder.currentRequestAttributes()
val locale = LocaleContextHolder.getLocale()
val contextMap = MDC.getCopyOfContextMap()
executor.execute {
SecurityContextHolder.setContext(securityContext)
RequestContextHolder.setRequestAttributes(requestAttributes)
LocaleContextHolder.setLocale(locale)
MDC.setContextMap(contextMap)
block.run()
}
}

override fun close() {
executor.shutdown()
}
}
}

@Configuration
internal class CoroutinesWebMvcConfigurer : WebMvcConfigurer {

@Autowired
private lateinit var coroutineContext: CoroutineContext

override fun addArgumentResolvers(resolvers: MutableList<HandlerMethodArgumentResolver>) {
resolvers.add(0, coroutineArgumentResolver(coroutineContext))
}

override fun addReturnValueHandlers(handlers: MutableList<HandlerMethodReturnValueHandler>) {
handlers.add(0, returnValueHandler())
}
}

private const val DEFERRED_RESULT = "deferred_result"

private fun <T> isContinuationClass(clazz: Class<T>) = Continuation::class.java.isAssignableFrom(clazz)
val Method?.isSuspend: Boolean get() = this?.kotlinFunction?.isSuspend ?: false

fun coroutineArgumentResolver(coroutineContext: CoroutineContext) =
object : HandlerMethodArgumentResolver {
override fun supportsParameter(parameter: MethodParameter) =
parameter.method.isSuspend && isContinuationClass(parameter.parameterType)

override fun resolveArgument(parameter: MethodParameter, mavContainer: ModelAndViewContainer,
webRequest: NativeWebRequest, binderFactory: WebDataBinderFactory) =
object : Continuation<Any> {
val deferredResult = DeferredResult<Any>()

override val context: CoroutineContext
get() = coroutineContext

override fun resumeWith(result: Result<Any>) {
if (result.isSuccess) {
deferredResult.setResult(result.getOrNull())
} else {
deferredResult.setErrorResult(result.exceptionOrNull())
}
}
}.apply {
mavContainer.model[DEFERRED_RESULT] = deferredResult
}
}

fun returnValueHandler() =
object: AsyncHandlerMethodReturnValueHandler {
private val delegate = DeferredResultMethodReturnValueHandler()

override fun supportsReturnType(returnType: MethodParameter): Boolean =
returnType.method.isSuspend

override fun handleReturnValue(returnValue: Any?, type: MethodParameter,
mavContainer: ModelAndViewContainer, webRequest: NativeWebRequest) {
val result = mavContainer.model[DEFERRED_RESULT] as DeferredResult<*>

return delegate.handleReturnValue(result, type, mavContainer, webRequest)
}

override fun isAsyncReturnValue(returnValue: Any, returnType: MethodParameter): Boolean =
returnValue === COROUTINE_SUSPENDED
}
}

This special implementation will take care of keeping the ThreadLocal values up to date in each coroutine scope, so you can continue processing your request without standing out of the servlet context values.

That was useful for me, and maybe it will be for you too, so good luck.

Leave a Reply

Your email address will not be published. Required fields are marked *

*

This site uses Akismet to reduce spam. Learn how your comment data is processed.