Allow configuring custom ThreadPoolTaskExecutor

Issue: SPR-12272

Author: rstoyanchev

spring-messaging/src/main/java/org/springframework/messaging/simp/config/ChannelRegistration.java

@@ -21,6 +21,7 @@
 import java.util.List;
 
 import org.springframework.messaging.support.ChannelInterceptor;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
 
 /**
  * A registration class for customizing the configuration for a
@@ -46,6 +47,17 @@ public TaskExecutorRegistration taskExecutor() {
 		return this.registration;
 	}
 
+	/**
+	 * Configure the thread pool backing this message channel using a custom
+	 * ThreadPoolTaskExecutor.
+	 */
+	public TaskExecutorRegistration taskExecutor(ThreadPoolTaskExecutor taskExecutor) {
+		if (this.registration == null) {
+			this.registration = new TaskExecutorRegistration(taskExecutor);
+		}
+		return this.registration;
+	}
+
 	/**
 	 * Configure interceptors for the message channel.
 	 */

spring-messaging/src/main/java/org/springframework/messaging/simp/config/TaskExecutorRegistration.java

@@ -26,6 +26,8 @@
  */
 public class TaskExecutorRegistration {
 
+	private ThreadPoolTaskExecutor taskExecutor;
+
 	private int corePoolSize = Runtime.getRuntime().availableProcessors() * 2;
 
 	private int maxPoolSize = Integer.MAX_VALUE;
@@ -35,6 +37,13 @@ public class TaskExecutorRegistration {
 	private int keepAliveSeconds = 60;
 
 
+	public TaskExecutorRegistration() {
+	}
+
+	public TaskExecutorRegistration(ThreadPoolTaskExecutor taskExecutor) {
+		this.taskExecutor = taskExecutor;
+	}
+
 	/**
 	 * Set the core pool size of the ThreadPoolExecutor.
 	 * <p><strong>NOTE:</strong> The core pool size is effectively the max pool size
@@ -93,7 +102,7 @@ public TaskExecutorRegistration keepAliveSeconds(int keepAliveSeconds) {
 	}
 
 	protected ThreadPoolTaskExecutor getTaskExecutor() {
-		ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
+		ThreadPoolTaskExecutor executor = (this.taskExecutor != null ? this.taskExecutor : new ThreadPoolTaskExecutor());
 		executor.setCorePoolSize(this.corePoolSize);
 		executor.setMaxPoolSize(this.maxPoolSize);
 		executor.setKeepAliveSeconds(this.keepAliveSeconds);

spring-messaging/src/main/java/org/springframework/messaging/support/Test.java

@@ -132,8 +132,8 @@ public void clientInboundChannelCustomized() {
 
 		assertEquals(2, channel.getInterceptors().size());
 
-		ThreadPoolTaskExecutor taskExecutor = this.customContext.getBean(
-				"clientInboundChannelExecutor", ThreadPoolTaskExecutor.class);
+		CustomThreadPoolTaskExecutor taskExecutor = this.customContext.getBean(
+				"clientInboundChannelExecutor", CustomThreadPoolTaskExecutor.class);
 
 		assertEquals(11, taskExecutor.getCorePoolSize());
 		assertEquals(12, taskExecutor.getMaxPoolSize());
@@ -489,7 +489,8 @@ static class CustomConfig extends AbstractMessageBrokerConfiguration {
 		@Override
 		protected void configureClientInboundChannel(ChannelRegistration registration) {
 			registration.setInterceptors(this.interceptor);
-			registration.taskExecutor().corePoolSize(11).maxPoolSize(12).keepAliveSeconds(13).queueCapacity(14);
+			registration.taskExecutor(new CustomThreadPoolTaskExecutor())
+					.corePoolSize(11).maxPoolSize(12).keepAliveSeconds(13).queueCapacity(14);
 		}
 
 		@Override
@@ -540,4 +541,7 @@ public void validate(Object target, Errors errors) {
 		}
 	}
 
+	private static class CustomThreadPoolTaskExecutor extends ThreadPoolTaskExecutor {
+	}
+
 }