Skip to content

Commit

Permalink
Simplify scope activation in AWS-SDK instrumentations.
Browse files Browse the repository at this point in the history 
Avoid storing the scope in the v1 attributes; instead close any active AWS-SDK span when the client
finishes its execution, like we do for v2. Also introduce a no-op scope when calling out from AWS to
Netty in async mode to make sure that it can't accidentally capture continuations which incorrectly
keep the AWS trace alive. (We close the AWS scope but there may be other scopes left on the stack.)
  • Loading branch information
@mcculls
mcculls committed Feb 23, 2023
1 parent 3ae5d9f commit 1b33ea5
Show file tree
Hide file tree
Showing 6 changed files with 69 additions and 100 deletions.
36 changes: 24 additions & 12 deletions ...11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AWSHttpClientInstrumentation.java
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
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.AWS_HTTP;
import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.DECORATE;
import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.SCOPE_CONTEXT_KEY;
import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.SPAN_CONTEXT_KEY;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;

import com.amazonaws.AmazonClientException;
Expand Down Expand Up @@ -49,14 +51,19 @@ public static class HttpClientAdvice {
public static void methodExit(
@Advice.Argument(value = 0, optional = true) final Request<?> request,
@Advice.Thrown final Throwable throwable) {
if (throwable != null) {
final AgentScope scope = request.getHandlerContext(SCOPE_CONTEXT_KEY);
if (scope != null) {
request.addHandlerContext(SCOPE_CONTEXT_KEY, null);
final AgentSpan span = scope.span();

final AgentScope scope = activeScope();
// check name in case TracingRequestHandler failed to activate the span
if (scope != null && AWS_HTTP.equals(scope.span().getSpanName())) {
scope.close();
}

if (throwable != null && request != 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);
scope.close();
span.finish();
}
}
Expand Down Expand Up @@ -92,14 +99,19 @@ public static class RequestExecutorAdvice {
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 && AWS_HTTP.equals(scope.span().getSpanName())) {
scope.close();
}

if (throwable != null) {
final AgentScope scope = request.getHandlerContext(SCOPE_CONTEXT_KEY);
if (scope != null) {
request.addHandlerContext(SCOPE_CONTEXT_KEY, null);
final AgentSpan span = scope.span();
final AgentSpan span = request.getHandlerContext(SPAN_CONTEXT_KEY);
if (span != null) {
request.addHandlerContext(SPAN_CONTEXT_KEY, null);
DECORATE.onError(span, throwable);
DECORATE.beforeFinish(span);
scope.close();
span.finish();
}
}
Expand Down
1 change: 1 addition & 0 deletions ...-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AwsSdkClientDecorator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
public class AwsSdkClientDecorator extends HttpClientDecorator<Request, Response>
implements AgentPropagation.Setter<Request<?>> {
public static final AwsSdkClientDecorator DECORATE = new AwsSdkClientDecorator();
public static final CharSequence AWS_HTTP = UTF8BytesString.create("aws.http");

private static final Pattern REQUEST_PATTERN = Pattern.compile("Request", Pattern.LITERAL);
private static final Pattern AMAZON_PATTERN = Pattern.compile("Amazon", Pattern.LITERAL);
Expand Down
14 changes: 8 additions & 6 deletions ...-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/OnErrorDecorator.java
View file
Original file line number Diff line number Diff line change
@@ -1,17 +1,19 @@
package datadog.trace.instrumentation.aws.v0;

import com.amazonaws.handlers.HandlerContextKey;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
import datadog.trace.bootstrap.instrumentation.decorator.BaseDecorator;

public class OnErrorDecorator extends BaseDecorator {
// aws1.x sdk doesn't have any truly async clients so we can store scope in request context safely
public static final HandlerContextKey<AgentScope> SCOPE_CONTEXT_KEY =
new HandlerContextKey<>("DatadogScope"); // same as TracingRequestHandler.SCOPE_CONTEXT_KEY

public static final HandlerContextKey<AgentSpan> SPAN_CONTEXT_KEY =
new HandlerContextKey<>("DatadogSpan"); // same as TracingRequestHandler.SPAN_CONTEXT_KEY

public static final OnErrorDecorator DECORATE = new OnErrorDecorator();
private static final CharSequence JAVA_AWS_SDK = UTF8BytesString.create("java-aws-sdk");
public static final CharSequence AWS_HTTP = UTF8BytesString.create("aws.http");

private static final CharSequence COMPONENT_NAME = UTF8BytesString.create("java-aws-sdk");

@Override
protected String[] instrumentationNames() {
Expand All @@ -25,6 +27,6 @@ protected CharSequence spanType() {

@Override
protected CharSequence component() {
return JAVA_AWS_SDK;
return COMPONENT_NAME;
}
}
63 changes: 18 additions & 45 deletions ...-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
package datadog.trace.instrumentation.aws.v0;

import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateNext;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
import static datadog.trace.instrumentation.aws.v0.AwsSdkClientDecorator.AWS_HTTP;
import static datadog.trace.instrumentation.aws.v0.AwsSdkClientDecorator.DECORATE;

import com.amazonaws.AmazonWebServiceRequest;
Expand All @@ -14,20 +13,15 @@
import com.amazonaws.handlers.RequestHandler2;
import datadog.trace.api.Config;
import datadog.trace.api.TracePropagationStyle;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Tracing Request Handler */
public class TracingRequestHandler extends RequestHandler2 {

// aws1.x sdk doesn't have any truly async clients so we can store scope in request context safely
public static final HandlerContextKey<AgentScope> SCOPE_CONTEXT_KEY =
new HandlerContextKey<>("DatadogScope"); // same as OnErrorDecorator.SCOPE_CONTEXT_KEY

private static final CharSequence AWS_HTTP = UTF8BytesString.create("aws.http");
public static final HandlerContextKey<AgentSpan> SPAN_CONTEXT_KEY =
new HandlerContextKey<>("DatadogSpan"); // same as OnErrorDecorator.SPAN_CONTEXT_KEY

private static final Logger log = LoggerFactory.getLogger(TracingRequestHandler.class);

Expand All @@ -38,61 +32,40 @@ public AmazonWebServiceRequest beforeMarshalling(final AmazonWebServiceRequest r

@Override
public void beforeRequest(final Request<?> request) {
boolean isPolling = isPollingRequest(request.getOriginalRequest());
if (isPolling) {
closePrevious(true);
}
final AgentSpan span = startSpan(AWS_HTTP);
DECORATE.afterStart(span);
DECORATE.onRequest(span, request);
if (isPolling) {
activateNext(span); // this scope will last until next poll
}
request.addHandlerContext(SCOPE_CONTEXT_KEY, activateSpan(span));
request.addHandlerContext(SPAN_CONTEXT_KEY, span);
if (Config.get().isAwsPropagationEnabled()) {
try {
propagate().inject(span, request, DECORATE, TracePropagationStyle.XRAY);
} catch (Throwable e) {
log.warn("Unable to inject trace header", e);
}
}
// This scope will be closed by AwsHttpClientInstrumentation
activateSpan(span);
}

@Override
public void afterResponse(final Request<?> request, final Response<?> response) {
final AgentScope scope = request.getHandlerContext(SCOPE_CONTEXT_KEY);
if (scope != null) {
request.addHandlerContext(SCOPE_CONTEXT_KEY, null);
DECORATE.onResponse(scope.span(), response);
DECORATE.beforeFinish(scope.span());
scope.close();
if (isPollingRequest(request.getOriginalRequest())) {
// will be finished on next poll
} else {
scope.span().finish();
}
final AgentSpan span = request.getHandlerContext(SPAN_CONTEXT_KEY);
if (span != null) {
request.addHandlerContext(SPAN_CONTEXT_KEY, null);
DECORATE.onResponse(span, response);
DECORATE.beforeFinish(span);
span.finish();
}
}

@Override
public void afterError(final Request<?> request, final Response<?> response, final Exception e) {
final AgentScope scope = request.getHandlerContext(SCOPE_CONTEXT_KEY);
if (scope != null) {
request.addHandlerContext(SCOPE_CONTEXT_KEY, null);
DECORATE.onError(scope.span(), e);
DECORATE.beforeFinish(scope.span());
scope.close();
if (isPollingRequest(request.getOriginalRequest())) {
// will be finished on next poll
} else {
scope.span().finish();
}
final AgentSpan span = request.getHandlerContext(SPAN_CONTEXT_KEY);
if (span != null) {
request.addHandlerContext(SPAN_CONTEXT_KEY, null);
DECORATE.onError(span, e);
DECORATE.beforeFinish(span);
span.finish();
}
}

private static boolean isPollingRequest(AmazonWebServiceRequest request) {
return null != request
&& "com.amazonaws.services.sqs.model.ReceiveMessageRequest"
.equals(request.getClass().getName());
}
}
30 changes: 15 additions & 15 deletions ...-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsHttpClientInstrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,10 @@
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.nameStartsWith;
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
import static datadog.trace.instrumentation.aws.v2.AwsSdkClientDecorator.AWS_HTTP;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.isPublic;

Expand Down Expand Up @@ -47,7 +50,7 @@ public void adviceTransformations(AdviceTransformation transformation) {
}

public static class AwsHttpClientAdvice {
// scope.close here doesn't actually close the span.
// scope.close here doesn't actually finish the span.

/**
* FIXME: This is a hack to prevent netty instrumentation from messing things up.
Expand All @@ -59,25 +62,22 @@ public static class AwsHttpClientAdvice {
* stored in channel attributes.
*/
@Advice.OnMethodEnter(suppress = Throwable.class)
public static boolean methodEnter(@Advice.This final Object thiz) {
if (thiz instanceof MakeAsyncHttpRequestStage) {
final AgentScope scope = activeScope();
if (scope != null) {
scope.close();
return true;
public static AgentScope methodEnter(@Advice.This final Object thiz) {
final AgentScope scope = activeScope();
// check name in case TracingExecutionInterceptor failed to activate the span
if (scope != null && AWS_HTTP.equals(scope.span().getSpanName())) {
if (thiz instanceof MakeAsyncHttpRequestStage) {
scope.close(); // close async legacy HTTP span to avoid Netty leak
} else {
return scope; // keep sync legacy HTTP span alive for duration of call
}
}
return false;
return activateSpan(noopSpan());
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void methodExit(@Advice.Enter final boolean scopeAlreadyClosed) {
if (!scopeAlreadyClosed) {
final AgentScope scope = activeScope();
if (scope != null) {
scope.close();
}
}
public static void methodExit(@Advice.Enter final AgentScope scope) {
scope.close();
}

/**
Expand Down
25 changes: 3 additions & 22 deletions ...k-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java
View file
Original file line number Diff line number Diff line change
@@ -1,20 +1,16 @@
package datadog.trace.instrumentation.aws.v2;

import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateNext;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
import static datadog.trace.instrumentation.aws.v2.AwsSdkClientDecorator.AWS_HTTP;
import static datadog.trace.instrumentation.aws.v2.AwsSdkClientDecorator.DECORATE;

import datadog.trace.api.Config;
import datadog.trace.api.TracePropagationStyle;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import software.amazon.awssdk.core.SdkRequest;
import software.amazon.awssdk.core.interceptor.Context;
import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
Expand All @@ -32,9 +28,6 @@ public class TracingExecutionInterceptor implements ExecutionInterceptor {
@Override
public void beforeExecution(
final Context.BeforeExecution context, final ExecutionAttributes executionAttributes) {
if (isPollingRequest(context.request())) {
closePrevious(true);
}
final AgentSpan span = startSpan(AWS_HTTP);
DECORATE.afterStart(span);
executionAttributes.putAttribute(SPAN_ATTRIBUTE, span);
Expand Down Expand Up @@ -70,11 +63,9 @@ public SdkHttpRequest modifyHttpRequest(
public void beforeTransmission(
final Context.BeforeTransmission context, final ExecutionAttributes executionAttributes) {
final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);

// This scope will be closed by AwsHttpClientInstrumentation since ExecutionInterceptor API
// doesn't provide a way to run code in the same thread after transmission has been scheduled.
final AgentScope scope = activateSpan(span);
scope.setAsyncPropagation(true);
// doesn't provide a way to run code in same thread after transmission has been scheduled.
activateSpan(span);
}

@Override
Expand All @@ -87,11 +78,7 @@ public void afterExecution(
DECORATE.onResponse(span, context.response());
DECORATE.onResponse(span, context.httpResponse());
DECORATE.beforeFinish(span);
if (isPollingRequest(context.request())) {
activateNext(span); // remains active until next poll or iteration scope timeout
} else {
span.finish();
}
span.finish();
}
}

Expand All @@ -107,12 +94,6 @@ public void onExecutionFailure(
}
}

private static boolean isPollingRequest(SdkRequest request) {
return null != request
&& "software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest"
.equals(request.getClass().getName());
}

public static void muzzleCheck() {
// Noop
}
Expand Down

0 comments on commit 1b33ea5

Please sign in to comment.