Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

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

Merged
merged 1 commit into from
Mar 17, 2024
Merged

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

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
88 changes: 57 additions & 31 deletions .../agent-builder/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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) {
Expand All @@ -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) {
Expand All @@ -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));
}
Expand All @@ -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
Expand Down Expand Up @@ -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) {
Expand All @@ -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);
Expand Down