Support micrometer context-propagation

[chickenchickenlove]

Motivation:

• Related Issue : Support micrometer-context-propagation #5145
• Armeria already support context-propagation to maintain RequestContext during executing Reactor code. How it requires maintenance.
• Reactor integrate micro-meter:context-propagation to do context-propagation during Flux, Mono officially. thus, it would be better to migrate from RequestContextHook to RequestContextPropagationHooks because it can reduce maintenance cost.

Modifications:

• Add new Hook for Reactor.
• Add new ThreadLocalAccessor for micro-meter:context-propagation to main RequestContext during executing Reactor code like Mono, Flux.
• Add new config enableContextPropagation to integrate micro-meter:context-propagation with spring-boot3.

Result:

• Closes Support micrometer-context-propagation #5145
• If user want to use micrometer:context-propagation to maintain RequestContext during executing Reactor code like Mono, Flux, just call RequestContextPropagationHook.enable().

[chickenchickenlove]

FYI, flow of micro-meter/context-propagation.

1. If we enable context-propagation, hook for context-propagation will be added Reactor's hooks.

2. if subscribe() are called, hook like captureThreadLocals() are called.

3. and then, ContextSnapshot are created to propagate context. in this time, accessor are used to capture state of ThreadLocals.

4. When Mono or Flux are executed in Scheduler's executor, scope are gotten by calling setThreadLocals(). Scope instance are returned, and it is concrete class of AutoClosable.

5. As you can see, Context-Propagation is restore thread local state at end of try ~ resource.

6. On the other hand, setThreadLocals() be about to restore Threadlocal state from ContextSnapshot before runnable.run().

IMHO, it is very similar to RequestContextHook on armeria.

[github-actions]

🔍 Build Scan® (commit: e80117c)

Job name	Status	Build Scan®

[chickenchickenlove]

I investigate that internal of Context-Propagation.

Case 1. publisher.subscribe(...). downstream -> upstream flow

• When publisher.subscribe(...) are called and publisher is instance of ContextWriteRestoringThreadLocals, it creates ContextSnapshot.Scope.in this sequence, read the Key-Value stored in ReactorContext, then use the ThreadLocalAccessor instance to store the value from ReactorContext into ThreadLocal. Additionally, keep the previously stored values in ThreadLocal as PreviousValues, and use this values to revert the state of ThreadLocal to its previous state when the ContextSnapshot.Scope ends. (FYI, ContextSnapshot.Scope implement AutoClosable).

Case 2. subscriber.onSubscribe(...). upstream -> downstream flow

• When subscriber.onSubscribe(...) are called and subscriber is instance of ContextWriteRestoringThreadLocal, it creates ContextSnapshot.Scope as well. it means that the values in ReactorContext will be stored to ThreadLocal at the start of subscriber.onSubscribe() and ThreadLocal will be revert to its previous state at the end of the function.

Case 3. subscription.request(...). downstream -> upstream flow

