Skip to content

Fix initialization of samplers for probe file #12067

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jpbempel wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Fix initialization of samplers for probe file #12067

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 12 additions & 53 deletions ...gent/agent-debugger/src/main/java/com/datadog/debugger/agent/ConfigurationFileLoader.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,30 +1,24 @@
package com.datadog.debugger.agent;

import com.datadog.debugger.probe.LogProbe;
import com.datadog.debugger.probe.MetricProbe;
import static com.datadog.debugger.probe.ProbeDefinitionDeserializer.deserializeLogProbe;
import static com.datadog.debugger.probe.ProbeDefinitionDeserializer.deserializeMetricProbe;
import static com.datadog.debugger.probe.ProbeDefinitionDeserializer.deserializeSpanDecorationProbe;
import static com.datadog.debugger.probe.ProbeDefinitionDeserializer.deserializeSpanProbe;
import static com.datadog.debugger.probe.ProbeDefinitionDeserializer.deserializeTriggerProbe;

import com.datadog.debugger.probe.ProbeDefinition;
import com.datadog.debugger.probe.SpanDecorationProbe;
import com.datadog.debugger.probe.SpanProbe;
import com.datadog.debugger.probe.TriggerProbe;
import com.datadog.debugger.util.MoshiHelper;
import com.squareup.moshi.JsonAdapter;
import com.squareup.moshi.JsonReader;
import com.squareup.moshi.JsonWriter;
import com.squareup.moshi.Moshi;
import com.squareup.moshi.Types;
import datadog.trace.util.SizeCheckedInputStream;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.lang.annotation.Annotation;
import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Type;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import okio.Okio;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
Expand All @@ -46,9 +40,7 @@ public static Configuration from(Path probeFilePath, long maxPayloadSize) {
}
} while (bytesRead > -1);
byte[] configContent = outputStream.toByteArray();
Moshi moshi = MoshiHelper.createMoshiConfigBuilder().add(new ProbeFileFactory()).build();
ParameterizedType type = Types.newParameterizedType(List.class, ProbeDefinition.class);
JsonAdapter<List<ProbeDefinition>> adapter = moshi.adapter(type);
JsonAdapter<List<ProbeDefinition>> adapter = new ProbeFileAdapter();
List<ProbeDefinition> probeDefinitions =
adapter.fromJson(
JsonReader.of(Okio.buffer(Okio.source(new ByteArrayInputStream(configContent)))));
Expand All @@ -59,40 +51,7 @@ public static Configuration from(Path probeFilePath, long maxPayloadSize) {
}
}

private static class ProbeFileFactory implements JsonAdapter.Factory {
@Override
public JsonAdapter<?> create(Type type, Set<? extends Annotation> annotations, Moshi moshi) {
if (Types.equals(type, Types.newParameterizedType(List.class, ProbeDefinition.class))) {
return new ProbeFileAdapter(
moshi.adapter(LogProbe.class),
moshi.adapter(MetricProbe.class),
moshi.adapter(SpanProbe.class),
moshi.adapter(SpanDecorationProbe.class),
moshi.adapter(TriggerProbe.class));
}
return null;
}
}

