Skip to content

Commit 8dfe6f4

Browse files
jpbempeljpbempel
committed
Fix initialization of samplers for probe file
When loading probes from file, initSamplers method is not called for deserialized probes. we centralize probe deserialization in ProbeDefinitionDeserializer class and add smoke test for this.
1 parent fc47338 commit 8dfe6f4

10 files changed

Lines changed: 261 additions & 79 deletions

File tree

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

Lines changed: 12 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -1,30 +1,24 @@
11
package com.datadog.debugger.agent;
22

3-
import com.datadog.debugger.probe.LogProbe;
4-
import com.datadog.debugger.probe.MetricProbe;
3+
import static com.datadog.debugger.probe.ProbeDefinitionDeserializer.deserializeLogProbe;
4+
import static com.datadog.debugger.probe.ProbeDefinitionDeserializer.deserializeMetricProbe;
5+
import static com.datadog.debugger.probe.ProbeDefinitionDeserializer.deserializeSpanDecorationProbe;
6+
import static com.datadog.debugger.probe.ProbeDefinitionDeserializer.deserializeSpanProbe;
7+
import static com.datadog.debugger.probe.ProbeDefinitionDeserializer.deserializeTriggerProbe;
8+
59
import com.datadog.debugger.probe.ProbeDefinition;
6-
import com.datadog.debugger.probe.SpanDecorationProbe;
7-
import com.datadog.debugger.probe.SpanProbe;
8-
import com.datadog.debugger.probe.TriggerProbe;
9-
import com.datadog.debugger.util.MoshiHelper;
1010
import com.squareup.moshi.JsonAdapter;
1111
import com.squareup.moshi.JsonReader;
1212
import com.squareup.moshi.JsonWriter;
13-
import com.squareup.moshi.Moshi;
14-
import com.squareup.moshi.Types;
1513
import datadog.trace.util.SizeCheckedInputStream;
1614
import java.io.ByteArrayInputStream;
1715
import java.io.ByteArrayOutputStream;
1816
import java.io.FileInputStream;
1917
import java.io.IOException;
2018
import java.io.InputStream;
21-
import java.lang.annotation.Annotation;
22-
import java.lang.reflect.ParameterizedType;
23-
import java.lang.reflect.Type;
2419
import java.nio.file.Path;
2520
import java.util.ArrayList;
2621
import java.util.List;
27-
import java.util.Set;
2822
import okio.Okio;
2923
import org.slf4j.Logger;
3024
import org.slf4j.LoggerFactory;
@@ -46,9 +40,7 @@ public static Configuration from(Path probeFilePath, long maxPayloadSize) {
4640
}
4741
} while (bytesRead > -1);
4842
byte[] configContent = outputStream.toByteArray();
49-
Moshi moshi = MoshiHelper.createMoshiConfigBuilder().add(new ProbeFileFactory()).build();
50-
ParameterizedType type = Types.newParameterizedType(List.class, ProbeDefinition.class);
51-
JsonAdapter<List<ProbeDefinition>> adapter = moshi.adapter(type);
43+
JsonAdapter<List<ProbeDefinition>> adapter = new ProbeFileAdapter();
5244
List<ProbeDefinition> probeDefinitions =
5345
adapter.fromJson(
5446
JsonReader.of(Okio.buffer(Okio.source(new ByteArrayInputStream(configContent)))));
@@ -59,40 +51,7 @@ public static Configuration from(Path probeFilePath, long maxPayloadSize) {
5951
}
6052
}
6153

