GH-9110: propagate Reactor context to ReactiveMessageHandler

Fixes: #9110

* Move context propagation utilities to the `IntegrationReactiveUtils`
* Capture context into message header in the `IntegrationReactiveUtils.adaptSubscribableChannelToPublisher()`
before `sink.tryEmitNext()`
* Restore the context from message header in the `flatMap()` for `ReactiveStreamsConsumer.reactiveMessageHandler`

**Auto-cherry-pick to `6.2.x`**

Author: artembilan

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

@@ -21,7 +21,6 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.locks.LockSupport;
 
-import io.micrometer.context.ContextSnapshotFactory;
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
 import reactor.core.Disposable;
@@ -31,17 +30,16 @@
 import reactor.core.publisher.Sinks;
 import reactor.core.scheduler.Scheduler;
 import reactor.core.scheduler.Schedulers;
-import reactor.util.context.Context;
 import reactor.util.context.ContextView;
 
 import org.springframework.core.log.LogMessage;
 import org.springframework.integration.IntegrationMessageHeaderAccessor;
 import org.springframework.integration.StaticMessageHeaderAccessor;
 import org.springframework.integration.support.MutableMessageBuilder;
+import org.springframework.integration.util.IntegrationReactiveUtils;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageDeliveryException;
 import org.springframework.util.Assert;
