Automatic context propagation with MDC using Project Reactor

build.gradle

// ...

dependencies {
    implementation 'io.micrometer:context-propagation:1.0.6'
    // ...
}

logback.xml

<configuration>

    <appender name="stdout" class="ch.qos.logback.core.ConsoleAppender">
        <encoder>
            <pattern>
                %d{HH:mm:ss.SSS} [%thread] [key=%X{key}] %-5level %logger{36} - %msg%n
            </pattern>
        </encoder>
    </appender>
    <logger name="reactor" level="info"/>

    <root level="info">
        <appender-ref ref="stdout"/>
    </root>

</configuration>

TestMdc.java

    @Test
    void testMDC() {
        Logger log = LoggerFactory.getLogger("test");

        Hooks.enableAutomaticContextPropagation();

        // To deal with the entire MDC (if we ensured no third-party code modifies it):
        // ContextRegistry.getInstance().registerThreadLocalAccessor(new MdcAccessor());

        // To deal with an individual key in the MDC:
        ContextRegistry.getInstance().registerThreadLocalAccessor(
            "key",
            () -> MDC.get("key"),
            value -> MDC.put("key", value),
            () -> MDC.remove("key"));

        MDC.put("key", "Hello");

        Mono.just("Delayed")
            .delayElement(Duration.ofMillis(10))
            .doOnNext(log::info)
            // It is vital to capture MDC contents into Reactor's context
            // as it is the source of truth for propagating ThreadLocals
            //.contextCapture() // Can be skipped since reactor-core:3.5.7
            .block();
    }

    // This implementation modifies the entire MDC.
    // If any other code (especially third-party libraries) interacts with the MDC
    // it would collide and the results will be corrupted.
    // The best strategy is to pick specific keys instead of working on the entire MDC map.
    static class MdcAccessor implements ThreadLocalAccessor<Map<String, String>> {

        static final String KEY = "mdc";

        @Override
        public Object key() {
            return KEY;
        }

        @Override
        public Map<String, String> getValue() {
            return MDC.getCopyOfContextMap();
        }

        @Override
        public void setValue(Map<String, String> value) {
            MDC.setContextMap(value);
        }

        @Override
        public void reset() {
            MDC.clear();
        }
    }

21:07:52.889 [parallel-1] [key=Hello] INFO test - Delayed

@chemicL Is it ok if we don't clear the MDC? I am using logbook + WebClient and sometimes MDC doesn't get restored before logbook logs the request.

I guess not, as the problem would still remain if thread is switched :/

@nakashkumar I'm not sure where the problem lies. Can you share anything about your setup? Is Observation from Micrometer in use? This mechanism from my gist would collide with it, as it clears the MDC entirely. You could replace the implementation to work on a single key instead of the entire map.

I created a small reproducer if you'd like to have a look.
Using a single key instead of the entire map has the same effect.

There's a lot of libraries at play in that example. I don't know logbook and have no knowledge of what it's doing. Is it suppose to only use what's in MDC? Perhaps you can remove logbook from the equation? I also see you are using micrometer's brave bridge. Perhaps the issue lies there (one of the handlers registered with brave deals with MDC). There are reports such as this one: micrometer-metrics/tracing#257 which look similar to what you are reporting. But first and foremost, you mustn't mix MDC accessors: either global or for the same key. If micrometer-tracing is dealing with e.g. "request-id" key from MDC and you register and accessor that also reads, sets, and removes "request-id", then one is going to overwrite the work of the other inevitably. So you shouldn't register your own accessor.

The only case when you can have more than one accessor for MDC is if they operate on independent keys, e.g.:

ContextRegistry.getInstance().registerThreadLocalAccessor(
	"KEY-A",
	() -> MDC.get("KEY-A"),
	value -> MDC.put("KEY-A", value),
	() -> MDC.remove("KEY-A"));

ContextRegistry.getInstance().registerThreadLocalAccessor(
	"KEY-B",
	() -> MDC.get("KEY-B"),
	value -> MDC.put("KEY-B", value),
	() -> MDC.remove("KEY-B"));

Thanks for taking the time to look into this. You are right, micrometer is already setting the thread locals from the reactor context using the context-propagation module. We shouldn't be introducing our own thread-local accessor for MDC.
It works for me if I log requests and responses in the context scope (the MDC is available then) and not in the netty channel handler.

Cześć @chemicL! I am here from your post on Spring Blog for context-propagation. I don't know if I got it well but when I use the micrometer-tracing it just populates spanId and traceId, I put some more keys into Context during contextWrite and these keys are used in logback.xml. But when debugging I never see the keys in the MDC.getCopyOfContextMap(). Should micrometer populate the MDC automatically with all the keys from context? Or should I add the MDC accessor as you did above?

@imosapatryk so you mean you have some key of your own, say "X" and when you do MDC.getCopyOfContextMap() then the map contains traceId and spanId, but not "X"? Have you registered an accessor for "X"? Like this:

ContextRegistry.getInstance().registerThreadLocalAccessor(
	"X",
	() -> MDC.get("X"),
	value -> MDC.put("X", value),
	() -> MDC.remove("X"));

Micrometer actually does not deal with MDC, it's the tracing instrumentation library, e.g. brave, for which a Micrometer bridge can be added, e.g. micrometer-tracing-bridge-brave. And it works with its particular keys (traceId, spanId), which are a derivative of instantiating a span into existence. So as far as I understand, yes, you do need the accessor.

This is exactly what I have done but I am bit worried how many times the "lambdas" for getting and setting are called (I debugged it and it is a lot)

Yes, it is quite a few, but it's required to guarantee proper cleanup and not leaking the state to other tasks on the event loop. You can measure the performance impact on the real scenario and decide if you want to refactor the code to directly interact with the ContextView and extract the identifies to explicitly pass them along to the log statements or if this automation is a price worth paying.

I'll do! 👍 Btw. I've started wondering if Hooks.enableAutomaticContextPropagation(); may have impact on reactive transaction propagation...

be more specific please :)

@leakin185
https://projectreactor.io/docs/core/release/api/reactor/core/publisher/Hooks.html#enableAutomaticContextPropagation--

Please make sure you're using the latest version :) reactor-core is now 3.5.10, and this feature was made available since 3.5.3. If it's missing, probably you're using an older version of the library.

I updated the gist to reflect the best practices.