Skip to content

Commit 621efbe

Browse files
amarzialidevflow.devflow-routing-intake
amarziali
and
devflow.devflow-routing-intake
authored
Use context swap for messaging instrumentations (#10753)
wip Changes migrate everything else add test for jakarta jms use full config name on the json file force kafka consumer scope to close in the dsm test Merge branch 'master' into andrea.marziali/swap-messaging Use DD_LEGACY_CONTEXT_MANAGER_ENABLED Update config print missed reverse rename rename Apply suggestions from code review Co-authored-by: Stuart McCulloch <stuart.mcculloch@datadoghq.com> Update dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterator.java Co-authored-by: Stuart McCulloch <stuart.mcculloch@datadoghq.com> Update dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingIterator.java Co-authored-by: Stuart McCulloch <stuart.mcculloch@datadoghq.com> suggestions Merge branch 'master' into andrea.marziali/swap-messaging Co-authored-by: devflow.devflow-routing-intake <devflow.devflow-routing-intake@kubernetes.us1.ddbuild.io>
1 parent c04d61b commit 621efbe

File tree

20 files changed

+291
-21
lines changed

20 files changed

+291
-21
lines changed

‎dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java‎

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateNext;
88
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious;
99
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
10+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext;
11+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
1012
import static datadog.trace.bootstrap.instrumentation.api.URIUtils.urlFileName;
1113
import static datadog.trace.instrumentation.aws.v1.sqs.MessageExtractAdapter.GETTER;
1214
import static datadog.trace.instrumentation.aws.v1.sqs.SqsDecorator.BROKER_DECORATE;
@@ -18,6 +20,7 @@
1820

1921
import com.amazonaws.services.sqs.model.Message;
2022
import datadog.trace.api.Config;
23+
import datadog.trace.api.InstrumenterConfig;
2124
import datadog.trace.api.datastreams.DataStreamsTags;
2225
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
2326
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
@@ -43,7 +46,14 @@ public boolean hasNext() {
4346
boolean moreMessages = delegate.hasNext();
4447
if (!moreMessages) {
4548
// no more messages, use this as a signal to close the last iteration scope
46-
closePrevious(true);
49+
if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
50+
closePrevious(true);
51+
} else {
52+
final AgentSpan previousSpan = spanFromContext(getRootContext().swap());
53+
if (previousSpan != null) {
54+
previousSpan.finishWithEndToEnd();
55+
}
56+
}
4757
}
4858
return moreMessages;
4959
}
@@ -57,7 +67,14 @@ public Message next() {
5767

5868
protected void startNewMessageSpan(Message message) {
5969
try {
60-
closePrevious(true);
70+
if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
71+
closePrevious(true);
72+
} else if (message == null) { // previous message span was the last
73+
final AgentSpan previousSpan = spanFromContext(getRootContext().swap());
74+
if (previousSpan != null) {
75+
previousSpan.finishWithEndToEnd();
76+
}
77+
}
6178
if (message != null) {
6279
AgentSpan queueSpan = null;
6380
if (batchContext == null) {
@@ -92,7 +109,14 @@ protected void startNewMessageSpan(Message message) {
92109

93110
CONSUMER_DECORATE.afterStart(span);
94111
CONSUMER_DECORATE.onConsume(span, queueUrl);
95-
activateNext(span);
112+
if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
113+
activateNext(span);
114+
} else {
115+
final AgentSpan previousSpan = spanFromContext(span.swap());
116+
if (previousSpan != null) {
117+
previousSpan.finishWithEndToEnd();
118+
}
119+
}
96120
if (queueSpan != null) {
97121
BROKER_DECORATE.beforeFinish(queueSpan);
98122
queueSpan.finish();

‎dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingListIterator.java‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
11
package datadog.trace.instrumentation.aws.v1.sqs;
22

33
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious;
4+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext;
5+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
46

57
import com.amazonaws.services.sqs.model.Message;
8+
import datadog.trace.api.InstrumenterConfig;
9+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
610
import java.util.ListIterator;
711

812
public class TracingListIterator extends TracingIterator<ListIterator<Message>>
@@ -17,7 +21,14 @@ public boolean hasPrevious() {
1721
boolean moreMessages = delegate.hasPrevious();
1822
if (!moreMessages) {
1923
// no more messages, use this as a signal to close the last iteration scope
20-
closePrevious(true);
24+
if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
25+
closePrevious(true);
26+
} else {
27+
final AgentSpan previousSpan = spanFromContext(getRootContext().swap());
28+
if (previousSpan != null) {
29+
previousSpan.finishWithEndToEnd();
30+
}
31+
}
2132
}
2233
return moreMessages;
2334
}

‎dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/test/groovy/SqsClientTest.groovy‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
22
import static java.nio.charset.StandardCharsets.UTF_8
3+
import datadog.trace.api.config.TraceInstrumentationConfig
34

45
import com.amazon.sqs.javamessaging.ProviderConfiguration
56
import com.amazon.sqs.javamessaging.SQSConnectionFactory
@@ -726,4 +727,10 @@ class SqsClientV1DataStreamsForkedTest extends SqsClientTest {
726727
}
727728
}
728729

729-
730+
class SqsClientV0ContextSwapForkedTest extends SqsClientV0Test {
731+
@Override
732+
protected void configurePreAgent() {
733+
super.configurePreAgent()
734+
injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false")
735+
}
736+
}

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

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateNext;
88
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious;
99
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
10+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext;
11+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
1012
import static datadog.trace.bootstrap.instrumentation.api.URIUtils.urlFileName;
1113
import static datadog.trace.instrumentation.aws.v2.sqs.MessageExtractAdapter.GETTER;
1214
import static datadog.trace.instrumentation.aws.v2.sqs.SqsDecorator.BROKER_DECORATE;
@@ -17,6 +19,7 @@
1719
import static java.util.concurrent.TimeUnit.MILLISECONDS;
1820

1921
import datadog.trace.api.Config;
22+
import datadog.trace.api.InstrumenterConfig;
2023
import datadog.trace.api.datastreams.DataStreamsTags;
2124
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
2225
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
@@ -45,7 +48,14 @@ public boolean hasNext() {
4548
boolean moreMessages = delegate.hasNext();
4649
if (!moreMessages) {
4750
// no more messages, use this as a signal to close the last iteration scope
48-
closePrevious(true);
51+
if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
52+
closePrevious(true);
53+
} else {
54+
final AgentSpan previousSpan = spanFromContext(getRootContext().swap());
55+
if (previousSpan != null) {
56+
previousSpan.finishWithEndToEnd();
57+
}
58+
}
4959
}
5060
return moreMessages;
5161
}
@@ -59,7 +69,14 @@ public Message next() {
5969

6070
protected void startNewMessageSpan(Message message) {
6171
try {
62-
closePrevious(true);
72+
if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
73+
closePrevious(true);
74+
} else if (message == null) { // previous message span was the last
75+
final AgentSpan previousSpan = spanFromContext(getRootContext().swap());
76+
if (previousSpan != null) {
77+
previousSpan.finishWithEndToEnd();
78+
}
79+
}
6380
if (message != null) {
6481
AgentSpan queueSpan = null;
6582
if (batchContext == null) {
@@ -94,7 +111,14 @@ protected void startNewMessageSpan(Message message) {
94111

95112
CONSUMER_DECORATE.afterStart(span);
96113
CONSUMER_DECORATE.onConsume(span, queueUrl, requestId);
97-
activateNext(span);
114+
if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
115+
activateNext(span);
116+
} else {
117+
final AgentSpan previousSpan = spanFromContext(span.swap());
118+
if (previousSpan != null) {
119+
previousSpan.finishWithEndToEnd();
120+
}
121+
}
98122
if (queueSpan != null) {
99123
BROKER_DECORATE.beforeFinish(queueSpan);
100124
queueSpan.finish();

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

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,11 @@
11
package datadog.trace.instrumentation.aws.v2.sqs;
22

33
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious;
4+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext;
5+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
46

7+
import datadog.trace.api.InstrumenterConfig;
8+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
59
import java.util.ListIterator;
610
import software.amazon.awssdk.services.sqs.model.Message;
711

@@ -17,7 +21,14 @@ public boolean hasPrevious() {
1721
boolean moreMessages = delegate.hasPrevious();
1822
if (!moreMessages) {
1923
// no more messages, use this as a signal to close the last iteration scope
20-
closePrevious(true);
24+
if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
25+
closePrevious(true);
26+
} else {
27+
final AgentSpan previousSpan = spanFromContext(getRootContext().swap());
28+
if (previousSpan != null) {
29+
previousSpan.finishWithEndToEnd();
30+
}
31+
}
2132
}
2233
return moreMessages;
2334
}

‎dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/test/groovy/SqsClientTest.groovy‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
22
import static java.nio.charset.StandardCharsets.UTF_8
3+
import datadog.trace.api.config.TraceInstrumentationConfig
34

45
import com.amazon.sqs.javamessaging.ProviderConfiguration
56
import com.amazon.sqs.javamessaging.SQSConnectionFactory
@@ -574,4 +575,12 @@ class SqsClientV1DataStreamsForkedTest extends SqsClientTest {
574575
}
575576
}
576577

578+
class SqsClientV0ContextSwapForkedTest extends SqsClientV0Test {
579+
@Override
580+
protected void configurePreAgent() {
581+
super.configurePreAgent()
582+
injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false")
583+
}
584+
}
585+
577586

‎dd-java-agent/instrumentation/jms/jakarta-jms-3.0/src/test/groovy/JMS2Test.groovy‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import com.google.common.io.Files
22
import datadog.trace.agent.test.InstrumentationSpecification
3+
import datadog.trace.api.config.TraceInstrumentationConfig
34
import datadog.trace.agent.test.asserts.ListWriterAssert
45
import datadog.trace.api.DDSpanTypes
56
import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags
@@ -254,3 +255,11 @@ class JMS2Test extends InstrumentationSpecification {
254255
}
255256
}
256257
}
258+
259+
class JMS2ContextSwapForkedTest extends JMS2Test {
260+
@Override
261+
protected void configurePreAgent() {
262+
super.configurePreAgent()
263+
injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false")
264+
}
265+
}

‎dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java‎

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@
99
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateNext;
1010
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious;
1111
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
12+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext;
13+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
1214
import static datadog.trace.instrumentation.jms.JMSDecorator.BROKER_DECORATE;
1315
import static datadog.trace.instrumentation.jms.JMSDecorator.CONSUMER_DECORATE;
1416
import static datadog.trace.instrumentation.jms.JMSDecorator.JMS_CONSUME;
@@ -24,6 +26,7 @@
2426

2527
import datadog.trace.agent.tooling.Instrumenter;
2628
import datadog.trace.api.Config;
29+
import datadog.trace.api.InstrumenterConfig;
2730
import datadog.trace.api.datastreams.DataStreamsContext;
2831
import datadog.trace.api.datastreams.DataStreamsTags;
2932
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
@@ -90,7 +93,14 @@ public static MessageConsumerState beforeReceive(@Advice.This final MessageConsu
9093
}
9194

9295
boolean finishSpan = consumerState.getSessionState().isAutoAcknowledge();
93-
closePrevious(finishSpan);
96+
if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
97+
closePrevious(finishSpan);
98+
} else {
99+
final AgentSpan previousSpan = spanFromContext(getRootContext().swap());
100+
if (previousSpan != null) {
101+
previousSpan.finishWithEndToEnd();
102+
}
103+
}
94104
if (finishSpan) {
95105
consumerState.finishTimeInQueueSpan(false);
96106
}
@@ -163,7 +173,14 @@ public static void afterReceive(
163173

164174
CONSUMER_DECORATE.onError(span, throwable);
165175

166-
activateNext(span); // scope is left open until next message or it times out
176+
if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
177+
activateNext(span); // scope is left open until next message or it times out
178+
} else {
179+
final AgentSpan previousSpan = spanFromContext(span.swap());
180+
if (previousSpan != null) {
181+
previousSpan.finishWithEndToEnd();
182+
}
183+
}
167184
JMSLogger.logIterationSpan(span);
168185

169186
SessionState sessionState = consumerState.getSessionState();
@@ -187,7 +204,15 @@ public static void beforeClose(@Advice.This final MessageConsumer consumer) {
187204
.get(consumer);
188205
if (null != consumerState) {
189206
boolean finishSpan = consumerState.getSessionState().isAutoAcknowledge();
190-
closePrevious(finishSpan);
207+
if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
208+
closePrevious(finishSpan);
209+
} else {
210+
final AgentSpan previousSpan = spanFromContext(getRootContext().swap());
211+
if (previousSpan != null) {
212+
CONSUMER_DECORATE.beforeFinish(previousSpan);
213+
previousSpan.finishWithEndToEnd();
214+
}
215+
}
191216
if (finishSpan) {
192217
consumerState.finishTimeInQueueSpan(true);
193218
}

‎dd-java-agent/instrumentation/jms/javax-jms-1.1/src/test/groovy/JMS1Test.groovy‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
1+
import static datadog.trace.api.config.TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED
2+
import static org.junit.jupiter.api.Assumptions.assumeTrue
3+
14
import datadog.trace.agent.test.asserts.ListWriterAssert
25
import datadog.trace.agent.test.asserts.TraceAssert
36
import datadog.trace.agent.test.naming.VersionedNamingTestBase
@@ -61,6 +64,10 @@ abstract class JMS1Test extends VersionedNamingTestBase {
6164

6265
abstract String operationForConsumer()
6366

67+
boolean testUnclosedScopeFinished() {
68+
true
69+
}
70+
6471
def setupSpec() {
6572
broker.start()
6673
final ActiveMQConnectionFactory connectionFactory = broker.createConnectionFactory()
@@ -164,6 +171,7 @@ abstract class JMS1Test extends VersionedNamingTestBase {
164171

165172
def "closing #destinationType session should close and finish any pending scopes"() {
166173
setup:
174+
assumeTrue(testUnclosedScopeFinished())
167175
def destination = destinationType.create(session)
168176
def localSession = connection.createSession(false, Session.AUTO_ACKNOWLEDGE)
169177
def producer = localSession.createProducer(destination)
@@ -1045,6 +1053,19 @@ class JMS1V0Test extends JMS1Test {
10451053
}
10461054
}
10471055
1056+
class JMSContextSwapForkedTest extends JMS1V0Test {
1057+
@Override
1058+
protected void configurePreAgent() {
1059+
injectSysConfig(LEGACY_CONTEXT_MANAGER_ENABLED, "false")
1060+
}
1061+
1062+
@Override
1063+
boolean testUnclosedScopeFinished() {
1064+
//TODO: This need to be removed when the Context manager will support it
1065+
false
1066+
}
1067+
}
1068+
10481069
class JMS1V1ForkedTest extends JMS1Test {
10491070
10501071
@Override

0 commit comments

Comments
 (0)