Spring Webflux + Coroutine에서 Context 관리하기 (downstream으로 context 변경 전파하기) - 2편

바로 이전 1편에서의 내용을 이어서 진행해보면

 

suspend fun 안에서는 아래의 방법으로 Reactor Context에 직접 접근하는 것이 가능하다. 이는 context 정보를 사용하기 위해 MDC를 거칠 필요가 없다는 것이다.

currentCoroutineContext()[ReactorContext]?.context

 

위에서 우리가 context-map이라는 key를 이용해 Map 객체를 context에 넣었던 것을 기억한다면, context-map 객체에 접근하기 위해 아래와 같이 활용할 수 있겠다는 생각이 들 수 있을 것이다.

suspend fun getContextMap(): MutableMap<String, String?>? =
    currentCoroutineContext()[ReactorContext]
        ?.context
        ?.get<MutableMap<String, String?>>("context-map")

우리는 이 메서드를 이용해 coroutine으로 동작하는 모든 메서드에서 Context 객체의 수정 가능한 Map에 접근(주소 참조)할 수 있게 된 것이다.

 

따라서 위에서 소개한 RequestContextInfoUtils를 살짝 수정하면 upstream에서 변경한 context 값을 downstream으로 전파할 수 있게 될 것이다.

object RequestContextInfoUtils {

    suspend fun getContextMap(): MutableMap<String, String?>? =
        currentCoroutineContext()[ReactorContext]
            ?.context
            ?.get<MutableMap<String, String?>>("context-map")

    suspend fun putToContext(key: String, value: String?) = getContextMap()?.put(key, value)

    suspend fun putToContexts(map: Map<String, String?>) = getContextMap()?.putAll(map)

    suspend inline fun getFromContext(key: String): String? = getContextMap()?.get(key)
    
    ...
    
    suspend fun getToken(): String? = getFromContext("token")
    
    suspend fun setToken(value: String?): String? = putToContext("token", value)
    
    ...
}

suspend fun test() {
    println(RequestContextInfoUtils.getToken())
    
    Mono.just("1")
        .publishOn(Schedulers.boundedElastic())
        .map {
            println(RequestContextInfoUtils.getToken())
            RequestContextInfoUtils.setToken("ABC2")
            println(RequestContextInfoUtils.getToken())
        }
        .awaitSingleOrNull()
        
    println(RequestContextInfoUtils.getToken())
}

// ABCKSNFIWNSKENFLSDKNFLSDKNF
// ABCKSNFIWNSKENFLSDKNFLSDKNF
// ABC2
// ABC2

 

물론 이 방식에도 단점은 존재한다. suspend fun에서만 사용할 수 있다는 점이다.

 

하지만 coroutine을 사용하게 되면 대부분의 코드는 suspend fun에서 수행하게 되고, MDC를 병행해서 사용하기 때문에 일반 메서드에서도 context 값에 접근은 가능하다는 사실을 잊어서는 안 된다.