Make IntegrationTaskSchedulerConfiguration virtual threads aware

Author: facewise

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/integration/IntegrationAutoConfiguration.java

@@ -32,16 +32,19 @@
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnSingleCandidate;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnThreading;
 import org.springframework.boot.autoconfigure.condition.SearchStrategy;
 import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
 import org.springframework.boot.autoconfigure.jmx.JmxAutoConfiguration;
 import org.springframework.boot.autoconfigure.jmx.JmxProperties;
 import org.springframework.boot.autoconfigure.rsocket.RSocketMessagingAutoConfiguration;
 import org.springframework.boot.autoconfigure.sql.init.OnDatabaseInitializationCondition;
 import org.springframework.boot.autoconfigure.task.TaskSchedulingAutoConfiguration;
+import org.springframework.boot.autoconfigure.thread.Threading;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.boot.context.properties.PropertyMapper;
 import org.springframework.boot.context.properties.source.MutuallyExclusiveConfigurationPropertiesException;
+import org.springframework.boot.task.SimpleAsyncTaskSchedulerBuilder;
 import org.springframework.boot.task.ThreadPoolTaskSchedulerBuilder;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Conditional;
@@ -65,6 +68,7 @@
 import org.springframework.messaging.rsocket.RSocketStrategies;
 import org.springframework.messaging.rsocket.annotation.support.RSocketMessageHandler;
 import org.springframework.scheduling.Trigger;
+import org.springframework.scheduling.concurrent.SimpleAsyncTaskScheduler;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
 import org.springframework.scheduling.support.CronTrigger;
 import org.springframework.scheduling.support.PeriodicTrigger;
@@ -79,6 +83,7 @@
  * @author Stephane Nicoll
  * @author Vedran Pavic
  * @author Madhura Bhave
+ * @author Yong-Hyun Kim
  * @since 1.1.0
  */
 @AutoConfiguration(after = { DataSourceAutoConfiguration.class, JmxAutoConfiguration.class,
@@ -169,15 +174,24 @@ private Trigger createPeriodicTrigger(Duration period, Duration initialDelay, bo
 	 * scheduling explicitly.
 	 */
 	@Configuration(proxyBeanMethods = false)
-	@ConditionalOnBean(ThreadPoolTaskSchedulerBuilder.class)
 	@ConditionalOnMissingBean(name = IntegrationContextUtils.TASK_SCHEDULER_BEAN_NAME)
 	protected static class IntegrationTaskSchedulerConfiguration {
 
 		@Bean(name = IntegrationContextUtils.TASK_SCHEDULER_BEAN_NAME)
+		@ConditionalOnBean(ThreadPoolTaskSchedulerBuilder.class)
+		@ConditionalOnThreading(Threading.PLATFORM)
 		public ThreadPoolTaskScheduler taskScheduler(ThreadPoolTaskSchedulerBuilder threadPoolTaskSchedulerBuilder) {
 			return threadPoolTaskSchedulerBuilder.build();
 		}
 
+		@Bean(name = IntegrationContextUtils.TASK_SCHEDULER_BEAN_NAME)
+		@ConditionalOnBean(SimpleAsyncTaskSchedulerBuilder.class)
+		@ConditionalOnThreading(Threading.VIRTUAL)
+		public SimpleAsyncTaskScheduler taskSchedulerVirtualThreads(
+				SimpleAsyncTaskSchedulerBuilder simpleAsyncTaskSchedulerBuilder) {
+			return simpleAsyncTaskSchedulerBuilder.build();
+		}
+
 	}
 
 	/**

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/integration/IntegrationAutoConfigurationTests.java

@@ -32,6 +32,8 @@
 import io.rsocket.transport.netty.client.TcpClientTransport;
 import org.assertj.core.api.InstanceOfAssertFactories;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.EnabledForJreRange;
+import org.junit.jupiter.api.condition.JRE;
 import reactor.core.publisher.Mono;
 
 import org.springframework.beans.DirectFieldAccessor;
@@ -54,9 +56,11 @@
 import org.springframework.boot.sql.init.DatabaseInitializationMode;
 import org.springframework.boot.sql.init.DatabaseInitializationSettings;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
+import org.springframework.boot.testsupport.assertj.SimpleAsyncTaskExecutorAssert;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.Primary;
+import org.springframework.core.task.SimpleAsyncTaskExecutor;
 import org.springframework.integration.annotation.IntegrationComponentScan;
 import org.springframework.integration.annotation.MessagingGateway;
 import org.springframework.integration.annotation.ServiceActivator;
@@ -84,6 +88,7 @@
 import org.springframework.messaging.rsocket.annotation.support.RSocketMessageHandler;
 import org.springframework.messaging.support.GenericMessage;
 import org.springframework.scheduling.TaskScheduler;
+import org.springframework.scheduling.concurrent.SimpleAsyncTaskScheduler;
 import org.springframework.scheduling.support.CronTrigger;
 import org.springframework.scheduling.support.PeriodicTrigger;
 import org.springframework.test.util.ReflectionTestUtils;
@@ -98,6 +103,7 @@
  * @author Artem Bilan
  * @author Stephane Nicoll
  * @author Vedran Pavic
+ * @author Yong-Hyun Kim
  */
 class IntegrationAutoConfigurationTests {
 
@@ -521,6 +527,18 @@ void integrationManagementInstrumentedWithObservation() {
 			});
 	}
 
+	@Test
+	@EnabledForJreRange(min = JRE.JAVA_21)
+	void integrationVirtualThreadsEnabled() {
+		this.contextRunner.withPropertyValues("spring.threads.virtual.enabled=true")
+			.run((context) -> assertThat(context).hasSingleBean(TaskScheduler.class)
+				.getBean(IntegrationContextUtils.TASK_SCHEDULER_BEAN_NAME, TaskScheduler.class)
+				.isInstanceOf(SimpleAsyncTaskScheduler.class)
+				.satisfies((taskScheduler) -> SimpleAsyncTaskExecutorAssert
+					.assertThat((SimpleAsyncTaskExecutor) taskScheduler)
+					.usesVirtualThreads()));
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	static class CustomMBeanExporter {