INT-4280: Add taskScheduler() to Java DSL

JIRA: https://jira.spring.io/browse/INT-4280

To let end-user easily to inject custom `TaskScheduler` to the endpoint
add `ConsumerEndpointSpec.taskScheduler()` option

* Address PR comments

spring-integration-core/src/main/java/org/springframework/integration/config/ConsumerEndpointFactoryBean.java

@@ -49,6 +49,7 @@ import org.springframework.messaging.PollableChannel;
 import org.springframework.messaging.SubscribableChannel;
 import org.springframework.messaging.core.BeanFactoryMessageChannelDestinationResolver;
 import org.springframework.messaging.core.DestinationResolver;
+import org.springframework.scheduling.TaskScheduler;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
 import org.springframework.util.StringUtils;
@@ -99,6 +100,8 @@ public class ConsumerEndpointFactoryBean
 
 	private volatile DestinationResolver<MessageChannel> channelResolver;
 
+	private TaskScheduler taskScheduler;
+
 	public void setHandler(MessageHandler handler) {
 		Assert.notNull(handler, "handler must not be null");
 		synchronized (this.handlerMonitor) {
@@ -160,6 +163,10 @@ public class ConsumerEndpointFactoryBean
 		this.adviceChain = adviceChain;
 	}
 
+	public void setTaskScheduler(TaskScheduler taskScheduler) {
+		this.taskScheduler = taskScheduler;
+	}
+
 	@Override
 	public void afterPropertiesSet() throws Exception {
 		if (this.beanName == null) {
@@ -296,6 +303,9 @@ public class ConsumerEndpointFactoryBean
 				phase = Integer.MAX_VALUE / 2;
 			}
 			this.endpoint.setPhase(phase);
+			if (this.taskScheduler != null) {
+				this.endpoint.setTaskScheduler(this.taskScheduler);
+			}
 			this.endpoint.afterPropertiesSet();
 			this.initialized = true;
 		}

spring-integration-core/src/main/java/org/springframework/integration/dsl/ConsumerEndpointSpec.java

@@ -30,9 +30,11 @@ import org.springframework.integration.router.AbstractMessageRouter;
 import org.springframework.integration.scheduling.PollerMetadata;
 import org.springframework.integration.transaction.TransactionInterceptorBuilder;
 import org.springframework.messaging.MessageHandler;
+import org.springframework.scheduling.TaskScheduler;
 import org.springframework.transaction.PlatformTransactionManager;
 import org.springframework.transaction.interceptor.DefaultTransactionAttribute;
 import org.springframework.transaction.interceptor.TransactionInterceptor;
+import org.springframework.util.Assert;
 
 import reactor.util.function.Tuple2;
 
@@ -78,6 +80,21 @@ public abstract class ConsumerEndpointSpec<S extends ConsumerEndpointSpec<S, H>,
 		return _this();
 	}
 
+	/**
+	 * Configure a {@link TaskScheduler} for scheduling tasks, for example in the
+	 * Polling Consumer. By default the global {@code ThreadPoolTaskScheduler} bean is used.
+	 * This configuration is useful when there are requirements to dedicate particular threads
+	 * for polling task, for example.
+	 * @param taskScheduler the {@link TaskScheduler} to use.
+	 * @return the endpoint spec.
+	 * @see org.springframework.integration.context.IntegrationContextUtils#getTaskScheduler
+	 */
+	public S taskScheduler(TaskScheduler taskScheduler) {
+		Assert.notNull(taskScheduler, "'taskScheduler' must not be null");
+		this.endpointFactoryBean.setTaskScheduler(taskScheduler);
+		return _this();
+	}
+
 	/**
 	 * Configure a list of {@link Advice} objects to be applied, in nested order, to the
 	 * endpoint's handler. The advice objects are applied only to the handler.

spring-integration-core/src/main/java/org/springframework/integration/dsl/CorrelationHandlerSpec.java

@@ -134,8 +134,10 @@ public abstract class CorrelationHandlerSpec<S extends CorrelationHandlerSpec<S,
 	 * @return the handler spec.
 	 * @see AbstractCorrelatingMessageHandler#setTaskScheduler(TaskScheduler)
 	 */
+	@Override
 	public S taskScheduler(TaskScheduler taskScheduler) {
-		Assert.notNull(taskScheduler, "'taskScheduler' must not be null.");
+		Assert.notNull(taskScheduler, "'taskScheduler' must not be null");
+		super.taskScheduler(taskScheduler);
 		this.handler.setTaskScheduler(taskScheduler);
 		return _this();
 	}

spring-integration-core/src/test/java/org/springframework/integration/dsl/flows/IntegrationFlowTests.java

@@ -26,8 +26,11 @@ import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 
 import org.aopalliance.aop.Advice;
 import org.aopalliance.intercept.MethodInterceptor;
@@ -427,6 +430,28 @@ public class IntegrationFlowTests {
 		assertEquals("foo", this.errorRecovererFlowGateway.testIt("foo"));
 	}
 
+	@Autowired
+	private MessageChannel dedicatedQueueChannel;
+
+	@Autowired
+	private SubscribableChannel dedicatedResults;
+
+	@Test
+	public void testDedicatedPollingThreadFlow() throws InterruptedException {
+		AtomicReference<String> threadNameReference = new AtomicReference<>();
+		CountDownLatch resultLatch = new CountDownLatch(1);
+		this.dedicatedResults.subscribe(m -> {
+			threadNameReference.set(Thread.currentThread().getName());
+			resultLatch.countDown();
+		});
+
+		this.dedicatedQueueChannel.send(new GenericMessage<>("foo"));
+
+		assertTrue(resultLatch.await(10, TimeUnit.SECONDS));
+
+		assertEquals("dedicatedTaskScheduler-1", threadNameReference.get());
+	}
+
 	@MessagingGateway
 	public interface ControlBusGateway {
 
@@ -721,6 +746,22 @@ public class IntegrationFlowTests {
 
 		}
 
+		@Bean
+		public IntegrationFlow dedicatedPollingThreadFlow() {
+			return IntegrationFlows.from(MessageChannels.queue("dedicatedQueueChannel"))
+					.bridge(e -> e
+							.poller(Pollers.fixedDelay(0).receiveTimeout(-1))
+							.taskScheduler(dedicatedTaskScheduler()))
+					.channel("dedicatedResults")
+					.get();
+		}
+
+
+		@Bean
+		public TaskScheduler dedicatedTaskScheduler() {
+			return new ThreadPoolTaskScheduler();
+		}
+
 	}
 
 	@MessagingGateway