Skip to content

Commit ffb48ae

Browse files
ygreeygree
authored
Fix message handler spans appear disconnected from the incoming SQS trace (#11511)
Add regression test for sync SQS receive response copy context loss Preserve SQS receive context across AWS SDK response rebuilding Store the receive queue URL before the SDK rebuilds ReceiveMessageResponse and propagate it through the response builder so messages() can still wrap the final list in TracingList. Also avoid wrapping messages during the SDK's internal MessageMD5ChecksumInterceptor pass to prevent creating consumer spans before user code actually consumes the messages. Merge branch 'master' into ygree/aws-sdk-sqs-bug Migrate SQS client tests to JUnit Document SQS checksum interceptor instrumentation add clarifying comments Merge branch 'master' into ygree/aws-sdk-sqs-bug Merge branch 'master' into ygree/aws-sdk-sqs-bug Co-authored-by: yury.gribkov <yury.gribkov@datadoghq.com>
1 parent a1a064b commit ffb48ae

9 files changed

Lines changed: 989 additions & 693 deletions

File tree

‎dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java‎

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import software.amazon.awssdk.core.interceptor.Context.BeforeTransmission;
2929
import software.amazon.awssdk.core.interceptor.Context.FailedExecution;
3030
import software.amazon.awssdk.core.interceptor.Context.ModifyHttpRequest;
31+
import software.amazon.awssdk.core.interceptor.Context.ModifyResponse;
3132
import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
3233
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
3334
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
@@ -111,6 +112,22 @@ public void beforeTransmission(
111112
}
112113
}
113114

115+
@Override
116+
public SdkResponse modifyResponse(
117+
final ModifyResponse context, final ExecutionAttributes executionAttributes) {
118+
final SdkResponse response = context.response();
119+
if (!AWS_LEGACY_TRACING && isPollingRequest(context.request()) && isPollingResponse(response)) {
120+
// Attach queueUrl before AWS SDK core rebuilds the response with
121+
// toBuilder().sdkHttpResponse(...).build(). afterExecution sees this pre-rebuild response,
122+
// not the final response returned to user code, so capturing queueUrl there is too late.
123+
context
124+
.request()
125+
.getValueForField("QueueUrl", String.class)
126+
.ifPresent(queueUrl -> responseQueueStore.put(response, queueUrl));
127+
}
128+
return response;
129+
}
130+
114131
@Override
115132
public void afterExecution(
116133
final AfterExecution context, final ExecutionAttributes executionAttributes) {
@@ -124,13 +141,6 @@ public void afterExecution(
124141
DECORATE.beforeFinish(span);
125142
span.finish();
126143
}
127-
if (!AWS_LEGACY_TRACING && isPollingResponse(context.response())) {
128-
// store queueUrl inside response for SqsReceiveResultInstrumentation
129-
context
130-
.request()
131-
.getValueForField("QueueUrl", String.class)
132-
.ifPresent(queueUrl -> responseQueueStore.put(context.response(), queueUrl));
133-
}
134144
}
135145

136146
@Override

‎dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsMd5ChecksumInterceptorInstrumentation.java‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
package datadog.trace.instrumentation.aws.v2.sqs;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
5+
6+
import datadog.trace.agent.tooling.Instrumenter;
7+
import net.bytebuddy.asm.Advice;
8+
9+
public final class SqsMd5ChecksumInterceptorInstrumentation
10+
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
11+
12+
@Override
13+
public String instrumentedType() {
14+
// The AWS SDK checksum interceptor reads ReceiveMessageResponse.messages() while finalizing
15+
// the response. Mark that internal access so we only wrap messages for application code.
16+
return "software.amazon.awssdk.services.sqs.internal.MessageMD5ChecksumInterceptor";
17+
}
18+
19+
@Override
20+
public void methodAdvice(MethodTransformer transformer) {
21+
transformer.applyAdvice(
22+
isMethod().and(named("afterExecution")), getClass().getName() + "$AfterExecutionAdvice");
23+
}
24+
25+
public static class AfterExecutionAdvice {
26+
@Advice.OnMethodEnter(suppress = Throwable.class)
27+
public static void onEnter() {
28+
SqsReceiveResponseInternalAccess.enter();
29+
}
30+
31+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
32+
public static void onExit() {
33+
SqsReceiveResponseInternalAccess.exit();
34+
}
35+
}
36+
}

