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 <grussell@vmware.com>

---------

Co-authored-by: Gary Russell <grussell@vmware.com>

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.Set;
 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 class IntegrationMessageHeaderAccessor extends MessageHeaderAccessor {
 		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.io.Closeable;
 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 final class StaticMessageHeaderAccessor {
 		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.publisher.Sinks;
 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 class FluxMessageChannel extends AbstractMessageChannel
 				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;