WebFlux Log Tracing

WebFlux를 처음 사용해보면서 로깅하는데 며칠을 애먹었다.

Spring MVC와는 다르게 레퍼런싱할 수 있는 자료들이 거의 없었다. 그래서 다른 분들의 시간을 단축시켜주고자 글을 적어본다.

(좀 더 좋은 자료가 있는 곳이 있다면 댓글로 공유부탁드립니다 🙇)

 

결론부터 말하자면,

WebFlux에서 MDC Context로 로깅을 '잘' 하려면 아래와 같이 HttpHandlerDecoratorFactory의 구현체를 하나 만들면 된다.

package com.tistory.mjin1220.decorator;

import org.slf4j.MDC;
import org.springframework.http.server.reactive.HttpHandler;
import org.springframework.http.server.reactive.HttpHandlerDecoratorFactory;
import org.springframework.stereotype.Component;
import reactor.util.context.Context;

import java.util.UUID;

@Component
public class CustomHttpHandlerDecorator implements HttpHandlerDecoratorFactory {

    private static final String MDC_KEY_TRACE_ID = "traceId";

    @Override
    public HttpHandler apply(HttpHandler httpHandler) {
        return (request, response) ->
                httpHandler.handle(request, response)
                           .contextWrite(context -> {
                               final String traceId = getTraceId();
                               MDC.put(MDC_KEY_TRACE_ID, traceId);
                               return Context.of(MDC_KEY_TRACE_ID, traceId);
                           });
    }

    private String getTraceId() {
        return UUID.randomUUID()
                   .toString();
    }
}

위와 같이 설정을 해주면 WebFilter에서 로직을 처리하다가 Exception이 발생하는 케이스에서도 커버가 된다.

WebFlux에서 AbstractErrorWebExceptionHandler를 이용해서 Exception을 처리하게 되는데, 이 때 WebFilter에서 아래와 같이 처리해주는 케이스로는 커버가 되지 않았다.

        return chain.filter(exchange)
                    .contextWrite(context -> {
                        final String traceId = getTraceId();
                        MDC.put(MDC_KEY_TRACE_ID, traceId);
                        return Context.of(MDC_KEY_TRACE_ID, traceId);
                    });

여기에서 내부적으로 WebFlux를 단일 ThreadLocal에서만 처리하는 경우는 드물다.

그러면 매번 컨텍스트 스위칭이 일어날 때마다 MDC에 있는 데이터들을 Context로 넘겨주는 것은 매우 귀찮은 일이다. (또는 Context에 있는 데이터들을 MDC로)

package com.tistory.mjin1220.config;

import lombok.RequiredArgsConstructor;
import org.reactivestreams.Subscription;
import org.slf4j.MDC;
import org.springframework.context.annotation.Configuration;
import reactor.core.CoreSubscriber;
import reactor.core.publisher.Hooks;
import reactor.core.publisher.Operators;
import reactor.util.context.Context;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @link https://www.novatec-gmbh.de/en/blog/how-can-the-mdc-context-be-used-in-the-reactive-spring-applications/
 */
@Configuration
public class MdcContextLifterConfig {

    public static final String MDC_CONTEXT_REACTOR_KEY = MdcContextLifterConfig.class.getName();

    @PostConstruct
    @SuppressWarnings({"unchecked", "rawtypes"})
    public void contextOperatorHook() {
        Hooks.onEachOperator(MDC_CONTEXT_REACTOR_KEY, Operators.lift((scannable, subscriber) -> new MdcContextLifter(subscriber)));
    }

    @PreDestroy
    public void cleanupHook() {
        Hooks.resetOnEachOperator(MDC_CONTEXT_REACTOR_KEY);
    }

    /**
     * Helper that copies the state of Reactor [Context] to MDC on the #onNext function.
     */
    @RequiredArgsConstructor
    public static class MdcContextLifter<T> implements CoreSubscriber<T> {

        private final CoreSubscriber<T> coreSubscriber;

        @Override
        public void onSubscribe(Subscription subscription) {
            coreSubscriber.onSubscribe(subscription);
        }

        @Override
        public void onNext(T t) {
            copyToMdc(coreSubscriber.currentContext());
            coreSubscriber.onNext(t);
        }

        @Override
        public void onError(Throwable throwable) {
            copyToMdc(coreSubscriber.currentContext());
            coreSubscriber.onError(throwable);
        }

        @Override
        public void onComplete() {
            coreSubscriber.onComplete();
        }

        @Override
        public Context currentContext() {
            return coreSubscriber.currentContext();
        }

        /**
         * Extension function for the Reactor [Context]. Copies the current context to the MDC, if context is empty clears the MDC.
         * State of the MDC after calling this method should be same as Reactor [Context] state.
         * One thread-local access only.
         */
        void copyToMdc(Context context) {
            if (context != null && !context.isEmpty()) {
                Map<String, String> map = context.stream()
                                                 .collect(Collectors.toMap(e -> e.getKey()
                                                                                 .toString(), e -> e.getValue()
                                                                                                    .toString()));

                MDC.setContextMap(map);
            } else {
                MDC.clear();
            }
        }
    }
}

※ Reference

• https://github.com/spring-projects/spring-framework/issues/26880
• https://techblog.woowahan.com/2667/