Allow an existing TaskExecutor to be configured in ChannelRegistration

This commit introduces a new method to configure an existing
TaskExecutor in ChannelRegistration. Contrary to
TaskExecutorRegistration, a ThreadPoolTaskExecutor is not necessary,
and it can't be further configured. This includes the thread name
prefix.

Closes gh-32081

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

@@ -21,6 +21,7 @@ import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Supplier;
 
 import org.springframework.beans.factory.BeanInitializationException;
 import org.springframework.beans.factory.annotation.Qualifier;
@@ -62,6 +63,7 @@ import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
+import org.springframework.util.CustomizableThreadCreator;
 import org.springframework.util.MimeTypeUtils;
 import org.springframework.util.PathMatcher;
 import org.springframework.util.StringUtils;
@@ -164,10 +166,8 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC
 
 	@Bean
 	public TaskExecutor clientInboundChannelExecutor() {
-		TaskExecutorRegistration reg = getClientInboundChannelRegistration().taskExecutor();
-		ThreadPoolTaskExecutor executor = reg.getTaskExecutor();
-		executor.setThreadNamePrefix("clientInboundChannel-");
-		return executor;
+		return getTaskExecutor(getClientInboundChannelRegistration(),
+				"clientInboundChannel-", this::defaultTaskExecutor);
 	}
 
 	protected final ChannelRegistration getClientInboundChannelRegistration() {
@@ -202,10 +202,8 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC
 
 	@Bean
 	public TaskExecutor clientOutboundChannelExecutor() {
-		TaskExecutorRegistration reg = getClientOutboundChannelRegistration().taskExecutor();
-		ThreadPoolTaskExecutor executor = reg.getTaskExecutor();
-		executor.setThreadNamePrefix("clientOutboundChannel-");
-		return executor;
+		return getTaskExecutor(getClientOutboundChannelRegistration(),
+				"clientOutboundChannel-", this::defaultTaskExecutor);
 	}
 
 	protected final ChannelRegistration getClientOutboundChannelRegistration() {
@@ -246,19 +244,31 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC
 
 		MessageBrokerRegistry registry = getBrokerRegistry(clientInboundChannel, clientOutboundChannel);
 		ChannelRegistration registration = registry.getBrokerChannelRegistration();
-		ThreadPoolTaskExecutor executor;
-		if (registration.hasTaskExecutor()) {
-			executor = registration.taskExecutor().getTaskExecutor();
-		}
-		else {
+		return getTaskExecutor(registration, "brokerChannel-", () -> {
 			// Should never be used
-			executor = new ThreadPoolTaskExecutor();
-			executor.setCorePoolSize(0);
-			executor.setMaxPoolSize(1);
-			executor.setQueueCapacity(0);
+			ThreadPoolTaskExecutor threadPoolTaskExecutor = new ThreadPoolTaskExecutor();
+			threadPoolTaskExecutor.setCorePoolSize(0);
+			threadPoolTaskExecutor.setMaxPoolSize(1);
+			threadPoolTaskExecutor.setQueueCapacity(0);
+			return threadPoolTaskExecutor;
+		});
+	}
+
+	private static TaskExecutor getTaskExecutor(ChannelRegistration registration,
+			String threadNamePrefix, Supplier<TaskExecutor> fallback) {
+
+		return registration.getTaskExecutor(fallback,
+				executor -> setThreadNamePrefix(executor, threadNamePrefix));
+	}
+
+	private TaskExecutor defaultTaskExecutor() {
+		return new TaskExecutorRegistration().getTaskExecutor();
+	}
+
+	private static void setThreadNamePrefix(TaskExecutor taskExecutor, String name) {
+		if (taskExecutor instanceof CustomizableThreadCreator ctc) {
+			ctc.setThreadNamePrefix(name);
 		}
-		executor.setThreadNamePrefix("brokerChannel-");
-		return executor;
 	}
 
 	/**

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2021 the original author or authors.
+ * Copyright 2002-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,7 +19,10 @@ package org.springframework.messaging.simp.config;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
 
+import org.springframework.core.task.TaskExecutor;
 import org.springframework.lang.Nullable;
 import org.springframework.messaging.support.ChannelInterceptor;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
@@ -29,6 +32,7 @@ import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
  * {@link org.springframework.messaging.MessageChannel}.
  *
  * @author Rossen Stoyanchev
+ * @author Stephane Nicoll
  * @since 4.0
  */
 public class ChannelRegistration {
@@ -36,6 +40,9 @@ public class ChannelRegistration {
 	@Nullable
 	private TaskExecutorRegistration registration;
 
+	@Nullable
+	private TaskExecutor executor;
+
 	private final List<ChannelInterceptor> interceptors = new ArrayList<>();
 
 
@@ -59,6 +66,18 @@ public class ChannelRegistration {
 		return this.registration;
 	}
 
+	/**
+	 * Configure the given {@link TaskExecutor} for this message channel,
+	 * taking precedence over a {@linkplain #taskExecutor() task executor
+	 * registration} if any.
+	 * @param taskExecutor the task executor to use
+	 * @since 6.1.4
+	 */
+	public ChannelRegistration executor(TaskExecutor taskExecutor) {
+		this.executor = taskExecutor;
+		return this;
+	}
+
 	/**
 	 * Configure the given interceptors for this message channel,
 	 * adding them to the channel's current list of interceptors.
@@ -71,13 +90,40 @@ public class ChannelRegistration {
 
 
 	protected boolean hasTaskExecutor() {
-		return (this.registration != null);
+		return (this.registration != null || this.executor != null);
 	}
 
 	protected boolean hasInterceptors() {
 		return !this.interceptors.isEmpty();
 	}
 
+	/**
+	 * Return the {@link TaskExecutor} to use. If no task executor has been
+	 * configured, the {@code fallback} supplier is used to provide a fallback
+	 * instance.
+	 * <p>
+	 * If the {@link TaskExecutor} to use is suitable for further customizations,
+	 * the {@code customizer} consumer is invoked.
+	 * @param fallback a supplier of a fallback task executor in case none is configured
+	 * @param customizer further customizations
+	 * @return the task executor to use
+	 */
+	protected TaskExecutor getTaskExecutor(Supplier<TaskExecutor> fallback, Consumer<TaskExecutor> customizer) {
+		if (this.executor != null) {
+			return this.executor;
+		}
+		else if (this.registration != null) {
+			ThreadPoolTaskExecutor registeredTaskExecutor = this.registration.getTaskExecutor();
+			customizer.accept(registeredTaskExecutor);
+			return registeredTaskExecutor;
+		}
+		else {
+			TaskExecutor taskExecutor = fallback.get();
+			customizer.accept(taskExecutor);
+			return taskExecutor;
+		}
+	}
+
 	protected List<ChannelInterceptor> getInterceptors() {
 		return this.interceptors;
 	}