62-
private static class ProbeFileFactory implements JsonAdapter.Factory {
63-
@Override
64-
public JsonAdapter<?> create(Type type, Set<? extends Annotation> annotations, Moshi moshi) {
65-
if (Types.equals(type, Types.newParameterizedType(List.class, ProbeDefinition.class))) {
66-
return new ProbeFileAdapter(
67-
moshi.adapter(LogProbe.class),
68-
moshi.adapter(MetricProbe.class),
69-
moshi.adapter(SpanProbe.class),
70-
moshi.adapter(SpanDecorationProbe.class),
71-
moshi.adapter(TriggerProbe.class));
72-
}
73-
return null;
74-
}
75-
}
76-
7754
private static class ProbeFileAdapter extends JsonAdapter<List<ProbeDefinition>> {
78-
private final JsonAdapter<LogProbe> logProbeAdapter;
79-
private final JsonAdapter<MetricProbe> metricProbeAdapter;
80-
private final JsonAdapter<SpanProbe> spanProbeAdapter;
81-
private final JsonAdapter<SpanDecorationProbe> spanDecorationProbeAdapter;
82-
private final JsonAdapter<TriggerProbe> triggerProbeAdapter;
83-
84-
public ProbeFileAdapter(
85-
JsonAdapter<LogProbe> logProbeAdapter,
86-
JsonAdapter<MetricProbe> metricProbeAdapter,
87-
JsonAdapter<SpanProbe> spanProbeAdapter,
88-
JsonAdapter<SpanDecorationProbe> spanDecorationProbeAdapter,
89-
JsonAdapter<TriggerProbe> triggerProbeAdapter) {
90-
this.logProbeAdapter = logProbeAdapter;
91-
this.metricProbeAdapter = metricProbeAdapter;
92-
this.spanProbeAdapter = spanProbeAdapter;
93-
this.spanDecorationProbeAdapter = spanDecorationProbeAdapter;
94-
this.triggerProbeAdapter = triggerProbeAdapter;
95-
}
9655

9756
@Override
9857
public List<ProbeDefinition> fromJson(JsonReader reader) throws IOException {
@@ -110,19 +69,19 @@ public List<ProbeDefinition> fromJson(JsonReader reader) throws IOException {
11069
String type = jsonPeekReader.nextString();
11170
switch (type) {
11271
case "LOG_PROBE":
113-
probeDefinitions.add(logProbeAdapter.fromJson(reader));
72+
probeDefinitions.add(deserializeLogProbe(reader));
11473
break;
11574
case "METRIC_PROBE":
116-
probeDefinitions.add(metricProbeAdapter.fromJson(reader));
75+
probeDefinitions.add(deserializeMetricProbe(reader));
11776
break;
11877
case "SPAN_PROBE":
119-
probeDefinitions.add(spanProbeAdapter.fromJson(reader));
78+
probeDefinitions.add(deserializeSpanProbe(reader));
12079
break;
12180
case "SPAN_DECORATION_PROBE":
122-
probeDefinitions.add(spanDecorationProbeAdapter.fromJson(reader));
81+
probeDefinitions.add(deserializeSpanDecorationProbe(reader));
12382
break;
12483
case "TRIGGER_PROBE":
125-
probeDefinitions.add(triggerProbeAdapter.fromJson(reader));
84+
probeDefinitions.add(deserializeTriggerProbe(reader));
12685
break;
12786
default:
12887
throw new RuntimeException("Unknown type: " + type);

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import com.datadog.debugger.agent.Configuration;
44
import com.datadog.debugger.util.MoshiHelper;
55
import com.squareup.moshi.JsonAdapter;
6+
import com.squareup.moshi.JsonReader;
67
import java.io.ByteArrayInputStream;
78
import java.io.IOException;
89
import okio.Okio;
@@ -29,22 +30,42 @@ public static MetricProbe deserializeMetricProbe(byte[] content) throws IOExcept
2930
return deserialize(METRIC_PROBE_JSON_ADAPTER, content);
3031
}
3132

33+
public static MetricProbe deserializeMetricProbe(JsonReader reader) throws IOException {
34+
return METRIC_PROBE_JSON_ADAPTER.fromJson(reader);
35+
}
36+
3237
public static LogProbe deserializeLogProbe(byte[] content) throws IOException {
3338
LogProbe logProbe = deserialize(LOG_PROBE_JSON_ADAPTER, content);
3439
logProbe.initSamplers();
3540
return logProbe;
3641
}
3742

43+
public static LogProbe deserializeLogProbe(JsonReader reader) throws IOException {
44+
LogProbe logProbe = LOG_PROBE_JSON_ADAPTER.fromJson(reader);
45+
logProbe.initSamplers();
46+
return logProbe;
47+
}
48+
3849
public static SpanProbe deserializeSpanProbe(byte[] content) throws IOException {
3950
return deserialize(SPAN_PROBE_JSON_ADAPTER, content);
4051
}
4152

53+
public static SpanProbe deserializeSpanProbe(JsonReader reader) throws IOException {
54+
return SPAN_PROBE_JSON_ADAPTER.fromJson(reader);
55+
}
56+
4257
public static TriggerProbe deserializeTriggerProbe(byte[] content) throws IOException {
4358
TriggerProbe triggerProbe = deserialize(TRIGGER_PROBE_JSON_ADAPTER, content);
4459
triggerProbe.initSamplers();
4560
return triggerProbe;
4661
}
4762

63+
public static TriggerProbe deserializeTriggerProbe(JsonReader jsonReader) throws IOException {
64+
TriggerProbe triggerProbe = TRIGGER_PROBE_JSON_ADAPTER.fromJson(jsonReader);
65+
triggerProbe.initSamplers();
66+
return triggerProbe;
67+
}
68+
4869
public static SpanDecorationProbe deserializeSpanDecorationProbe(byte[] content)
4970
throws IOException {
5071
SpanDecorationProbe spanDecorationProbe =
@@ -53,6 +74,14 @@ public static SpanDecorationProbe deserializeSpanDecorationProbe(byte[] content)
5374
return spanDecorationProbe;
5475
}
5576

77+
public static SpanDecorationProbe deserializeSpanDecorationProbe(JsonReader jsonReader)
78+
throws IOException {
79+
SpanDecorationProbe spanDecorationProbe =
80+
SPAN_DECORATION_PROBE_JSON_ADAPTER.fromJson(jsonReader);
81+
spanDecorationProbe.initSamplers();
82+
return spanDecorationProbe;
83+
}
84+
5685
private static <T> T deserialize(JsonAdapter<T> adapter, byte[] content) throws IOException {
5786
return adapter.fromJson(Okio.buffer(Okio.source(new ByteArrayInputStream(content))));
5887
}

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

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import com.datadog.debugger.probe.ProbeDefinition;
88
import com.datadog.debugger.probe.SpanDecorationProbe;
99
import com.datadog.debugger.probe.SpanProbe;
10+
import com.datadog.debugger.probe.TriggerProbe;
1011
import java.nio.file.Path;
1112
import java.nio.file.Paths;
1213
import java.util.List;
@@ -21,11 +22,12 @@ public void load() throws Exception {
2122
Configuration configuration = ConfigurationFileLoader.from(probeFilePath, 1024 * 1024);
2223
assertNotNull(configuration);
2324
List<ProbeDefinition> definitions = configuration.getDefinitions();
24-
assertEquals(5, definitions.size());
25-
assertInstanceOf(MetricProbe.class, definitions.get(0));
26-
assertInstanceOf(LogProbe.class, definitions.get(1));
25+
assertEquals(6, definitions.size());
26+
assertInstanceOf(TriggerProbe.class, definitions.get(0));
27+
assertInstanceOf(MetricProbe.class, definitions.get(1));
2728
assertInstanceOf(LogProbe.class, definitions.get(2));
28-
assertInstanceOf(SpanProbe.class, definitions.get(3));
29-
assertInstanceOf(SpanDecorationProbe.class, definitions.get(4));
29+
assertInstanceOf(LogProbe.class, definitions.get(3));
30+
assertInstanceOf(SpanProbe.class, definitions.get(4));
31+
assertInstanceOf(SpanDecorationProbe.class, definitions.get(5));
3032
}
3133
}

‎dd-java-agent/agent-debugger/src/test/resources/test_probe_file.json‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,5 +87,19 @@
8787
},
8888
"tags": [],
8989
"evaluateAt": "EXIT"
90+
},
91+
{
92+
"id": "70b55d06-f9fa-403b-a329-4f2f960aed02",
93+
"version": 0,
94+
"type": "TRIGGER_PROBE",
95+
"language": "java",
96+
"where": {
97+
"typeName": "MetadataClientUtils",
98+
"methodName": "listTableWithContinuation"
99+
},
100+
"tags": [],
101+
"evaluateAt": "EXIT",
102+
"sessionId": "session",
103+
"sampling": { "snapshotsPerSecond": 100 }
90104
}
91105
]

