Re-use helpers across all instrumenters in a module (#6822)

DataDog · Mar 17, 2024 · f674b8b · f674b8b
1 parent 36e924e
commit f674b8b
Showing 1 changed file with 57 additions and 31 deletions.
diff --git a/.../agent-builder/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java b/.../agent-builder/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java
@@ -50,15 +50,22 @@ public final class CombiningTransformerBuilder
       new HashMap<>();
 
   private final AgentBuilder agentBuilder;
+  private int nextSupplementaryId;
 
   private final List<MatchRecorder> matchers = new ArrayList<>();
   private final BitSet knownTypesMask;
   private AdviceStack[] transformers;
-  private int nextSupplementaryId;
+
+  // module defined matchers and transformers, shared across members
+  private ElementMatcher<? super MethodDescription> ignoredMethods;
+  private ElementMatcher<ClassLoader> classLoaderMatcher;
+  private Map<String, String> contextStore;
+  private AgentBuilder.Transformer contextRequestRewriter;
+  private HelperTransformer helperTransformer;
+  private MuzzleCheck muzzle;
 
   // temporary buffer for collecting advice; reset for each instrumenter
   private final List<AgentBuilder.Transformer> advice = new ArrayList<>();
-  private ElementMatcher<? super MethodDescription> ignoredMethods;
 
   /**
    * Post processor to be applied to instrumenter advices if they implement {@link
@@ -68,22 +75,49 @@ public final class CombiningTransformerBuilder
 
   public CombiningTransformerBuilder(AgentBuilder agentBuilder, int maxInstrumentationId) {
     this.agentBuilder = agentBuilder;
-    int maxInstrumentationCount = maxInstrumentationId + 1;
-    this.knownTypesMask = new BitSet(maxInstrumentationCount);
-    this.transformers = new AdviceStack[maxInstrumentationCount];
-    this.nextSupplementaryId = maxInstrumentationId + 1;
+    int transformationCount = maxInstrumentationId + 1;
+    this.nextSupplementaryId = transformationCount;
+    this.knownTypesMask = new BitSet(transformationCount);
+    this.transformers = new AdviceStack[transformationCount];
   }
 
+  /** Builds matchers and transformers for an instrumentation module and its members. */
   public void applyInstrumentation(InstrumenterModule module) {
     if (module.isEnabled()) {
       InstrumenterState.registerInstrumentation(module);
+      prepareInstrumentation(module);
       for (Instrumenter member : module.typeInstrumentations()) {
-        buildInstrumentation(module, member);
+        buildTypeInstrumentation(module, member);
       }
     }
   }
 
-  private void buildInstrumentation(InstrumenterModule module, Instrumenter member) {
+  /** Prepares shared matchers and transformers defined by an instrumentation module. */
+  private void prepareInstrumentation(InstrumenterModule module) {
+    ignoredMethods = module.methodIgnoreMatcher();
+    classLoaderMatcher = module.classLoaderMatcher();
+    contextStore = module.contextStore();
+
+    contextRequestRewriter =
+        !contextStore.isEmpty()
+            ? new VisitingTransformer(
+                new FieldBackedContextRequestRewriter(contextStore, module.name()))
+            : null;
+
+    String[] helperClassNames = module.helperClassNames();
+    if (module.injectHelperDependencies()) {
+      helperClassNames = HelperScanner.withClassDependencies(helperClassNames);
+    }
+    helperTransformer =
+        helperClassNames.length > 0
+            ? new HelperTransformer(module.getClass().getSimpleName(), helperClassNames)
+            : null;
+
+    muzzle = new MuzzleCheck(module);
+  }
+
+  /** Builds a type-specific transformer, controlled by one or more matchers. */
+  private void buildTypeInstrumentation(InstrumenterModule module, Instrumenter member) {
 
     int id = module.instrumentationId();
     if (module != member) {
@@ -96,11 +130,11 @@ private void buildInstrumentation(InstrumenterModule module, Instrumenter member
       }
     }
 
-    buildInstrumentationMatcher(module, member, id);
-    buildInstrumentationAdvice(module, member, id);
+    buildTypeMatcher(member, id);
+    buildTypeAdvice(member, id);
   }
 
-  private void buildInstrumentationMatcher(InstrumenterModule module, Instrumenter member, int id) {
+  private void buildTypeMatcher(Instrumenter member, int id) {
 
     if (member instanceof Instrumenter.ForSingleType
         || member instanceof Instrumenter.ForKnownTypes) {
@@ -124,7 +158,6 @@ private void buildInstrumentationMatcher(InstrumenterModule module, Instrumenter
       matchers.add(new MatchRecorder.ForHierarchy(id, (Instrumenter.ForTypeHierarchy) member));
     }
 
-    ElementMatcher<ClassLoader> classLoaderMatcher = module.classLoaderMatcher();
     if (classLoaderMatcher != ANY_CLASS_LOADER) {
       matchers.add(new MatchRecorder.NarrowLocation(id, classLoaderMatcher));
     }
@@ -135,42 +168,35 @@ private void buildInstrumentationMatcher(InstrumenterModule module, Instrumenter
               id, ((Instrumenter.WithTypeStructure) member).structureMatcher()));
     }
 
-    matchers.add(new MatchRecorder.NarrowLocation(id, new MuzzleCheck(module)));
+    matchers.add(new MatchRecorder.NarrowLocation(id, muzzle));
   }
 
-  private void buildInstrumentationAdvice(InstrumenterModule module, Instrumenter member, int id) {
+  private void buildTypeAdvice(Instrumenter member, int id) {
 
     postProcessor =
         member instanceof WithPostProcessor ? ((WithPostProcessor) member).postProcessor() : null;
 
-    String[] helperClassNames = module.helperClassNames();
-    if (module.injectHelperDependencies()) {
-      helperClassNames = HelperScanner.withClassDependencies(helperClassNames);
-    }
-    if (helperClassNames.length > 0) {
-      advice.add(new HelperTransformer(module.getClass().getSimpleName(), helperClassNames));
+    if (null != helperTransformer) {
+      advice.add(helperTransformer);
     }
 
-    Map<String, String> contextStore = module.contextStore();
-    if (!contextStore.isEmpty()) {
+    if (null != contextRequestRewriter) {
+      registerContextStoreInjection(member, contextStore);
       // rewrite context store access to call FieldBackedContextStores with assigned store-id
-      advice.add(
-          new VisitingTransformer(
-              new FieldBackedContextRequestRewriter(contextStore, module.name())));
-
-      registerContextStoreInjection(module, member, contextStore);
+      advice.add(contextRequestRewriter);
     }
 
-    ignoredMethods = module.methodIgnoreMatcher();
     if (member instanceof Instrumenter.HasTypeAdvice) {
       ((Instrumenter.HasTypeAdvice) member).typeAdvice(this);
     }
     if (member instanceof Instrumenter.HasMethodAdvice) {
       ((Instrumenter.HasMethodAdvice) member).methodAdvice(this);
     }
+
+    // record the advice collected for this type instrumentation
     transformers[id] = new AdviceStack(advice);
 
-    advice.clear();
+    advice.clear(); // reset for next type instrumentation
   }
 
   @Override
@@ -203,7 +229,7 @@ private void applyContextStoreInjection() {
 
   /** Tracks which class-loader matchers are associated with each store request. */
   private void registerContextStoreInjection(
-      InstrumenterModule module, Instrumenter member, Map<String, String> contextStore) {
+      Instrumenter member, Map<String, String> contextStore) {
     ElementMatcher<ClassLoader> activation;
 
     if (member instanceof Instrumenter.ForBootstrap) {
@@ -219,7 +245,7 @@ private void registerContextStoreInjection(
       activation = ANY_CLASS_LOADER;
     }
 
-    activation = requireBoth(activation, module.classLoaderMatcher());
+    activation = requireBoth(activation, classLoaderMatcher);
 
     for (Map.Entry<String, String> storeEntry : contextStore.entrySet()) {
       ElementMatcher<ClassLoader> oldActivation = contextStoreInjection.get(storeEntry);