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)
  • Loading branch information
mcculls committed Jan 2, 2025
1 parent c24db36 commit 8c31aee
Show file tree
Hide file tree
Showing 4 changed files with 108 additions and 98 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,46 @@
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;

@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();
}
}
}
}
}

0 comments on commit 8c31aee

Please sign in to comment.