Skip to content

Commit

Permalink
Example of how to group related instrumentations under a single modul…
Browse files Browse the repository at this point in the history
…e (AwsSdkModule) (#8141)
  • Loading branch information
mcculls authored Jan 2, 2025
1 parent a9dc3be commit ace20ff
Show file tree
Hide file tree
Showing 5 changed files with 138 additions and 101 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,7 @@
import com.amazonaws.AmazonClientException;
import com.amazonaws.Request;
import com.amazonaws.handlers.RequestHandler2;
import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.Instrumenter;
import datadog.trace.agent.tooling.InstrumenterModule;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
Expand All @@ -22,26 +20,14 @@
* {@link AmazonClientException} (for example an error thrown by another handler). In these cases
* {@link RequestHandler2#afterError} is not called.
*/
@AutoService(InstrumenterModule.class)
public class AWSHttpClientInstrumentation extends InstrumenterModule.Tracing
public class AWSHttpClientInstrumentation
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {

public AWSHttpClientInstrumentation() {
super("aws-sdk");
}

@Override
public String instrumentedType() {
return "com.amazonaws.http.AmazonHttpClient";
}

@Override
public String[] helperClassNames() {
return new String[] {
packageName + ".OnErrorDecorator", packageName + ".AwsNameCache",
};
}

@Override
public void methodAdvice(MethodTransformer transformer) {
transformer.applyAdvice(
Expand Down Expand Up @@ -74,57 +60,4 @@ public static void methodExit(
}
}
}

/**
* Due to a change in the AmazonHttpClient class, this instrumentation is needed to support newer
* versions. The above class should cover older versions.
*/
@AutoService(InstrumenterModule.class)
public static final class RequestExecutorInstrumentation extends AWSHttpClientInstrumentation {

@Override
public String instrumentedType() {
return "com.amazonaws.http.AmazonHttpClient$RequestExecutor";
}

@Override
public String[] helperClassNames() {
return new String[] {
packageName + ".OnErrorDecorator", packageName + ".AwsNameCache",
};
}

@Override
public void methodAdvice(MethodTransformer transformer) {
transformer.applyAdvice(
isMethod().and(named("doExecute")),
RequestExecutorInstrumentation.class.getName() + "$RequestExecutorAdvice");
}

public static class RequestExecutorAdvice {
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void methodExit(
@Advice.FieldValue("request") final Request<?> request,
@Advice.Thrown final Throwable throwable) {

final AgentScope scope = activeScope();
// check name in case TracingRequestHandler failed to activate the span
if (scope != null
&& (AwsNameCache.spanName(request).equals(scope.span().getSpanName())
|| scope.span() instanceof AgentTracer.NoopAgentSpan)) {
scope.close();
}

if (throwable != null) {
final AgentSpan span = request.getHandlerContext(SPAN_CONTEXT_KEY);
if (span != null) {
request.addHandlerContext(SPAN_CONTEXT_KEY, null);
DECORATE.onError(span, throwable);
DECORATE.beforeFinish(span);
span.finish();
}
}
}
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
package datadog.trace.instrumentation.aws.v0;

import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.Instrumenter;
import datadog.trace.agent.tooling.InstrumenterModule;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

/** Groups the instrumentations for AWS SDK 1.11.0+. */
@AutoService(InstrumenterModule.class)
public final class AwsSdkModule extends InstrumenterModule.Tracing {

public AwsSdkModule() {
super("aws-sdk");
}

@Override
public String[] helperClassNames() {
return new String[] {
packageName + ".AwsSdkClientDecorator",
packageName + ".GetterAccess",
packageName + ".GetterAccess$1",
packageName + ".TracingRequestHandler",
packageName + ".AwsNameCache",
packageName + ".OnErrorDecorator",
};
}

@Override
public Map<String, String> contextStore() {
Map<String, String> map = new java.util.HashMap<>();
map.put("com.amazonaws.services.sqs.model.ReceiveMessageResult", "java.lang.String");
map.put(
"com.amazonaws.AmazonWebServiceRequest",
"datadog.trace.bootstrap.instrumentation.api.AgentSpan");
return map;
}

@Override
public List<Instrumenter> typeInstrumentations() {
return Arrays.asList(
new AWSHttpClientInstrumentation(),
new RequestExecutorInstrumentation(),
new HandlerChainFactoryInstrumentation());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -4,52 +4,23 @@
import static net.bytebuddy.matcher.ElementMatchers.isMethod;

import com.amazonaws.handlers.RequestHandler2;
import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.Instrumenter;
import datadog.trace.agent.tooling.InstrumenterModule;
import datadog.trace.bootstrap.InstrumentationContext;
import java.util.List;
import java.util.Map;
import net.bytebuddy.asm.Advice;

/**
* This instrumentation might work with versions before 1.11.0, but this was the first version that
* is tested. It could possibly be extended earlier.
*/
@AutoService(InstrumenterModule.class)
public final class HandlerChainFactoryInstrumentation extends InstrumenterModule.Tracing
public final class HandlerChainFactoryInstrumentation
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {

public HandlerChainFactoryInstrumentation() {
super("aws-sdk");
}

@Override
public String instrumentedType() {
return "com.amazonaws.handlers.HandlerChainFactory";
}

@Override
public String[] helperClassNames() {
return new String[] {
packageName + ".AwsSdkClientDecorator",
packageName + ".GetterAccess",
packageName + ".GetterAccess$1",
packageName + ".TracingRequestHandler",
packageName + ".AwsNameCache",
};
}

@Override
public Map<String, String> contextStore() {
Map<String, String> map = new java.util.HashMap<>();
map.put("com.amazonaws.services.sqs.model.ReceiveMessageResult", "java.lang.String");
map.put(
"com.amazonaws.AmazonWebServiceRequest",
"datadog.trace.bootstrap.instrumentation.api.AgentSpan");
return map;
}

@Override
public void methodAdvice(MethodTransformer transformer) {
transformer.applyAdvice(
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
package datadog.trace.instrumentation.aws.v0;

import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.DECORATE;
import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.SPAN_CONTEXT_KEY;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;

import com.amazonaws.Request;
import datadog.trace.agent.tooling.Instrumenter;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
import net.bytebuddy.asm.Advice;

/**
* Due to a change in the AmazonHttpClient class, this instrumentation is needed to support newer
* versions. The {@link AWSHttpClientInstrumentation} class should cover older versions.
*/
public final class RequestExecutorInstrumentation
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {

@Override
public String instrumentedType() {
return "com.amazonaws.http.AmazonHttpClient$RequestExecutor";
}

@Override
public void methodAdvice(MethodTransformer transformer) {
transformer.applyAdvice(
isMethod().and(named("doExecute")),
RequestExecutorInstrumentation.class.getName() + "$RequestExecutorAdvice");
}

public static class RequestExecutorAdvice {
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void methodExit(
@Advice.FieldValue("request") final Request<?> request,
@Advice.Thrown final Throwable throwable) {

final AgentScope scope = activeScope();
// check name in case TracingRequestHandler failed to activate the span
if (scope != null
&& (AwsNameCache.spanName(request).equals(scope.span().getSpanName())
|| scope.span() instanceof AgentTracer.NoopAgentSpan)) {
scope.close();
}

if (throwable != null) {
final AgentSpan span = request.getHandlerContext(SPAN_CONTEXT_KEY);
if (span != null) {
request.addHandlerContext(SPAN_CONTEXT_KEY, null);
DECORATE.onError(span, throwable);
DECORATE.beforeFinish(span);
span.finish();
}
}
}
}
}
32 changes: 29 additions & 3 deletions docs/how_instrumentations_work.md
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,7 @@ At this point the instrumentation should override the method `muzzleDirective()`

## Instrumentation classes

The Instrumentation class is where the Instrumentation begins. It will:
The Instrumentation class is where the instrumentation begins. It will:

1. Use Matchers to choose target types (i.e., classes)
2. From only those target types, use Matchers to select the members (i.e., methods) to instrument.
Expand All @@ -110,8 +110,9 @@ The Instrumentation class is where the Instrumentation begins. It will:
Instrumentation classes:

1. Must be annotated with `@AutoService(InstrumenterModule.class)`
2. Should extend one of the six abstract TargetSystem `InstrumenterModule` classes
3. Should implement one of the `Instrumenter` interfaces
2. Should be declared in a file that ends with `Instrumentation.java`
3. Should extend one of the six abstract TargetSystem `InstrumenterModule` classes
4. Should implement one of the `Instrumenter` interfaces

For example:

Expand All @@ -136,6 +137,31 @@ public class RabbitChannelInstrumentation extends InstrumenterModule.Tracing
| `InstrumenterModule.`[`Usm`](https://github.com/DataDog/dd-trace-java/blob/82a3400cd210f4051b92fe1a86cd1b64a17e005e/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java#L273) | |
| [`InstrumenterModule`](https://github.com/DataDog/dd-trace-java/blob/82a3400cd210f4051b92fe1a86cd1b64a17e005e/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java) | Avoid extending `InstrumenterModule` directly. When no other TargetGroup is applicable we generally default to `InstrumenterModule.Tracing` |

### Grouping Instrumentations

Related instrumentations may be grouped under a single `InstrumenterModule` to share common details
such as integration name, helpers, context store use, and optional `classLoaderMatcher()`.

Module classes:

1. Must be annotated with `@AutoService(InstrumenterModule.class)`
2. Should be declared in a file that ends with `Module.java`
3. Should extend one of the six abstract TargetSystem `InstrumenterModule` classes
4. Should have a `typeInstrumentations()` method that returns the instrumentations in the group
5. Should NOT implement one of the `Instrumenter` interfaces

> [!WARNING]
> Grouped instrumentations must NOT be annotated with `@AutoService(InstrumenterModule.class)
> and must NOT extend any of the six abstract TargetSystem `InstrumenterModule` classes
Existing instrumentations can be grouped under a new module, assuming they share the same integration name.

For each member instrumentation:
1. Remove `@AutoService(InstrumenterModule.class)`
2. Remove `extends InstrumenterModule...`
3. Move the list of helpers to the module, merging as necessary
4. Move the context store map to the module, merging as necessary

### Type Matching

Instrumentation classes should implement an
Expand Down

0 comments on commit ace20ff

Please sign in to comment.