Add log4j MDC log tags in webflux

By | May 11, 2023
0 Comment
Share the joy
  •  
  •  
  •  
  •  
  •  
  •  

Webflux is async. So log4j with MDC tags may not work correctly natively. So adding below to copy MDC context in each operator changes. Hooks is some advanced usage for customizing action in operators.

1. MdcContextLifterConfiguration

@Configuration
public class MdcContextLifterConfiguration {

    private static final String MDC_CONTEXT_REACTOR_KEY =
        MdcContextLifterConfiguration.class.getName();

    @PostConstruct
    private void contextOperatorHook() {
        Hooks.onEachOperator(
            MDC_CONTEXT_REACTOR_KEY,
            Operators.lift((scannable, coreSubscriber) -> new MdcContextLifter<>(coreSubscriber)));
    }

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

2. MdcContextLifter

public class MdcContextLifter<T> implements CoreSubscriber<T> {

    CoreSubscriber<T> coreSubscriber;

    public MdcContextLifter(CoreSubscriber<T> coreSubscriber) {
        this.coreSubscriber = coreSubscriber;
    }

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

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

    @Override
    public void onError(Throwable err) {
        coreSubscriber.onError(err);
    }

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

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

    private void copyToMdc(Context context) {
        if (!context.isEmpty()) {
            // TODO filter MDC context to ignore unused context
            Map<String, String> map =
                context.stream()
                    .collect(Collectors.toMap(e -> e.getKey().toString(), e -> e.getValue().toString()));
            MDC.setContextMap(map);
        } else {
            MDC.clear();
        }
    }
}