From e1f8be0bfe196e7d4929f2a8357f39b0d52ab91f Mon Sep 17 00:00:00 2001
From: Oliver Drotbohm <oliver.drotbohm@broadcom.com>
Date: Thu, 23 May 2024 17:22:40 +0200
Subject: [PATCH] GH-637 - Enable TraceContext propagation across asynchronous
 threads.

We're now registering a customizers for both the SimpleAsyncTaskExecutor (used for virtual threads) and the ThreadPoolTaskExecutor to register a ContextPropagatingTaskDecorator.
---
 .../ModuleObservabilityAutoConfiguration.java   | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/spring-modulith-observability/src/main/java/org/springframework/modulith/observability/autoconfigure/ModuleObservabilityAutoConfiguration.java b/spring-modulith-observability/src/main/java/org/springframework/modulith/observability/autoconfigure/ModuleObservabilityAutoConfiguration.java
index a4e406fd1..2023d3d5e 100644
--- a/spring-modulith-observability/src/main/java/org/springframework/modulith/observability/autoconfigure/ModuleObservabilityAutoConfiguration.java
+++ b/spring-modulith-observability/src/main/java/org/springframework/modulith/observability/autoconfigure/ModuleObservabilityAutoConfiguration.java
@@ -28,8 +28,13 @@
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnThreading;
+import org.springframework.boot.autoconfigure.thread.Threading;
+import org.springframework.boot.task.SimpleAsyncTaskExecutorCustomizer;
+import org.springframework.boot.task.ThreadPoolTaskExecutorCustomizer;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.core.task.support.ContextPropagatingTaskDecorator;
 import org.springframework.modulith.observability.ModuleEventListener;
 import org.springframework.modulith.observability.ModuleTracingBeanPostProcessor;
 import org.springframework.modulith.runtime.ApplicationModulesRuntime;
@@ -53,6 +58,18 @@ static ModuleEventListener tracingModuleEventListener(ApplicationModulesRuntime
 		return new ModuleEventListener(runtime, () -> tracer.getObject());
 	}
 
+	@Bean
+	@ConditionalOnThreading(Threading.VIRTUAL)
+	SimpleAsyncTaskExecutorCustomizer simpleAsyncTaskExecutorCustomizer() {
+		return executor -> executor.setTaskDecorator(new ContextPropagatingTaskDecorator());
+	}
+
+	@Bean
+	@ConditionalOnThreading(Threading.PLATFORM)
+	ThreadPoolTaskExecutorCustomizer threadPoolTaskExecutorCustomizer() {
+		return executor -> executor.setTaskDecorator(new ContextPropagatingTaskDecorator());
+	}
+
 	/**
 	 * Brave-specific auto configuration.
 	 *