Supporting context-propagation ContextSnapshot restore with an operator

[simonbasle]

The goal of this feature request is to support context-propagation library, and more specifically the restoration of a ContextSnapshot into ThreadLocals, in a more integrated fashion.

Namely, we would like to somehow wrap user-provided "functions" (or actually any functional interface required in the Flux/Mono API) so that each restores the thread locals from the Reactor Context (and clears after itself).

The onus can be put on the users themselves by providing some facility to perform the wrapping, or can be implicit.

Considered Alternatives

Broad alternatives

1. explicit wrapping: a FunctionalWrappers instance is exposed to the user which must wrap each of her functions using said instance. one instance per ContextView/subscription
2. implicit wrapping within scope: define boundaries within which operators could discover a flag, instructing them to wrap any function passed as input when Subscriber is created
3. spec/builder approach: expose what looks like a subset of the Flux API, except it just defines a blueprint that will be applied at subscription time with additional wrapping of the user-provided functions

[simonbasle]

(1) Explicit approach

Explored in #3147

Expected API

The API is close to transformDeferredContextual, except it doesn't expose the ContextView but rather a Wrappers instance capable of wrapping any functional type.

interface Wrappers {
    <T,R> Function<T, R> function(Function<T, R> original);
//and many others...

The one instance created by the operator would use context-propagation API when wrapping functions, but that's transparently done. User needs to call out where and what to wrap though:

aStringFlux.withFunctionalWrappers(
    source, wrapper -> source.map(wrapper.function(
        v -> v.length()
    ))
)

Caveats

The wrapping resets thread locals the moment the function has finished. This could be hard to grasp for some users, eg. when wrapping a Function<T, Flux<R>> for a flatMap. In that case, the threadlocals wouldn't be visible at subscription time in the returned inner Flux, since it would be subscribed to by flatMap AFTER the function has returned.

This approach also cannot work if the "scoped" chain of operators is built by third-parties that are not aware of the Wrappers instance.

Possible variant

Rather than a dedicated operator, use transformDeferredContextual and a factory method that provides the necessary BiFunction. Multiple factories can thus be exposed for additional use cases.

aStringFlux.transformDeferredContextual(
    ContextPropagation.withWrappers(
        (source, wrapper) -> source.map(wrapper.function(
            v -> v.length()
        ))
    )
);

User writes a BiFunction<Flux<T>, Wrapper, Flux<R>>, factory method turns that into a BiFunction<Flux<T>, ContextView, Flux<R>> which is directly workable in transformDeferredContextual.

Or add a variant with a configurable resource:

• First parameter is a Function to get a RESOURCE from a ContextView
• Second parameter is still a BiFunction, but this time it offers the RESOURCE rather than the ContextView as second input

Like in the following:

asStringFlux.transformDeferredContextual(
    ContextPropagation::wrappers, //instantiate a `Wrappers` 
    //can be also written as a more verbose Function: ctx -> ContextPropagation.wrappers(ctx),

    //separately, define the transformation. except now we get a Wrappers resource
    (source, w) -> source.map(w.function(
        v -> v.length
    ))
);

[simonbasle]

(2) Scope with Implicit Flag approach

Where to put the Flag and how to define its range

One place is the Reactor Context. It then boils down to putting the flag inside context towards the "end" of the sub-chain and removing it towards the beginning.

Another idea would be to consider this flag as "Publisher-level configuration". There is currently no support for this so it implies A LOT of work, with a lot of potential caveats/blind spots.

Expected API

We could introduce a new operator or use the transformDeferredContextual name, except not with a BiFunction:

flux
    .transformDeferredContextual(f ->
        f.map(v -> v)
    );

The contextual suffix without a Context parameter hints at the implicit behavior.

Another option is to have a pair of scope-bounding operators:

flux
    .startScope()
    .map(v -> v)
    .endScope();

In both cases, if Context-based this would be similar to doing the following:

flux
    .contextWrite(c -> c.delete(SCOPED_KEY)
    .map(v -> v)
    .contextWrite(c -> c.put(SCOPED_KEY, SOMETHING);

Behind the scenes, transformDeferredContextual(Function) is more of an alias for transformDeferred since it doesn't really need the context: contextWrite operators will take care of things.

Note that Publisher-level configuration would likely need similar pairing of operators as startScope() and endScope().

Caveats

• for the scope's flag to have any effect, each operator needs to be aware of it and implement some behavior in onSubscribe to accommodate it (similar to doOnDiscard support)
• arbitrary support: within the scope there's no API restriction YET some operators used there might simply not support the feature. there's no way to warn users other than via documentation

In the Context-based approach:

• cognitive dissonance: from the perspective of the user, the scope is in order START-op1-op2-op3-END. but in reality it is reverted, more like REMOVE_FLAG-op1-op2-op3-PUT_FLAG
• virality: like Context, the flag would "propagate" inside inner chains, like eg. in a flatMap. is that desirable? can it be reasoned about by users?
• one needs to be careful not to blindly remove the flag towards the beginning of the scope: what if we're in a nested scope (which is possible since we don't constraint which operators can be called within the scope)?
• with pair of operators, in addition to nesting, there's now the possibility that one uses endScope() but not startScope(). while it might seem like an innocent mistake, it is actually a big one: endScope() would need to be the one that actually puts the flag in the context (because it reads like it should be at the bottom of the chain, and we're using Context which flows backwards). this trick means that forgetting startScope() would make the flag active all the way up to the top of the chain

In the Publisher-based approach:

• if the flag lives on an arbitrary parent Publisher, then it cannot be visible from inner publishers like eg. in a flatMap (which might make this easier to reason about)
• the config flag should be inherited from the parent, or visible in the parent, which implies it is always retained until Publisher#subscribe
• if as little as one operator which doesn't support exposing this flag (or its parent's flag) is injected in the chain, it can break the whole thing
• one needs to be careful not to blindly remove the flag towards the end of the scope: what if we're in a nested scope (which is possible since we don't constraint which operators can be called within the scope)?
• forgetting the endScope() means the rest of the chain sees the flag, but that's rather more acceptable than when it leaks backwards

[simonbasle]

(3) Sub-API approach

Expected API

???

Caveat

Need to come up with a signature that the type system understands rather than ending up with something that users will need to massage into the type system.

Consumer<Spec<T, R>> approach doesn't work because of this. the compiler can't know which pseudo-operators the user will invoke, so it cannot know in advance the output type <R> of the spec.

🚫 This is a blocker, I wasn't able to find an elegant way to supporting this with anything that changes the return type (which includes map, flatMap, actually most operators but side effects and filter...).

Benefits

1. constrained: we'd only expose equivalents to operators that we want to support, and users cannot invoke any other operator within the "scope"
2. best compromise between explicit and implicit: users are explicit about the operators they want to be context aware, but the wrapping of their functions is implicit (they don't need to call any wrapper method and can focus on building the chain)

[simonbasle]

(4) Implicit restoration support in a couple of selected operators

Focusing eg. on tap and handle could be a way to offer 0-config option.

Explored in #3180

Caveat

• Very limited set of operators supporting this: how to communicate to the users that no, we won't add implicit support into more operators. (this is something that comes up regularly with the context-exposing variants: "can I have myFavoriteOperator overload with an exposed ContextView?")
• arbitrary support: within the scope there's no API restriction YET some operators used there might simply not support the feature. there's no way to warn users other than via documentation

[chemicL]

This has been completed with #3180 and later evolved into automatic context propagation with #3335 and following improvements.