Skip to content

Commit 324a983

Browse files
amarzialiamarziali
authored
Merge branch 'master' into andrea.marziali/separate-context-tracking
2 parents 100b7c8 + cc12228 commit 324a983

43 files changed

Lines changed: 1565 additions & 544 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/JsonSnapshotSerializer.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,16 @@
99
import datadog.trace.api.ProcessTags;
1010
import datadog.trace.bootstrap.debugger.CapturedContext;
1111
import datadog.trace.bootstrap.debugger.DebuggerContext;
12+
import java.time.Duration;
13+
import java.time.temporal.ChronoUnit;
1214

1315
/** Serializes snapshots in Json using Moshi */
1416
public class JsonSnapshotSerializer implements DebuggerContext.ValueSerializer {
1517
private static final JsonAdapter<IntakeRequest> ADAPTER =
16-
MoshiHelper.createMoshiSnapshot().adapter(IntakeRequest.class);
18+
MoshiHelper.createMoshiSnapshot(
19+
Duration.of(
20+
Config.get().getDynamicInstrumentationCaptureTimeout(), ChronoUnit.MILLIS))
21+
.adapter(IntakeRequest.class);
1722
private static final JsonAdapter<CapturedContext.CapturedValue> VALUE_ADAPTER =
1823
new MoshiSnapshotHelper.CapturedValueAdapter();
1924

‎dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/util/MoshiHelper.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import com.squareup.moshi.Types;
1212
import datadog.trace.bootstrap.debugger.el.DebuggerScript;
1313
import java.lang.reflect.ParameterizedType;
14+
import java.time.Duration;
1415
import java.util.Map;
1516

1617
/** Helper for creating Moshi instances with the right adapters depending on the context */
@@ -32,9 +33,9 @@ public static Moshi.Builder createMoshiConfigBuilder() {
3233
.add(ProbeDefinition.Tag[].class, new ProbeDefinition.TagAdapter());
3334
}
3435

35-
public static Moshi createMoshiSnapshot() {
36+
public static Moshi createMoshiSnapshot(Duration captureTimeOut) {
3637
return new Moshi.Builder()
37-
.add(new MoshiSnapshotHelper.SnapshotJsonFactory())
38+
.add(new MoshiSnapshotHelper.SnapshotJsonFactory(captureTimeOut))
3839
.add(
3940
DebuggerScript.class,
4041
new ProbeCondition.ProbeConditionJsonAdapter()) // ProbeDetails in Snapshot

‎dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/util/MoshiSnapshotHelper.java‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import com.squareup.moshi.JsonWriter;
77
import com.squareup.moshi.Moshi;
88
import com.squareup.moshi.Types;
9+
import datadog.trace.api.Config;
910
import datadog.trace.bootstrap.debugger.CapturedContext;
1011
import datadog.trace.bootstrap.debugger.CapturedStackFrame;
1112
import datadog.trace.bootstrap.debugger.Limits;
@@ -62,14 +63,20 @@ public class MoshiSnapshotHelper {
6263
public static final String LOCATION = "location";
6364
public static final String MESSAGE = "message";
6465
public static final String STACKTRACE = "stacktrace";
65-
private static final Duration TIME_OUT = Duration.ofMillis(200);
6666

6767
public static class SnapshotJsonFactory implements JsonAdapter.Factory {
68+
private final Duration captureTimeOut;
69+
70+
public SnapshotJsonFactory(Duration captureTimeOut) {
71+
this.captureTimeOut = captureTimeOut;
72+
}
73+
6874
@Override
6975
public JsonAdapter<?> create(Type type, Set<? extends Annotation> set, Moshi moshi) {
7076
if (Types.equals(type, Snapshot.Captures.class)) {
7177
return new CapturesAdapter(
7278
moshi,
79+
captureTimeOut,
7380
new CapturedContextAdapter(
7481
moshi, new CapturedValueAdapter(), new CapturedThrowableAdapter(moshi)));
7582
}
@@ -88,12 +95,15 @@ public JsonAdapter<?> create(Type type, Set<? extends Annotation> set, Moshi mos
8895
}
8996

9097
public static class CapturesAdapter extends JsonAdapter<Snapshot.Captures> {
98+
protected final Duration captureTimeOut;
9199
protected final JsonAdapter<CapturedContext> capturedContextAdapter;
92100
protected final JsonAdapter<Map<Integer, CapturedContext>> linesAdapter;
93101
protected final JsonAdapter<List<CapturedContext.CapturedThrowable>> caughtExceptionsAdapter;
94102

95-
public CapturesAdapter(Moshi moshi, JsonAdapter<CapturedContext> capturedContextAdapter) {
103+
public CapturesAdapter(
104+
Moshi moshi, Duration captureTimeout, JsonAdapter<CapturedContext> capturedContextAdapter) {
96105
this.capturedContextAdapter = capturedContextAdapter;
106+
this.captureTimeOut = captureTimeout;
97107
linesAdapter =
98108
moshi.adapter(
99109
Types.newParameterizedType(Map.class, Integer.class, CapturedContext.class));
@@ -109,7 +119,7 @@ public void toJson(JsonWriter jsonWriter, Snapshot.Captures captures) throws IOE
109119
return;
110120
}
111121
jsonWriter.setTag(
112-
TimeoutChecker.class, new TimeoutChecker(System.currentTimeMillis(), TIME_OUT));
122+
TimeoutChecker.class, new TimeoutChecker(System.currentTimeMillis(), captureTimeOut));
113123
jsonWriter.beginObject();
114124
jsonWriter.name(ENTRY);
115125
capturedContextAdapter.toJson(jsonWriter, captures.getEntry());
@@ -149,7 +159,9 @@ public void toJson(JsonWriter jsonWriter, CapturedContext capturedContext) throw
149159
}
150160
TimeoutChecker timeoutChecker = jsonWriter.tag(TimeoutChecker.class);
151161
if (timeoutChecker == null) {
152-
timeoutChecker = new TimeoutChecker(TIME_OUT);
162+
Duration timeout =
163+
Duration.of(Config.get().getDynamicInstrumentationCaptureTimeout(), ChronoUnit.MILLIS);
164+
timeoutChecker = new TimeoutChecker(timeout);
153165
}
154166
// need to 'freeze' the context before serializing it
155167
capturedContext.freeze(timeoutChecker);

‎dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/SnapshotSerializationTest.java‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -789,7 +789,10 @@ public void length255() throws IOException {
789789
@Test
790790
public void capturesAdapterNull() {
791791
MoshiSnapshotHelper.CapturesAdapter capturesAdapter =
792-
new MoshiSnapshotHelper.CapturesAdapter(MoshiHelper.createMoshiSnapshot(), null);
792+
new MoshiSnapshotHelper.CapturesAdapter(
793+
MoshiSnapshotTestHelper.createMoshiSnapshot(),
794+
Duration.of(5, ChronoUnit.SECONDS),
795+
null);
793796
Assertions.assertEquals("null", capturesAdapter.toJson(null));
794797
}
795798

@@ -921,7 +924,9 @@ public void fieldCount20() throws IOException {
921924
public void timeOut() throws IOException {
922925
DebuggerContext.initValueSerializer(
923926
new TimeoutSnapshotSerializer(Duration.of(150, ChronoUnit.MILLIS)));
924-
JsonAdapter<Snapshot> adapter = createSnapshotAdapter();
927+
JsonAdapter<Snapshot> adapter =
928+
MoshiHelper.createMoshiSnapshot(Duration.of(100, ChronoUnit.MILLIS))
929+
.adapter(Snapshot.class);
925930
Snapshot snapshot = createSnapshot();
926931
CapturedContext context = new CapturedContext();
927932
CapturedContext.CapturedValue arg1 = CapturedContext.CapturedValue.of("arg1", "int", 42);

‎dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/util/MoshiSnapshotTestHelper.java‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@
3939
import java.io.IOException;
4040
import java.lang.annotation.Annotation;
4141
import java.lang.reflect.Type;
42+
import java.time.Duration;
43+
import java.time.temporal.ChronoUnit;
4244
import java.util.ArrayList;
4345
import java.util.Arrays;
4446
import java.util.Collections;
@@ -90,11 +92,18 @@ public static String getValue(CapturedContext.CapturedValue capturedValue) {
9092
}
9193

9294
public static class SnapshotJsonFactory implements JsonAdapter.Factory {
95+
private final Duration captureTimeOut;
96+
97+
public SnapshotJsonFactory(Duration captureTimeout) {
98+
this.captureTimeOut = captureTimeout;
99+
}
100+
93101
@Override
94102
public JsonAdapter<?> create(Type type, Set<? extends Annotation> set, Moshi moshi) {
95103
if (Types.equals(type, Snapshot.Captures.class)) {
96104
return new MoshiSnapshotTestHelper.CapturesAdapter(
97105
moshi,
106+
captureTimeOut,
98107
new CapturedContextAdapter(
99108
moshi, new CapturedValueAdapter(), new CapturedThrowableAdapter(moshi)));
100109
}
@@ -113,8 +122,9 @@ public JsonAdapter<?> create(Type type, Set<? extends Annotation> set, Moshi mos
113122
}
114123

115124
public static Moshi createMoshiSnapshot() {
125+
// By default, for tests we have a capture timeout of 5 seconds
116126
return new Moshi.Builder()
117-
.add(new MoshiSnapshotTestHelper.SnapshotJsonFactory())
127+
.add(new MoshiSnapshotTestHelper.SnapshotJsonFactory(Duration.of(5, ChronoUnit.SECONDS)))
118128
.add(
119129
DebuggerScript.class,
120130
new ProbeCondition.ProbeConditionJsonAdapter()) // ProbeDetails in Snapshot
@@ -152,8 +162,9 @@ private static String primitiveArrayToString(Object obj) {
152162

153163
private static class CapturesAdapter extends MoshiSnapshotHelper.CapturesAdapter {
154164

155-
public CapturesAdapter(Moshi moshi, JsonAdapter<CapturedContext> capturedContextAdapter) {
156-
super(moshi, capturedContextAdapter);
165+
public CapturesAdapter(
166+
Moshi moshi, Duration captureTimeout, JsonAdapter<CapturedContext> capturedContextAdapter) {
167+
super(moshi, captureTimeout, capturedContextAdapter);
157168
}
158169

159170
@Override

‎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
}

0 commit comments

Comments
 (0)