Add ContextPropagation runtime util + captureContext operator (reactor#3145)

This commit introduces a `ContextPropagation` runtime-detection utility
class as well as a `contextCapture()` operator.

If context-propagation isn't on the classpath, this operator is NO-OP.

If context-propagation is on the classpath however, the operator will
use it to capture relevant ThreadLocals during the subscription phase
(at the point a contextWrite would be effected) and store it in the
ContextView visible from upstream of the operator.

Co-authored-by: Rossen Stoyanchev <[email protected]>

Author: simonbasle

docs/asciidoc/advanced-contextPropagation.adoc

@@ -0,0 +1,40 @@
+[[context.propagation]]
+= Context-Propagation Support
+
+Since 3.5.0, Reactor-Core embeds support for the `io.micrometer:context-propagation` SPI.
+This library is intended as a means to easily adapt between various implementations of the concept of a Context, of which
+`ContextView`/`Context` is an example, and between `ThreadLocal` variables as well.
+
+`ReactorContextAccessor` allows the Context-Propagation library to understand Reactor `Context` and `Contextview`.
+It implements the SPI and is loaded via `java.util.ServiceLoader`.
+No user action is required, other than having a dependency on both reactor-core and `io.micrometer:context-propagation`. The `ReactorContextAccessor` class is public but shouldn't generally be accessed by user code.
+
+On top of that, Reactor-Core 3.5.0 also introduces the `contextCapture` operator that transparently deals with `ContextSnapshot`s if the library is available at runtime, for users' convenience.
+
+== `contextCapture` Operator
+This operator can be used when one needs to capture `ThreadLocal` value(s) at subscription time and reflect these values in the Reactor `Context` for the benefit of upstream operators.
+It relies on the `context-propagation` library and notably the registered `ThreadLocalAccessor`(s) to discover relevant ThreadLocal values.
+
+This is a convenient alternative to `contextWrite` which uses the `context-propagation` API to obtain a `ContextSnapshot` and then uses that snapshot to populate the Reactor `Context`.
+
+As a result, if there were any ThreadLocal values during subscription phase, for which there is a registered `ThreadLocalAccessor`, their values would now be stored in the Reactor `Context` and visible
+at runtime in upstream operators.
+
+====
+[source,java]
+----
+//assuming TL is known to Context-Propagation as key TLKEY.
+static final ThreadLocal<String> TL = new ThreadLocal<>();
+
+//in the main thread, TL is set to "HELLO"
+TL.set("HELLO");
+
+Mono.deferContextual(ctx ->
+  Mono.delay(Duration.ofSeconds(1))
+      //we're now in another thread, TL is not set
+      .map(v -> "delayed ctx[" + TLKEY + "]=" + ctx.getOrDefault(TLKEY, "not found") + ", TL=" + TL.get())
+)
+.contextCapture()
+.block(); // returns "delayed ctx[TLKEY]=HELLO, TL=null"
+----
+====

docs/asciidoc/advancedFeatures.adoc

@@ -11,6 +11,7 @@ This chapter covers advanced features and concepts of Reactor, including the fol
 * <<scheduler-factory>>
 * <<hooks>>
 * <<context>>
+* <<context.propagation>>
 * <<null-safety>>
 * <<cleanup>>
 
@@ -841,15 +842,6 @@ TIP: In order to read from the `Context` without misleading users into thinking
 while data is running through the pipeline, only the `ContextView` is exposed by the operators above.
 In case one needs to use one of the remaining APIs that still require a `Context`, one can use `Context.of(contextView)` for conversion.
 
-[[context.propagation]]
-=== Micrometer Context-Propagation Support
-Since 3.5.0, Reactor-Core embeds basic support for the `io.micrometer:context-propagation` SPI.
-This library is intended as a mean to easily adapt between various implementations of the concept of a Context, of which
-`ContextView`/`Context` is an example, and between `ThreadLocal` variables as well.
-
-`ReactorContextAccessor` is one implementation of this SPI that is loaded via `ServiceLoader`. No user action is required,
-other than depending on reactor-core and `context-propagation` (the class is public but shouldn't generally be accessed by user code).
-
 === Simple `Context` Examples
 
 The examples in this section are meant as ways to better understand some of the caveats of
@@ -1088,6 +1080,8 @@ public void contextForLibraryReactivePut() {
 ----
 ====
 
+include::advanced-contextPropagation.adoc[leveloffset=1]
+
 [[cleanup]]
 == Dealing with Objects that Need Cleanup
 

reactor-core/src/main/java/reactor/core/publisher/ContextPropagation.java

@@ -0,0 +1,120 @@
+/*
+ * Copyright (c) 2022 VMware Inc. or its affiliates, All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package reactor.core.publisher;
+
+import java.util.function.Function;
+import java.util.function.Predicate;
+
+import io.micrometer.context.ContextRegistry;
+import io.micrometer.context.ContextSnapshot;
+
+import reactor.util.annotation.Nullable;
+import reactor.util.context.Context;
+
+/**
+ * Utility private class to detect if the <a href="https://github.com/micrometer-metrics/context-propagation">context-propagation library</a> is on the classpath and to offer
+ * ContextSnapshot support to {@link Flux} and {@link Mono}.
+ *
+ * @author Simon Baslé
+ */
+final class ContextPropagation {
+
+	static final boolean isContextPropagationAvailable;
+
+	static final Predicate<Object> PREDICATE_TRUE = v -> true;
+	static final Function<Context, Context> NO_OP = c -> c;
+	static final Function<Context, Context> WITH_GLOBAL_REGISTRY_NO_PREDICATE;
+
+	static {
+		boolean contextPropagation;
+		try {
+			Class.forName("io.micrometer.context.ContextRegistry", false, ContextPropagation.class.getClassLoader());
+			contextPropagation = true;
+		}
+		catch (Throwable t) {
+			contextPropagation = false;
+		}
+		isContextPropagationAvailable = contextPropagation;
+		if (contextPropagation) {
+			WITH_GLOBAL_REGISTRY_NO_PREDICATE = new ContextCaptureFunction(PREDICATE_TRUE, ContextRegistry.getInstance());
+		}
+		else {
+			WITH_GLOBAL_REGISTRY_NO_PREDICATE = NO_OP;
+		}
+	}
+
+	/**
+	 * Is Micrometer {@code context-propagation} API on the classpath?
+	 *
+	 * @return true if context-propagation is available at runtime, false otherwise
+	 */
+	static boolean isContextPropagationAvailable() {
+		return isContextPropagationAvailable;
+	}
+
+	/**
+	 * Create a support function that takes a snapshot of thread locals and merges them with the
+	 * provided {@link Context}, resulting in a new {@link Context} which includes entries
+	 * captured from threadLocals by the Context-Propagation API.
+	 *
+	 * @return the {@link Context} augmented with captured entries
+	 */
+	public static Function<Context, Context> contextCapture() {
+		if (!isContextPropagationAvailable) {
+			return NO_OP;
+		}
+		return WITH_GLOBAL_REGISTRY_NO_PREDICATE;
+	}
+
+	/**
+	 * Create a support function that takes a snapshot of thread locals and merges them with the
+	 * provided {@link Context}, resulting in a new {@link Context} which includes entries
+	 * captured from threadLocals by the Context-Propagation API.
+	 * <p>
+	 * The provided {@link Predicate} is used on keys associated to said thread locals
+	 * by the Context-Propagation API to filter which entries should be captured in the
+	 * first place.
+	 *
+	 * @param captureKeyPredicate a {@link Predicate} used on keys to determine if each entry
+	 * should be injected into the new {@link Context}
+	 * @return a {@link Function} augmenting {@link Context} with captured entries
+	 */
+	public static Function<Context, Context> contextCapture(Predicate<Object> captureKeyPredicate) {
+		if (!isContextPropagationAvailable) {
+			return NO_OP;
+		}
+		return new ContextCaptureFunction(captureKeyPredicate, null);
+	}
+
+	//the Function indirection allows tests to directly assert code in this class rather than static methods
+	static final class ContextCaptureFunction implements Function<Context, Context> {
+
+		final Predicate<Object> capturePredicate;
+		final ContextRegistry registry;
+
+		ContextCaptureFunction(Predicate<Object> capturePredicate, @Nullable ContextRegistry registry) {
+			this.capturePredicate = capturePredicate;
+			this.registry = registry != null ? registry : ContextRegistry.getInstance();
+		}
+
+		@Override
+		public Context apply(Context target) {
+			return ContextSnapshot.captureAllUsing(capturePredicate, this.registry).updateContext(target);
+		}
+	}
+
+}

reactor-core/src/main/java/reactor/core/publisher/Flux.java

@@ -4029,6 +4029,27 @@ public final Flux<T> concatWith(Publisher<? extends T> other) {
 		return concat(this, other);
 	}
 
+	/**
+	 * If <a href="https://github.com/micrometer-metrics/context-propagation">context-propagation library</a>
+	 * is on the classpath, this is a convenience shortcut to capture thread local values during the
+	 * subscription phase and put them in the {@link Context} that is visible upstream of this operator.
+	 * <p>
+	 * As a result this operator should generally be used as close as possible to the end of
+	 * the chain / subscription point.
+	 * <p>
+	 * If context-propagation is not available at runtime, this operator simply returns the current {@link Flux}
+	 * instance.
+	 *
+	 * @return a new {@link Flux} where context-propagation API has been used to capture entries and
+	 * inject them into the {@link Context}
+	 */
+	public final Flux<T> contextCapture() {
+		if (!ContextPropagation.isContextPropagationAvailable()) {
+			return this;
+		}
+		return onAssembly(new FluxContextWrite<>(this, ContextPropagation.contextCapture()));
+	}
+
 	/**
 	 * Enrich the {@link Context} visible from downstream for the benefit of upstream
 	 * operators, by making all values from the provided {@link ContextView} visible on top

reactor-core/src/main/java/reactor/core/publisher/Mono.java

@@ -2227,6 +2227,27 @@ public final Flux<T> concatWith(Publisher<? extends T> other) {
 		return Flux.concat(this, other);
 	}
 
+	/**
+	 * If <a href="https://github.com/micrometer-metrics/context-propagation">context-propagation library</a>
+	 * is on the classpath, this is a convenience shortcut to capture thread local values during the
+	 * subscription phase and put them in the {@link Context} that is visible upstream of this operator.
+	 * <p>
+	 * As a result this operator should generally be used as close as possible to the end of
+	 * the chain / subscription point.
+	 * <p>
+	 * If context-propagation is not available at runtime, this operator simply returns the current {@link Mono}
+	 * instance.
+	 *
+	 * @return a new {@link Flux} where context-propagation API has been used to capture entries and
+	 * inject them into the {@link Context}
+	 */
+	public final Mono<T> contextCapture() {
+		if (!ContextPropagation.isContextPropagationAvailable()) {
+			return this;
+		}
+		return onAssembly(new MonoContextWrite<>(this, ContextPropagation.contextCapture()));
+	}
+
 	/**
 	 * Enrich the {@link Context} visible from downstream for the benefit of upstream
 	 * operators, by making all values from the provided {@link ContextView} visible on top

reactor-core/src/test/java/reactor/core/publisher/ContextPropagationNotThereSmokeTest.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) 2022 VMware Inc. or its affiliates, All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package reactor.core.publisher;
+
+import java.util.function.Predicate;
+
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * For tests that actually assert things when context-propagation is available, see
+ * {@code withMicrometer} test set.
+ * @author Simon Baslé
+ */
+class ContextPropagationNotThereSmokeTest {
+
+	@Test
+	void contextPropagationIsNotAvailable() {
+		assertThat(ContextPropagation.isContextPropagationAvailable()).isFalse();
+	}
+
+	@Test
+	void contextCaptureIsNoOp() {
+		assertThat(ContextPropagation.contextCapture()).as("without predicate").isSameAs(ContextPropagation.NO_OP);
+		assertThat(ContextPropagation.contextCapture(v -> true)).as("with predicate").isSameAs(ContextPropagation.NO_OP);
+	}
+
+	@Test
+	void contextCaptureFluxApiIsNoOp() {
+		Flux<Integer> source = Flux.empty();
+		assertThat(source.contextCapture()).isSameAs(source);
+	}
+
+	@Test
+	void contextCaptureMonoApiIsNoOp() {
+		Mono<Integer> source = Mono.empty();
+		assertThat(source.contextCapture()).isSameAs(source);
+	}
+
+}