Merge pull request #39611 from BenchmarkingBuffalo

* pr/39611:
  Polish "Configure suitable TaskExecutor for WebSocket"
  Configure suitable TaskExecutor for WebSocket

Closes gh-39611

Author: mhalbritter

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2023 the original author or authors.
+ * Copyright 2012-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.
@@ -17,6 +17,7 @@
 package org.springframework.boot.autoconfigure.websocket.servlet;
 
 import java.util.List;
+import java.util.Map;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 
@@ -28,14 +29,17 @@
 import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplication;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplication.Type;
 import org.springframework.boot.autoconfigure.jackson.JacksonAutoConfiguration;
+import org.springframework.boot.autoconfigure.task.TaskExecutionAutoConfiguration;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.core.task.AsyncTaskExecutor;
 import org.springframework.messaging.converter.ByteArrayMessageConverter;
 import org.springframework.messaging.converter.DefaultContentTypeResolver;
 import org.springframework.messaging.converter.MappingJackson2MessageConverter;
 import org.springframework.messaging.converter.MessageConverter;
 import org.springframework.messaging.converter.StringMessageConverter;
 import org.springframework.messaging.simp.config.AbstractMessageBrokerConfiguration;
+import org.springframework.messaging.simp.config.ChannelRegistration;
 import org.springframework.util.MimeTypeUtils;
 import org.springframework.web.socket.config.annotation.DelegatingWebSocketMessageBrokerConfiguration;
 import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;
@@ -44,6 +48,8 @@
  * {@link EnableAutoConfiguration Auto-configuration} for WebSocket-based messaging.
  *
  * @author Andy Wilkinson
+ * @author Lasse Wulff
+ * @author Moritz Halbritter
  * @since 1.3.0
  */
 @AutoConfiguration(after = JacksonAutoConfiguration.class)
@@ -58,8 +64,19 @@ static class WebSocketMessageConverterConfiguration implements WebSocketMessageB
 
 		private final ObjectMapper objectMapper;
 
-		WebSocketMessageConverterConfiguration(ObjectMapper objectMapper) {
+		private final AsyncTaskExecutor executor;
+
+		WebSocketMessageConverterConfiguration(ObjectMapper objectMapper,
+				Map<String, AsyncTaskExecutor> taskExecutors) {
 			this.objectMapper = objectMapper;
+			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
@@ -74,6 +91,20 @@ public boolean configureMessageConverters(List<MessageConverter> messageConverte
 			return false;
 		}
 
+		@Override
+		public void configureClientInboundChannel(ChannelRegistration registration) {
+			if (this.executor != null) {
+				registration.executor(this.executor);
+			}
+		}
+
+		@Override
+		public void configureClientOutboundChannel(ChannelRegistration registration) {
+			if (this.executor != null) {
+				registration.executor(this.executor);
+			}
+		}
+
 		@Bean
 		static LazyInitializationExcludeFilter eagerStompWebSocketHandlerMapping() {
 			return (name, definition, type) -> name.equals("stompWebSocketHandlerMapping");

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2023 the original author or authors.
+ * Copyright 2012-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,8 +19,10 @@
 import java.lang.reflect.Type;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
@@ -35,6 +37,7 @@
 import org.springframework.boot.LazyInitializationBeanFactoryPostProcessor;
 import org.springframework.boot.autoconfigure.ImportAutoConfiguration;
 import org.springframework.boot.autoconfigure.jackson.JacksonAutoConfiguration;
+import org.springframework.boot.autoconfigure.task.TaskExecutionAutoConfiguration;
 import org.springframework.boot.autoconfigure.web.servlet.DispatcherServletAutoConfiguration;
 import org.springframework.boot.autoconfigure.web.servlet.ServletWebServerFactoryAutoConfiguration;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
@@ -43,17 +46,22 @@
 import org.springframework.boot.web.servlet.context.AnnotationConfigServletWebServerApplicationContext;
 import org.springframework.context.annotation.Bean;
 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;
 import org.springframework.messaging.simp.annotation.SubscribeMapping;
+import org.springframework.messaging.simp.config.ChannelRegistration;
 import org.springframework.messaging.simp.config.MessageBrokerRegistry;
 import org.springframework.messaging.simp.stomp.StompCommand;
 import org.springframework.messaging.simp.stomp.StompFrameHandler;
 import org.springframework.messaging.simp.stomp.StompHeaders;
 import org.springframework.messaging.simp.stomp.StompSession;
 import org.springframework.messaging.simp.stomp.StompSessionHandler;
 import org.springframework.messaging.simp.stomp.StompSessionHandlerAdapter;
+import org.springframework.security.util.FieldUtils;
 import org.springframework.stereotype.Controller;
 import org.springframework.web.client.RestTemplate;
 import org.springframework.web.socket.client.standard.StandardWebSocketClient;
@@ -75,6 +83,7 @@
  * Tests for {@link WebSocketMessagingAutoConfiguration}.
  *
  * @author Andy Wilkinson
+ * @author Lasse Wulff
  */
 class WebSocketMessagingAutoConfigurationTests {
 
@@ -129,10 +138,34 @@ void customizedConverterTypesMatchDefaultConverterTypes() {
 		}
 	}
 
+	@Test
+	void predefinedThreadExecutorIsSelectedForInboundChannel() throws Throwable {
+		AsyncTaskExecutor expectedExecutor = new SimpleAsyncTaskExecutor();
+		ChannelRegistration registration = new ChannelRegistration();
+		WebSocketMessagingAutoConfiguration.WebSocketMessageConverterConfiguration configuration = new WebSocketMessagingAutoConfiguration.WebSocketMessageConverterConfiguration(
+				new ObjectMapper(),
+				Map.of(TaskExecutionAutoConfiguration.APPLICATION_TASK_EXECUTOR_BEAN_NAME, expectedExecutor));
+		configuration.configureClientInboundChannel(registration);
+		TaskExecutor executor = (TaskExecutor) FieldUtils.getFieldValue(registration, "executor");
+		assertThat(executor).isEqualTo(expectedExecutor);
+	}
+
+	@Test
+	void predefinedThreadExecutorIsSelectedForOutboundChannel() throws Throwable {
+		AsyncTaskExecutor expectedExecutor = new SimpleAsyncTaskExecutor();
+		ChannelRegistration registration = new ChannelRegistration();
+		WebSocketMessagingAutoConfiguration.WebSocketMessageConverterConfiguration configuration = new WebSocketMessagingAutoConfiguration.WebSocketMessageConverterConfiguration(
+				new ObjectMapper(),
+				Map.of(TaskExecutionAutoConfiguration.APPLICATION_TASK_EXECUTOR_BEAN_NAME, expectedExecutor));
+		configuration.configureClientOutboundChannel(registration);
+		TaskExecutor executor = (TaskExecutor) FieldUtils.getFieldValue(registration, "executor");
+		assertThat(executor).isEqualTo(expectedExecutor);
+	}
+
 	private List<MessageConverter> getCustomizedConverters() {
 		List<MessageConverter> customizedConverters = new ArrayList<>();
 		WebSocketMessagingAutoConfiguration.WebSocketMessageConverterConfiguration configuration = new WebSocketMessagingAutoConfiguration.WebSocketMessageConverterConfiguration(
-				new ObjectMapper());
+				new ObjectMapper(), Collections.emptyMap());
 		configuration.configureMessageConverters(customizedConverters);
 		return customizedConverters;
 	}