Merge pull request #1261 from DataDog/mar-kolya/speedup-context-injection-matchers

speedup context injection matchers

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

@@ -7,6 +7,7 @@ import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.none;
 
+import datadog.trace.agent.tooling.context.FieldBackedProvider;
 import datadog.trace.api.Config;
 import java.lang.instrument.Instrumentation;
 import java.util.ArrayList;
@@ -53,6 +54,8 @@ public class AgentInstaller {
 
     addByteBuddyRawSetting();
 
+    FieldBackedProvider.resetContextMatchers();
+
     AgentBuilder agentBuilder =
         new AgentBuilder.Default()
             .disableClassFormatChanges()

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedProvider.java

@@ -2,7 +2,6 @@ package datadog.trace.agent.tooling.context;
 
 import static datadog.trace.agent.tooling.ClassLoaderMatcher.BOOTSTRAP_CLASSLOADER;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.DDElementMatchers.safeHasSuperType;
-import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 
 import datadog.trace.agent.tooling.HelperInjector;
@@ -20,6 +19,7 @@ import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.Map;
 import java.util.Set;
@@ -341,12 +341,46 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
     };
   }
 
+  /*
+  Set of pairs (context holder, context class) for which we have matchers installed.
+  We use this to make sure we do not install matchers repeatedly for cases when same
+  context class is used by multiple instrumentations.
+   */
+  private static final Set<Map.Entry<String, String>> INSTALLED_CONTEXT_MATCHERS = new HashSet<>();
+
+  /** Clear set that prevents multiple matchers for same context class */
+  public static void resetContextMatchers() {
+    synchronized (INSTALLED_CONTEXT_MATCHERS) {
+      INSTALLED_CONTEXT_MATCHERS.clear();
+    }
+  }
+
   @Override
   public AgentBuilder.Identified.Extendable additionalInstrumentation(
       AgentBuilder.Identified.Extendable builder) {
 
     if (fieldInjectionEnabled) {
       for (final Map.Entry<String, String> entry : instrumenter.contextStore().entrySet()) {
+        /*
+         * For each context store defined in a current instrumentation we create an agent builder
+         * that injects necessary fields.
+         * Note: this synchronization should not have any impact on performance
+         * since this is done when agent builder is being made, it doesn't affect actual
+         * class transformation.
+         */
+        synchronized (INSTALLED_CONTEXT_MATCHERS) {
+          // FIXME: This makes an assumption that class loader matchers for instrumenters that use
+          // same context classes should be the same - which seems reasonable, but is not checked.
+          // Addressing this properly requires some notion of 'compound intrumenters' which we
+          // currently do not have.
+          if (INSTALLED_CONTEXT_MATCHERS.contains(entry)) {
+            log.debug("Skipping builder for {} {}", instrumenter.getClass().getName(), entry);
+            continue;
+          }
+
+          log.debug("Making builder for {} {}", instrumenter.getClass().getName(), entry);
+          INSTALLED_CONTEXT_MATCHERS.add(entry);
+
           /*
            * For each context store defined in a current instrumentation we create an agent builder
            * that injects necessary fields.
@@ -359,9 +393,9 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
                   .transform(AgentBuilder.Transformer.NoOp.INSTANCE);
 
           /*
-         * We inject helpers here as well as when instrumentation is applied to ensure that helpers
+           * We inject helpers here as well as when instrumentation is applied to ensure that
-         * are present even if instrumented classes are not loaded, but classes with state fields
+           * helpers are present even if instrumented classes are not loaded, but classes with state
-         * added are loaded (e.g. sun.net.www.protocol.https.HttpsURLConnectionImpl).
+           * fields added are loaded (e.g. sun.net.www.protocol.https.HttpsURLConnectionImpl).
            */
           builder = injectHelpersIntoBootstrapClassloader(builder);
 
@@ -371,6 +405,7 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
                       getFieldInjectionVisitor(entry.getKey(), entry.getValue())));
         }
       }
+    }
     return builder;
   }