-import org.springframework.util.ClassUtils;
 
 /**
  * The {@link AbstractMessageChannel} implementation for the
@@ -56,9 +54,6 @@
 public class FluxMessageChannel extends AbstractMessageChannel
 		implements Publisher<Message<?>>, ReactiveStreamsSubscribableChannel {
 
-	private static final boolean isContextPropagationPresent = ClassUtils.isPresent(
-			"io.micrometer.context.ContextSnapshot", FluxMessageChannel.class.getClassLoader());
-
 	private final Scheduler scheduler = Schedulers.boundedElastic();
 
 	private final Sinks.Many<Message<?>> sink = Sinks.many().multicast().onBackpressureBuffer(1, false);
@@ -91,8 +86,8 @@ protected boolean doSend(Message<?> message, long timeout) {
 
 	private boolean tryEmitMessage(Message<?> message) {
 		Message<?> messageToEmit = message;
-		if (isContextPropagationPresent) {
-			ContextView contextView = ContextSnapshotHelper.captureContext();
+		if (IntegrationReactiveUtils.isContextPropagationPresent) {
+			ContextView contextView = IntegrationReactiveUtils.captureReactorContext();
 			if (!contextView.isEmpty()) {
 				messageToEmit = MutableMessageBuilder.fromMessage(message)
 						.setHeader(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT, contextView)
@@ -196,14 +191,4 @@ public void destroy() {
 		super.destroy();
 	}
 
-	private static final class ContextSnapshotHelper {
-
-		private static final ContextSnapshotFactory CONTEXT_SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build();
-
-		static ContextView captureContext() {
-			return CONTEXT_SNAPSHOT_FACTORY.captureAll().updateContext(Context.empty());
-		}
-
-	}
-
 }

spring-integration-core/src/main/java/org/springframework/integration/endpoint/ReactiveStreamsConsumer.java

@@ -19,7 +19,6 @@
 import java.util.function.Consumer;
 import java.util.function.Function;
 
-import io.micrometer.context.ContextSnapshotFactory;
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
@@ -31,6 +30,7 @@
 
 import org.springframework.context.Lifecycle;
 import org.springframework.integration.IntegrationMessageHeaderAccessor;
+import org.springframework.integration.StaticMessageHeaderAccessor;
 import org.springframework.integration.channel.ChannelUtils;
 import org.springframework.integration.channel.NullChannel;
 import org.springframework.integration.core.MessageProducer;
@@ -44,7 +44,6 @@
 import org.springframework.messaging.MessageHandler;
 import org.springframework.messaging.ReactiveMessageHandler;
 import org.springframework.util.Assert;
-import org.springframework.util.ClassUtils;
 import org.springframework.util.ErrorHandler;
 
 /**
@@ -58,9 +57,6 @@
  */
 public class ReactiveStreamsConsumer extends AbstractEndpoint implements IntegrationConsumer {
 
-	private static final boolean isContextPropagationPresent = ClassUtils.isPresent(
-			"io.micrometer.context.ContextSnapshot", ReactiveStreamsConsumer.class.getClassLoader());
-
 	private final MessageChannel inputChannel;
 
 	private final Publisher<Message<Object>> publisher;
@@ -188,7 +184,9 @@ protected void doStart() {
 		if (this.reactiveMessageHandler != null) {
 			this.subscription =
 					fluxFromChannel
-							.flatMap(this.reactiveMessageHandler::handleMessage)
+							.flatMap((message) ->
+									this.reactiveMessageHandler.handleMessage(message)
+											.contextWrite(StaticMessageHeaderAccessor.getReactorContext(message)))
 							.onErrorContinue((ex, data) -> this.errorHandler.handleError(ex))
 							.subscribe();
 		}
@@ -300,7 +298,7 @@ protected void hookOnSubscribe(Subscription subscription) {
 		protected void hookOnNext(Message<?> message) {
 			Message<?> messageToDeliver = message;
 
-			if (isContextPropagationPresent) {
+			if (IntegrationReactiveUtils.isContextPropagationPresent) {
 				ContextView reactorContext = message.getHeaders()
 						.get(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT, ContextView.class);
 
@@ -310,7 +308,7 @@ protected void hookOnNext(Message<?> message) {
 									.removeHeader(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT)
 									.build();
 
-					try (AutoCloseable scope = ContextSnapshotHelper.setContext(reactorContext)) {
+					try (AutoCloseable scope = IntegrationReactiveUtils.setThreadLocalsFromReactorContext(reactorContext)) {
 						this.delegate.onNext(messageToDeliver);
 					}
 					catch (Exception ex) {
@@ -335,14 +333,4 @@ protected void hookOnComplete() {
 
 	}
 
-	private static final class ContextSnapshotHelper {
-
-		private static final ContextSnapshotFactory CONTEXT_SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build();
-
-		static AutoCloseable setContext(ContextView context) {
-			return CONTEXT_SNAPSHOT_FACTORY.setThreadLocalsFrom(context);
-		}
-
-	}
-
 }

spring-integration-core/src/main/java/org/springframework/integration/util/IntegrationReactiveUtils.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2020-2022 the original author or authors.
+ * Copyright 2020-2024 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.
@@ -19,24 +19,30 @@
 import java.time.Duration;
 import java.util.concurrent.locks.LockSupport;
 
+import io.micrometer.context.ContextSnapshotFactory;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.Sinks;
 import reactor.core.scheduler.Schedulers;
+import reactor.util.context.Context;
+import reactor.util.context.ContextView;
 import reactor.util.retry.Retry;
 
+import org.springframework.integration.IntegrationMessageHeaderAccessor;
 import org.springframework.integration.StaticMessageHeaderAccessor;
 import org.springframework.integration.acks.AckUtils;
 import org.springframework.integration.core.MessageSource;
+import org.springframework.integration.support.MutableMessageBuilder;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageChannel;
 import org.springframework.messaging.MessageHandler;
 import org.springframework.messaging.MessagingException;
 import org.springframework.messaging.PollableChannel;
 import org.springframework.messaging.SubscribableChannel;
+import org.springframework.util.ClassUtils;
 
 /**
  * Utilities for adapting integration components to/from reactive types.
@@ -60,9 +66,40 @@ public final class IntegrationReactiveUtils {
 	 */
 	public static final Duration DEFAULT_DELAY_WHEN_EMPTY = Duration.ofSeconds(1);
 
+	/**
+	 * The indicator that {@code io.micrometer:context-propagation} library is on classpath.
+	 * @since 6.2.5
+	 */
+	public static final boolean isContextPropagationPresent = ClassUtils.isPresent(
+			"io.micrometer.context.ContextSnapshot", IntegrationReactiveUtils.class.getClassLoader());
+
+	private static final ContextSnapshotFactory CONTEXT_SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build();
+
 	private IntegrationReactiveUtils() {
 	}
 
+	/**
+	 * Capture a Reactor {@link ContextView} from the current thread local state
+	 * according to the {@link ContextSnapshotFactory} logic.
+	 * @return the Reactor {@link ContextView} from the current thread local state.
+	 * @since 6.2.5
+	 */
+	public static ContextView captureReactorContext() {
+		return CONTEXT_SNAPSHOT_FACTORY.captureAll().updateContext(Context.empty());
+	}
+
+	/**
+	 * Populate thread local variables from the provided Reactor {@link ContextView}
+	 * according to the {@link ContextSnapshotFactory} logic.
+	 * @param context the Reactor {@link ContextView} to populate from.
+	 * @return the {@link io.micrometer.context.ContextSnapshot.Scope} as a {@link AutoCloseable}
+	 * to not pollute the target classpath. Can be cast if necessary.
+	 * @since 6.2.5
+	 */
+	public static AutoCloseable setThreadLocalsFromReactorContext(ContextView context) {
+		return CONTEXT_SNAPSHOT_FACTORY.setThreadLocalsFrom(context);
+	}
+
 	/**
 	 * Wrap a provided {@link MessageSource} into a {@link Flux} for pulling the on demand.
 	 * When {@link MessageSource#receive()} returns {@code null}, the source {@link Mono}
@@ -137,8 +174,17 @@ private static <T> Flux<Message<T>> adaptSubscribableChannelToPublisher(Subscrib
 		return Flux.defer(() -> {
 			Sinks.Many<Message<T>> sink = Sinks.many().unicast().onBackpressureError();
 			MessageHandler messageHandler = (message) -> {
+				Message<?> messageToEmit = message;
+				if (IntegrationReactiveUtils.isContextPropagationPresent) {
+					ContextView contextView = IntegrationReactiveUtils.captureReactorContext();
+					if (!contextView.isEmpty()) {
+						messageToEmit = MutableMessageBuilder.fromMessage(message)
+								.setHeader(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT, contextView)
+								.build();
+					}
+				}
 				while (true) {
-					switch (sink.tryEmitNext((Message<T>) message)) {
+					switch (sink.tryEmitNext((Message<T>) messageToEmit)) {
 						case FAIL_NON_SERIALIZED:
 						case FAIL_OVERFLOW:
 							LockSupport.parkNanos(1000); // NOSONAR

spring-integration-core/src/test/java/org/springframework/integration/support/management/observation/IntegrationObservabilityZipkinTests.java

@@ -18,25 +18,31 @@
 
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 
 import io.micrometer.common.KeyValues;
 import io.micrometer.core.tck.MeterRegistryAssert;
+import io.micrometer.observation.Observation;
 import io.micrometer.observation.ObservationRegistry;
 import io.micrometer.tracing.Span;
 import io.micrometer.tracing.test.SampleTestRunner;
 import io.micrometer.tracing.test.simple.SpansAssert;
+import reactor.core.publisher.Mono;
 
 import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.integration.annotation.BridgeTo;
 import org.springframework.integration.annotation.EndpointId;
 import org.springframework.integration.annotation.Poller;
 import org.springframework.integration.annotation.ServiceActivator;
 import org.springframework.integration.channel.NullChannel;
+import org.springframework.integration.channel.PublishSubscribeChannel;
 import org.springframework.integration.channel.QueueChannel;
 import org.springframework.integration.config.EnableIntegration;
 import org.springframework.integration.config.EnableIntegrationManagement;
+import org.springframework.integration.dsl.IntegrationFlow;
 import org.springframework.integration.gateway.MessagingGatewaySupport;
 import org.springframework.integration.handler.BridgeHandler;
 import org.springframework.integration.handler.advice.HandleMessageAdvice;
@@ -46,6 +52,7 @@
 import org.springframework.messaging.support.GenericMessage;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.awaitility.Awaitility.await;
 
 /**
  * @author Artem Bilan
@@ -93,6 +100,8 @@ public SampleTestRunnerConsumer yourCode() {
 				assertThat(receive).isNull();
 
 				assertThat(configuration.observedHandlerLatch.await(10, TimeUnit.SECONDS)).isTrue();
+
+				await().untilAsserted(() -> assertThat(configuration.observationReference.get()).isNotNull());
 			}
 
 			SpansAssert.assertThat(bb.getFinishedSpans())
@@ -110,7 +119,7 @@ public SampleTestRunnerConsumer yourCode() {
 							.hasTag(IntegrationObservation.ProducerTags.COMPONENT_NAME.asString(), "queueChannel")
 							.hasTag(IntegrationObservation.ProducerTags.COMPONENT_TYPE.asString(), "producer")
 							.hasKindEqualTo(Span.Kind.PRODUCER))
-					.hasSize(3);
+					.hasSize(4);
 
 			MeterRegistryAssert.assertThat(getMeterRegistry())
 					.hasTimerWithNameAndTags("spring.integration.handler",
@@ -125,7 +134,7 @@ public SampleTestRunnerConsumer yourCode() {
 	@EnableIntegration
 	@EnableIntegrationManagement(
 			observationPatterns = {
-					"${spring.integration.management.observation-patterns:testInboundGateway,skippedObservationInboundGateway,queueChannel,observedEndpoint}",
+					"${spring.integration.management.observation-patterns:testInboundGateway,skippedObservationInboundGateway,queueChannel,observedEndpoint,publishSubscribeChannel}",
 					"${spring.integration.management.observation-patterns:}"
 			})
 	public static class ObservationIntegrationTestConfiguration {
@@ -166,8 +175,10 @@ TestMessagingGatewaySupport skippedObservationInboundGateway() {
 		@ServiceActivator(inputChannel = "queueChannel",
 				poller = @Poller(fixedDelay = "100"),
 				adviceChain = "observedHandlerAdvice")
-		BridgeHandler bridgeHandler() {
-			return new BridgeHandler();
+		BridgeHandler bridgeHandler(PublishSubscribeChannel publishSubscribeChannel) {
+			BridgeHandler bridgeHandler = new BridgeHandler();
+			bridgeHandler.setOutputChannel(publishSubscribeChannel);
+			return bridgeHandler;
 		}
 
 		@Bean
@@ -182,6 +193,26 @@ HandleMessageAdvice observedHandlerAdvice() {
 			};
 		}
 
+		@Bean
+		@BridgeTo
+		PublishSubscribeChannel publishSubscribeChannel() {
+			return new PublishSubscribeChannel();
+		}
+
+		AtomicReference<Observation> observationReference = new AtomicReference<>();
+
+		@Bean
+		IntegrationFlow handleReactiveFlow(PublishSubscribeChannel publishSubscribeChannel,
+				ObservationRegistry observationRegistry) {
+
+			return IntegrationFlow.from(publishSubscribeChannel)
+					.handleReactive(m ->
+							Mono.just("Hi There")
+									.doOnSuccess(val ->
+											observationReference.set(observationRegistry.getCurrentObservation()))
+									.then());
+		}
+
 	}
 
 	private static class TestMessagingGatewaySupport extends MessagingGatewaySupport {