Add support for context propagation in task execution

Signed-off-by: Sandipan <[email protected]>

Author: sbera13

core/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/task/TaskExecutionProperties.java

@@ -44,6 +44,11 @@ public class TaskExecutionProperties {
 	 */
 	private Mode mode = Mode.AUTO;
 
+	/**
+	 * Indicates whether to register ContextPropagatingTaskDecorator bean
+	 */
+	private boolean propagateContext = false;
+
 	/**
 	 * Prefix to use for the names of newly created threads.
 	 */
@@ -69,6 +74,14 @@ public void setMode(Mode mode) {
 		this.mode = mode;
 	}
 
+	public boolean getPropagateContext() {
+		return this.propagateContext;
+	}
+
+	public void setPropagateContext(boolean propagateContext) {
+		this.propagateContext = propagateContext;
+	}
+
 	public String getThreadNamePrefix() {
 		return this.threadNamePrefix;
 	}

core/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/task/TaskExecutorConfigurations.java

@@ -75,17 +75,39 @@ static class TaskExecutorConfiguration {
 
 		@Bean(TaskExecutionAutoConfiguration.APPLICATION_TASK_EXECUTOR_BEAN_NAME)
 		@ConditionalOnThreading(Threading.VIRTUAL)
-		SimpleAsyncTaskExecutor applicationTaskExecutorVirtualThreads(SimpleAsyncTaskExecutorBuilder builder) {
-			return builder.build();
+		SimpleAsyncTaskExecutor applicationTaskExecutorVirtualThreads(SimpleAsyncTaskExecutorBuilder builder,
+																	  ObjectProvider<TaskDecorator> decorators) {
+			TaskDecorator decorator = decorators.getIfAvailable();
+			SimpleAsyncTaskExecutor executor = builder.build();
+			if (decorator != null) {
+				executor.setTaskDecorator(decorator);
+			}
+			return executor;
 		}
 
 		@Bean(TaskExecutionAutoConfiguration.APPLICATION_TASK_EXECUTOR_BEAN_NAME)
 		@Lazy
 		@ConditionalOnThreading(Threading.PLATFORM)
-		ThreadPoolTaskExecutor applicationTaskExecutor(ThreadPoolTaskExecutorBuilder threadPoolTaskExecutorBuilder) {
-			return threadPoolTaskExecutorBuilder.build();
+		ThreadPoolTaskExecutor applicationTaskExecutor( ThreadPoolTaskExecutorBuilder threadPoolTaskExecutorBuilder,
+			    										ObjectProvider<TaskDecorator> decorators ) {
+			TaskDecorator decorator = decorators.getIfAvailable();
+			ThreadPoolTaskExecutor executor = builder.build();
+			if (decorator != null) {
+				executor.setTaskDecorator(decorator);
+			}
+			return executor;
 		}
 
+		@Bean
+		@ConditionalOnProperty(
+				prefix = "spring.task.execution",
+				name = "propagate-context",
+				havingValue = "true",
+				matchIfMissing = false
+		)
+		TaskDecorator contextPropagatingTaskDecorator() {
+			return new ContextPropagatingTaskDecorator();
+		}
 	}
 
 	@Configuration(proxyBeanMethods = false)