diff --git a/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java b/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java index 1fe1a07845d..c6540098cdc 100644 --- a/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java +++ b/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java @@ -204,7 +204,10 @@ public static ClassFileTransformer installBytebuddyAgent( int installedCount = 0; for (InstrumenterModule module : instrumenterModules) { - if (!module.isApplicable(enabledSystems)) { + final boolean hasGeneralPurposeAdvices = + module instanceof Instrumenter.HasGeneralPurposeAdvices; + final boolean isApplicable = module.isApplicable(enabledSystems); + if (!hasGeneralPurposeAdvices && !module.isApplicable(enabledSystems)) { if (DEBUG) { log.debug("Not applicable - instrumentation.class={}", module.getClass().getName()); } @@ -214,7 +217,7 @@ public static ClassFileTransformer installBytebuddyAgent( log.debug("Loading - instrumentation.class={}", module.getClass().getName()); } try { - transformerBuilder.applyInstrumentation(module); + transformerBuilder.applyInstrumentation(module, isApplicable); installedCount++; } catch (Exception | LinkageError e) { log.error("Failed to load - instrumentation.class={}", module.getClass().getName(), e); diff --git a/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java b/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java index 21a0b4c6e47..65fabc06632 100644 --- a/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java +++ b/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java @@ -7,6 +7,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.declaresAnnotation; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf; +import static java.util.Collections.emptySet; import static net.bytebuddy.matcher.ElementMatchers.not; import datadog.trace.agent.tooling.bytebuddy.ExceptionHandlers; @@ -25,6 +26,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Set; import net.bytebuddy.agent.builder.AgentBuilder; import net.bytebuddy.asm.Advice; import net.bytebuddy.asm.AsmVisitorWrapper; @@ -75,6 +77,7 @@ public final class CombiningTransformerBuilder private HelperTransformer helperTransformer; private Advice.PostProcessor.Factory postProcessor; private MuzzleCheck muzzle; + private Set generalPurposeAdviceClasses = null; // temporary buffer for collecting advice; reset for each instrumenter private final List advice = new ArrayList<>(); @@ -92,7 +95,8 @@ public CombiningTransformerBuilder( } /** Builds matchers and transformers for an instrumentation module and its members. */ - public void applyInstrumentation(InstrumenterModule module) { + public void applyInstrumentation( + InstrumenterModule module, boolean isModuleApplicableOnTargetSystems) { if (module.isEnabled()) { int instrumentationId = instrumenterIndex.instrumentationId(module); if (instrumentationId < 0) { @@ -100,7 +104,7 @@ public void applyInstrumentation(InstrumenterModule module) { instrumentationId = nextRuntimeInstrumentationId++; } InstrumenterState.registerInstrumentation(module, instrumentationId); - prepareInstrumentation(module, instrumentationId); + prepareInstrumentation(module, instrumentationId, isModuleApplicableOnTargetSystems); for (Instrumenter member : module.typeInstrumentations()) { buildTypeInstrumentation(member); } @@ -108,7 +112,8 @@ public void applyInstrumentation(InstrumenterModule module) { } /** Prepares shared matchers and transformers defined by an instrumentation module. */ - private void prepareInstrumentation(InstrumenterModule module, int instrumentationId) { + private void prepareInstrumentation( + InstrumenterModule module, int instrumentationId, boolean isModuleApplicableOnTargetSystems) { ignoredMethods = module.methodIgnoreMatcher(); classLoaderMatcher = module.classLoaderMatcher(); contextStore = module.contextStore(); @@ -137,6 +142,16 @@ private void prepareInstrumentation(InstrumenterModule module, int instrumentati postProcessor = module.postProcessor(); muzzle = new MuzzleCheck(module, instrumentationId); + + if (!isModuleApplicableOnTargetSystems) { + if (module instanceof Instrumenter.HasGeneralPurposeAdvices) { + generalPurposeAdviceClasses = + ((Instrumenter.HasGeneralPurposeAdvices) module).generalPurposeAdviceClasses(); + if (generalPurposeAdviceClasses == null) { + this.generalPurposeAdviceClasses = emptySet(); + } + } + } } /** Builds a type-specific transformer, controlled by one or more matchers. */ @@ -239,7 +254,10 @@ public void applyAdvice(Instrumenter.TransformingAdvice typeAdvice) { } @Override - public void applyAdvice(ElementMatcher matcher, String adviceClass) { + public void applyAdvice( + ElementMatcher matcher, + String adviceClass, + String... additionalAdviceClasses) { Advice.WithCustomMapping customMapping = Advice.withCustomMapping(); if (postProcessor != null) { customMapping = customMapping.with(postProcessor); @@ -254,7 +272,18 @@ public void applyAdvice(ElementMatcher matcher, Strin } else { forAdvice = forAdvice.include(adviceLoader); } - advice.add(forAdvice.advice(not(ignoredMethods).and(matcher), adviceClass)); + if (generalPurposeAdviceClasses == null || generalPurposeAdviceClasses.contains(adviceClass)) { + forAdvice = forAdvice.advice(not(ignoredMethods).and(matcher), adviceClass); + } + if (additionalAdviceClasses != null) { + for (String adviceClassName : additionalAdviceClasses) { + if (generalPurposeAdviceClasses == null + || generalPurposeAdviceClasses.contains(adviceClassName)) { + forAdvice = forAdvice.advice(not(ignoredMethods).and(matcher), adviceClassName); + } + } + } + advice.add(forAdvice); } public ClassFileTransformer installOn(Instrumentation instrumentation) { diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java index a92fb326bd4..3d991cdfa08 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java @@ -5,6 +5,7 @@ import java.security.ProtectionDomain; import java.util.Collection; +import java.util.Set; import net.bytebuddy.asm.AsmVisitorWrapper; import net.bytebuddy.description.method.MethodDescription; import net.bytebuddy.description.type.TypeDescription; @@ -90,6 +91,10 @@ interface HasTypeAdvice extends Instrumenter { void typeAdvice(TypeTransformer transformer); } + interface HasGeneralPurposeAdvices extends Instrumenter { + Set generalPurposeAdviceClasses(); + } + /** Instrumentation that provides advice specific to one or more methods. */ interface HasMethodAdvice extends Instrumenter { /** @@ -110,7 +115,10 @@ default void applyAdvice(AsmVisitorWrapper typeVisitor) { /** Applies method advice from an instrumentation that {@link HasMethodAdvice}. */ interface MethodTransformer { - void applyAdvice(ElementMatcher matcher, String adviceClass); + void applyAdvice( + ElementMatcher matcher, + String adviceClass, + String... additionalAdviceClasses); } /** Contributes a transformation step to the dynamic type builder. */ diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java index 62ddf7c32fe..3dc62f69442 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java @@ -327,4 +327,15 @@ public boolean isApplicable(Set enabledSystems) { return enabledSystems.contains(TargetSystem.CIVISIBILITY); } } + + public abstract static class ContextTracking extends InstrumenterModule { + public ContextTracking(String instrumentationName, String... additionalNames) { + super(instrumentationName, additionalNames); + } + + @Override + public boolean isApplicable(Set enabledSystems) { + return true; + } + } } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerator.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerator.java index a2d5a035b1c..a1490670e51 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerator.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerator.java @@ -1,5 +1,7 @@ package datadog.trace.agent.tooling.muzzle; +import static java.util.Arrays.asList; + import datadog.trace.agent.tooling.AdviceShader; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; @@ -7,7 +9,6 @@ import java.io.IOException; import java.nio.file.Files; import java.util.ArrayList; -import java.util.Arrays; import java.util.Collections; import java.util.HashSet; import java.util.LinkedHashMap; @@ -85,7 +86,13 @@ private static Reference[] generateReferences( final Set referenceSources = new HashSet<>(); final Map references = new LinkedHashMap<>(); final Set adviceClasses = new HashSet<>(); - instrumenter.methodAdvice((matcher, adviceClass) -> adviceClasses.add(adviceClass)); + instrumenter.methodAdvice( + (matcher, adviceClass, otherAdviceClasses) -> { + adviceClasses.add(adviceClass); + if (otherAdviceClasses != null) { + adviceClasses.addAll(asList(otherAdviceClasses)); + } + }); ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader(); for (String adviceClass : adviceClasses) { if (referenceSources.add(adviceClass)) { @@ -107,7 +114,7 @@ private static Reference[] generateReferences( /** This code is generated in a separate side-class. */ private static byte[] generateMuzzleClass(InstrumenterModule module) { - Set ignoredClassNames = new HashSet<>(Arrays.asList(module.muzzleIgnoredClassNames())); + Set ignoredClassNames = new HashSet<>(asList(module.muzzleIgnoredClassNames())); AdviceShader adviceShader = AdviceShader.with(module.adviceShading()); List references = new ArrayList<>(); diff --git a/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java b/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java index ba3e5ad848f..0e351763438 100644 --- a/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java @@ -9,6 +9,7 @@ import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE; import static datadog.trace.instrumentation.tomcat.TomcatDecorator.DD_PARENT_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.tomcat.TomcatDecorator.DECORATE; +import static java.util.Collections.singleton; import static java.util.Collections.singletonMap; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; @@ -23,10 +24,12 @@ import datadog.trace.api.gateway.Flow; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter; import java.util.Arrays; import java.util.Collection; import java.util.Map; +import java.util.Set; import net.bytebuddy.asm.Advice; import org.apache.catalina.connector.CoyoteAdapter; import org.apache.catalina.connector.Request; @@ -34,7 +37,10 @@ @AutoService(InstrumenterModule.class) public final class TomcatServerInstrumentation extends InstrumenterModule.Tracing - implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice, ExcludeFilterProvider { + implements Instrumenter.ForSingleType, + Instrumenter.HasMethodAdvice, + Instrumenter.HasGeneralPurposeAdvices, + ExcludeFilterProvider { public TomcatServerInstrumentation() { super("tomcat"); @@ -84,6 +90,8 @@ public void methodAdvice(MethodTransformer transformer) { named("service") .and(takesArgument(0, named("org.apache.coyote.Request"))) .and(takesArgument(1, named("org.apache.coyote.Response"))), + TomcatServerInstrumentation.class.getName() + + "$ContextTrackingAdvice", // context tracking must be applied first TomcatServerInstrumentation.class.getName() + "$ServiceAdvice"); transformer.applyAdvice( named("postParseRequest") @@ -111,21 +119,41 @@ public void methodAdvice(MethodTransformer transformer) { "org.apache.tomcat.util.net.NioBlockingSelector$BlockPoller")); } - public static class ServiceAdvice { + @Override + public Set generalPurposeAdviceClasses() { + return singleton(packageName + ".TomcatServerInstrumentation$ContextTrackingAdvice"); + } + public static class ContextTrackingAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static ContextScope onService(@Advice.Argument(0) org.apache.coyote.Request req) { - + public static void extractParent( + @Advice.Argument(0) org.apache.coyote.Request req, + @Advice.Local("parentScope") ContextScope parentScope) { Object existingCtx = req.getAttribute(DD_CONTEXT_ATTRIBUTE); if (existingCtx instanceof Context) { // Request already gone through initial processing, so just attach the context. - return ((Context) existingCtx).attach(); + parentScope = ((Context) existingCtx).attach(); + } else { + final Context parentContext = DECORATE.extract(req); + req.setAttribute(DD_PARENT_CONTEXT_ATTRIBUTE, parentContext); + parentScope = parentContext.attach(); } + } + + @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) + public static void closeScope(@Advice.Local("parentScope") ContextScope scope) { + scope.close(); + } + } - final Context parentContext = DECORATE.extract(req); - req.setAttribute(DD_PARENT_CONTEXT_ATTRIBUTE, parentContext); - final Context context = DECORATE.startSpan(req, parentContext); - final ContextScope scope = context.attach(); + public static class ServiceAdvice { + + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onService( + @Advice.Argument(0) org.apache.coyote.Request req, + @Advice.Local("serverScope") ContextScope serverScope) { + final Context context = DECORATE.startSpan(req, Java8BytecodeBridge.getCurrentContext()); + serverScope = context.attach(); // This span is finished when Request.recycle() is called by RequestInstrumentation. final AgentSpan span = spanFromContext(context); @@ -134,12 +162,11 @@ public static ContextScope onService(@Advice.Argument(0) org.apache.coyote.Reque req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); - return scope; } @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) - public static void closeScope(@Advice.Enter final ContextScope scope) { - scope.close(); + public static void closeScope(@Advice.Local("serverScope") ContextScope serverScope) { + serverScope.close(); } private void muzzleCheck(CoyoteAdapter adapter, Request request, Response response)