• When subscription.request(...) are called and subscriptionis instance ofContextWriteRestoringThreadLocal, it creates ContextSnapshot.scope` as well.

ContextWriteRestoringThreadLocals operator are integrated when both [Mono|Flux]#contextCapture() and [Mono|Flux]#contextWrite() are called. It means that Reactor Context are essential to propagate RequestContext to each ThreadLocal during Reactor Operations.

In practice, it works like this:

1. Reads all values accessible by the KEY of ThreadLocalAccessor in the Reactor Context.
2. Stores the read values in the Threadlocal by using ThreadLocalAccessor. At this time, the values that previously existed in the Threadllocal are maintained in a Map called PreviousValues.
3. When the Scope ends, it restores the PreviousValues back to the ThreadLocal.

It means that Reactor Context for writing, ThreadLocals for reading.

[chickenchickenlove]

@trustin nim, There are also major differences in the test.
I would like to inform you about them.

The test utility function addCallbacks() should be change.

private static <T> Mono<T> addCallbacks(Mono<T> mono, ClientRequestContext ctx) {
        return mono.doFirst(() -> assertThat(ctxExists(ctx)).isTrue())
                   .doOnSubscribe(s -> assertThat(ctxExists(ctx)).isTrue())
                   .doOnRequest(l -> assertThat(ctxExists(ctx)).isTrue())
                   .doOnNext(foo -> assertThat(ctxExists(ctx)).isTrue())
                   .doOnSuccess(t -> assertThat(ctxExists(ctx)).isTrue())
                   .doOnEach(s -> assertThat(ctxExists(ctx)).isTrue())
                   .doOnError(t -> assertThat(ctxExists(ctx)).isTrue())
                   .doAfterTerminate(() -> assertThat(ctxExists(ctx)).isTrue())
                   // I added contextWrite(...)
                   .contextWrite(Context.of(RequestContextAccessor.getInstance().key(), ctx));
        // doOnCancel and doFinally do not have context because we cannot add a hook to the cancel.
    }

contextWrite(Context.of(RequestContextAccessor.getInstance().key(), ctx)); are added. As we know, micro-meter:context-propagation require Reactor Context to propagate context during reactor operations. thus, i added this method.

// Before : StepVerifier.create(mono1)
// After : Add initiali Reactor Context  to StepVerifier. 
StepVerifier.create(mono1, initialReactorContext(ctx))
                    .expectSubscriptionMatches(s -> ctxExists(ctx))
                    .expectNextMatches(s -> ctxExists(ctx) && "baz".equals(s))
                    .verifyComplete();

In previous test code, StepVerifier.create(mono1) is enough. however, it is not enough to micro-meter:context-propagation.
StepVerifier create DefaultVerifySubscriber to subscribe Flux|Mono and valid result is correct. however, DefaultVerifySubscriber has only empty Reactor Context. it means that .expectSubscriptionMatches(s -> ctxExists(ctx)) should be failed.

micro-meter:context-propagation is used to read the values from the Reactor Context and restore the state of ThreadLocal. however, since the DefaultStepVerifierSubscriber has an Empty Reactor Context, the RequestContext stored in ThreadLocal will become Null.

Thus, initial Reactor Context should be include to StepVerifier as well.

[trustin]

Could you fix the build failures before getting reviews?

[minwoox]

These look great! Thanks a lot for completing this PR. 👍 👍 👍

[trustin]

@minwoox IIRC you had some comments related with request context hooks you wanted @chickenchickenlove to address. Could you leave some comment about that?

[minwoox]

Let's just use RequestContextUtil.getAndSet(value);.
value.push() internally invokes a hook but the closeable never be closed because we discard the result of value.push().

armeria/core/src/main/java/com/linecorp/armeria/internal/common/RequestContextUtil.java

Line 192 in 87f2318

final SafeCloseable closeable = invokeHook(current);

armeria/core/src/main/java/com/linecorp/armeria/internal/common/RequestContextUtil.java

Line 197 in 87f2318

closeable.close();

This is a limitation of the current design of Mircometer context-proparation.
Because we cannot call the closeable, we should avoid invoking the hook altogether and inform users of the limitation of this API.

[chickenchickenlove]

@minwoox nim, thanks for your comments.
RequestContextUtil.getAndSet(value) seems not get SafeClosable instance.
I make new commit to apply your comments.
Please take another look, when you have free time 🙇‍♂️

[minwoox]

Still looks good. Thanks!

[minwoox]

ditto let's just use RequestContextUtil.getAndSet(previousValue);

[minwoox]

@minwoox IIRC you had some comments related with request context hooks you wanted @chickenchickenlove to address. Could you leave some comment about that?

Thanks! I left my opinion here: #5577 (comment)

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 74.05%. Comparing base (b8eb810) to head (e80117c).
Report is 194 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main    #5577      +/-   ##
============================================
+ Coverage     73.95%   74.05%   +0.09%     
- Complexity    20115    21242    +1127     
============================================
  Files          1730     1850     +120     
  Lines         74161    78540    +4379     
  Branches       9465    10020     +555     
============================================
+ Hits          54847    58164    +3317     
- Misses        14837    15680     +843     
- Partials       4477     4696     +219     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.