Add virtual thread option for ThreadPoolTaskExecutorBuilder/`Thread…

…PoolTaskSchedulerBuilder` See spring-projects/spring-framework#33807
quaff · Nov 18, 2024 · d046f95 · d046f95
1 parent 3792772
commit d046f95
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 22 deletions.
diff --git a/...pring-boot/src/main/java/org/springframework/boot/task/ThreadPoolTaskExecutorBuilder.java b/...pring-boot/src/main/java/org/springframework/boot/task/ThreadPoolTaskExecutorBuilder.java
@@ -65,6 +65,8 @@ public class ThreadPoolTaskExecutorBuilder {
 
 	private final TaskDecorator taskDecorator;
 
+	private final Boolean virtualThreads;
+
 	private final Set<ThreadPoolTaskExecutorCustomizer> customizers;
 
 	public ThreadPoolTaskExecutorBuilder() {
@@ -78,13 +80,14 @@ public ThreadPoolTaskExecutorBuilder() {
 		this.awaitTerminationPeriod = null;
 		this.threadNamePrefix = null;
 		this.taskDecorator = null;
+		this.virtualThreads = null;
 		this.customizers = null;
 	}
 
 	private ThreadPoolTaskExecutorBuilder(Integer queueCapacity, Integer corePoolSize, Integer maxPoolSize,
 			Boolean allowCoreThreadTimeOut, Duration keepAlive, Boolean acceptTasksAfterContextClose,
 			Boolean awaitTermination, Duration awaitTerminationPeriod, String threadNamePrefix,
-			TaskDecorator taskDecorator, Set<ThreadPoolTaskExecutorCustomizer> customizers) {
+			TaskDecorator taskDecorator, Boolean virtualThreads, Set<ThreadPoolTaskExecutorCustomizer> customizers) {
 		this.queueCapacity = queueCapacity;
 		this.corePoolSize = corePoolSize;
 		this.maxPoolSize = maxPoolSize;
@@ -95,6 +98,7 @@ private ThreadPoolTaskExecutorBuilder(Integer queueCapacity, Integer corePoolSiz
 		this.awaitTerminationPeriod = awaitTerminationPeriod;
 		this.threadNamePrefix = threadNamePrefix;
 		this.taskDecorator = taskDecorator;
+		this.virtualThreads = virtualThreads;
 		this.customizers = customizers;
 	}
 
@@ -107,7 +111,8 @@ private ThreadPoolTaskExecutorBuilder(Integer queueCapacity, Integer corePoolSiz
 	public ThreadPoolTaskExecutorBuilder queueCapacity(int queueCapacity) {
 		return new ThreadPoolTaskExecutorBuilder(queueCapacity, this.corePoolSize, this.maxPoolSize,
 				this.allowCoreThreadTimeOut, this.keepAlive, this.acceptTasksAfterContextClose, this.awaitTermination,
-				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.customizers);
+				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.virtualThreads,
+				this.customizers);
 	}
 
 	/**
@@ -122,7 +127,8 @@ public ThreadPoolTaskExecutorBuilder queueCapacity(int queueCapacity) {
 	public ThreadPoolTaskExecutorBuilder corePoolSize(int corePoolSize) {
 		return new ThreadPoolTaskExecutorBuilder(this.queueCapacity, corePoolSize, this.maxPoolSize,
 				this.allowCoreThreadTimeOut, this.keepAlive, this.acceptTasksAfterContextClose, this.awaitTermination,
-				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.customizers);
+				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.virtualThreads,
+				this.customizers);
 	}
 
 	/**
@@ -137,7 +143,8 @@ public ThreadPoolTaskExecutorBuilder corePoolSize(int corePoolSize) {
 	public ThreadPoolTaskExecutorBuilder maxPoolSize(int maxPoolSize) {
 		return new ThreadPoolTaskExecutorBuilder(this.queueCapacity, this.corePoolSize, maxPoolSize,
 				this.allowCoreThreadTimeOut, this.keepAlive, this.acceptTasksAfterContextClose, this.awaitTermination,
-				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.customizers);
+				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.virtualThreads,
+				this.customizers);
 	}
 
 	/**
@@ -149,7 +156,8 @@ public ThreadPoolTaskExecutorBuilder maxPoolSize(int maxPoolSize) {
 	public ThreadPoolTaskExecutorBuilder allowCoreThreadTimeOut(boolean allowCoreThreadTimeOut) {
 		return new ThreadPoolTaskExecutorBuilder(this.queueCapacity, this.corePoolSize, this.maxPoolSize,
 				allowCoreThreadTimeOut, this.keepAlive, this.acceptTasksAfterContextClose, this.awaitTermination,
-				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.customizers);
+				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.virtualThreads,
+				this.customizers);
 	}
 
 	/**
@@ -160,7 +168,8 @@ public ThreadPoolTaskExecutorBuilder allowCoreThreadTimeOut(boolean allowCoreThr
 	public ThreadPoolTaskExecutorBuilder keepAlive(Duration keepAlive) {
 		return new ThreadPoolTaskExecutorBuilder(this.queueCapacity, this.corePoolSize, this.maxPoolSize,
 				this.allowCoreThreadTimeOut, keepAlive, this.acceptTasksAfterContextClose, this.awaitTermination,
-				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.customizers);
+				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.virtualThreads,
+				this.customizers);
 	}
 
 	/**
@@ -174,7 +183,8 @@ public ThreadPoolTaskExecutorBuilder keepAlive(Duration keepAlive) {
 	public ThreadPoolTaskExecutorBuilder acceptTasksAfterContextClose(boolean acceptTasksAfterContextClose) {
 		return new ThreadPoolTaskExecutorBuilder(this.queueCapacity, this.corePoolSize, this.maxPoolSize,
 				this.allowCoreThreadTimeOut, this.keepAlive, acceptTasksAfterContextClose, this.awaitTermination,
-				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.customizers);
+				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.virtualThreads,
+				this.customizers);
 	}
 
 	/**
@@ -188,7 +198,8 @@ public ThreadPoolTaskExecutorBuilder acceptTasksAfterContextClose(boolean accept
 	public ThreadPoolTaskExecutorBuilder awaitTermination(boolean awaitTermination) {
 		return new ThreadPoolTaskExecutorBuilder(this.queueCapacity, this.corePoolSize, this.maxPoolSize,
 				this.allowCoreThreadTimeOut, this.keepAlive, this.acceptTasksAfterContextClose, awaitTermination,
-				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.customizers);
+				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.virtualThreads,
+				this.customizers);
 	}
 
 	/**
@@ -203,7 +214,8 @@ public ThreadPoolTaskExecutorBuilder awaitTermination(boolean awaitTermination)
 	public ThreadPoolTaskExecutorBuilder awaitTerminationPeriod(Duration awaitTerminationPeriod) {
 		return new ThreadPoolTaskExecutorBuilder(this.queueCapacity, this.corePoolSize, this.maxPoolSize,
 				this.allowCoreThreadTimeOut, this.keepAlive, this.acceptTasksAfterContextClose, this.awaitTermination,
-				awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.customizers);
+				awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.virtualThreads,
+				this.customizers);
 	}
 
 	/**
@@ -214,7 +226,8 @@ public ThreadPoolTaskExecutorBuilder awaitTerminationPeriod(Duration awaitTermin
 	public ThreadPoolTaskExecutorBuilder threadNamePrefix(String threadNamePrefix) {
 		return new ThreadPoolTaskExecutorBuilder(this.queueCapacity, this.corePoolSize, this.maxPoolSize,
 				this.allowCoreThreadTimeOut, this.keepAlive, this.acceptTasksAfterContextClose, this.awaitTermination,
-				this.awaitTerminationPeriod, threadNamePrefix, this.taskDecorator, this.customizers);
+				this.awaitTerminationPeriod, threadNamePrefix, this.taskDecorator, this.virtualThreads,
+				this.customizers);
 	}
 
 	/**
@@ -225,7 +238,20 @@ public ThreadPoolTaskExecutorBuilder threadNamePrefix(String threadNamePrefix) {
 	public ThreadPoolTaskExecutorBuilder taskDecorator(TaskDecorator taskDecorator) {
 		return new ThreadPoolTaskExecutorBuilder(this.queueCapacity, this.corePoolSize, this.maxPoolSize,
 				this.allowCoreThreadTimeOut, this.keepAlive, this.acceptTasksAfterContextClose, this.awaitTermination,
-				this.awaitTerminationPeriod, this.threadNamePrefix, taskDecorator, this.customizers);
+				this.awaitTerminationPeriod, this.threadNamePrefix, taskDecorator, this.virtualThreads,
+				this.customizers);
+	}
+
+	/**
+	 * Specify whether to use virtual threads instead of platform threads.
+	 * @param virtualThreads whether to use virtual threads instead of platform threads
+	 * @return a new builder instance
+	 */
+	public ThreadPoolTaskExecutorBuilder virtualThreads(boolean virtualThreads) {
+		return new ThreadPoolTaskExecutorBuilder(this.queueCapacity, this.corePoolSize, this.maxPoolSize,
+				this.allowCoreThreadTimeOut, this.keepAlive, this.acceptTasksAfterContextClose, this.awaitTermination,
+				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, virtualThreads,
+				this.customizers);
 	}
 
 	/**
@@ -255,7 +281,8 @@ public ThreadPoolTaskExecutorBuilder customizers(Iterable<? extends ThreadPoolTa
 		Assert.notNull(customizers, "Customizers must not be null");
 		return new ThreadPoolTaskExecutorBuilder(this.queueCapacity, this.corePoolSize, this.maxPoolSize,
 				this.allowCoreThreadTimeOut, this.keepAlive, this.acceptTasksAfterContextClose, this.awaitTermination,
-				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, append(null, customizers));
+				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.virtualThreads,
+				append(null, customizers));
 	}
 
 	/**
@@ -284,7 +311,7 @@ public ThreadPoolTaskExecutorBuilder additionalCustomizers(
 		Assert.notNull(customizers, "Customizers must not be null");
 		return new ThreadPoolTaskExecutorBuilder(this.queueCapacity, this.corePoolSize, this.maxPoolSize,
 				this.allowCoreThreadTimeOut, this.keepAlive, this.acceptTasksAfterContextClose, this.awaitTermination,
-				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator,
+				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator, this.virtualThreads,
 				append(this.customizers, customizers));
 	}
 
@@ -332,6 +359,7 @@ public <T extends ThreadPoolTaskExecutor> T configure(T taskExecutor) {
 		map.from(this.awaitTerminationPeriod).as(Duration::toMillis).to(taskExecutor::setAwaitTerminationMillis);
 		map.from(this.threadNamePrefix).whenHasText().to(taskExecutor::setThreadNamePrefix);
 		map.from(this.taskDecorator).to(taskExecutor::setTaskDecorator);
+		map.from(this.virtualThreads).to(taskExecutor::setVirtualThreads);
 		if (!CollectionUtils.isEmpty(this.customizers)) {
 			this.customizers.forEach((customizer) -> customizer.customize(taskExecutor));
 		}

diff --git a/...ring-boot/src/main/java/org/springframework/boot/task/ThreadPoolTaskSchedulerBuilder.java b/...ring-boot/src/main/java/org/springframework/boot/task/ThreadPoolTaskSchedulerBuilder.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.
@@ -36,6 +36,7 @@
  * bean and can be injected whenever a {@link ThreadPoolTaskScheduler} is needed.
  *
  * @author Stephane Nicoll
+ * @author Yanming Zhou
  * @since 3.2.0
  */
 public class ThreadPoolTaskSchedulerBuilder {
@@ -48,22 +49,27 @@ public class ThreadPoolTaskSchedulerBuilder {
 
 	private final String threadNamePrefix;
 
+	private final Boolean virtualThreads;
+
 	private final Set<ThreadPoolTaskSchedulerCustomizer> customizers;
 
 	public ThreadPoolTaskSchedulerBuilder() {
 		this.poolSize = null;
 		this.awaitTermination = null;
 		this.awaitTerminationPeriod = null;
 		this.threadNamePrefix = null;
+		this.virtualThreads = null;
 		this.customizers = null;
 	}
 
 	public ThreadPoolTaskSchedulerBuilder(Integer poolSize, Boolean awaitTermination, Duration awaitTerminationPeriod,
-			String threadNamePrefix, Set<ThreadPoolTaskSchedulerCustomizer> taskSchedulerCustomizers) {
+			String threadNamePrefix, Boolean virtualThreads,
+			Set<ThreadPoolTaskSchedulerCustomizer> taskSchedulerCustomizers) {
 		this.poolSize = poolSize;
 		this.awaitTermination = awaitTermination;
 		this.awaitTerminationPeriod = awaitTerminationPeriod;
 		this.threadNamePrefix = threadNamePrefix;
+		this.virtualThreads = virtualThreads;
 		this.customizers = taskSchedulerCustomizers;
 	}
 
@@ -74,7 +80,7 @@ public ThreadPoolTaskSchedulerBuilder(Integer poolSize, Boolean awaitTermination
 	 */
 	public ThreadPoolTaskSchedulerBuilder poolSize(int poolSize) {
 		return new ThreadPoolTaskSchedulerBuilder(poolSize, this.awaitTermination, this.awaitTerminationPeriod,
-				this.threadNamePrefix, this.customizers);
+				this.threadNamePrefix, this.virtualThreads, this.customizers);
 	}
 
 	/**
@@ -87,7 +93,7 @@ public ThreadPoolTaskSchedulerBuilder poolSize(int poolSize) {
 	 */
 	public ThreadPoolTaskSchedulerBuilder awaitTermination(boolean awaitTermination) {
 		return new ThreadPoolTaskSchedulerBuilder(this.poolSize, awaitTermination, this.awaitTerminationPeriod,
-				this.threadNamePrefix, this.customizers);
+				this.threadNamePrefix, this.virtualThreads, this.customizers);
 	}
 
 	/**
@@ -101,7 +107,7 @@ public ThreadPoolTaskSchedulerBuilder awaitTermination(boolean awaitTermination)
 	 */
 	public ThreadPoolTaskSchedulerBuilder awaitTerminationPeriod(Duration awaitTerminationPeriod) {
 		return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, awaitTerminationPeriod,
-				this.threadNamePrefix, this.customizers);
+				this.threadNamePrefix, this.virtualThreads, this.customizers);
 	}
 
 	/**
@@ -111,7 +117,17 @@ public ThreadPoolTaskSchedulerBuilder awaitTerminationPeriod(Duration awaitTermi
 	 */
 	public ThreadPoolTaskSchedulerBuilder threadNamePrefix(String threadNamePrefix) {
 		return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, this.awaitTerminationPeriod,
-				threadNamePrefix, this.customizers);
+				threadNamePrefix, this.virtualThreads, this.customizers);
+	}
+
+	/**
+	 * Specify whether to use virtual threads instead of platform threads.
+	 * @param virtualThreads whether to use virtual threads instead of platform threads
+	 * @return a new builder instance
+	 */
+	public ThreadPoolTaskSchedulerBuilder virtualThreads(boolean virtualThreads) {
+		return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, this.awaitTerminationPeriod,
+				this.threadNamePrefix, virtualThreads, this.customizers);
 	}
 
 	/**
@@ -143,7 +159,7 @@ public ThreadPoolTaskSchedulerBuilder customizers(
 			Iterable<? extends ThreadPoolTaskSchedulerCustomizer> customizers) {
 		Assert.notNull(customizers, "Customizers must not be null");
 		return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, this.awaitTerminationPeriod,
-				this.threadNamePrefix, append(null, customizers));
+				this.threadNamePrefix, this.virtualThreads, append(null, customizers));
 	}
 
 	/**
@@ -173,7 +189,7 @@ public ThreadPoolTaskSchedulerBuilder additionalCustomizers(
 			Iterable<? extends ThreadPoolTaskSchedulerCustomizer> customizers) {
 		Assert.notNull(customizers, "Customizers must not be null");
 		return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, this.awaitTerminationPeriod,
-				this.threadNamePrefix, append(this.customizers, customizers));
+				this.threadNamePrefix, this.virtualThreads, append(this.customizers, customizers));
 	}
 
 	/**
@@ -199,6 +215,7 @@ public <T extends ThreadPoolTaskScheduler> T configure(T taskScheduler) {
 		map.from(this.awaitTermination).to(taskScheduler::setWaitForTasksToCompleteOnShutdown);
 		map.from(this.awaitTerminationPeriod).asInt(Duration::getSeconds).to(taskScheduler::setAwaitTerminationSeconds);
 		map.from(this.threadNamePrefix).to(taskScheduler::setThreadNamePrefix);
+		map.from(this.virtualThreads).to(taskScheduler::setVirtualThreads);
 		if (!CollectionUtils.isEmpty(this.customizers)) {
 			this.customizers.forEach((customizer) -> customizer.customize(taskScheduler));
 		}

diff --git a/...-boot/src/test/java/org/springframework/boot/task/ThreadPoolTaskExecutorBuilderTests.java b/...-boot/src/test/java/org/springframework/boot/task/ThreadPoolTaskExecutorBuilderTests.java
@@ -89,6 +89,12 @@ void taskDecoratorShouldApply() {
 		assertThat(executor).extracting("taskDecorator").isSameAs(taskDecorator);
 	}
 
+	@Test
+	void virtualThreadsShouldApply() {
+		ThreadPoolTaskExecutor executor = this.builder.virtualThreads(true).build();
+		assertThat(executor).hasFieldOrPropertyWithValue("virtualThreads", true);
+	}
+
 	@Test
 	void customizersWhenCustomizersAreNullShouldThrowException() {
 		assertThatIllegalArgumentException()

diff --git a/...boot/src/test/java/org/springframework/boot/task/ThreadPoolTaskSchedulerBuilderTests.java b/...boot/src/test/java/org/springframework/boot/task/ThreadPoolTaskSchedulerBuilderTests.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.
@@ -34,6 +34,7 @@
  * Tests for {@link ThreadPoolTaskSchedulerBuilder}.
  *
  * @author Stephane Nicoll
+ * @author Yanming Zhou
  */
 class ThreadPoolTaskSchedulerBuilderTests {
 
@@ -64,6 +65,12 @@ void threadNamePrefixShouldApply() {
 		assertThat(scheduler.getThreadNamePrefix()).isEqualTo("test-");
 	}
 
+	@Test
+	void virtualThreadsShouldApply() {
+		ThreadPoolTaskScheduler scheduler = this.builder.virtualThreads(true).build();
+		assertThat(scheduler).hasFieldOrPropertyWithValue("virtualThreads", true);
+	}
+
 	@Test
 	void customizersWhenCustomizersAreNullShouldThrowException() {
 		assertThatIllegalArgumentException()