Skip to content

Commit

Permalink
GH-637 - Enable TraceContext propagation across asynchronous threads.
Browse files Browse the repository at this point in the history
We're now registering a ThreadPoolTaskScheduler and customize the resources created to propagate the Trace context into threads created by it. This is necessary to connect the tracing spans created for events handled asynchronously.
  • Loading branch information
odrotbohm committed May 23, 2024
1 parent 324e824 commit 4e7dd5f
Showing 1 changed file with 47 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -22,22 +22,36 @@
import brave.handler.MutableSpan;
import brave.handler.SpanHandler;
import brave.propagation.TraceContext;
import io.micrometer.context.ContextExecutorService;
import io.micrometer.context.ContextScheduledExecutorService;
import io.micrometer.context.ContextSnapshot;
import io.micrometer.context.ContextSnapshotFactory;
import io.micrometer.tracing.Tracer;

import java.util.concurrent.ExecutorService;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadFactory;
import java.util.function.Supplier;

import org.springframework.beans.factory.ObjectProvider;
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
import org.springframework.boot.autoconfigure.AutoConfigureBefore;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.autoconfigure.task.TaskExecutionAutoConfiguration;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.modulith.observability.ModuleEventListener;
import org.springframework.modulith.observability.ModuleTracingBeanPostProcessor;
import org.springframework.modulith.runtime.ApplicationModulesRuntime;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;

/**
* @author Oliver Drotbohm
*/
@Configuration(proxyBeanMethods = false)
@AutoConfigureBefore(TaskExecutionAutoConfiguration.class)
@ConditionalOnProperty(name = "management.tracing.enabled", havingValue = "true", matchIfMissing = true)
class ModuleObservabilityAutoConfiguration {

Expand All @@ -53,6 +67,39 @@ static ModuleEventListener tracingModuleEventListener(ApplicationModulesRuntime
return new ModuleEventListener(runtime, () -> tracer.getObject());
}

/**
* Custom override of default {@link ThreadPoolTaskScheduler} to make sure asynchronous method invocations get the
* {@link io.micrometer.tracing.handler.TracingObservationHandler.TracingContext} forwarded into threads spawned for
* those methods. <em>The name of the bean is important for it to be picked up by the async invocation
* infrastructure!</em>
*/
@Bean(name = "taskExecutor", destroyMethod = "shutdown")
ThreadPoolTaskScheduler threadPoolTaskScheduler() {

ThreadPoolTaskScheduler threadPoolTaskScheduler = new ThreadPoolTaskScheduler() {

private static final long serialVersionUID = -3935299327010101697L;
private final Supplier<ContextSnapshot> captureAll = () -> ContextSnapshotFactory.builder().build().captureAll();

@Override
protected ExecutorService initializeExecutor(ThreadFactory threadFactory,
RejectedExecutionHandler rejectedExecutionHandler) {

return ContextExecutorService.wrap(super.initializeExecutor(threadFactory, rejectedExecutionHandler),
captureAll);
}

@Override
public ScheduledExecutorService getScheduledExecutor() throws IllegalStateException {
return ContextScheduledExecutorService.wrap(super.getScheduledExecutor(), captureAll);
}
};

threadPoolTaskScheduler.initialize();

return threadPoolTaskScheduler;
}

/**
* Brave-specific auto configuration.
*
Expand Down

0 comments on commit 4e7dd5f

Please sign in to comment.