Merge pull request #5026 from DataDog/mcculls/memoized-type-matching

Memoized type matching, to restore previous approach use -Ddd.resolver.cache.config=NO_MEMOS

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

@@ -7,6 +7,7 @@
 import datadog.trace.agent.tooling.bytebuddy.DDOutlinePoolStrategy;
 import datadog.trace.agent.tooling.bytebuddy.SharedTypePools;
 import datadog.trace.agent.tooling.bytebuddy.matcher.DDElementMatchers;
+import datadog.trace.agent.tooling.bytebuddy.memoize.MemoizedMatchers;
 import datadog.trace.agent.tooling.usm.UsmExtractorImpl;
 import datadog.trace.agent.tooling.usm.UsmMessageFactoryImpl;
 import datadog.trace.api.InstrumenterConfig;
@@ -101,7 +102,11 @@ public static ClassFileTransformer installBytebuddyAgent(
       DDCachingPoolStrategy.registerAsSupplier();
     }
 
-    DDElementMatchers.registerAsSupplier();
+    if (InstrumenterConfig.get().isResolverMemoizingEnabled()) {
+      MemoizedMatchers.registerAsSupplier();
+    } else {
+      DDElementMatchers.registerAsSupplier();
+    }
 
     if (enabledSystems.contains(Instrumenter.TargetSystem.USM)) {
       UsmMessageFactoryImpl.registerAsSupplier();

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/matcher/ClassLoaderMatchers.java

@@ -89,7 +89,7 @@ public static ElementMatcher.Junction<ClassLoader> hasClassNamedOneOf(
     return new ElementMatcher.Junction.Disjunction<>(matchers);
   }
 
-  public static void reset() {
+  public static void resetState() {
     hasClassCache.clear();
   }
 

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/memoize/HasContextField.java

@@ -0,0 +1,61 @@
+package datadog.trace.agent.tooling.bytebuddy.memoize;
+
+import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
+import java.util.BitSet;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+/** Matches types that should have context-store fields injected. */
+final class HasContextField extends ElementMatcher.Junction.ForNonNullValues<TypeDescription> {
+  private final ElementMatcher<TypeDescription> storeMatcher;
+  private final ElementMatcher<TypeDescription> skipMatcher;
+  private final BitSet skippableStoreIds = new BitSet();
+
+  HasContextField(ElementMatcher<TypeDescription> storeMatcher) {
+    this(storeMatcher, null);
+  }
+
+  HasContextField(
+      ElementMatcher<TypeDescription> storeMatcher, ElementMatcher<TypeDescription> skipMatcher) {
+    this.storeMatcher = storeMatcher;
+    this.skipMatcher = skipMatcher;
+  }
+
+  @Override
+  protected boolean doMatch(TypeDescription target) {
+    // match first type in the hierarchy which injects this store, that isn't skipped
+    return storeMatcher.matches(target)
+        && (null == skipMatcher || !skipMatcher.matches(target))
+        && (!storeMatcher.matches(target.getSuperClass().asErasure()));
+  }
+
+  void maybeSkip(int contextStoreId) {
+    skippableStoreIds.set(contextStoreId);
+  }
+
+  boolean hasSuperStore(TypeDescription target, BitSet weakStoreIds) {
+    TypeDescription superTarget = target.getSuperClass().asErasure();
+    boolean hasSuperStore = storeMatcher.matches(superTarget);
+    if (hasSuperStore) {
+      // report if super-class was skipped from field-injection
+      if (null != skipMatcher && skipMatcher.matches(superTarget)) {
+        weakStoreIds.or(skippableStoreIds);
+      }
+    }
+    return hasSuperStore;
+  }
+
+  /** Matches types that would have had fields injected, but were explicitly excluded. */
+  static final class Skip extends ElementMatcher.Junction.ForNonNullValues<TypeDescription> {
+    private final ExcludeFilter.ExcludeType excludeType;
+
+    Skip(ExcludeFilter.ExcludeType excludeType) {
+      this.excludeType = excludeType;
+    }
+
+    @Override
+    protected boolean doMatch(TypeDescription target) {
+      return ExcludeFilter.exclude(excludeType, target.getName());
+    }
+  }
+}

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/memoize/HasSuperMethod.java

@@ -0,0 +1,83 @@
+package datadog.trace.agent.tooling.bytebuddy.memoize;
+
+import static net.bytebuddy.matcher.ElementMatchers.hasSignature;
+
+import java.util.HashSet;
+import java.util.Set;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDefinition;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.description.type.TypeList;
+import net.bytebuddy.matcher.ElementMatcher;
+
+/**
+ * Matches methods that are either a direct match or have a super-class method with the same
+ * signature that matches. Uses a memoized type matcher to reduce the potential search space.
+ */
+final class HasSuperMethod extends ElementMatcher.Junction.ForNonNullValues<MethodDescription> {
+  private final ElementMatcher<TypeDescription> typeMatcher;
+  private final ElementMatcher<? super MethodDescription> methodMatcher;
+
+  HasSuperMethod(
+      ElementMatcher<TypeDescription> typeMatcher,
+      ElementMatcher<? super MethodDescription> methodMatcher) {
+    this.typeMatcher = typeMatcher;
+    this.methodMatcher = methodMatcher;
+  }
+
+  @Override
+  protected boolean doMatch(MethodDescription target) {
+    if (target.isConstructor()) {
+      return false;
+    }
+
+    TypeDefinition type = target.getDeclaringType();
+    if (!typeMatcher.matches(type.asErasure())) {
+      return false; // no further matches recorded in hierarchy
+    } else if (methodMatcher.matches(target)) {
+      return true; // direct match, no need to check hierarchy
+    }
+
+    // need to search hierarchy; use expected signature to filter candidate methods
+    ElementMatcher<MethodDescription> signatureMatcher = hasSignature(target.asSignatureToken());
+    Set<String> visited = new HashSet<>();
+
+    if (interfaceMatches(type.getInterfaces(), signatureMatcher, visited)) {
+      return true;
+    }
+
+    type = type.getSuperClass();
+    while (null != type && typeMatcher.matches(type.asErasure())) {
+      for (MethodDescription method : type.getDeclaredMethods()) {
+        if (signatureMatcher.matches(method) && methodMatcher.matches(method)) {
+          return true;
+        }
+      }
+      if (interfaceMatches(type.getInterfaces(), signatureMatcher, visited)) {
+        return true;
+      }
+      type = type.getSuperClass();
+    }
+    return false;
+  }
+
+  private boolean interfaceMatches(
+      TypeList.Generic interfaces,
+      ElementMatcher<MethodDescription> signatureMatcher,
+      Set<String> visited) {
+    for (TypeDefinition type : interfaces) {
+      if (!visited.add(type.getTypeName()) || !typeMatcher.matches(type.asErasure())) {
+        continue; // skip if already visited or that part of the hierarchy doesn't match
+      }
+      for (MethodDescription method : type.getDeclaredMethods()) {
+        if (signatureMatcher.matches(method) && methodMatcher.matches(method)) {
+          return true;
+        }
+      }
+      if (interfaceMatches(type.getInterfaces(), signatureMatcher, visited)) {
+        return true;
+      }
+    }
+    return false;
+  }
+}

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/memoize/MemoizedMatchers.java

@@ -0,0 +1,130 @@
+package datadog.trace.agent.tooling.bytebuddy.memoize;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.agent.tooling.bytebuddy.memoize.Memoizer.MatcherKind.ANNOTATION;
+import static datadog.trace.agent.tooling.bytebuddy.memoize.Memoizer.MatcherKind.CLASS;
+import static datadog.trace.agent.tooling.bytebuddy.memoize.Memoizer.MatcherKind.FIELD;
+import static datadog.trace.agent.tooling.bytebuddy.memoize.Memoizer.MatcherKind.INTERFACE;
+import static datadog.trace.agent.tooling.bytebuddy.memoize.Memoizer.MatcherKind.METHOD;
+import static datadog.trace.agent.tooling.bytebuddy.memoize.Memoizer.MatcherKind.TYPE;
+import static datadog.trace.bootstrap.FieldBackedContextStores.getContextStoreId;
+
+import datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers;
+import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
+import java.util.BitSet;
+import java.util.HashMap;
+import java.util.Map;
+import net.bytebuddy.description.NamedElement;
+import net.bytebuddy.description.field.FieldDescription;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+/** Supplies memoized matchers. */
+public final class MemoizedMatchers implements HierarchyMatchers.Supplier {
+  public static void registerAsSupplier() {
+    HierarchyMatchers.registerIfAbsent(new MemoizedMatchers());
+    Memoizer.resetState();
+  }
+
+  @Override
+  public ElementMatcher.Junction<TypeDescription> declaresAnnotation(
+      ElementMatcher.Junction<? super NamedElement> matcher) {
+    return Memoizer.prepare(ANNOTATION, matcher, false);
+  }
+
+  @Override
+  public ElementMatcher.Junction<TypeDescription> declaresField(
+      ElementMatcher.Junction<? super FieldDescription> matcher) {
+    return Memoizer.prepare(FIELD, matcher, false);
+  }
+
+  @Override
+  public ElementMatcher.Junction<TypeDescription> declaresMethod(
+      ElementMatcher.Junction<? super MethodDescription> matcher) {
+    return Memoizer.prepare(METHOD, matcher, false);
+  }
+
+  @Override
+  public ElementMatcher.Junction<TypeDescription> concreteClass() {
+    return Memoizer.isConcrete;
+  }
+
+  @Override
+  public ElementMatcher.Junction<TypeDescription> extendsClass(
+      ElementMatcher.Junction<? super TypeDescription> matcher) {
+    return Memoizer.prepare(CLASS, matcher, true);
+  }
+
+  @Override
+  public ElementMatcher.Junction<TypeDescription> implementsInterface(
+      ElementMatcher.Junction<? super TypeDescription> matcher) {
+    return Memoizer.isClass.and(Memoizer.prepare(INTERFACE, matcher, true));
+  }
+
+  @Override
+  public ElementMatcher.Junction<TypeDescription> hasInterface(
+      ElementMatcher.Junction<? super TypeDescription> matcher) {
+    return Memoizer.prepare(INTERFACE, matcher, true);
+  }
+
+  @Override
+  public ElementMatcher.Junction<TypeDescription> hasSuperType(
+      ElementMatcher.Junction<? super TypeDescription> matcher) {
+    return Memoizer.isClass.and(Memoizer.prepare(TYPE, matcher, true));
+  }
+
+  @Override
+  public ElementMatcher.Junction<MethodDescription> hasSuperMethod(
+      ElementMatcher.Junction<? super MethodDescription> matcher) {
+    return new HasSuperMethod(Memoizer.prepare(METHOD, matcher, true), matcher);
+  }
+
+  /** Keeps track of which context-field matchers we've supplied so far. */
+  private static final Map<String, HasContextField> contextFields = new HashMap<>();
+
+  @Override
+  public ElementMatcher.Junction<TypeDescription> declaresContextField(
+      String keyType, String contextType) {
+
+    // is there a chance a type might match, but be excluded (skipped) from field-injection?
+    ExcludeFilter.ExcludeType excludeType = ExcludeFilter.ExcludeType.fromFieldType(keyType);
+
+    HasContextField contextField = contextFields.get(keyType);
+    if (null == contextField) {
+      ElementMatcher<TypeDescription> storeMatcher = hasSuperType(named(keyType));
+      if (null != excludeType) {
+        ElementMatcher<TypeDescription> skipMatcher =
+            Memoizer.prepare(CLASS, new HasContextField.Skip(excludeType), true);
+        contextField = new HasContextField(storeMatcher, skipMatcher);
+      } else {
+        contextField = new HasContextField(storeMatcher);
+      }
+      contextFields.put(keyType, contextField);
+    }
+
+    if (null != excludeType) {
+      // record that this store may be skipped from field-injection
+      contextField.maybeSkip(getContextStoreId(keyType, contextType));
+    }
+
+    return contextField;
+  }
+
+  /**
+   * Returns {@code true} if ancestors of the given type have any context-stores.
+   *
+   * <p>Stores which were skipped from field-injection in the super-class hierarchy are recorded in
+   * {@code weakStoreIds} so the appropriate weak-map delegation can be applied when injecting code
+   * to handle context-store requests.
+   */
+  public static boolean hasSuperStores(TypeDescription target, BitSet weakStoreIds) {
+    boolean hasSuperStores = false;
+    for (HasContextField contextField : contextFields.values()) {
+      if (contextField.hasSuperStore(target, weakStoreIds)) {
+        hasSuperStores = true;
+      }
+    }
+    return hasSuperStores;
+  }
+}