private static class ProbeFileAdapter extends JsonAdapter<List<ProbeDefinition>> {
private final JsonAdapter<LogProbe> logProbeAdapter;
private final JsonAdapter<MetricProbe> metricProbeAdapter;
private final JsonAdapter<SpanProbe> spanProbeAdapter;
private final JsonAdapter<SpanDecorationProbe> spanDecorationProbeAdapter;
private final JsonAdapter<TriggerProbe> triggerProbeAdapter;

public ProbeFileAdapter(
JsonAdapter<LogProbe> logProbeAdapter,
JsonAdapter<MetricProbe> metricProbeAdapter,
JsonAdapter<SpanProbe> spanProbeAdapter,
JsonAdapter<SpanDecorationProbe> spanDecorationProbeAdapter,
JsonAdapter<TriggerProbe> triggerProbeAdapter) {
this.logProbeAdapter = logProbeAdapter;
this.metricProbeAdapter = metricProbeAdapter;
this.spanProbeAdapter = spanProbeAdapter;
this.spanDecorationProbeAdapter = spanDecorationProbeAdapter;
this.triggerProbeAdapter = triggerProbeAdapter;
}

@Override
public List<ProbeDefinition> fromJson(JsonReader reader) throws IOException {
Expand All @@ -110,19 +69,19 @@ public List<ProbeDefinition> fromJson(JsonReader reader) throws IOException {
String type = jsonPeekReader.nextString();
switch (type) {
case "LOG_PROBE":
probeDefinitions.add(logProbeAdapter.fromJson(reader));
probeDefinitions.add(deserializeLogProbe(reader));
break;
case "METRIC_PROBE":
probeDefinitions.add(metricProbeAdapter.fromJson(reader));
probeDefinitions.add(deserializeMetricProbe(reader));
break;
case "SPAN_PROBE":
probeDefinitions.add(spanProbeAdapter.fromJson(reader));
probeDefinitions.add(deserializeSpanProbe(reader));
break;
case "SPAN_DECORATION_PROBE":
probeDefinitions.add(spanDecorationProbeAdapter.fromJson(reader));
probeDefinitions.add(deserializeSpanDecorationProbe(reader));
break;
case "TRIGGER_PROBE":
probeDefinitions.add(triggerProbeAdapter.fromJson(reader));
probeDefinitions.add(deserializeTriggerProbe(reader));
break;
default:
throw new RuntimeException("Unknown type: " + type);
Expand Down
29 changes: 29 additions & 0 deletions .../agent-debugger/src/main/java/com/datadog/debugger/probe/ProbeDefinitionDeserializer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import com.datadog.debugger.agent.Configuration;
import com.datadog.debugger.util.MoshiHelper;
import com.squareup.moshi.JsonAdapter;
import com.squareup.moshi.JsonReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import okio.Okio;
Expand All @@ -29,22 +30,42 @@ public static MetricProbe deserializeMetricProbe(byte[] content) throws IOExcept
return deserialize(METRIC_PROBE_JSON_ADAPTER, content);
}

public static MetricProbe deserializeMetricProbe(JsonReader reader) throws IOException {
return METRIC_PROBE_JSON_ADAPTER.fromJson(reader);
}

public static LogProbe deserializeLogProbe(byte[] content) throws IOException {
LogProbe logProbe = deserialize(LOG_PROBE_JSON_ADAPTER, content);
logProbe.initSamplers();
return logProbe;
}

public static LogProbe deserializeLogProbe(JsonReader reader) throws IOException {
LogProbe logProbe = LOG_PROBE_JSON_ADAPTER.fromJson(reader);
logProbe.initSamplers();
return logProbe;
}

public static SpanProbe deserializeSpanProbe(byte[] content) throws IOException {
return deserialize(SPAN_PROBE_JSON_ADAPTER, content);
}

public static SpanProbe deserializeSpanProbe(JsonReader reader) throws IOException {
return SPAN_PROBE_JSON_ADAPTER.fromJson(reader);
}

public static TriggerProbe deserializeTriggerProbe(byte[] content) throws IOException {
TriggerProbe triggerProbe = deserialize(TRIGGER_PROBE_JSON_ADAPTER, content);
triggerProbe.initSamplers();
return triggerProbe;
}

public static TriggerProbe deserializeTriggerProbe(JsonReader jsonReader) throws IOException {
TriggerProbe triggerProbe = TRIGGER_PROBE_JSON_ADAPTER.fromJson(jsonReader);
triggerProbe.initSamplers();
return triggerProbe;
}

public static SpanDecorationProbe deserializeSpanDecorationProbe(byte[] content)
throws IOException {
SpanDecorationProbe spanDecorationProbe =
Expand All @@ -53,6 +74,14 @@ public static SpanDecorationProbe deserializeSpanDecorationProbe(byte[] content)
return spanDecorationProbe;
}

public static SpanDecorationProbe deserializeSpanDecorationProbe(JsonReader jsonReader)
throws IOException {
SpanDecorationProbe spanDecorationProbe =
SPAN_DECORATION_PROBE_JSON_ADAPTER.fromJson(jsonReader);
spanDecorationProbe.initSamplers();
return spanDecorationProbe;
}

private static <T> T deserialize(JsonAdapter<T> adapter, byte[] content) throws IOException {
return adapter.fromJson(Okio.buffer(Okio.source(new ByteArrayInputStream(content))));
}
Expand Down
12 changes: 7 additions & 5 deletions .../agent-debugger/src/test/java/com/datadog/debugger/agent/ConfigurationFileLoaderTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import com.datadog.debugger.probe.ProbeDefinition;
import com.datadog.debugger.probe.SpanDecorationProbe;
import com.datadog.debugger.probe.SpanProbe;
import com.datadog.debugger.probe.TriggerProbe;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;
Expand All @@ -21,11 +22,12 @@ public void load() throws Exception {
Configuration configuration = ConfigurationFileLoader.from(probeFilePath, 1024 * 1024);
assertNotNull(configuration);
List<ProbeDefinition> definitions = configuration.getDefinitions();
assertEquals(5, definitions.size());
assertInstanceOf(MetricProbe.class, definitions.get(0));
assertInstanceOf(LogProbe.class, definitions.get(1));
assertEquals(6, definitions.size());
assertInstanceOf(TriggerProbe.class, definitions.get(0));

@chatgpt-codex-connector chatgpt-codex-connector Bot Jul 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Fix expected probe ordering in loader test

The probe-file fixture still starts with a LOG_PROBE and appends the new TRIGGER_PROBE at the end, and ConfigurationFileLoader appends definitions in JSON array order. This assertion therefore fails immediately with a LogProbe at index 0, so the updated test blocks the module test suite before it can verify the sampler initialization change.

Useful? React with 👍 / 👎.

@jpbempel jpbempel Jul 24, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no, Configuration::getDefinitions add definitions in a different order

assertInstanceOf(MetricProbe.class, definitions.get(1));
assertInstanceOf(LogProbe.class, definitions.get(2));
assertInstanceOf(SpanProbe.class, definitions.get(3));
assertInstanceOf(SpanDecorationProbe.class, definitions.get(4));
assertInstanceOf(LogProbe.class, definitions.get(3));
assertInstanceOf(SpanProbe.class, definitions.get(4));
assertInstanceOf(SpanDecorationProbe.class, definitions.get(5));
}
}
14 changes: 14 additions & 0 deletions dd-java-agent/agent-debugger/src/test/resources/test_probe_file.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,5 +87,19 @@
},
"tags": [],
"evaluateAt": "EXIT"
},
{
"id": "70b55d06-f9fa-403b-a329-4f2f960aed02",
"version": 0,
"type": "TRIGGER_PROBE",
"language": "java",
"where": {
"typeName": "MetadataClientUtils",
"methodName": "listTableWithContinuation"
},
"tags": [],
"evaluateAt": "EXIT",
"sessionId": "session",
"sampling": { "snapshotsPerSecond": 100 }
}
]
23 changes: 13 additions & 10 deletions ...tests/debugger-integration-tests/src/test/java/datadog/smoketest/BaseIntegrationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import com.squareup.moshi.Moshi;
import com.squareup.moshi.Types;
import datadog.trace.bootstrap.debugger.CapturedContext;
import datadog.trace.bootstrap.debugger.ProbeId;
import datadog.trace.bootstrap.debugger.ProbeRateLimiter;
import datadog.trace.test.agent.decoder.DecodedMessage;
import datadog.trace.test.agent.decoder.DecodedTrace;
Expand Down Expand Up @@ -332,23 +333,25 @@ protected void registerProbeStatusListener(Consumer<ProbeStatus> listener) {
probeStatusListeners.add(listener);
}

