[Kotlin] MDCContext 동작 원리 간단히 알아보기

이전에 Sl4fj MDC로 요청 구분하기에서 withContext를 통해 다른 스레드로 전환이 발생할 때
MDCContext를 사용하면 된다고 언급했었다.

withContext(Dispatchers.IO + MDCContext()) { ... }

MDCContext()를 더해주는 것만으로 이전 context를 이어서 사용할 수 있다는 건 굉장한 이점이다.
어떤 과정으로 context를 이어서 사용할 수 있는지 간단히 알아보자.

아래는 MDCContext 스펙이다.

일단 MDCContext()를 호출해주는 순간 생성자 내 필드인 contextMap는 기본값인 MDC.getCopyOfContextMap()를 갖는다.
MDC.getCopyOfContextMap()는 현재 MDC 정보를 그대로 복사한다. 즉, contextMap는 현재 context 정보를 갖게된다.

public class MDCContext(
    public val contextMap: MDCContextMap = MDC.getCopyOfContextMap()
) : ThreadContextElement<MDCContextMap>, AbstractCoroutineContextElement(Key) {

    override fun updateThreadContext(context: CoroutineContext): MDCContextMap { ... }
    override fun restoreThreadContext(context: CoroutineContext, oldState: MDCContextMap) { ... }
    ...
}

updateThreadContext(), restoreThreadContext()는 오버라이드된걸 보니 MDCContext에서 최초로 정의된게 아니다.
MDCContext가 상속받는 ThreadContextElement interface를 따라가보면 명세를 확인할 수 있다.

ThreadContextElement

public interface ThreadContextElement<S> : CoroutineContext.Element {
    /**
     * 현재 스레드의 context를 업데이트
     * 기존 context를 응답으로 반환하고, 응답값은 restoreThreadContext에서 복원 용도로 사용
     * (응답값을 oldState이라 부름)
     */
    public fun updateThreadContext(context: CoroutineContext): S

    /**
     * 현재 스레드의 context를 복원
     * updateThreadContext 응답인 oldState를 param으로 받아 복원함
     */
    public fun restoreThreadContext(context: CoroutineContext, oldState: S)
}

updateThreadContext

context를 업데이트하는 함수로, 코루틴이 시작(or 재시작)될 때 호출된다.
특이한 점은 응답으로 S 타입의 oldState를 반환하는데 이는 아래에서 설명할 복원 함수 restoreThreadContext에서 사용된다.

restoreThreadContext

context를 복원하는 함수로, 코루틴이 중단될 때 호출된다.
파라미터로 oldState를 받는데 이는 updateThreadContext의 응답으로, 타입이 S로 같은걸 알 수 있다.

그니까 updateThreadContext에서 복원할 정보(oldState)를 남겨뒀다가 코루틴이 중단되면 다른 코루틴이 들어와서 사용할 수 있게 restoreThreadContext를 호출해서 복원해둔다. 그러면 다른 코루틴의 context를 침범하지 않으면서 동작할 수 있다.

Reentrancy and thread-safety

Correct implementations of this interface must expect that calls to restoreThreadContext may happen in parallel to the subsequent updateThreadContext and restoreThreadContext operations. See CopyableThreadContextElement for advanced interleaving details.

All implementations of ThreadContextElement should be thread-safe and guard their internal mutable state within an element accordingly.

그래서인지 ThreadContextElement 명세에서 Reentrancy and thread-safety 하다고 표현하고 있다.

MDCContext에서의 구현

ThreadContextElement 구현체 MDCContext는 각 함수를 어떻게 구현하고 있는지 알아보자.

updateThreadContext

override fun updateThreadContext(context: CoroutineContext): MDCContextMap {
    val oldState = MDC.getCopyOfContextMap() // 현재 context 복원 용도로 기록
    setCurrent(contextMap) // 업데이트
    return oldState
}

• contextMap는 MDCContext의 생성자 내 필드로, MDCContext 초기화 시점에 context를 저장한다.

명세 그대로 현재 코루틴의 context를 oldState에 저장해둔다. 그리고 그걸 반환한다.

참고로 setCurrent는 아래와 같이 정의돼 있다. 현재 스레드의 MDC에 contextMap을 업데이트하는 모습이다.

private fun setCurrent(contextMap: MDCContextMap) {
    if (contextMap == null) {
        MDC.clear()
    } else {
        MDC.setContextMap(contextMap)
    }
}

• (참고) MDC.setContextMap(null)을 실행하면 예외를 던진다. 그래서인지 null 분기 처리가 돼있다.

restoreThreadContext

override fun restoreThreadContext(context: CoroutineContext, oldState: MDCContextMap) {
    setCurrent(oldState)
}

updateThreadContext 응답인 oldState를 param으로 받아 MDC를 복구하는 모습이다.
역시 명세와 같다.

withContext 활용

기존에 아래처럼 withContext를 사용하고 있었다면, 사용하는 곳마다 + MDCContext()를 해줘야한다.

// as-is
withContext(Dispatchers.IO) { ... }
withContext(Dispatchers.Default) { ... }

// to-be
withContext(Dispatchers.IO + MDCContext()) { ... }
withContext(Dispatchers.Default + MDCContext()) { ... }

추가로 작성한 로직에 withContext를 써야할 때, MDCContext()를 실수로 빠뜨릴 수도 있다.
그래서 아래처럼 정의해두고 사용하면 간편하다. withContext를 withMDCContext로만 교체해주면 된다.

suspend fun <T> withMDContext(
    context: CoroutineContext,
    block: suspend CoroutineScope.() -> T
): T {
    return withContext(context + MDCContext()) {
        block()
    }
}

// 활용
withMDCContext(Dispatchers.IO) { ... }
withMDCContext(Dispatchers.Default) { ... }

References

• kotlinlang.org: MDCContext