From 1ca5b9c988feba093101e9c43790fe49e03b8dc4 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Fri, 18 Sep 2026 17:18:06 +0200 Subject: [PATCH] Reduce SQS body propagation parsing overhead --- .../aws-java/aws-java-sqs-1.0/build.gradle | 2 + .../aws/v1/sqs/MessageExtractAdapter.java | 44 +++++- .../aws/v1/sqs/MessageExtractAdapterTest.java | 138 ++++++++++++++++++ 3 files changed, 179 insertions(+), 5 deletions(-) create mode 100644 dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/test/java/datadog/trace/instrumentation/aws/v1/sqs/MessageExtractAdapterTest.java diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/build.gradle b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/build.gradle index ab37cfcac66..067cc222cf2 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/build.gradle @@ -16,6 +16,8 @@ addTestSuiteForDir('latestDepTest', 'test') addTestSuiteExtendingForDir('latestDepForkedTest', 'latestDepTest', 'test') dependencies { + testImplementation libs.bundles.junit5 + compileOnly group: 'com.amazonaws', name: 'aws-java-sdk-sqs', version: '1.11.0' compileOnly group: 'com.amazonaws', name: 'amazon-sqs-java-messaging-lib', version: '1.0.0' diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/MessageExtractAdapter.java b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/MessageExtractAdapter.java index 0d97bdbf0f9..da37db8de49 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/MessageExtractAdapter.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/MessageExtractAdapter.java @@ -4,6 +4,8 @@ import com.amazonaws.services.sqs.model.Message; import com.amazonaws.services.sqs.model.MessageAttributeValue; +import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.core.JsonToken; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; import datadog.trace.api.Config; @@ -66,11 +68,25 @@ public void forEachKey(Message carrier, AgentPropagation.KeyClassifier classifie public void forEachKeyInBody(String body, AgentPropagation.KeyClassifier classifier) throws IOException { - // Parse the JSON string into a JsonNode - JsonNode rootNode = MAPPER.readTree(body); - - // Navigate to MessageAttributes._datadog - JsonNode messageAttributes = rootNode.path("MessageAttributes").path("_datadog"); + JsonNode messageAttributes = null; + try (JsonParser parser = MAPPER.getFactory().createParser(body)) { + if (parser.nextToken() == JsonToken.START_OBJECT) { + while (parser.nextToken() == JsonToken.FIELD_NAME) { + String name = parser.getCurrentName(); + parser.nextToken(); + if ("MessageAttributes".equals(name)) { + messageAttributes = readDatadogAttribute(parser); + } else { + parser.skipChildren(); + } + } + } else { + MAPPER.readTree(parser); + } + } + if (messageAttributes == null) { + return; + } // Extract Value and Type String value = messageAttributes.path("Value").asText(); @@ -83,6 +99,24 @@ public void forEachKeyInBody(String body, AgentPropagation.KeyClassifier classif } } + private static JsonNode readDatadogAttribute(JsonParser parser) throws IOException { + JsonNode attribute = null; + if (parser.getCurrentToken() == JsonToken.START_OBJECT) { + while (parser.nextToken() == JsonToken.FIELD_NAME) { + String name = parser.getCurrentName(); + parser.nextToken(); + if ("_datadog".equals(name)) { + attribute = MAPPER.readTree(parser); + } else { + parser.skipChildren(); + } + } + } else { + parser.skipChildren(); + } + return attribute; + } + public long extractTimeInQueueStart(final Message carrier) { try { Map systemAttributes = carrier.getAttributes(); diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/test/java/datadog/trace/instrumentation/aws/v1/sqs/MessageExtractAdapterTest.java b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/test/java/datadog/trace/instrumentation/aws/v1/sqs/MessageExtractAdapterTest.java new file mode 100644 index 00000000000..b7a0395fe75 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/test/java/datadog/trace/instrumentation/aws/v1/sqs/MessageExtractAdapterTest.java @@ -0,0 +1,138 @@ +package datadog.trace.instrumentation.aws.v1.sqs; + +import static java.nio.charset.StandardCharsets.UTF_8; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; +import datadog.trace.bootstrap.instrumentation.messaging.DatadogAttributeParser; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Base64; +import java.util.List; +import org.junit.jupiter.api.Test; + +class MessageExtractAdapterTest { + private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final String CONTEXT = + "{\"x-datadog-trace-id\":\"123\",\"x-datadog-parent-id\":\"456\"}"; + + private static String attribute(String type, String value) throws IOException { + return "{\"Type\":" + + MAPPER.writeValueAsString(type) + + ",\"Value\":" + + MAPPER.writeValueAsString(value) + + "}"; + } + + private static List extract(String body, boolean baseline) throws IOException { + List result = new ArrayList<>(); + AgentPropagation.KeyClassifier classifier = + (key, value) -> { + result.add(key + "=" + value); + return true; + }; + if (baseline) { + JsonNode node = MAPPER.readTree(body).path("MessageAttributes").path("_datadog"); + String value = node.path("Value").asText(); + String type = node.path("Type").asText(); + if ("String".equals(type)) { + DatadogAttributeParser.forEachProperty(classifier, value); + } else if ("Binary".equals(type)) { + DatadogAttributeParser.forEachProperty( + classifier, ByteBuffer.wrap(Base64.getDecoder().decode(value))); + } + } else { + MessageExtractAdapter.GETTER.forEachKeyInBody(body, classifier); + } + return result; + } + + @Test + void extractsStringAndBinaryWithEitherFieldOrder() throws IOException { + for (String type : new String[] {"String", "Binary"}) { + String value = + "Binary".equals(type) + ? Base64.getEncoder().encodeToString(CONTEXT.getBytes(UTF_8)) + : CONTEXT; + String attr = attribute(type, value); + String reversed = + "{\"Value\":" + MAPPER.writeValueAsString(value) + ",\"Type\":\"" + type + "\"}"; + for (String a : new String[] {attr, reversed}) { + String body = + "{\"Message\":\"ignored\",\"MessageAttributes\":{\"other\":{\"nested\":[1,2]},\"_datadog\":" + + a + + "},\"tail\":[{},true]}"; + assertEquals(extract(body, true), extract(body, false)); + assertEquals(2, extract(body, false).size()); + } + } + } + + @Test + void preservesLastDuplicateAtEveryLevel() throws IOException { + String attr = attribute("String", CONTEXT); + String attrs = "{\"_datadog\":" + attr + "}"; + String[] bodies = { + "{\"MessageAttributes\":" + attrs + ",\"MessageAttributes\":{}}", + "{\"MessageAttributes\":{},\"MessageAttributes\":" + attrs + "}", + "{\"MessageAttributes\":" + attrs + ",\"MessageAttributes\":null}", + "{\"MessageAttributes\":{\"_datadog\":" + attr + ",\"_datadog\":null}}", + "{\"MessageAttributes\":{\"_datadog\":null,\"_datadog\":" + attr + "}}", + "{\"MessageAttributes\":{\"_datadog\":{\"Type\":\"Binary\",\"Type\":\"String\",\"Value\":false,\"Value\":" + + MAPPER.writeValueAsString(CONTEXT) + + "}}}", + "{\"MessageAttributes\":{" + "\"_datadog\":" + attr + ",\"other\":1}}" + }; + for (String body : bodies) assertEquals(extract(body, true), extract(body, false), body); + } + + @Test + void ignoresUnrelatedFieldsAndPreservesNodeConversions() throws IOException { + String attr = attribute("String", CONTEXT); + for (String value : new String[] {"null", "true", "123", "1.25", "[]", "{}", "\"\""}) { + String[] bodies = { + "{\"MessageAttributes\":" + value + "}", + "{\"MessageAttributes\":{\"_datadog\":" + value + "}}", + "{\"MessageAttributes\":{\"_datadog\":{\"Type\":\"String\",\"Value\":" + value + "}}}", + "{\"nested\":{\"MessageAttributes\":{\"_datadog\":" + attr + "}},\"other\":" + value + "}" + }; + for (String body : bodies) assertEquals(extract(body, true), extract(body, false), body); + } + for (String body : new String[] {"{}", "[]", "null", "42", "\"plain text\""}) { + assertEquals(extract(body, true), extract(body, false), body); + } + } + + @Test + void validatesTheEnvelopeBeforePublishingContext() throws IOException { + String prefix = "{\"MessageAttributes\":{\"_datadog\":" + attribute("String", CONTEXT) + "}"; + for (String suffix : + new String[] {"", ",\"tail\":[1,]}", ",\"tail\":\"\\uZZZZ\"}", ",\"tail\":{\"x\":}}"}) { + String body = prefix + suffix; + List observed = new ArrayList<>(); + assertThrows(IOException.class, () -> extract(body, true)); + assertThrows( + IOException.class, + () -> + MessageExtractAdapter.GETTER.forEachKeyInBody( + body, + (key, value) -> { + observed.add(key); + return true; + })); + assertTrue(observed.isEmpty()); + } + } + + @Test + void preservesIgnoringTrailingRootValues() throws IOException { + String body = + "{\"MessageAttributes\":{\"_datadog\":" + attribute("String", CONTEXT) + "}} trailing"; + assertEquals(extract(body, true), extract(body, false)); + } +}