MdcContextBridge

MdcContextBridge.java

import org.reactivestreams.Subscription;
import org.slf4j.MDC;
import reactor.core.CoreSubscriber;
import reactor.core.publisher.Hooks;
import reactor.core.publisher.Operators;
import reactor.util.context.Context;

import java.util.Map;
import java.util.stream.Collectors;

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

    private final CoreSubscriber<T> actual;

    public MdcContextBridge(CoreSubscriber<T> actual) {
        this.actual = actual;
    }

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

    @Override
    public void onSubscribe(Subscription s) {
        actual.onSubscribe(s);
    }

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

    @Override
    public void onError(Throwable t) {
        actual.onError(t);
    }

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

    public static void copyToMdc(Context context) {
        if (!context.isEmpty()) {
            Map<String, String> map = context.stream()
                    .collect(Collectors.toMap(e -> String.valueOf(e.getKey()), e -> String.valueOf(e.getValue())));
            MDC.setContextMap(map);
        } else {
            MDC.clear();
        }
    }

    public static void install() {
        Hooks.onEachOperator(MDC_CONTEXT_BRIDGE_KEY, Operators.lift((__, sub) -> new MdcContextBridge<>(sub)));
    }

    public static void reset() {
        Hooks.resetOnEachOperator(MDC_CONTEXT_BRIDGE_KEY);
    }

    public static final String MDC_CONTEXT_BRIDGE_KEY = MdcContextBridge.class.getName();
}