‎dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsModule.java‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55
import datadog.trace.agent.tooling.InstrumenterModule;
66
import datadog.trace.api.InstrumenterConfig;
77
import java.util.ArrayList;
8-
import java.util.Collections;
98
import java.util.List;
109
import java.util.Map;
1110

@@ -24,6 +23,7 @@ public String[] helperClassNames() {
2423
"datadog.trace.instrumentation.aws.v2.sqs.MessageAttributeInjector",
2524
"datadog.trace.instrumentation.aws.v2.sqs.MessageExtractAdapter",
2625
"datadog.trace.instrumentation.aws.v2.sqs.SqsDecorator",
26+
"datadog.trace.instrumentation.aws.v2.sqs.SqsReceiveResponseInternalAccess",
2727
"datadog.trace.instrumentation.aws.v2.sqs.TracingIterator",
2828
"datadog.trace.instrumentation.aws.v2.sqs.TracingList",
2929
"datadog.trace.instrumentation.aws.v2.sqs.TracingListIterator"
@@ -32,17 +32,25 @@ public String[] helperClassNames() {
3232

3333
@Override
3434
public Map<String, String> contextStore() {
35-
return Collections.singletonMap(
35+
Map<String, String> contextStore = new java.util.HashMap<>();
36+
contextStore.put(
3637
"software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse", "java.lang.String");
38+
contextStore.put(
39+
"software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse$BuilderImpl",
40+
"java.lang.String");
41+
return contextStore;
3742
}
3843

3944
@Override
4045
public List<Instrumenter> typeInstrumentations() {
41-
final List<Instrumenter> ret = new ArrayList<>(4);
46+
final List<Instrumenter> ret = new ArrayList<>(6);
4247
ret.add(new SqsClientInstrumentation());
4348
ret.add(new SqsReceiveRequestInstrumentation());
4449
// we don't need to instrument messages when we're doing legacy AWS-SDK tracing
4550
if (!InstrumenterConfig.get().isLegacyInstrumentationEnabled(false, "aws-sdk")) {
51+
ret.add(new SqsMd5ChecksumInterceptorInstrumentation());
52+
ret.add(new SqsReceiveResponseBuilderInstrumentation());
53+
ret.add(new SqsReceiveResponseBuilderImplInstrumentation());
4654
ret.add(new SqsReceiveResultInstrumentation());
4755
}
4856
return ret;

‎dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsReceiveResponseBuilderImplInstrumentation.java‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
package datadog.trace.instrumentation.aws.v2.sqs;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
5+
import static net.bytebuddy.matcher.ElementMatchers.returns;
6+
import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments;
7+
8+
import datadog.trace.agent.tooling.Instrumenter;
9+
import datadog.trace.bootstrap.ContextStore;
10+
import datadog.trace.bootstrap.InstrumentationContext;
11+
import net.bytebuddy.asm.Advice;
12+
import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse;
13+
14+
public final class SqsReceiveResponseBuilderImplInstrumentation
15+
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
16+
17+
private static final String BUILDER_IMPL =
18+
"software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse$BuilderImpl";
19+
20+
@Override
21+
public String instrumentedType() {
22+
return BUILDER_IMPL;
23+
}
24+
25+
@Override
26+
public void methodAdvice(MethodTransformer transformer) {
27+
transformer.applyAdvice(
28+
isMethod()
29+
.and(named("build"))
30+
.and(takesNoArguments())
31+
.and(
32+
returns(named("software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse"))),
33+
getClass().getName() + "$BuildAdvice");
34+
}
35+
36+
public static class BuildAdvice {
37+
@Advice.OnMethodExit(suppress = Throwable.class)
38+
public static void onExit(
39+
@Advice.This Object builder, @Advice.Return ReceiveMessageResponse response) {
40+
if (response == null) {
41+
return;
42+
}
43+
ContextStore<Object, String> builderStore =
44+
InstrumentationContext.get(BUILDER_IMPL, "java.lang.String");
45+
String queueUrl = builderStore.get(builder);
46+
if (queueUrl != null) {
47+
// Complete the handoff from the pre-rebuild response to the final response user code
48+
// sees.
49+
InstrumentationContext.get(ReceiveMessageResponse.class, String.class)
50+
.put(response, queueUrl);
51+
}
52+
}
53+
}
54+
}

‎dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsReceiveResponseBuilderInstrumentation.java‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
package datadog.trace.instrumentation.aws.v2.sqs;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
5+
import static net.bytebuddy.matcher.ElementMatchers.returns;
6+
import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments;
7+
8+
import datadog.trace.agent.tooling.Instrumenter;
9+
import datadog.trace.bootstrap.InstrumentationContext;
10+
import net.bytebuddy.asm.Advice;
11+
import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse;
12+
13+
public final class SqsReceiveResponseBuilderInstrumentation
14+
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
15+
16+
private static final String BUILDER_IMPL =
17+
"software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse$BuilderImpl";
18+
19+
@Override
20+
public String instrumentedType() {
21+
return "software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse";
22+
}
23+
24+
@Override
25+
public void methodAdvice(MethodTransformer transformer) {
26+
transformer.applyAdvice(
27+
isMethod()
28+
.and(named("toBuilder"))
29+
.and(takesNoArguments())
30+
.and(
31+
returns(
32+
named(
33+
"software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse$Builder"))),
34+
getClass().getName() + "$ToBuilderAdvice");
35+
}
36+
37+
public static class ToBuilderAdvice {
38+
@Advice.OnMethodExit(suppress = Throwable.class)
39+
public static void onExit(
40+
@Advice.This ReceiveMessageResponse response, @Advice.Return Object builder) {
41+
if (builder == null) {
42+
return;
43+
}
44+
String queueUrl =
45+
InstrumentationContext.get(ReceiveMessageResponse.class, String.class).get(response);
46+
if (queueUrl != null) {
47+
// AWS SDK core finalizes modeled responses through response.toBuilder().build().
48+
// Carry queueUrl onto the builder so it survives that response instance change.
49+
InstrumentationContext.get(BUILDER_IMPL, "java.lang.String").put(builder, queueUrl);
50+
}
51+
}
52+
}
53+
}

‎dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsReceiveResponseInternalAccess.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
package datadog.trace.instrumentation.aws.v2.sqs;
2+
3+
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
4+
5+
public final class SqsReceiveResponseInternalAccess {
6+
7+
private SqsReceiveResponseInternalAccess() {}
8+
9+
public static void enter() {
10+
CallDepthThreadLocalMap.incrementCallDepth(SqsReceiveResponseInternalAccess.class);
11+
}
12+
13+
public static void exit() {
14+
CallDepthThreadLocalMap.decrementCallDepth(SqsReceiveResponseInternalAccess.class);
15+
}
16+
17+
public static boolean active() {
18+
return CallDepthThreadLocalMap.getCallDepth(SqsReceiveResponseInternalAccess.class) > 0;
19+
}
20+
}

‎dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsReceiveResultInstrumentation.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,11 @@ public static class GetMessagesAdvice {
2929
public static void onExit(
3030
@Advice.This ReceiveMessageResponse result,
3131
@Advice.Return(readOnly = false) List<Message> messages) {
32+
if (SqsReceiveResponseInternalAccess.active()) {
33+
// AWS SDK's MD5 checksum interceptor calls messages() during afterExecution. That should
34+
// not create consumer spans; those belong around application message processing.
35+
return;
36+
}
3237
if (messages != null && !messages.isEmpty() && !(messages instanceof TracingList)) {
3338
String queueUrl =
3439
InstrumentationContext.get(ReceiveMessageResponse.class, String.class).get(result);

0 commit comments

Comments
 (0)