Skip to content

Replace ThreadPoolExecutor context-store with ClassValue cache #12105

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gh-worker-dd-mergequeue-cf854d merged 1 commit into from
Jul 30, 2026
Merged

Replace ThreadPoolExecutor context-store with ClassValue cache #12105

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 9 additions & 18 deletions ...trap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/TPEHelper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import datadog.context.ContextScope;
import datadog.trace.api.GenericClassValue;
import datadog.trace.api.InstrumenterConfig;
import datadog.trace.api.Platform;
import datadog.trace.bootstrap.ContextStore;
import java.util.Set;
import java.util.concurrent.ThreadPoolExecutor;
Expand Down Expand Up @@ -51,25 +52,15 @@ public static boolean useWrapping(Runnable task) {
return useWrapping || task instanceof Wrapper || (task != null && WRAP.get(task.getClass()));
}

public static void setPropagate(
ContextStore<ThreadPoolExecutor, Boolean> contextStore, ThreadPoolExecutor executor) {
if (executor == null || contextStore == null || contextStore.get(executor) != null) {
return;
}
String executorType = executor.getClass().getName();
if (excludedClasses.contains(executorType)) {
contextStore.put(executor, Boolean.FALSE);
} else {
contextStore.put(executor, Boolean.TRUE);
}
}
private static final ClassValue<Boolean> PROPAGATE =
GenericClassValue.of(input -> !excludedClasses.contains(input.getName()));

public static boolean shouldPropagate(
ContextStore<ThreadPoolExecutor, Boolean> contextStore, ThreadPoolExecutor executor) {
if (executor == null || contextStore == null) {
return false;
}
return Boolean.TRUE.equals(contextStore.get(executor));
public static boolean shouldPropagate(ThreadPoolExecutor executor) {
// avoid tracking threads when building native images as it confuses the scanner
// (we still want instrumentation applied, so tracking works in the built image)
return !Platform.isNativeImageBuilder()
&& executor != null
&& PROPAGATE.get(executor.getClass());
}

public static void capture(ContextStore<Runnable, State> contextStore, Runnable task) {
Expand Down
2 changes: 0 additions & 2 deletions .../src/main/java/datadog/trace/instrumentation/java/concurrent/executor/ExecutorModule.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.EXECUTOR;
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE;
import static datadog.trace.instrumentation.java.concurrent.ConcurrentInstrumentationNames.EXECUTOR_INSTRUMENTATION_NAME;
import static datadog.trace.instrumentation.java.concurrent.executor.ThreadPoolExecutorInstrumentation.TPE;
import static java.util.Collections.singleton;

import com.google.auto.service.AutoService;
Expand Down Expand Up @@ -41,7 +40,6 @@ public Map<String, String> contextStore() {
contextStore.put("java.util.concurrent.RunnableFuture", State.class.getName());
// TODO get rid of this
contextStore.put("java.lang.Runnable", State.class.getName());
contextStore.put(TPE, Boolean.class.getName());
return Collections.unmodifiableMap(contextStore);
}

Expand Down
36 changes: 8 additions & 28 deletions ...dog/trace/instrumentation/java/concurrent/executor/ThreadPoolExecutorInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE;
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE_FUTURE;
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.exclude;
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.not;
Expand All @@ -16,7 +15,6 @@

import datadog.context.ContextScope;
import datadog.trace.agent.tooling.Instrumenter;
import datadog.trace.api.Platform;
import datadog.trace.bootstrap.InstrumentationContext;
import datadog.trace.bootstrap.instrumentation.java.concurrent.QueueTimerHelper;
import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
Expand Down Expand Up @@ -57,7 +55,6 @@ public final class ThreadPoolExecutorInstrumentation
implements Instrumenter.ForBootstrap,
Instrumenter.ForTypeHierarchy,
Instrumenter.HasMethodAdvice {
static final String TPE = "java.util.concurrent.ThreadPoolExecutor";

// executors which do their own wrapping before calling super,
// leading to double wrapping, once at the child level and once
Expand All @@ -75,12 +72,11 @@ public String hierarchyMarkerType() {
@Override
public ElementMatcher<TypeDescription> hierarchyMatcher() {
return not(named("java.util.concurrent.ScheduledThreadPoolExecutor"))
.and(extendsClass(named(TPE)));
.and(extendsClass(named("java.util.concurrent.ThreadPoolExecutor")));
}

@Override
public void methodAdvice(MethodTransformer transformer) {
transformer.applyAdvice(isConstructor(), getClass().getName() + "$Init");
transformer.applyAdvice(
named("execute")
.and(isMethod())
Expand All @@ -102,25 +98,12 @@ public void methodAdvice(MethodTransformer transformer) {
getClass().getName() + "$Remove");
}

public static final class Init {
@Advice.OnMethodExit(suppress = Throwable.class)
public static void decideWrapping(@Advice.This final ThreadPoolExecutor zis) {
// avoid tracking threads when building native images as it confuses the scanner
// (we still want instrumentation applied, so tracking works in the built image)
if (!Platform.isNativeImageBuilder()) {
TPEHelper.setPropagate(
InstrumentationContext.get(ThreadPoolExecutor.class, Boolean.class), zis);
}
}
}

public static final class Execute {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static void capture(
@Advice.This final ThreadPoolExecutor tpe,
@Advice.Argument(readOnly = false, value = 0) Runnable task) {
if (TPEHelper.shouldPropagate(
InstrumentationContext.get(ThreadPoolExecutor.class, Boolean.class), tpe)) {
if (TPEHelper.shouldPropagate(tpe)) {
if (TPEHelper.useWrapping(task)) {
task = Wrapper.wrap(task);
} else {
Expand Down Expand Up @@ -155,10 +138,9 @@ public static void capture(
public static final class BeforeExecute {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static ContextScope beforeExecuteEnter(
@Advice.This final ThreadPoolExecutor zis,
@Advice.This final ThreadPoolExecutor tpe,
@Advice.Argument(readOnly = false, value = 1) Runnable task) {
if (TPEHelper.shouldPropagate(
InstrumentationContext.get(ThreadPoolExecutor.class, Boolean.class), zis)) {
if (TPEHelper.shouldPropagate(tpe)) {
if (TPEHelper.useWrapping(task)) {
task = Wrapper.unwrap(task);
} else {
Expand All @@ -181,10 +163,9 @@ public static void beforeExecuteExit(
public static final class AfterExecute {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static ContextScope afterExecuteEnter(
@Advice.This final ThreadPoolExecutor zis,
@Advice.This final ThreadPoolExecutor tpe,
@Advice.Argument(readOnly = false, value = 0) Runnable task) {
if (TPEHelper.shouldPropagate(
InstrumentationContext.get(ThreadPoolExecutor.class, Boolean.class), zis)) {
if (TPEHelper.shouldPropagate(tpe)) {
if (TPEHelper.useWrapping(task)) {
task = Wrapper.unwrap(task);
} else {
Expand All @@ -206,10 +187,9 @@ public static void afterExecuteExit(
public static final class Remove {
@Advice.OnMethodExit(suppress = Throwable.class)
public static void remove(
@Advice.This final ThreadPoolExecutor zis,
@Advice.This final ThreadPoolExecutor tpe,
@Advice.Return(readOnly = false) Runnable removed) {
if (TPEHelper.shouldPropagate(
InstrumentationContext.get(ThreadPoolExecutor.class, Boolean.class), zis)) {
if (TPEHelper.shouldPropagate(tpe)) {
if (TPEHelper.useWrapping(removed)) {
if (removed instanceof Wrapper) {
Wrapper<?> wrapper = ((Wrapper<?>) removed);
Expand Down