Propagate Reactor context over headers (#8591)

* Propagate Reactor context over headers

When we do something like `Flux.from(Publisher)`
and don't compose it with the one involved in the `Subscriber` context,
we lose this context.

* Provide a mechanism to propagate a Reactor context over message header
produce within that context.
* Restore this context in the `FluxMessageChannel` for a new publisher
we use in this channel

**Cherry-pick to `6.0.x`**

* Fix language in docs

Co-authored-by: Gary Russell <[email protected]>

---------

Co-authored-by: Gary Russell <[email protected]>

Author: artembilan

spring-integration-core/src/main/java/org/springframework/integration/IntegrationMessageHeaderAccessor.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2019 the original author or authors.
+ * Copyright 2013-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -25,6 +25,8 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiFunction;
 
+import reactor.util.context.ContextView;
+
 import org.springframework.integration.acks.AcknowledgmentCallback;
 import org.springframework.lang.Nullable;
 import org.springframework.messaging.Message;
@@ -77,6 +79,12 @@ public class IntegrationMessageHeaderAccessor extends MessageHeaderAccessor {
 	 */
 	public static final String SOURCE_DATA = "sourceData";
 
+	/**
+	 * Raw source message.
+	 */
+	public static final String REACTOR_CONTEXT = "reactorContext";
+
+
 	private static final BiFunction<String, String, String> TYPE_VERIFY_MESSAGE_FUNCTION =
 			(name, trailer) -> "The '" + name + trailer;
 
@@ -175,6 +183,16 @@ public <T> T getSourceData() {
 		return (T) getHeader(SOURCE_DATA);
 	}
 
+	/**
+	 * Get a {@link ContextView} header if present.
+	 * @return the {@link ContextView} header if present.
+	 * @since 6.0.5
+	 */
+	@Nullable
+	public ContextView getReactorContext() {
+		return getHeader(REACTOR_CONTEXT, ContextView.class);
+	}
+
 	@SuppressWarnings("unchecked")
 	@Nullable
 	public <T> T getHeader(String key, Class<T> type) {

spring-integration-core/src/main/java/org/springframework/integration/StaticMessageHeaderAccessor.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2017-2020 the original author or authors.
+ * Copyright 2017-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -20,6 +20,9 @@
 import java.util.UUID;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import reactor.util.context.Context;
+import reactor.util.context.ContextView;
+
 import org.springframework.integration.acks.AcknowledgmentCallback;
 import org.springframework.integration.acks.SimpleAcknowledgment;
 import org.springframework.lang.Nullable;
@@ -120,4 +123,19 @@ public static <T> T getSourceData(Message<?> message) {
 		return (T) message.getHeaders().get(IntegrationMessageHeaderAccessor.SOURCE_DATA);
 	}
 
+	/**
+	 * Get a {@link ContextView} header if present.
+	 * @param message the message to get a header from.
+	 * @return the {@link ContextView} header if present.
+	 * @since 6.0.5
+	 */
+	public static ContextView getReactorContext(Message<?> message) {
+		ContextView reactorContext = message.getHeaders()
+				.get(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT, ContextView.class);
+		if (reactorContext == null) {
+			reactorContext = Context.empty();
+		}
+		return reactorContext;
+	}
+
 }

spring-integration-core/src/main/java/org/springframework/integration/channel/FluxMessageChannel.java

@@ -30,6 +30,10 @@
 import reactor.core.scheduler.Scheduler;
 import reactor.core.scheduler.Schedulers;
 
+import org.springframework.core.log.LogMessage;
+import org.springframework.integration.IntegrationMessageHeaderAccessor;
+import org.springframework.integration.StaticMessageHeaderAccessor;
+import org.springframework.integration.support.MessageBuilder;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageDeliveryException;
 import org.springframework.util.Assert;
@@ -111,22 +115,35 @@ public void subscribeTo(Publisher<? extends Message<?>> publisher) {
 				Flux.from(publisher)
 						.delaySubscription(this.subscribedSignal.asFlux().filter(Boolean::booleanValue).next())
 						.publishOn(this.scheduler)
-						.handle((message, synchronousSink) -> {
-							try {
-								if (!send(message)) {
-									logger.warn(new MessageDeliveryException(message,
-											"Failed to send message to channel '" + this),
-											"Message was not delivered");
-								}
-							}
-							catch (Exception ex) {
-								logger.warn(ex, () -> "Error during processing event: " + message);
-							}
-						})
+						.flatMap((message) ->
+								Mono.just(message)
+										.handle((messageToHandle, sink) -> sendReactiveMessage(messageToHandle))
+										.contextWrite(StaticMessageHeaderAccessor.getReactorContext(message)))
 						.contextCapture()
 						.subscribe());
 	}
 
+	private void sendReactiveMessage(Message<?> message) {
+		Message<?> messageToSend = message;
+		// We have just restored Reactor context, so no need in a header anymore.
+		if (messageToSend.getHeaders().containsKey(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT)) {
+			messageToSend =
+					MessageBuilder.fromMessage(message)
+							.removeHeader(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT)
+							.build();
+		}
+		try {
+			if (!send(messageToSend)) {
+				logger.warn(
+						new MessageDeliveryException(messageToSend, "Failed to send message to channel '" + this),
+						"Message was not delivered");
+			}
+		}
+		catch (Exception ex) {
+			logger.warn(ex, LogMessage.format("Error during processing event: %s", messageToSend));
+		}
+	}
+
 	@Override
 	public void destroy() {
 		this.active = false;

spring-integration-rsocket/src/main/java/org/springframework/integration/rsocket/inbound/RSocketInboundGateway.java

@@ -22,13 +22,15 @@
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.util.context.ContextView;
 
 import org.springframework.core.ReactiveAdapter;
 import org.springframework.core.ResolvableType;
 import org.springframework.core.codec.Decoder;
 import org.springframework.core.codec.Encoder;
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
+import org.springframework.integration.IntegrationMessageHeaderAccessor;
 import org.springframework.integration.gateway.MessagingGatewaySupport;
 import org.springframework.integration.rsocket.AbstractRSocketConnector;
 import org.springframework.integration.rsocket.ClientRSocketConnector;
@@ -151,7 +153,7 @@ public void setRequestElementClass(Class<?> requestElementClass) {
 	/**
 	 * Specify the type of payload to be generated when the inbound RSocket request
 	 * content is read by the converters/encoders.
-	 * By default this value is null which means at runtime any "text" Content-Type will
+	 * By default, this value is null which means at runtime any "text" Content-Type will
 	 * result in String while all others default to {@code byte[].class}.
 	 * @param requestElementType The payload type.
 	 */
@@ -212,9 +214,22 @@ public Mono<Void> handleMessage(Message<?> requestMessage) {
 		}
 		else {
 			return requestMono
-					.doOnNext(this::send)
-					.then();
+					.flatMap((message) ->
+							Mono.deferContextual((context) ->
+									Mono.just(message)
+											.handle((messageToSend, sink) ->
+													send(messageWithReactorContextIfAny(messageToSend, context)))));
+		}
+	}
+
+	private Message<?> messageWithReactorContextIfAny(Message<?> message, ContextView context) {
+		if (!context.isEmpty()) {
+			return getMessageBuilderFactory()
+					.fromMessage(message)
+					.setHeader(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT, context)
+					.build();
 		}
+		return message;
 	}
 
 	private Mono<Message<?>> decodeRequestMessage(Message<?> requestMessage) {

spring-integration-webflux/src/main/java/org/springframework/integration/webflux/inbound/WebFluxInboundEndpoint.java

@@ -30,6 +30,7 @@
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.util.context.ContextView;
 import reactor.util.function.Tuple2;
 
 import org.springframework.core.ReactiveAdapter;
@@ -47,6 +48,7 @@
 import org.springframework.http.codec.ServerCodecConfigurer;
 import org.springframework.http.server.reactive.ServerHttpRequest;
 import org.springframework.http.server.reactive.ServerHttpResponse;
+import org.springframework.integration.IntegrationMessageHeaderAccessor;
 import org.springframework.integration.expression.ExpressionEvalMap;
 import org.springframework.integration.http.HttpHeaders;
 import org.springframework.integration.http.inbound.BaseHttpInboundEndpoint;
@@ -99,7 +101,7 @@ public WebFluxInboundEndpoint(boolean expectReply) {
 
 	/**
 	 * A {@link ServerCodecConfigurer} for the request readers and response writers.
-	 * By default the {@link ServerCodecConfigurer#create()} factory is used.
+	 * By default, the {@link ServerCodecConfigurer#create()} factory is used.
 	 * @param codecConfigurer the {@link ServerCodecConfigurer} to use.
 	 */
 	public void setCodecConfigurer(ServerCodecConfigurer codecConfigurer) {
@@ -133,9 +135,9 @@ public String getComponentType() {
 
 	@Override
 	public Mono<Void> handle(ServerWebExchange exchange) {
-		return Mono.defer(() -> {
+		return Mono.deferContextual((context) -> {
 			if (isRunning()) {
-				return doHandle(exchange);
+				return doHandle(exchange, context);
 			}
 			else {
 				return Mono.error(new ResponseStatusException(HttpStatus.SERVICE_UNAVAILABLE, "Endpoint is stopped"))
@@ -144,13 +146,13 @@ public Mono<Void> handle(ServerWebExchange exchange) {
 		});
 	}
 
-	private Mono<Void> doHandle(ServerWebExchange exchange) {
+	private Mono<Void> doHandle(ServerWebExchange exchange, ContextView context) {
 		return extractRequestBody(exchange)
 				.doOnSubscribe(s -> this.activeCount.incrementAndGet())
 				.map(body ->
 						new RequestEntity<>(body, exchange.getRequest().getHeaders(),
 								exchange.getRequest().getMethod(), exchange.getRequest().getURI()))
-				.flatMap(entity -> buildMessage(entity, exchange))
+				.flatMap(entity -> buildMessage(entity, exchange, context))
 				.flatMap(requestTuple -> {
 					if (isExpectReply()) {
 						return sendAndReceiveMessageReactive(requestTuple.getT1())
@@ -253,7 +255,7 @@ private Mono<?> readRequestBody(ServerWebExchange exchange, MediaType contentTyp
 	}
 
 	private Mono<Tuple2<Message<Object>, RequestEntity<?>>> buildMessage(RequestEntity<?> httpEntity,
-			ServerWebExchange exchange) {
+			ServerWebExchange exchange, ContextView context) {
 
 		ServerHttpRequest request = exchange.getRequest();
 		MultiValueMap<String, String> requestParams = request.getQueryParams();
@@ -283,6 +285,10 @@ private Mono<Tuple2<Message<Object>, RequestEntity<?>>> buildMessage(RequestEnti
 		AbstractIntegrationMessageBuilder<Object> messageBuilder =
 				prepareRequestMessageBuilder(request, payload, headers);
 
+		if (!context.isEmpty()) {
+			messageBuilder.setHeader(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT, context);
+		}
+
 		return exchange.getPrincipal()
 				.map(principal -> messageBuilder.setHeader(HttpHeaders.USER_PRINCIPAL, principal))
 				.defaultIfEmpty(messageBuilder)

src/reference/asciidoc/reactive-streams.adoc

@@ -105,7 +105,7 @@ The subscription for this `Mono` is done using `Schedulers.boundedElastic()` to
 When the message source returns `null` (no data to pull), the `Mono` is turned into a `repeatWhenEmpty()` state with a `delay` for a subsequent re-subscription based on a `IntegrationReactiveUtils.DELAY_WHEN_EMPTY_KEY` `Duration` entry from the subscriber context.
 By default, it is 1 second.
 If the `MessageSource` produces messages with a `IntegrationMessageHeaderAccessor.ACKNOWLEDGMENT_CALLBACK` information in the headers, it is acknowledged (if necessary) in the `doOnSuccess()` of the original `Mono` and rejected in the `doOnError()` if the downstream flow throws a `MessagingException` with the failed message to reject.
-This `ReactiveMessageSourceProducer` could be used for any use-case when a a polling channel adapter's features should be turned into a reactive, on demand solution for any existing `MessageSource<?>` implementation.
+This `ReactiveMessageSourceProducer` could be used for any use-case when a polling channel adapter's features should be turned into a reactive, on demand solution for any existing `MessageSource<?>` implementation.
 
 === Splitter and Aggregator
 
@@ -332,3 +332,53 @@ Currently, Spring Integration provides channel adapter (or gateway) implementati
 The <<./redis.adoc#redis-stream-outbound,Redis Stream Channel Adapters>> are also reactive and uses `ReactiveStreamOperations` from Spring Data.
 More reactive channel adapters are coming, for example for Apache Kafka in <<./kafka.adoc#kafka,Kafka>> based on the `ReactiveKafkaProducerTemplate` and `ReactiveKafkaConsumerTemplate` from https://spring.io/projects/spring-kafka[Spring for Apache Kafka] etc.
 For many other non-reactive channel adapters thread pools are recommended to avoid blocking during reactive stream processing.
+
+[[context-propagation]]
+=== Reactive to Imperative Context Propagation
+
+When the https://github.com/micrometer-metrics/context-propagation[Context Propagation] library is on the classpath, the Project Reactor can take `ThreadLocal` values (e.g. https://micrometer.io/docs/observation[Micrometer Observation] or `SecurityContextHolder`) and store them into a `Subscriber` context.
+The opposite operation is also possible, when we need to populate a logging MDC for tracing or let services we call from the reactive stream to restore an observation from the scope.
+See more information in Project Reactor https://projectreactor.io/docs/core/release/reference/#context.propagation[documentation] about its special operators for context propagation.
+The storing and restoring context works smoothly if our whole solution is a single reactive stream composition since a `Subscriber` context is visible from downstream up to the beginning of the composition(`Flux` or `Mono`).
+But, if the application switches between different `Flux` instances or into imperative processing and back, then the context tied to the `Subscriber` might not be available.
+For such a use case, Spring Integration provides an additional capability (starting with version `6.0.5`) to store a Reactor `ContextView` into the `IntegrationMessageHeaderAccessor.REACTOR_CONTEXT` message header produced from the reactive stream, e.g. when we perform direct `send()` operation.
+This header is used then in the `FluxMessageChannel.subscribeTo()` to restore a Reactor context for the `Message` that this channel is going to emit.
+Currently, this header is populated from the `WebFluxInboundEndpoint` and `RSocketInboundGateway` components, but can be used in any solution where reactive to imperative integration is performed.
+The logic to populate this header is like this:
+
+====
+[source, java]
+----
+return requestMono
+        .flatMap((message) ->
+                Mono.deferContextual((context) ->
+                        Mono.just(message)
+                                .handle((messageToSend, sink) ->
+                                        send(messageWithReactorContextIfAny(messageToSend, context)))));
+...
+
+private Message<?> messageWithReactorContextIfAny(Message<?> message, ContextView context) {
+    if (!context.isEmpty()) {
+        return getMessageBuilderFactory()
+                .fromMessage(message)
+                .setHeader(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT, context)
+                .build();
+    }
+    return message;
+}
+----
+====
+
+Note, that we still need to use a `handle()` operator to make Reactor restore `ThreadLocal` values from the context.
+Even if it is sent as a header, the framework cannot make an assumption if it is going to be to restore onto `ThreadLocal` values downstream.
+
+To restore the context from a `Message` on the other `Flux` or `Mono` composition, this logic can be performed:
+
+====
+[source, java]
+----
+Mono.just(message)
+        .handle((messageToHandle, sink) -> ...)
+        .contextWrite(StaticMessageHeaderAccessor.getReactorContext(message)));
+----
+====