protected AtomicBoolean registerCheckReceivedInstalledEmitting() {
protected AtomicBoolean registerCheckReceivedInstalledEmitting(ProbeId probeId) {
AtomicBoolean received = new AtomicBoolean();
AtomicBoolean installed = new AtomicBoolean();
AtomicBoolean emitting = new AtomicBoolean();
AtomicBoolean result = new AtomicBoolean();
registerProbeStatusListener(
probeStatus -> {
if (probeStatus.getDiagnostics().getStatus() == ProbeStatus.Status.RECEIVED) {
received.set(true);
}
if (probeStatus.getDiagnostics().getStatus() == ProbeStatus.Status.INSTALLED) {
installed.set(true);
}
if (probeStatus.getDiagnostics().getStatus() == ProbeStatus.Status.EMITTING) {
emitting.set(true);
if (probeStatus.getDiagnostics().getProbeId().equals(probeId)) {
if (probeStatus.getDiagnostics().getStatus() == ProbeStatus.Status.RECEIVED) {
received.set(true);
}
if (probeStatus.getDiagnostics().getStatus() == ProbeStatus.Status.INSTALLED) {
installed.set(true);
}
if (probeStatus.getDiagnostics().getStatus() == ProbeStatus.Status.EMITTING) {
emitting.set(true);
}
result.set(received.get() && installed.get() && emitting.get());
}
result.set(received.get() && installed.get() && emitting.get());
});
return result;
}
Expand Down
14 changes: 7 additions & 7 deletions .../debugger-integration-tests/src/test/java/datadog/smoketest/LogProbesIntegrationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ void testFullMethod() throws Exception {
assertNull(snapshot.getCaptures().getReturn().getCapturedThrowable());
snapshotReceived.set(true);
});
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(PROBE_ID);
processRequests(
() -> snapshotReceived.get() && statusResult.get(),
() ->
Expand Down Expand Up @@ -132,7 +132,7 @@ void testFullMethodWithCondition() throws Exception {
assertFullMethodCaptureArgs(snapshot.getCaptures().getEntry());
snapshotReceived.set(true);
});
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(PROBE_ID);
processRequests(
() -> snapshotReceived.get() && statusResult.get(),
() ->
Expand Down Expand Up @@ -166,7 +166,7 @@ void testFullMethodWithConditionAtExit() throws Exception {
assertFullMethodCaptureArgs(snapshot.getCaptures().getReturn());
snapshotReceived.set(true);
});
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(PROBE_ID);
processRequests(
() -> snapshotReceived.get() && statusResult.get(),
() ->
Expand Down Expand Up @@ -199,7 +199,7 @@ void testFullMethodWithConditionFailed() throws Exception {
"Cannot find symbol: noarg", snapshot.getEvaluationErrors().get(0).getMessage());
snapshotReceived.set(true);
});
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(PROBE_ID);
processRequests(
() -> snapshotReceived.get() && statusResult.get(),
() ->
Expand Down Expand Up @@ -236,7 +236,7 @@ void testFullMethodWithLogTemplate() throws Exception {
assertEquals(PROBE_ID.getId(), snapshot.getProbe().getId());
snapshotReceived.set(true);
});
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(PROBE_ID);
processRequests(
() -> snapshotReceived.get() && correctLogMessage.get() && statusResult.get(),
() ->
Expand Down Expand Up @@ -288,7 +288,7 @@ void testFullMethodWithCaptureExpressions() throws Exception {
assertNull(key2Value.getNotCapturedReason());
snapshotReceived.set(true);
});
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(PROBE_ID);
processRequests(
() -> snapshotReceived.get() && statusResult.get(),
() ->
Expand Down Expand Up @@ -371,7 +371,7 @@ void testLineProbe() throws Exception {
assertNull(capturedContext.getCapturedThrowable());
snapshotReceived.set(true);
});
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(LINE_PROBE_ID1);
processRequests(
() -> snapshotReceived.get() && statusResult.get(),
() ->
Expand Down
4 changes: 2 additions & 2 deletions ...bugger-integration-tests/src/test/java/datadog/smoketest/MetricProbesIntegrationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,7 +112,7 @@ private void doMethodMetric(
targetProcess = createProcessBuilder(logFilePath, METHOD_NAME, EXPECTED_UPLOADS).start();
String msgExpected = String.format(expectedMsgFormat, metricName, PROBE_ID.getId());
assertNotNull(retrieveStatsdMessage(msgExpected));
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(PROBE_ID);
processRequests(
statusResult::get, () -> String.format("timeout statusResult=%s", statusResult.get()));
}
Expand Down Expand Up @@ -219,7 +219,7 @@ private void doLineMetric(
targetProcess = createProcessBuilder(logFilePath, METHOD_NAME, EXPECTED_UPLOADS).start();
String msgExpected = String.format(expectedMsgFormat, metricName, PROBE_ID.getId());
assertNotNull(retrieveStatsdMessage(msgExpected));
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting(PROBE_ID);
processRequests(
statusResult::get, () -> String.format("timeout statusResult=%s", statusResult.get()));
}
Expand Down
Loading
Loading