Memoized type matching #5026
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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()); | ||
|
There was a problem hiding this comment. This is not a comment on this PR, but more of an idea. I assume that this isn't hit that often now with memoization, but would it make sense to let the exclude filter build up a trie and use that instead of the hash set lookup and string comparisons? There was a problem hiding this comment. Good idea, worth looking into |
||
| } | ||
| } | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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; | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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; | ||
| } | ||
| } |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This seems less exact than the weak store id bit sets in
ShouldInjectFieldsStatebecause of themaybeSkip, if I read the code correctly. Is there a possibility that we use a weak store unnecessary, or are could we just add redundant code that never gets used?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes this is potentially less exact - let me add some tests to confirm when it would make a difference, I can then make it more exact based on that.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It turns out the
java-completablefuturetests already demonstrate the difference - and the new approach is actually more deterministic. For each context class that can be skipped/excluded we collect store ids associated with that context class (ie. the value side of all the context-value associations) and therefore might need to be 'skipped' when that excluded class is a super-class - by delegating to the weak map instead.When we do find a context class that should be excluded and it's a super-class of the class we're injecting then we notify the field-injection code that super-class requests involving those stores should delegate to the weak map. That part of the injected code will only be used if someone tries to push/get values to those stores. If a store isn't used then it just means we injected an unused case block.
The old approach was more lazy in that it only recorded a store as needing weak map delegation at the time the match was made. Unfortunately variations in class-loading order can make this non-deterministic - for example with the
java-completablefuturetests theUniCompletionclass is processed before its superclasses (since both are not yet loaded at matching time) so it doesn't realize that it should delegateRunnable/ForkJoinTasksuper-class requests to the weak map. This changes if theCompletionclass is forcibly loaded first, because it now sees the skipped class before processingUniCompletion.In practice this doesn't make a difference because we add a last-resort
catchblock to handle unexpected requests and thiscatchblock ends up being used by the old approach. Whereas with the new approach we always add those weak map case blocks (and avoid triggering the exceptionalcatchblock)There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks @mcculls for the detailed explanation. Determinism FTW!