‎dd-smoke-tests/debugger-integration-tests/src/test/java/datadog/smoketest/BaseIntegrationTest.java‎

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import com.squareup.moshi.Moshi;
2222
import com.squareup.moshi.Types;
2323
import datadog.trace.bootstrap.debugger.CapturedContext;
24+
import datadog.trace.bootstrap.debugger.ProbeId;
2425
import datadog.trace.bootstrap.debugger.ProbeRateLimiter;
2526
import datadog.trace.test.agent.decoder.DecodedMessage;
2627
import datadog.trace.test.agent.decoder.DecodedTrace;
@@ -332,23 +333,25 @@ protected void registerProbeStatusListener(Consumer<ProbeStatus> listener) {
332333
probeStatusListeners.add(listener);
333334
}
334335

335-
protected AtomicBoolean registerCheckReceivedInstalledEmitting() {
336+
protected AtomicBoolean registerCheckReceivedInstalledEmitting(ProbeId probeId) {
336337
AtomicBoolean received = new AtomicBoolean();
337338
AtomicBoolean installed = new AtomicBoolean();
338339
AtomicBoolean emitting = new AtomicBoolean();
339340
AtomicBoolean result = new AtomicBoolean();
340341
registerProbeStatusListener(
341342
probeStatus -> {
342-
if (probeStatus.getDiagnostics().getStatus() == ProbeStatus.Status.RECEIVED) {
343-
received.set(true);
344-
}
345-
if (probeStatus.getDiagnostics().getStatus() == ProbeStatus.Status.INSTALLED) {
346-
installed.set(true);
347-
}
348-
if (probeStatus.getDiagnostics().getStatus() == ProbeStatus.Status.EMITTING) {
349-
emitting.set(true);
343+
if (probeStatus.getDiagnostics().getProbeId().equals(probeId)) {
344+
if (probeStatus.getDiagnostics().getStatus() == ProbeStatus.Status.RECEIVED) {
345+
received.set(true);
346+
}
347+
if (probeStatus.getDiagnostics().getStatus() == ProbeStatus.Status.INSTALLED) {
348+
installed.set(true);
349+
}
350+
if (probeStatus.getDiagnostics().getStatus() == ProbeStatus.Status.EMITTING) {
351+
emitting.set(true);
352+
}
353+
result.set(received.get() && installed.get() && emitting.get());
350354
}
351-
result.set(received.get() && installed.get() && emitting.get());
352355
});
353356
return result;
354357
}

‎dd-smoke-tests/debugger-integration-tests/src/test/java/datadog/smoketest/LogProbesIntegrationTest.java‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ void testFullMethod() throws Exception {
100100
assertNull(snapshot.getCaptures().getReturn().getCapturedThrowable());
101101
snapshotReceived.set(true);
102102
});
103-
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
103+
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(PROBE_ID);
104104
processRequests(
105105
() -> snapshotReceived.get() && statusResult.get(),
106106
() ->
@@ -132,7 +132,7 @@ void testFullMethodWithCondition() throws Exception {
132132
assertFullMethodCaptureArgs(snapshot.getCaptures().getEntry());
133133
snapshotReceived.set(true);
134134
});
135-
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
135+
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(PROBE_ID);
136136
processRequests(
137137
() -> snapshotReceived.get() && statusResult.get(),
138138
() ->
@@ -166,7 +166,7 @@ void testFullMethodWithConditionAtExit() throws Exception {
166166
assertFullMethodCaptureArgs(snapshot.getCaptures().getReturn());
167167
snapshotReceived.set(true);
168168
});
169-
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
169+
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(PROBE_ID);
170170
processRequests(
171171
() -> snapshotReceived.get() && statusResult.get(),
172172
() ->
@@ -199,7 +199,7 @@ void testFullMethodWithConditionFailed() throws Exception {
199199
"Cannot find symbol: noarg", snapshot.getEvaluationErrors().get(0).getMessage());
200200
snapshotReceived.set(true);
201201
});
202-
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
202+
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(PROBE_ID);
203203
processRequests(
204204
() -> snapshotReceived.get() && statusResult.get(),
205205
() ->
@@ -236,7 +236,7 @@ void testFullMethodWithLogTemplate() throws Exception {
236236
assertEquals(PROBE_ID.getId(), snapshot.getProbe().getId());
237237
snapshotReceived.set(true);
238238
});
239-
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
239+
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(PROBE_ID);
240240
processRequests(
241241
() -> snapshotReceived.get() && correctLogMessage.get() && statusResult.get(),
242242
() ->
@@ -288,7 +288,7 @@ void testFullMethodWithCaptureExpressions() throws Exception {
288288
assertNull(key2Value.getNotCapturedReason());
289289
snapshotReceived.set(true);
290290
});
291-
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
291+
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(PROBE_ID);
292292
processRequests(
293293
() -> snapshotReceived.get() && statusResult.get(),
294294
() ->
@@ -371,7 +371,7 @@ void testLineProbe() throws Exception {
371371
assertNull(capturedContext.getCapturedThrowable());
372372
snapshotReceived.set(true);
373373
});
374-
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
374+
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(LINE_PROBE_ID1);
375375
processRequests(
376376
() -> snapshotReceived.get() && statusResult.get(),
377377
() ->

‎dd-smoke-tests/debugger-integration-tests/src/test/java/datadog/smoketest/MetricProbesIntegrationTest.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ private void doMethodMetric(
112112
targetProcess = createProcessBuilder(logFilePath, METHOD_NAME, EXPECTED_UPLOADS).start();
113113
String msgExpected = String.format(expectedMsgFormat, metricName, PROBE_ID.getId());
114114
assertNotNull(retrieveStatsdMessage(msgExpected));
115-
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
115+
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(PROBE_ID);
116116
processRequests(
117117
statusResult::get, () -> String.format("timeout statusResult=%s", statusResult.get()));
118118
}
@@ -219,7 +219,7 @@ private void doLineMetric(
219219
targetProcess = createProcessBuilder(logFilePath, METHOD_NAME, EXPECTED_UPLOADS).start();
220220
String msgExpected = String.format(expectedMsgFormat, metricName, PROBE_ID.getId());
221221
assertNotNull(retrieveStatsdMessage(msgExpected));
222-
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
222+
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(PROBE_ID);
223223
processRequests(
224224
statusResult::get, () -> String.format("timeout statusResult=%s", statusResult.get()));
225225
}

0 commit comments

Comments
 (0)