MDC와 TaskDecorator를 통한 TraceId 전파

비동기 작업에서 동일한 traceId를 추적할 수 있도록 MDC와 TaskDecorator를 사용하는 방법을 설명합니다.

기본적으로 MDC 필터는 다음과 같이 사용됩니다

MDC(Mapped Diagnostic Context)는 로깅할 때 스레드별로 고유한 진단 정보를 저장하기 위한 메커니즘입니다. Filter를 통해 DispatcherServlet에 요청을 전달하기 전에 MDC에 고유한 UUID를 저장하여, 이후의 로깅에서 해당 요청에 대한 정보를 traceId를 통해 일관되게 추적할 수 있게 합니다.

MDC 필터 예시

다음은 Filter를 상속받은 MDC 필터 클래스의 일부분입니다. doFilter() 메서드에서는 고유한 traceId를 생성하고 MDC에 저장하여 각 요청마다 추적 가능한 ID를 부여합니다.

public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
        throws IOException, ServletException {
    // traceId를 UUID로 생성하여 설정
    String traceId = UUID.randomUUID().toString().substring(0, 8);
    MDC.put("traceId", traceId);
    
    try {
        chain.doFilter(request, response);
    } finally {
        MDC.clear();  // 요청 처리 완료 후 MDC 초기화
    }
}

문제점: ThreadLocal 기반의 MDC 한계

MDC는 ThreadLocal을 사용하여 각 스레드에 독립적인 진단 정보를 저장합니다. 하지만 비동기 작업이나 이벤트 호출에서 다른 스레드가 사용되면 MDC에 저장된 traceId가 전파되지 않으므로, 새로운 스레드는 기존 traceId를 알 수 없습니다. 이로 인해 비동기 처리 중 요청이 어떤 traceId와 연관된 것인지 추적할 수 없는 문제가 발생하게 됩니다.

해결 방법: TaskDecorator를 사용한 MDC 전파

비동기 환경에서도 동일한 traceId를 추적할 수 있도록, 비동기 작업을 시작할 때마다 MDC 정보를 복사하여 새로운 스레드에 설정해주는 TaskDecorator를 사용합니다. Spring의 TaskDecorator는 ThreadPoolTaskExecutor의 스레드가 실행할 Runnable을 감싸서 비동기 작업에서 MDC 정보가 일관되게 유지되도록 설정할 수 있습니다.

MdcTaskDecorator 클래스

아래는 TaskDecorator를 구현한 MdcTaskDecorator 클래스입니다. decorate() 메서드에서 현재 스레드의 MDC 정보를 복사해 두고, 비동기 스레드가 실행될 때 이를 설정하여 traceId가 유지되도록 합니다.

public class MdcTaskDecorator implements TaskDecorator {

    @Override
    public Runnable decorate(Runnable runnable) {
        // 현재 스레드의 MDC 컨텍스트 복사
        Map<String, String> contextMap = MDC.getCopyOfContextMap();

        return () -> {
            try {
                // 비동기 스레드에 MDC 컨텍스트 설정
                if (contextMap != null) {
                    MDC.setContextMap(contextMap);
                }
                runnable.run();
            } finally {
                MDC.clear();  // 작업 완료 후 MDC 초기화
            }
        };
    }
}

왜 MdcTaskDecorator로 해결이 가능한가?

MdcTaskDecorator는 현재 스레드의 MDC 정보를 복사하여, 새로운 스레드에서도 동일한 traceId를 사용할 수 있게 해줍니다. MDC는 ThreadLocal을 사용하여 스레드 단위로 정보를 저장하므로, 비동기 작업에서 새로운 스레드가 시작될 때 원래의 traceId가 자동으로 전파되지 않습니다. 이 문제를 해결하기 위해 MdcTaskDecorator는 다음과 같은 방식으로 동작합니다:

1. 현재 스레드의 MDC 정보를 복사: MDC.getCopyOfContextMap()을 사용하여 현재 스레드의 MDC 데이터를 Map<String, String> 형태로 복사합니다. 이 복사된 Map은 ThreadLocal이 아니기 때문에, 다른 스레드에 안전하게 전달할 수 있습니다.

2. 비동기 스레드에서 복사된 정보를 설정: 비동기 작업이 실행될 때 decorate() 메서드가 반환한 Runnable을 실행하면서, MDC.setContextMap(contextMap)을 호출해 복사된 MDC 정보를 비동기 스레드의 MDC에 설정합니다. 이렇게 함으로써 새로운 스레드에서도 동일한 traceId를 유지할 수 있습니다.

3. 작업 완료 후 MDC 초기화: 비동기 작업이 끝나면 MDC.clear()를 호출하여 스레드 로컬에 저장된 MDC 데이터를 초기화하고, 메모리 누수를 방지합니다.

이 과정을 통해 비동기 작업에서도 동일한 MDC 컨텍스트를 유지할 수 있어, traceId를 기반으로 한 로그 추적이 일관되게 이루어집니다.

MdcTaskDecorator의 호출 흐름

전체적인 MdcTaskDecorator 호출 흐름은 다음과 같습니다:

1. 기존 쓰레드: 요청을 처리하던 기존 쓰레드가 비동기 작업을 실행하기 위해 Runnable을 생성하고 스레드 풀에 제출합니다.
2. 비동기 쓰레드 생성 요청: 기존 쓰레드가 ThreadPoolTaskExecutor를 통해 비동기 작업을 요청합니다. 이때 ThreadPoolTaskExecutor는 TaskDecorator가 설정되어 있으면 해당 TaskDecorator를 사용해 Runnable을 꾸며줍니다.
3. MdcTaskDecorator의 decorate() 호출: ThreadPoolTaskExecutor는 decorate() 메서드를 호출하여 Runnable을 감쌉니다. 이 과정에서 현재 스레드의 MDC 컨텍스트 맵을 복사합니다.
4. 비동기 스레드에서 MDC 적용: 비동기 스레드가 Runnable을 실행하기 직전에, 복사된 MDC 컨텍스트가 새 스레드의 MDC에 설정됩니다. 이로 인해 비동기 스레드에서도 기존 쓰레드와 동일한 traceId 등 MDC 정보를 사용할 수 있게 됩니다.

이 흐름을 통해 MdcTaskDecorator는 비동기 환경에서도 MDC 정보를 안전하게 전파하며, 로그 추적이 일관되게 유지되도록 합니다.