test

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);

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<String> generalPurposeAdviceClasses = null;
 
   // temporary buffer for collecting advice; reset for each instrumenter
   private final List<AgentBuilder.Transformer> advice = new ArrayList<>();
@@ -92,23 +95,25 @@ 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) {
         // this is a non-indexed instrumentation configured at runtime
         instrumentationId = nextRuntimeInstrumentationId++;
       }
       InstrumenterState.registerInstrumentation(module, instrumentationId);
-      prepareInstrumentation(module, instrumentationId);
+      prepareInstrumentation(module, instrumentationId, isModuleApplicableOnTargetSystems);
       for (Instrumenter member : module.typeInstrumentations()) {
         buildTypeInstrumentation(member);
       }
     }
   }
 
   /** 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<? super MethodDescription> matcher, String adviceClass) {
+  public void applyAdvice(
+      ElementMatcher<? super MethodDescription> 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<? super MethodDescription> 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) {

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<String> 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<? super MethodDescription> matcher, String adviceClass);
+    void applyAdvice(
+        ElementMatcher<? super MethodDescription> matcher,
+        String adviceClass,
+        String... additionalAdviceClasses);
   }
 
   /** Contributes a transformation step to the dynamic type builder. */

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java

@@ -327,4 +327,15 @@ public boolean isApplicable(Set<TargetSystem> 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<TargetSystem> enabledSystems) {
+      return true;
+    }
+  }
 }

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerator.java

@@ -1,13 +1,14 @@
 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;
 import java.io.File;
 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<String> referenceSources = new HashSet<>();
     final Map<String, Reference> references = new LinkedHashMap<>();
     final Set<String> 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<String> ignoredClassNames = new HashSet<>(Arrays.asList(module.muzzleIgnoredClassNames()));
+    Set<String> ignoredClassNames = new HashSet<>(asList(module.muzzleIgnoredClassNames()));
     AdviceShader adviceShader = AdviceShader.with(module.adviceShading());
 
     List<Reference> references = new ArrayList<>();

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,18 +24,23 @@
 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;
 import org.apache.catalina.connector.Response;
 
 @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<String> 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)