GH-3869: Add ContextHolderRequestHandlerAdvice

Fixes #3869

* Move `ContextHolderRequestHandlerAdvice` to the `core` module for more general purposes
* Add `ContextHolderRequestHandlerAdviceTests`
* Rework `DelegatingSessionFactoryTests` to rely on the `ContextHolderRequestHandlerAdvice`.
This allows us to remove unnecessary XML configuration for this test class
* Document the feature

Author: adel-haidar

spring-integration-core/src/main/java/org/springframework/integration/handler/advice/ContextHolderRequestHandlerAdvice.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright 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.
+ * 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 org.springframework.integration.handler.advice;
+
+import java.util.function.Consumer;
+import java.util.function.Function;
+
+import org.springframework.messaging.Message;
+import org.springframework.util.Assert;
+
+/**
+ * An {@link AbstractRequestHandlerAdvice} implementation to store and reset
+ * a value into/from some context (e.g. {@link ThreadLocal}) against a request message.
+ * The context is populated before {@code callback.execute()} and reset after.
+ *
+ * @author Adel Haidar
+ * @author Artem Bilan
+ *
+ * @since 6.1
+ */
+public class ContextHolderRequestHandlerAdvice extends AbstractRequestHandlerAdvice {
+
+	public final Function<Message<?>, Object> valueProvider;
+
+	public final Consumer<Object> contextSetHook;
+
+	public final Runnable contextClearHook;
+
+	private boolean enableContextPropagation;
+
+	/**
+	 * Construct an instance based on the provided hooks.
+	 * @param valueProvider The key provider function.
+	 * @param contextSetHook The context set hook consumer.
+	 * @param contextClearHook The context clear hook consumer.
+	 */
+	public ContextHolderRequestHandlerAdvice(Function<Message<?>, Object> valueProvider,
+			Consumer<Object> contextSetHook, Runnable contextClearHook) {
+
+		Assert.notNull(valueProvider, "'valueProvider' must not be null");
+		Assert.notNull(contextSetHook, "'contextSetHook' must not be null");
+		Assert.notNull(contextClearHook, "'contextClearHook' must not be null");
+		this.valueProvider = valueProvider;
+		this.contextSetHook = contextSetHook;
+		this.contextClearHook = contextClearHook;
+	}
+
+	@Override
+	protected Object doInvoke(ExecutionCallback callback, Object target, Message<?> message) {
+		Object value = this.valueProvider.apply(message);
+		logger.trace(() -> "Setting context value to: " + value + " from message: " + message);
+		try {
+			this.contextSetHook.accept(value);
+			return callback.execute();
+		}
+		finally {
+			this.contextClearHook.run();
+		}
+	}
+
+}

spring-integration-core/src/test/java/org/springframework/integration/handler/advice/ContextHolderRequestHandlerAdviceTests.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright 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.
+ * 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 org.springframework.integration.handler.advice;
+
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.junit.jupiter.api.Test;
+
+import org.springframework.aop.framework.ProxyFactoryBean;
+import org.springframework.messaging.MessageHandler;
+import org.springframework.messaging.support.GenericMessage;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * @author Artem Bilan
+ *
+ * @since 6.1
+ */
+public class ContextHolderRequestHandlerAdviceTests {
+
+	@Test
+	void contextHolderRequestHandlerAdviceInAction() {
+		AtomicReference<Object> context = new AtomicReference<>();
+
+		AtomicReference<Object> valueFromHandler = new AtomicReference<>();
+		MessageHandler testHandler = message -> valueFromHandler.set(context.get());
+
+		String testContextValue = "test data";
+
+		ContextHolderRequestHandlerAdvice contextHolderRequestHandlerAdvice =
+				new ContextHolderRequestHandlerAdvice(m -> testContextValue, context::set, () -> context.set(null));
+
+		ProxyFactoryBean fb = new ProxyFactoryBean();
+		fb.setTarget(testHandler);
+		fb.addAdvice(contextHolderRequestHandlerAdvice);
+		testHandler = (MessageHandler) fb.getObject();
+
+		testHandler.handleMessage(new GenericMessage<>(""));
+
+		assertThat(valueFromHandler.get()).isEqualTo(testContextValue);
+		assertThat(context.get()).isNull();
+	}
+
+}

spring-integration-file/src/main/java/org/springframework/integration/file/remote/session/DelegatingSessionFactory.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2015-2021 the original author or authors.
+ * Copyright 2015-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.
@@ -28,6 +28,7 @@
  * @param <F> the target system file type.
  *
  * @author Gary Russell
+ *
  * @since 4.2
  *
  */
