Polish "Configure suitable TaskExecutor for WebSocket"

See gh-39611

Author: mhalbritter

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/data/jpa/JpaRepositoriesAutoConfiguration.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2024 the original author or authors.
+ * Copyright 2012-2022 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.
@@ -67,7 +67,6 @@
  * @author Josh Long
  * @author Scott Frederick
  * @author Stefano Cordio
- * @author Lasse Wulff
  * @since 1.0.0
  * @see EnableJpaRepositories
  */
@@ -85,14 +84,20 @@ public class JpaRepositoriesAutoConfiguration {
 	public EntityManagerFactoryBuilderCustomizer entityManagerFactoryBootstrapExecutorCustomizer(
 			Map<String, AsyncTaskExecutor> taskExecutors) {
 		return (builder) -> {
-			AsyncTaskExecutor bootstrapExecutor = TaskExecutionAutoConfiguration
-				.determineAsyncTaskExecutor(taskExecutors);
+			AsyncTaskExecutor bootstrapExecutor = determineBootstrapExecutor(taskExecutors);
 			if (bootstrapExecutor != null) {
 				builder.setBootstrapExecutor(bootstrapExecutor);
 			}
 		};
 	}
 
+	private AsyncTaskExecutor determineBootstrapExecutor(Map<String, AsyncTaskExecutor> taskExecutors) {
+		if (taskExecutors.size() == 1) {
+			return taskExecutors.values().iterator().next();
+		}
+		return taskExecutors.get(TaskExecutionAutoConfiguration.APPLICATION_TASK_EXECUTOR_BEAN_NAME);
+	}
+
 	private static final class BootstrapExecutorCondition extends AnyNestedCondition {
 
 		BootstrapExecutorCondition() {

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/task/TaskExecutionAutoConfiguration.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2024 the original author or authors.
+ * Copyright 2012-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.
@@ -16,14 +16,11 @@
 
 package org.springframework.boot.autoconfigure.task;
 
-import java.util.Map;
-
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.context.annotation.Import;
-import org.springframework.core.task.AsyncTaskExecutor;
 import org.springframework.core.task.TaskExecutor;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
 
@@ -33,7 +30,6 @@
  * @author Stephane Nicoll
  * @author Camille Vienot
  * @author Moritz Halbritter
- * @author Lasse Wulff
  * @since 2.1.0
  */
 @ConditionalOnClass(ThreadPoolTaskExecutor.class)
@@ -50,11 +46,4 @@ public class TaskExecutionAutoConfiguration {
 	 */
 	public static final String APPLICATION_TASK_EXECUTOR_BEAN_NAME = "applicationTaskExecutor";
 
-	public static AsyncTaskExecutor determineAsyncTaskExecutor(Map<String, AsyncTaskExecutor> taskExecutors) {
-		if (taskExecutors.size() == 1) {
-			return taskExecutors.values().iterator().next();
-		}
-		return taskExecutors.get(APPLICATION_TASK_EXECUTOR_BEAN_NAME);
-	}
-
 }

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/websocket/servlet/WebSocketMessagingAutoConfiguration.java

@@ -49,6 +49,7 @@
  *
  * @author Andy Wilkinson
  * @author Lasse Wulff
+ * @author Moritz Halbritter
  * @since 1.3.0
  */
 @AutoConfiguration(after = JacksonAutoConfiguration.class)
@@ -68,7 +69,14 @@ static class WebSocketMessageConverterConfiguration implements WebSocketMessageB
 		WebSocketMessageConverterConfiguration(ObjectMapper objectMapper,
 				Map<String, AsyncTaskExecutor> taskExecutors) {
 			this.objectMapper = objectMapper;
-			this.executor = TaskExecutionAutoConfiguration.determineAsyncTaskExecutor(taskExecutors);
+			this.executor = determineAsyncTaskExecutor(taskExecutors);
+		}
+
+		private static AsyncTaskExecutor determineAsyncTaskExecutor(Map<String, AsyncTaskExecutor> taskExecutors) {
+			if (taskExecutors.size() == 1) {
+				return taskExecutors.values().iterator().next();
+			}
+			return taskExecutors.get(TaskExecutionAutoConfiguration.APPLICATION_TASK_EXECUTOR_BEAN_NAME);
 		}
 
 		@Override
@@ -85,12 +93,16 @@ public boolean configureMessageConverters(List<MessageConverter> messageConverte
 
 		@Override
 		public void configureClientInboundChannel(ChannelRegistration registration) {
-			registration.executor(this.executor);
+			if (this.executor != null) {
+				registration.executor(this.executor);
+			}
 		}
 
 		@Override
 		public void configureClientOutboundChannel(ChannelRegistration registration) {
-			registration.executor(this.executor);
+			if (this.executor != null) {
+				registration.executor(this.executor);
+			}
 		}
 
 		@Bean

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/websocket/servlet/WebSocketMessagingAutoConfigurationTests.java

@@ -48,6 +48,7 @@
 import org.springframework.context.annotation.Configuration;
 import org.springframework.core.task.AsyncTaskExecutor;
 import org.springframework.core.task.SimpleAsyncTaskExecutor;
+import org.springframework.core.task.TaskExecutor;
 import org.springframework.messaging.converter.CompositeMessageConverter;
 import org.springframework.messaging.converter.MessageConverter;
 import org.springframework.messaging.converter.SimpleMessageConverter;
@@ -144,11 +145,9 @@ void predefinedThreadExecutorIsSelectedForInboundChannel() throws Throwable {
 		WebSocketMessagingAutoConfiguration.WebSocketMessageConverterConfiguration configuration = new WebSocketMessagingAutoConfiguration.WebSocketMessageConverterConfiguration(
 				new ObjectMapper(),
 				Map.of(TaskExecutionAutoConfiguration.APPLICATION_TASK_EXECUTOR_BEAN_NAME, expectedExecutor));
-
 		configuration.configureClientInboundChannel(registration);
-
-		AsyncTaskExecutor mappedExecutor = (AsyncTaskExecutor) FieldUtils.getFieldValue(registration, "executor");
-		assertThat(mappedExecutor).isEqualTo(expectedExecutor);
+		TaskExecutor executor = (TaskExecutor) FieldUtils.getFieldValue(registration, "executor");
+		assertThat(executor).isEqualTo(expectedExecutor);
 	}
 
 	@Test
@@ -158,11 +157,9 @@ void predefinedThreadExecutorIsSelectedForOutboundChannel() throws Throwable {
 		WebSocketMessagingAutoConfiguration.WebSocketMessageConverterConfiguration configuration = new WebSocketMessagingAutoConfiguration.WebSocketMessageConverterConfiguration(
 				new ObjectMapper(),
 				Map.of(TaskExecutionAutoConfiguration.APPLICATION_TASK_EXECUTOR_BEAN_NAME, expectedExecutor));
-
 		configuration.configureClientOutboundChannel(registration);
-
-		AsyncTaskExecutor mappedExecutor = (AsyncTaskExecutor) FieldUtils.getFieldValue(registration, "executor");
-		assertThat(mappedExecutor).isEqualTo(expectedExecutor);
+		TaskExecutor executor = (TaskExecutor) FieldUtils.getFieldValue(registration, "executor");
+		assertThat(executor).isEqualTo(expectedExecutor);
 	}
 
 	private List<MessageConverter> getCustomizedConverters() {