Fix message handler spans appear disconnected from the incoming SQS trace #11511
There are no files selected for viewing
This file contains hidden or 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 hidden or 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,36 @@ | ||
| package datadog.trace.instrumentation.aws.v2.sqs; | ||
|
|
||
| import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; | ||
| import static net.bytebuddy.matcher.ElementMatchers.isMethod; | ||
|
|
||
| import datadog.trace.agent.tooling.Instrumenter; | ||
| import net.bytebuddy.asm.Advice; | ||
|
|
||
| public final class SqsMd5ChecksumInterceptorInstrumentation | ||
| implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { | ||
|
|
||
| @Override | ||
| public String instrumentedType() { | ||
| // The AWS SDK checksum interceptor reads ReceiveMessageResponse.messages() while finalizing | ||
| // the response. Mark that internal access so we only wrap messages for application code. | ||
| return "software.amazon.awssdk.services.sqs.internal.MessageMD5ChecksumInterceptor"; | ||
| } | ||
|
|
||
| @Override | ||
| public void methodAdvice(MethodTransformer transformer) { | ||
| transformer.applyAdvice( | ||
| isMethod().and(named("afterExecution")), getClass().getName() + "$AfterExecutionAdvice"); | ||
| } | ||
|
|
||
| public static class AfterExecutionAdvice { | ||
| @Advice.OnMethodEnter(suppress = Throwable.class) | ||
| public static void onEnter() { | ||
| SqsReceiveResponseInternalAccess.enter(); | ||
|
Contributor
There was a problem hiding this comment. don't we need to register SqsReceiveResponseInternalAccess as a helper for this advice ?
Contributor
Author
There was a problem hiding this comment. Correct. It's registered in the module: https://github.com/DataDog/dd-trace-java/pull/11511/changes#diff-11edc0f6db02d9c1bf0fcfb56354d561b9ee7b9cf8e69398c36642bc9e852541R26 |
||
| } | ||
|
|
||
| @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) | ||
| public static void onExit() { | ||
| SqsReceiveResponseInternalAccess.exit(); | ||
| } | ||
| } | ||
| } | ||
This file contains hidden or 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 hidden or 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,54 @@ | ||
| package datadog.trace.instrumentation.aws.v2.sqs; | ||
|
|
||
| import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; | ||
| import static net.bytebuddy.matcher.ElementMatchers.isMethod; | ||
| import static net.bytebuddy.matcher.ElementMatchers.returns; | ||
| import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; | ||
|
|
||
| import datadog.trace.agent.tooling.Instrumenter; | ||
| import datadog.trace.bootstrap.ContextStore; | ||
| import datadog.trace.bootstrap.InstrumentationContext; | ||
| import net.bytebuddy.asm.Advice; | ||
| import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse; | ||
|
|
||
| public final class SqsReceiveResponseBuilderImplInstrumentation | ||
| implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { | ||
|
|
||
| private static final String BUILDER_IMPL = | ||
| "software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse$BuilderImpl"; | ||
|
|
||
| @Override | ||
| public String instrumentedType() { | ||
| return BUILDER_IMPL; | ||
| } | ||
|
|
||
| @Override | ||
| public void methodAdvice(MethodTransformer transformer) { | ||
| transformer.applyAdvice( | ||
| isMethod() | ||
| .and(named("build")) | ||
| .and(takesNoArguments()) | ||
| .and( | ||
| returns(named("software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse"))), | ||
| getClass().getName() + "$BuildAdvice"); | ||
| } | ||
|
|
||
| public static class BuildAdvice { | ||
| @Advice.OnMethodExit(suppress = Throwable.class) | ||
| public static void onExit( | ||
| @Advice.This Object builder, @Advice.Return ReceiveMessageResponse response) { | ||
| if (response == null) { | ||
| return; | ||
| } | ||
| ContextStore<Object, String> builderStore = | ||
| InstrumentationContext.get(BUILDER_IMPL, "java.lang.String"); | ||
| String queueUrl = builderStore.get(builder); | ||
| if (queueUrl != null) { | ||
| // Complete the handoff from the pre-rebuild response to the final response user code | ||
| // sees. | ||
| InstrumentationContext.get(ReceiveMessageResponse.class, String.class) | ||
| .put(response, queueUrl); | ||
| } | ||
| } | ||
| } | ||
| } |
This file contains hidden or 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,53 @@ | ||
| package datadog.trace.instrumentation.aws.v2.sqs; | ||
|
|
||
| import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; | ||
| import static net.bytebuddy.matcher.ElementMatchers.isMethod; | ||
| import static net.bytebuddy.matcher.ElementMatchers.returns; | ||
| import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; | ||
|
|
||
| import datadog.trace.agent.tooling.Instrumenter; | ||
| import datadog.trace.bootstrap.InstrumentationContext; | ||
| import net.bytebuddy.asm.Advice; | ||
| import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse; | ||
|
|
||
| public final class SqsReceiveResponseBuilderInstrumentation | ||
| implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { | ||
|
|
||
| private static final String BUILDER_IMPL = | ||
| "software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse$BuilderImpl"; | ||
|
|
||
| @Override | ||
| public String instrumentedType() { | ||
| return "software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse"; | ||
| } | ||
|
|
||
| @Override | ||
| public void methodAdvice(MethodTransformer transformer) { | ||
| transformer.applyAdvice( | ||
| isMethod() | ||
| .and(named("toBuilder")) | ||
| .and(takesNoArguments()) | ||
| .and( | ||
| returns( | ||
| named( | ||
| "software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse$Builder"))), | ||
| getClass().getName() + "$ToBuilderAdvice"); | ||
| } | ||
|
|
||
| public static class ToBuilderAdvice { | ||
| @Advice.OnMethodExit(suppress = Throwable.class) | ||
| public static void onExit( | ||
| @Advice.This ReceiveMessageResponse response, @Advice.Return Object builder) { | ||
| if (builder == null) { | ||
| return; | ||
| } | ||
| String queueUrl = | ||
| InstrumentationContext.get(ReceiveMessageResponse.class, String.class).get(response); | ||
| if (queueUrl != null) { | ||
| // AWS SDK core finalizes modeled responses through response.toBuilder().build(). | ||
| // Carry queueUrl onto the builder so it survives that response instance change. | ||
| InstrumentationContext.get(BUILDER_IMPL, "java.lang.String").put(builder, queueUrl); | ||
| } | ||
| } | ||
| } | ||
| } |
This file contains hidden or 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,20 @@ | ||
| package datadog.trace.instrumentation.aws.v2.sqs; | ||
|
|
||
| import datadog.trace.bootstrap.CallDepthThreadLocalMap; | ||
|
|
||
| public final class SqsReceiveResponseInternalAccess { | ||
|
|
||
| private SqsReceiveResponseInternalAccess() {} | ||
|
|
||
| public static void enter() { | ||
| CallDepthThreadLocalMap.incrementCallDepth(SqsReceiveResponseInternalAccess.class); | ||
| } | ||
|
|
||
| public static void exit() { | ||
| CallDepthThreadLocalMap.decrementCallDepth(SqsReceiveResponseInternalAccess.class); | ||
| } | ||
|
|
||
| public static boolean active() { | ||
| return CallDepthThreadLocalMap.getCallDepth(SqsReceiveResponseInternalAccess.class) > 0; | ||
| } | ||
| } |
This file contains hidden or 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
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.
I think a comment why we need to instrument this internal class would be nice
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.
Yep, this is part of the commit message. It's needed to avoid wrapping messages during the SDK's internal MessageMD5ChecksumInterceptor pass to prevent creating consumer spans
before user code actually consumes the messages.
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.
yeah but when reading the code (in the editor, after merging), you don't go and look at the commit messages in the blame to understand what's going on, which is why I think having a comment in the code explaining things is better
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.
I agree with including this in the code; I just didn't want to run CI again just for a single comment line :)