@@ -52,7 +53,7 @@ public DelegatingSessionFactory(Map<Object, SessionFactory<F>> factories, Sessio
 	 * @param factoryLocator the factory.
 	 */
 	public DelegatingSessionFactory(SessionFactoryLocator<F> factoryLocator) {
-		Assert.notNull(factoryLocator, "'factoryFactory' cannot be null");
+		Assert.notNull(factoryLocator, "'factoryLocator' cannot be null");
 		this.factoryLocator = factoryLocator;
 	}
 

spring-integration-file/src/test/java/org/springframework/integration/file/remote/session/DelegatingSessionFactoryTests.java

@@ -26,11 +26,12 @@
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
-import org.springframework.context.annotation.ImportResource;
 import org.springframework.integration.annotation.ServiceActivator;
+import org.springframework.integration.channel.QueueChannel;
 import org.springframework.integration.config.EnableIntegration;
 import org.springframework.integration.file.remote.AbstractFileInfo;
 import org.springframework.integration.file.remote.gateway.AbstractRemoteFileOutboundGateway;
+import org.springframework.integration.handler.advice.ContextHolderRequestHandlerAdvice;
 import org.springframework.integration.test.util.TestUtils;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageChannel;
@@ -102,8 +103,6 @@ public void testFlow() throws Exception {
 	}
 
 	@Configuration
-	@ImportResource(
-			"classpath:/org/springframework/integration/file/remote/session/delegating-session-factory-context.xml")
 	@EnableIntegration
 	public static class Config {
 
@@ -132,11 +131,21 @@ public SessionFactoryLocator<String> sessionFactoryLocator() {
 			return new DefaultSessionFactoryLocator<>(factories, bar);
 		}
 
-		@ServiceActivator(inputChannel = "c1")
+		@Bean
+		QueueChannel out() {
+			return new QueueChannel();
+		}
+
+		@Bean
+		ContextHolderRequestHandlerAdvice contextHolderRequestHandlerAdvice(DelegatingSessionFactory<String> dsf) {
+			return new ContextHolderRequestHandlerAdvice(Message::getPayload, dsf::setThreadKey, dsf::clearThreadKey);
+		}
+
+		@ServiceActivator(inputChannel = "in", adviceChain = "contextHolderRequestHandlerAdvice")
 		@Bean
 		MessageHandler handler() {
 			AbstractRemoteFileOutboundGateway<String> gateway =
-					new AbstractRemoteFileOutboundGateway<String>(dsf(), "ls", "payload") {
+					new AbstractRemoteFileOutboundGateway<>(dsf(), "ls", "payload") {
 
 						@Override
 						protected boolean isDirectory(String file) {
@@ -172,8 +181,9 @@ protected List<AbstractFileInfo<String>> asFileInfoList(Collection<String> files
 						protected String enhanceNameWithSubDirectory(String file, String directory) {
 							return null;
 						}
+
 					};
-			gateway.setOutputChannelName("c2");
+			gateway.setOutputChannelName("out");
 			gateway.setOptions("-1");
 			return gateway;
 		}

spring-integration-file/src/test/java/org/springframework/integration/file/remote/session/delegating-session-factory-context.xml

@@ -231,33 +231,10 @@ private static final class SharedSSLFTPSClient extends FTPSClient {
 
 Version 4.2 introduced the `DelegatingSessionFactory`, which allows the selection of the actual session factory at runtime.
 Prior to invoking the FTP endpoint, call `setThreadKey()` on the factory to associate a key with the current thread.
-That key is then used to lookup the actual session factory to be used.
+That key is then used to look up the actual session factory to be used.
 You can clear the key by calling `clearThreadKey()` after use.
 
-We added convenience methods so that you can easily do use a delegating session factory from a message flow.
-
-The following example shows how to declare a delegating session factory:
-
-====
-[source, xml]
-----
-<bean id="dsf" class="org.springframework.integration.file.remote.session.DelegatingSessionFactory">
-    <constructor-arg>
-        <bean class="o.s.i.file.remote.session.DefaultSessionFactoryLocator">
-            <!-- delegate factories here -->
-        </bean>
-    </constructor-arg>
-</bean>
-
-<int:service-activator input-channel="in" output-channel="c1"
-        expression="@dsf.setThreadKey(#root, headers['factoryToUse'])" />
-
-<int-ftp:outbound-gateway request-channel="c1" reply-channel="c2" ... />
-
-<int:service-activator input-channel="c2" output-channel="out"
-        expression="@dsf.clearThreadKey(#root)" />
-----
-====
+See <<./handler-advice.adoc#context-holder-advice, Context Holder Advice>> for more information how this factory can be used together with a `ContextHolderRequestHandlerAdvice`.
 
 IMPORTANT: When you use session caching (see <<ftp-session-caching>>), each of the delegates should be cached.
 You cannot cache the `DelegatingSessionFactory` itself.

src/reference/asciidoc/ftp.adoc

@@ -57,6 +57,9 @@ In addition to providing the general mechanism to apply AOP advice classes, Spri
 * `RateLimiterRequestHandlerAdvice` (described in <<rate-limiter-advice>>)
 * `CacheRequestHandlerAdvice` (described in <<cache-advice>>)
 * `ReactiveRequestHandlerAdvice` (described in <<reactive-advice>>)
+* `ContextHolderRequestHandlerAdvice` (described in <<context-holder-advice>>)
+
+[[expression-advice]]
 
 [[retry-advice]]
 ===== Retry Advice
@@ -571,6 +574,44 @@ The `message` argument is the request message for the message handler and can be
 The `mono` argument is the result of this message handler's `handleRequestMessage()` method implementation.
 A nested `Mono.transform()` can also be called from this function to apply, for example, a https://spring.io/projects/spring-cloud-circuitbreaker[Reactive Circuit Breaker].
 
+[[context-holder-advice]]
+==== Context Holder Advice
+
+Starting with version 6.1, the `ContextHolderRequestHandlerAdvice` has been introduced.
+This advice takes some value from the request message as and stores it in the context holder.
+The value is clear from the context when an execution is finished on the target `MessageHandler`.
+The best way to think about this advice is similar to the programming flow where we store some value into a `ThreadLocal`, get access to it from the target call and then clean up the `ThreadLocal` after execution.
+The `ContextHolderRequestHandlerAdvice` requires these constructor arguments: a `Function<Message<?>, Object>` as a value provider, `Consumer<Object>` as a context set callback and `Runnable` as a context clean up hook.
+
+Following is a sample how a `ContextHolderRequestHandlerAdvice` can be used in combination with a `o.s.i.file.remote.session.DelegatingSessionFactory`:
+
+====
+[source, java]
+----
+@Bean
+DelegatingSessionFactory<?> dsf(SessionFactory<?> one, SessionFactory<?> two) {
+    return new DelegatingSessionFactory<>(Map.of("one", one, "two", two), null);
+}
+
+@Bean
+ContextHolderRequestHandlerAdvice contextHolderRequestHandlerAdvice(DelegatingSessionFactory<String> dsf) {
+    return new ContextHolderRequestHandlerAdvice(message -> message.getHeaders().get("FACTORY_KEY"),
+                                      dsf::setThreadKey, dsf::clearThreadKey);
+}
+
+@ServiceActivator(inputChannel = "in", adviceChain = "contextHolderRequestHandlerAdvice")
+FtpOutboundGateway ftpOutboundGateway(DelegatingSessionFactory<?> sessionFactory) {
+	return new FtpOutboundGateway(sessionFactory, "ls", "payload");
+}
+----
+====
+
+And it is just enough to send a message to the `in` channel with a `FACTORY_KEY` header set to either `one` or `two`.
+The `ContextHolderRequestHandlerAdvice` sets the value from that header into a `DelegatingSessionFactory` via its `setThreadKey`.
+Then when `FtpOutboundGateway` executes an `ls` command a proper delegating `SessionFactory` is chosen from the `DelegatingSessionFactory` according to the value in its `ThreadLocal`.
+When the result is produced from the `FtpOutboundGateway`, a `ThreadLocal` value in the `DelegatingSessionFactory` is cleared according to the `clearThreadKey()` call from the `ContextHolderRequestHandlerAdvice`.
+See <<./ftp.adoc#ftp-dsf,Delegating Session Factory>> for more information.
+
 [[custom-advice]]
 ==== Custom Advice Classes
 

src/reference/asciidoc/handler-advice.adoc

@@ -23,6 +23,11 @@ In general the project has been moved to the latest dependency versions.
 The Zip Spring Integration Extension project has been migrated as the `spring-integration-zip` module.
 See <<./zip.adoc#zip,Zip Support>>  for more information.
 
+[[x6.1-context-holder-advice]]
+==== `ContextHolderRequestHandlerAdvice`
+
+The `ContextHolderRequestHandlerAdvice` allows to store a value from a request message into some context around `MessageHandler` execution.
+See <<./handler-advice.adoc#context-holder-advice, Context Holder Advice>> for more information.
 
 [[x6.1-general]]
 === General Changes