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,6 +341,20 @@ 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) {
@@ -350,25 +364,46 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
         /*
          * 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.
          */
-        builder =
+        synchronized (INSTALLED_CONTEXT_MATCHERS) {
-            builder
+          // FIXME: This makes an assumption that class loader matchers for instrumenters that use
-                .type(safeHasSuperType(named(entry.getKey())), instrumenter.classLoaderMatcher())
+          // same context classes should be the same - which seems reasonable, but is not checked.
-                .and(safeToInjectFieldsMatcher())
+          // Addressing this properly requires some notion of 'compound intrumenters' which we
-                .and(Default.NOT_DECORATOR_MATCHER)
+          // currently do not have.
-                .transform(AgentBuilder.Transformer.NoOp.INSTANCE);
+          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);
-         * We inject helpers here as well as when instrumentation is applied to ensure that helpers
+          INSTALLED_CONTEXT_MATCHERS.add(entry);
-         * are present even if instrumented classes are not loaded, but classes with state fields
-         * added are loaded (e.g. sun.net.www.protocol.https.HttpsURLConnectionImpl).
-         */
-        builder = injectHelpersIntoBootstrapClassloader(builder);
 
-        builder =
+          /*
-            builder.transform(
+           * For each context store defined in a current instrumentation we create an agent builder
-                getTransformerForASMVisitor(
+           * that injects necessary fields.
-                    getFieldInjectionVisitor(entry.getKey(), entry.getValue())));
+           */
+          builder =
+              builder
+                  .type(safeHasSuperType(named(entry.getKey())), instrumenter.classLoaderMatcher())
+                  .and(safeToInjectFieldsMatcher())
+                  .and(Default.NOT_DECORATOR_MATCHER)
+                  .transform(AgentBuilder.Transformer.NoOp.INSTANCE);
+
+          /*
+           * We inject helpers here as well as when instrumentation is applied to ensure that
+           * helpers are present even if instrumented classes are not loaded, but classes with state
+           * fields added are loaded (e.g. sun.net.www.protocol.https.HttpsURLConnectionImpl).
+           */
+          builder = injectHelpersIntoBootstrapClassloader(builder);
+
+          builder =
+              builder.transform(
+                  getTransformerForASMVisitor(
+                      getFieldInjectionVisitor(entry.getKey(), entry.getValue())));
+        }
       }
     }
     return builder;