From 7f8721dc93f82fffcda11561a59a9f82df915787 Mon Sep 17 00:00:00 2001 From: Dylan Yang Date: Thu, 5 Sep 2024 16:41:18 -0400 Subject: [PATCH 01/16] inject step functions trace context --- .../aws/v2/AwsSdkClientDecorator.java | 36 +++++++++++++++++++ .../aws/v2/TracingExecutionInterceptor.java | 12 +++++++ 2 files changed, 48 insertions(+) diff --git a/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java b/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java index b8ed62ba99f..70dd161ab7b 100644 --- a/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java +++ b/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java @@ -6,6 +6,9 @@ import static datadog.trace.core.datastreams.TagsProcessor.TOPIC_TAG; import static datadog.trace.core.datastreams.TagsProcessor.TYPE_TAG; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; import datadog.trace.api.Config; import datadog.trace.api.DDTags; import datadog.trace.api.cache.DDCache; @@ -23,6 +26,8 @@ import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString; import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator; import datadog.trace.core.datastreams.TagsProcessor; +import java.lang.reflect.Field; +import java.lang.reflect.Method; import java.net.URI; import java.time.Instant; import java.util.Collections; @@ -398,6 +403,37 @@ public AgentSpan onSdkResponse( return span; } + public static void injectTraceToStepFunctionInput(SdkRequest request, AgentSpan span) { + Class clazz = request.getClass(); + if ((clazz.getSimpleName().equals("StartExecutionRequest")) + || clazz.getSimpleName().equals("StartSyncExecutionRequest")) { + try { + Method method = clazz.getMethod("input"); + String input = (String) method.invoke(request); + + ObjectMapper mapper = new ObjectMapper(); + JsonNode inputJsonNode = mapper.readTree(input); + if (inputJsonNode.isObject()) { + ObjectNode traceContext = mapper.createObjectNode(); + traceContext.put("x-datadog-trace-id", span.getTraceId().toString()); + traceContext.put("x-datadog-parent-id", String.valueOf(span.getSpanId())); + traceContext.put("x-datadog-sampling-priority", span.getSamplingPriority().toString()); + traceContext.put("x-datadog-tags", mapper.writeValueAsString(span.getTags())); + + ObjectNode ddInputNode = (ObjectNode) inputJsonNode; + ddInputNode.put("_datadog", traceContext); + } + + Field inputField = clazz.getDeclaredField("input"); + inputField.setAccessible(true); + inputField.set(request, mapper.writeValueAsString(inputJsonNode)); + + } catch (Throwable e) { + // Failed to inject trace context + } + } + } + @Override protected String[] instrumentationNames() { return new String[] {"aws-sdk"}; diff --git a/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java b/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java index 6234af389cc..fc586cb2a14 100644 --- a/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java @@ -63,6 +63,18 @@ public void afterMarshalling( } } + @Override + public SdkRequest modifyRequest( + Context.ModifyRequest context, ExecutionAttributes executionAttributes) { + SdkRequest sdkRequest = context.request(); + final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + if (span != null) { + DECORATE.injectTraceToStepFunctionInput(sdkRequest, span); + } + + return sdkRequest; + } + @Override public SdkHttpRequest modifyHttpRequest( Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) { From 7ecfa8cb94b5888a9a55e4e9ab11182bd22e58cc Mon Sep 17 00:00:00 2001 From: Dylan Yang Date: Fri, 13 Sep 2024 14:17:11 -0400 Subject: [PATCH 02/16] remove previous logic --- .../aws/v2/AwsSdkClientDecorator.java | 36 ------------------- .../aws/v2/TracingExecutionInterceptor.java | 12 ------- 2 files changed, 48 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java b/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java index 70dd161ab7b..b8ed62ba99f 100644 --- a/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java +++ b/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java @@ -6,9 +6,6 @@ import static datadog.trace.core.datastreams.TagsProcessor.TOPIC_TAG; import static datadog.trace.core.datastreams.TagsProcessor.TYPE_TAG; -import com.fasterxml.jackson.databind.JsonNode; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.fasterxml.jackson.databind.node.ObjectNode; import datadog.trace.api.Config; import datadog.trace.api.DDTags; import datadog.trace.api.cache.DDCache; @@ -26,8 +23,6 @@ import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString; import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator; import datadog.trace.core.datastreams.TagsProcessor; -import java.lang.reflect.Field; -import java.lang.reflect.Method; import java.net.URI; import java.time.Instant; import java.util.Collections; @@ -403,37 +398,6 @@ public AgentSpan onSdkResponse( return span; } - public static void injectTraceToStepFunctionInput(SdkRequest request, AgentSpan span) { - Class clazz = request.getClass(); - if ((clazz.getSimpleName().equals("StartExecutionRequest")) - || clazz.getSimpleName().equals("StartSyncExecutionRequest")) { - try { - Method method = clazz.getMethod("input"); - String input = (String) method.invoke(request); - - ObjectMapper mapper = new ObjectMapper(); - JsonNode inputJsonNode = mapper.readTree(input); - if (inputJsonNode.isObject()) { - ObjectNode traceContext = mapper.createObjectNode(); - traceContext.put("x-datadog-trace-id", span.getTraceId().toString()); - traceContext.put("x-datadog-parent-id", String.valueOf(span.getSpanId())); - traceContext.put("x-datadog-sampling-priority", span.getSamplingPriority().toString()); - traceContext.put("x-datadog-tags", mapper.writeValueAsString(span.getTags())); - - ObjectNode ddInputNode = (ObjectNode) inputJsonNode; - ddInputNode.put("_datadog", traceContext); - } - - Field inputField = clazz.getDeclaredField("input"); - inputField.setAccessible(true); - inputField.set(request, mapper.writeValueAsString(inputJsonNode)); - - } catch (Throwable e) { - // Failed to inject trace context - } - } - } - @Override protected String[] instrumentationNames() { return new String[] {"aws-sdk"}; diff --git a/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java b/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java index fc586cb2a14..6234af389cc 100644 --- a/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java @@ -63,18 +63,6 @@ public void afterMarshalling( } } - @Override - public SdkRequest modifyRequest( - Context.ModifyRequest context, ExecutionAttributes executionAttributes) { - SdkRequest sdkRequest = context.request(); - final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); - if (span != null) { - DECORATE.injectTraceToStepFunctionInput(sdkRequest, span); - } - - return sdkRequest; - } - @Override public SdkHttpRequest modifyHttpRequest( Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) { From fa36832ab004141daa0257c86127e3830f08c329 Mon Sep 17 00:00:00 2001 From: Dylan Yang Date: Mon, 16 Sep 2024 13:43:05 -0400 Subject: [PATCH 03/16] aws sfn instrumentation --- .../aws-java-sfn-2.0/build.gradle | 29 ++++ .../aws/v2/sfn/InputAttributeInjector.java | 44 ++++++ .../aws/v2/sfn/SfnClientInstrumentation.java | 50 +++++++ .../aws/v2/sfn/SfnInterceptor.java | 55 +++++++ .../src/test/groovy/SfnClientTest.groovy | 137 ++++++++++++++++++ settings.gradle | 1 + 6 files changed, 316 insertions(+) create mode 100644 dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle create mode 100644 dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java create mode 100644 dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnClientInstrumentation.java create mode 100644 dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java create mode 100644 dd-java-agent/instrumentation/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle new file mode 100644 index 00000000000..0205b243c17 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle @@ -0,0 +1,29 @@ +muzzle { + pass { + group = "software.amazon.awssdk" + module = "sfn" + versions = "[2.0.0,)" + assertInverse = true + } +} + +apply from: "$rootDir/gradle/java.gradle" + +addTestSuiteForDir('latestDepTest', 'test') +addTestSuiteExtendingForDir('latestDepForkedTest', 'latestDepTest', 'test') + +dependencies { + compileOnly group: 'software.amazon.awssdk', name: 'sfn', version: '2.27.2' + + // Include httpclient instrumentation for testing because it is a dependency for aws-sdk. + testImplementation project(':dd-java-agent:instrumentation:apache-httpclient-4') + testImplementation project(':dd-java-agent:instrumentation:aws-java-sdk-2.2') + testImplementation 'software.amazon.awssdk:sfn:2.27.2' + testImplementation 'org.testcontainers:localstack:1.19.7' + + latestDepTestImplementation group: 'software.amazon.awssdk', name: 'sfn', version: '+' +} + +tasks.withType(Test).configureEach { + usesService(testcontainersLimit) +} diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java new file mode 100644 index 00000000000..d9f30b28429 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java @@ -0,0 +1,44 @@ +package datadog.trace.instrumentation.aws.v2.sfn; + +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; + +public class InputAttributeInjector { + public static String buildTraceContext(AgentSpan span) { + // Extract span tags + StringBuilder spanTagsJSON = new StringBuilder(); + spanTagsJSON.append("{"); + span.getTags() + .forEach( + (tagKey, tagValue) -> + spanTagsJSON + .append("\"") + .append(tagKey) + .append("\":\"") + .append(tagValue) + .append("\",")); + spanTagsJSON.setLength(spanTagsJSON.length() - 1); // remove trailing comma + spanTagsJSON.append("}"); + + // Build DD trace context object + String ddTraceContextJSON = + String.format( + "\"_datadog\": { \"x-datadog-trace-id\": \"%s\",\"x-datadog-parent-id\":\"%s\", \"x-datadog-tags\": %s }", + span.getTraceId().toString(), span.getSpanId(), spanTagsJSON); + + return ddTraceContextJSON; + } + + public static StringBuilder getModifiedInput(String request, String ddTraceContextJSON) { + StringBuilder modifiedInput = new StringBuilder(request); + int startPos = modifiedInput.indexOf("{"); + int endPos = modifiedInput.lastIndexOf("}"); + String inputContent = modifiedInput.substring(startPos + 1, endPos); + if (inputContent.isEmpty()) { + modifiedInput.insert(endPos, ddTraceContextJSON); + } else { + modifiedInput.insert( + endPos, String.format(", %s", ddTraceContextJSON)); // prepend comma to existing input + } + return modifiedInput; + } +} diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnClientInstrumentation.java b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnClientInstrumentation.java new file mode 100644 index 00000000000..b89b8aa9dbe --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnClientInstrumentation.java @@ -0,0 +1,50 @@ +package datadog.trace.instrumentation.aws.v2.sfn; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.isMethod; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import java.util.List; +import net.bytebuddy.asm.Advice; +import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; + +/** AWS SDK v2 Step Function instrumentation */ +@AutoService(InstrumenterModule.class) +public final class SfnClientInstrumentation extends InstrumenterModule.Tracing + implements Instrumenter.ForSingleType { + + public SfnClientInstrumentation() { + super("sfn", "aws-sdk"); + } + + @Override + public String instrumentedType() { + return "software.amazon.awssdk.core.client.builder.SdkDefaultClientBuilder"; + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isMethod().and(named("resolveExecutionInterceptors")), + SfnClientInstrumentation.class.getName() + "$AwsSfnBuilderAdvice"); + } + + @Override + public String[] helperClassNames() { + return new String[] {packageName + ".SfnInterceptor", packageName + ".InputAttributeInjector"}; + } + + public static class AwsSfnBuilderAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + public static void addHandler(@Advice.Return final List interceptors) { + for (ExecutionInterceptor interceptor : interceptors) { + if (interceptor instanceof SfnInterceptor) { + return; + } + } + interceptors.add(new SfnInterceptor()); + } + } +} diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java new file mode 100644 index 00000000000..e9d05bbe835 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java @@ -0,0 +1,55 @@ +package datadog.trace.instrumentation.aws.v2.sfn; + +import datadog.trace.bootstrap.InstanceStore; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import software.amazon.awssdk.core.SdkRequest; +import software.amazon.awssdk.core.interceptor.Context; +import software.amazon.awssdk.core.interceptor.ExecutionAttribute; +import software.amazon.awssdk.core.interceptor.ExecutionAttributes; +import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; +import software.amazon.awssdk.services.sfn.model.StartExecutionRequest; +import software.amazon.awssdk.services.sfn.model.StartSyncExecutionRequest; + +public class SfnInterceptor implements ExecutionInterceptor { + + public static final ExecutionAttribute SPAN_ATTRIBUTE = + InstanceStore.of(ExecutionAttribute.class) + .putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan")); + + public SfnInterceptor() {} + + @Override + public SdkRequest modifyRequest( + Context.ModifyRequest context, ExecutionAttributes executionAttributes) { + final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + // StartExecutionRequest + if (context.request() instanceof StartExecutionRequest) { + StartExecutionRequest request = (StartExecutionRequest) context.request(); + if (request.input() == null) { + return request; + } + String ddTraceContextJSON = InputAttributeInjector.buildTraceContext(span); + // Inject the trace context into the Step Function input + StringBuilder modifiedInput = + InputAttributeInjector.getModifiedInput(request.input(), ddTraceContextJSON); + + return request.toBuilder().input(modifiedInput.toString()).build(); + } + + // StartSyncExecutionRequest + if (context.request() instanceof StartSyncExecutionRequest) { + StartSyncExecutionRequest request = (StartSyncExecutionRequest) context.request(); + if (request.input() == null) { + return request; + } + String ddTraceContextJSON = InputAttributeInjector.buildTraceContext(span); + // Inject the trace context into the Step Function input + StringBuilder modifiedInput = + InputAttributeInjector.getModifiedInput(request.input(), ddTraceContextJSON); + + return request.toBuilder().input(modifiedInput.toString()).build(); + } + + return context.request(); + } +} diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy new file mode 100644 index 00000000000..ad5605005ce --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy @@ -0,0 +1,137 @@ +import datadog.trace.agent.test.naming.VersionedNamingTestBase +import datadog.trace.agent.test.utils.TraceUtils +import datadog.trace.api.DDSpanTypes +import datadog.trace.bootstrap.instrumentation.api.Tags +import groovy.json.JsonSlurper +import org.testcontainers.containers.GenericContainer +import org.testcontainers.utility.DockerImageName +import software.amazon.awssdk.services.sfn.SfnClient +import software.amazon.awssdk.services.sfn.model.StartExecutionResponse +import software.amazon.awssdk.regions.Region +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials +import spock.lang.Shared + +import java.time.Duration + +import static datadog.trace.agent.test.utils.TraceUtils.basicSpan + + +abstract class SfnClientTest extends VersionedNamingTestBase { + static final LOCALSTACK = new GenericContainer(DockerImageName.parse("localstack/localstack")) + .withExposedPorts(4566) + .withEnv("SERVICES", "stepfunctions") + .withReuse(true) + .withStartupTimeout(Duration.ofSeconds(120)) + + @Shared SfnClient sfnClient + + @Shared String testStateMachineARN + + def setupSpec() { + LOCALSTACK.start() + def endPoint = "http://" + LOCALSTACK.getHost() + ":" + LOCALSTACK.getMappedPort(4566) + sfnClient = SfnClient.builder() + .endpointOverride(URI.create(endPoint)) + .region(Region.US_EAST_1) + .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test", "test"))) + .build() + + def response = sfnClient.createStateMachine { builder -> + builder.name("testStateMachine") + .definition("{\"StartAt\": \"HelloWorld\", \"States\": {\"HelloWorld\": {\"Type\": \"Pass\", \"End\": true}}}") + .build() + } + testStateMachineARN = response.stateMachineArn() + } + + def cleanupSpec() { + LOCALSTACK.stop() + } + + def "Step Functions span is created"() { + when: + StartExecutionResponse response + TraceUtils.runUnderTrace('parent', { + response = sfnClient.startExecution { builder -> + builder.stateMachineArn(testStateMachineARN) + .input("{\"key\": \"value\"}") + .build() + } + }) + + def endPoint = "http://" + LOCALSTACK.getHost() + ":" + LOCALSTACK.getMappedPort(4566) + + + then: + assertTraces(1) { + trace(2) { + basicSpan(it, "parent") + span { + serviceName "java-aws-sdk" + operationName "aws.http" + resourceName "Sfn.StartExecution" + spanType DDSpanTypes.HTTP_CLIENT + errored false + measured true + childOf(span(0)) + tags { + "$Tags.COMPONENT" "java-aws-sdk" + "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT + "$Tags.HTTP_URL" endPoint+'/' + "$Tags.HTTP_METHOD" "POST" + "$Tags.HTTP_STATUS" 200 + "$Tags.PEER_PORT" LOCALSTACK.getMappedPort(4566) + "$Tags.PEER_HOSTNAME" LOCALSTACK.getHost() + "aws.service" "Sfn" + "aws.operation" "StartExecution" + "aws.agent" "java-aws-sdk" + "aws.requestId" response.responseMetadata().requestId() + "aws_service" "Sfn" + defaultTags() + } + } + } + } + } + + def "Trace context is injected to Step Functions input"() { + when: + StartExecutionResponse response + TraceUtils.runUnderTrace('parent', { + response = sfnClient.startExecution { builder -> + builder.stateMachineArn(testStateMachineARN) + .input("{\"key\": \"value\"}") + .build() + } + }) + + then: + def execution = sfnClient.describeExecution { builder -> + builder.executionArn(response.executionArn()) + .build() + } + def input = new JsonSlurper().parseText(execution.input()) + input["key"] == "value" + input["_datadog"]["x-datadog-trace-id"] != null + input["_datadog"]["x-datadog-parent-id"] != null + input["_datadog"]["x-datadog-tags"] != null + } +} + +class SfnClientV0Test extends SfnClientTest { + @Override + int version() { + 0 + } + + @Override + String service() { + return null + } + + @Override + String operation() { + return null + } +} diff --git a/settings.gradle b/settings.gradle index 279384630b7..fe306785c31 100644 --- a/settings.gradle +++ b/settings.gradle @@ -184,6 +184,7 @@ include ':dd-java-agent:instrumentation:aws-java-sns-1.0' include ':dd-java-agent:instrumentation:aws-java-sns-2.0' include ':dd-java-agent:instrumentation:aws-java-sqs-1.0' include ':dd-java-agent:instrumentation:aws-java-sqs-2.0' +include ':dd-java-agent:instrumentation:aws-java-sfn-2.0' include ':dd-java-agent:instrumentation:aws-lambda-handler' include ':dd-java-agent:instrumentation:axis-2' include ':dd-java-agent:instrumentation:axway-api' From eb642b0d9705c3016a57d1a71aa1b6f00967c437 Mon Sep 17 00:00:00 2001 From: Dylan Yang Date: Mon, 16 Sep 2024 15:28:41 -0400 Subject: [PATCH 04/16] muzzle --- dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle index 0205b243c17..da35dbeacad 100644 --- a/dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle @@ -2,7 +2,7 @@ muzzle { pass { group = "software.amazon.awssdk" module = "sfn" - versions = "[2.0.0,)" + versions = "[2.15.35,)" assertInverse = true } } @@ -13,7 +13,7 @@ addTestSuiteForDir('latestDepTest', 'test') addTestSuiteExtendingForDir('latestDepForkedTest', 'latestDepTest', 'test') dependencies { - compileOnly group: 'software.amazon.awssdk', name: 'sfn', version: '2.27.2' + compileOnly group: 'software.amazon.awssdk', name: 'sfn', version: '2.15.35' // Include httpclient instrumentation for testing because it is a dependency for aws-sdk. testImplementation project(':dd-java-agent:instrumentation:apache-httpclient-4') From 0778f3d29b5d83caf1e2f3dc6a03b90027e2d72c Mon Sep 17 00:00:00 2001 From: Dylan Yang Date: Tue, 17 Sep 2024 10:26:56 -0400 Subject: [PATCH 05/16] single char --- .../instrumentation/aws/v2/sfn/InputAttributeInjector.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java index d9f30b28429..9209e5ca517 100644 --- a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java @@ -6,7 +6,7 @@ public class InputAttributeInjector { public static String buildTraceContext(AgentSpan span) { // Extract span tags StringBuilder spanTagsJSON = new StringBuilder(); - spanTagsJSON.append("{"); + spanTagsJSON.append('{'); span.getTags() .forEach( (tagKey, tagValue) -> @@ -17,7 +17,7 @@ public static String buildTraceContext(AgentSpan span) { .append(tagValue) .append("\",")); spanTagsJSON.setLength(spanTagsJSON.length() - 1); // remove trailing comma - spanTagsJSON.append("}"); + spanTagsJSON.append('}'); // Build DD trace context object String ddTraceContextJSON = From 421856b3fab61816beddcca4d088fd2151f33825 Mon Sep 17 00:00:00 2001 From: Dylan Yang Date: Tue, 8 Oct 2024 11:20:25 -0400 Subject: [PATCH 06/16] update muzzle --- dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle index da35dbeacad..90796d3ff32 100644 --- a/dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle @@ -18,8 +18,8 @@ dependencies { // Include httpclient instrumentation for testing because it is a dependency for aws-sdk. testImplementation project(':dd-java-agent:instrumentation:apache-httpclient-4') testImplementation project(':dd-java-agent:instrumentation:aws-java-sdk-2.2') - testImplementation 'software.amazon.awssdk:sfn:2.27.2' - testImplementation 'org.testcontainers:localstack:1.19.7' + testImplementation 'software.amazon.awssdk:sfn:2.15.35' + testImplementation libs.testcontainers latestDepTestImplementation group: 'software.amazon.awssdk', name: 'sfn', version: '+' } From 895341a14cb334bfbe434163f35664ecfc4971a2 Mon Sep 17 00:00:00 2001 From: Dylan Yang Date: Tue, 8 Oct 2024 11:22:32 -0400 Subject: [PATCH 07/16] refactor input attr injector --- .../aws/v2/sfn/InputAttributeInjector.java | 8 ++++---- .../trace/instrumentation/aws/v2/sfn/SfnInterceptor.java | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java index 9209e5ca517..6ec9b63ab57 100644 --- a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java @@ -11,7 +11,7 @@ public static String buildTraceContext(AgentSpan span) { .forEach( (tagKey, tagValue) -> spanTagsJSON - .append("\"") + .append('"') .append(tagKey) .append("\":\"") .append(tagValue) @@ -28,7 +28,7 @@ public static String buildTraceContext(AgentSpan span) { return ddTraceContextJSON; } - public static StringBuilder getModifiedInput(String request, String ddTraceContextJSON) { + public static String getModifiedInput(String request, String ddTraceContextJSON) { StringBuilder modifiedInput = new StringBuilder(request); int startPos = modifiedInput.indexOf("{"); int endPos = modifiedInput.lastIndexOf("}"); @@ -37,8 +37,8 @@ public static StringBuilder getModifiedInput(String request, String ddTraceConte modifiedInput.insert(endPos, ddTraceContextJSON); } else { modifiedInput.insert( - endPos, String.format(", %s", ddTraceContextJSON)); // prepend comma to existing input + endPos, ",".concat(ddTraceContextJSON)); // prepend comma to existing input } - return modifiedInput; + return modifiedInput.toString(); } } diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java index e9d05bbe835..d4e2fe71861 100644 --- a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java @@ -30,10 +30,10 @@ public SdkRequest modifyRequest( } String ddTraceContextJSON = InputAttributeInjector.buildTraceContext(span); // Inject the trace context into the Step Function input - StringBuilder modifiedInput = + String modifiedInput = InputAttributeInjector.getModifiedInput(request.input(), ddTraceContextJSON); - return request.toBuilder().input(modifiedInput.toString()).build(); + return request.toBuilder().input(modifiedInput).build(); } // StartSyncExecutionRequest @@ -44,10 +44,10 @@ public SdkRequest modifyRequest( } String ddTraceContextJSON = InputAttributeInjector.buildTraceContext(span); // Inject the trace context into the Step Function input - StringBuilder modifiedInput = + String modifiedInput = InputAttributeInjector.getModifiedInput(request.input(), ddTraceContextJSON); - return request.toBuilder().input(modifiedInput.toString()).build(); + return request.toBuilder().input(modifiedInput).build(); } return context.request(); From 9d76dec62a9898cb741df00a4d8a8a5e10f8d0c6 Mon Sep 17 00:00:00 2001 From: Dylan Yang Date: Wed, 9 Oct 2024 10:27:40 -0400 Subject: [PATCH 08/16] refactor interceptor --- .../aws/v2/sfn/SfnInterceptor.java | 32 ++++++++++++------- 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java index d4e2fe71861..4538da4deff 100644 --- a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java @@ -28,12 +28,7 @@ public SdkRequest modifyRequest( if (request.input() == null) { return request; } - String ddTraceContextJSON = InputAttributeInjector.buildTraceContext(span); - // Inject the trace context into the Step Function input - String modifiedInput = - InputAttributeInjector.getModifiedInput(request.input(), ddTraceContextJSON); - - return request.toBuilder().input(modifiedInput).build(); + return injectTraceContext(span, request); } // StartSyncExecutionRequest @@ -42,14 +37,27 @@ public SdkRequest modifyRequest( if (request.input() == null) { return request; } - String ddTraceContextJSON = InputAttributeInjector.buildTraceContext(span); - // Inject the trace context into the Step Function input - String modifiedInput = - InputAttributeInjector.getModifiedInput(request.input(), ddTraceContextJSON); - - return request.toBuilder().input(modifiedInput).build(); + return injectTraceContext(span, request); } return context.request(); } + + private SdkRequest injectTraceContext(AgentSpan span, StartExecutionRequest request) { + String ddTraceContextJSON = InputAttributeInjector.buildTraceContext(span); + // Inject the trace context into the Step Function input + String modifiedInput = + InputAttributeInjector.getModifiedInput(request.input(), ddTraceContextJSON); + + return request.toBuilder().input(modifiedInput).build(); + } + + private SdkRequest injectTraceContext(AgentSpan span, StartSyncExecutionRequest request) { + String ddTraceContextJSON = InputAttributeInjector.buildTraceContext(span); + // Inject the trace context into the Step Function input + String modifiedInput = + InputAttributeInjector.getModifiedInput(request.input(), ddTraceContextJSON); + + return request.toBuilder().input(modifiedInput).build(); + } } From aba111908a2d03def78eb7e556ce02f5042f2527 Mon Sep 17 00:00:00 2001 From: Dylan Yang Date: Wed, 9 Oct 2024 10:29:21 -0400 Subject: [PATCH 09/16] cleanup test --- .../src/test/groovy/SfnClientTest.groovy | 36 +++++++++---------- 1 file changed, 17 insertions(+), 19 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy index ad5605005ce..c70268988ee 100644 --- a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy @@ -18,19 +18,19 @@ import static datadog.trace.agent.test.utils.TraceUtils.basicSpan abstract class SfnClientTest extends VersionedNamingTestBase { - static final LOCALSTACK = new GenericContainer(DockerImageName.parse("localstack/localstack")) - .withExposedPorts(4566) - .withEnv("SERVICES", "stepfunctions") - .withReuse(true) - .withStartupTimeout(Duration.ofSeconds(120)) - + @Shared GenericContainer localStack @Shared SfnClient sfnClient - @Shared String testStateMachineARN + @Shared Object endPoint def setupSpec() { - LOCALSTACK.start() - def endPoint = "http://" + LOCALSTACK.getHost() + ":" + LOCALSTACK.getMappedPort(4566) + localStack = new GenericContainer(DockerImageName.parse("localstack/localstack")) + .withExposedPorts(4566) + .withEnv("SERVICES", "stepfunctions") + .withReuse(true) + .withStartupTimeout(Duration.ofSeconds(120)) + localStack.start() + endPoint = "http://" + localStack.getHost() + ":" + localStack.getMappedPort(4566) sfnClient = SfnClient.builder() .endpointOverride(URI.create(endPoint)) .region(Region.US_EAST_1) @@ -46,7 +46,8 @@ abstract class SfnClientTest extends VersionedNamingTestBase { } def cleanupSpec() { - LOCALSTACK.stop() + sfnClient.close() + localStack.stop() } def "Step Functions span is created"() { @@ -60,16 +61,13 @@ abstract class SfnClientTest extends VersionedNamingTestBase { } }) - def endPoint = "http://" + LOCALSTACK.getHost() + ":" + LOCALSTACK.getMappedPort(4566) - - then: assertTraces(1) { trace(2) { basicSpan(it, "parent") span { - serviceName "java-aws-sdk" - operationName "aws.http" + serviceName service() + operationName operation() resourceName "Sfn.StartExecution" spanType DDSpanTypes.HTTP_CLIENT errored false @@ -81,8 +79,8 @@ abstract class SfnClientTest extends VersionedNamingTestBase { "$Tags.HTTP_URL" endPoint+'/' "$Tags.HTTP_METHOD" "POST" "$Tags.HTTP_STATUS" 200 - "$Tags.PEER_PORT" LOCALSTACK.getMappedPort(4566) - "$Tags.PEER_HOSTNAME" LOCALSTACK.getHost() + "$Tags.PEER_PORT" localStack.getMappedPort(4566) + "$Tags.PEER_HOSTNAME" localStack.getHost() "aws.service" "Sfn" "aws.operation" "StartExecution" "aws.agent" "java-aws-sdk" @@ -127,11 +125,11 @@ class SfnClientV0Test extends SfnClientTest { @Override String service() { - return null + return "java-aws-sdk" } @Override String operation() { - return null + return "aws.http" } } From 370df6d8c800d31241ea104465d0015c5f71e59b Mon Sep 17 00:00:00 2001 From: Dylan Yang Date: Thu, 24 Oct 2024 11:14:17 -0400 Subject: [PATCH 10/16] use JsonBuffer --- .../aws/v2/sfn/InputAttributeInjector.java | 36 ++++++++++--------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java index 6ec9b63ab57..9df6549229f 100644 --- a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java @@ -1,31 +1,33 @@ package datadog.trace.instrumentation.aws.v2.sfn; +import datadog.trace.bootstrap.JsonBuffer; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; public class InputAttributeInjector { public static String buildTraceContext(AgentSpan span) { // Extract span tags - StringBuilder spanTagsJSON = new StringBuilder(); - spanTagsJSON.append('{'); + JsonBuffer spanTagsJSON = new JsonBuffer(); + spanTagsJSON.beginObject(); span.getTags() - .forEach( - (tagKey, tagValue) -> - spanTagsJSON - .append('"') - .append(tagKey) - .append("\":\"") - .append(tagValue) - .append("\",")); - spanTagsJSON.setLength(spanTagsJSON.length() - 1); // remove trailing comma - spanTagsJSON.append('}'); + .forEach((tagKey, tagValue) -> spanTagsJSON.name(tagKey).value(tagValue.toString())); + spanTagsJSON.endObject(); // Build DD trace context object - String ddTraceContextJSON = - String.format( - "\"_datadog\": { \"x-datadog-trace-id\": \"%s\",\"x-datadog-parent-id\":\"%s\", \"x-datadog-tags\": %s }", - span.getTraceId().toString(), span.getSpanId(), spanTagsJSON); + JsonBuffer ddTraceContextJSON = new JsonBuffer(); + ddTraceContextJSON + .beginObject() + .name("_datadog") + .beginObject() + .name("x-datadog-trace-id") + .value(span.getTraceId().toString()) + .name("x-datadog-parent-id") + .value(String.valueOf(span.getSpanId())) + .name("x-datadog-tags") + .value(spanTagsJSON) + .endObject() + .endObject(); - return ddTraceContextJSON; + return ddTraceContextJSON.toString(); } public static String getModifiedInput(String request, String ddTraceContextJSON) { From a644a8b26ed490679e4721f611800ef0c952daa1 Mon Sep 17 00:00:00 2001 From: Dylan Yang Date: Thu, 24 Oct 2024 11:15:45 -0400 Subject: [PATCH 11/16] catch exceptions --- .../instrumentation/aws/v2/sfn/SfnInterceptor.java | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java index 4538da4deff..e0acf28ca3c 100644 --- a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java @@ -21,6 +21,15 @@ public SfnInterceptor() {} @Override public SdkRequest modifyRequest( Context.ModifyRequest context, ExecutionAttributes executionAttributes) { + try { + return modifyRequestImpl(context, executionAttributes); + } catch (Exception e) { + return context.request(); + } + } + + public SdkRequest modifyRequestImpl( + Context.ModifyRequest context, ExecutionAttributes executionAttributes) { final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); // StartExecutionRequest if (context.request() instanceof StartExecutionRequest) { @@ -45,7 +54,7 @@ public SdkRequest modifyRequest( private SdkRequest injectTraceContext(AgentSpan span, StartExecutionRequest request) { String ddTraceContextJSON = InputAttributeInjector.buildTraceContext(span); - // Inject the trace context into the Step Function input + // Inject the trace context into the StartExecutionRequest input String modifiedInput = InputAttributeInjector.getModifiedInput(request.input(), ddTraceContextJSON); @@ -54,7 +63,7 @@ private SdkRequest injectTraceContext(AgentSpan span, StartExecutionRequest requ private SdkRequest injectTraceContext(AgentSpan span, StartSyncExecutionRequest request) { String ddTraceContextJSON = InputAttributeInjector.buildTraceContext(span); - // Inject the trace context into the Step Function input + // Inject the trace context into the StartSyncExecutionRequest input String modifiedInput = InputAttributeInjector.getModifiedInput(request.input(), ddTraceContextJSON); From 973cd79b40c0bf788837ca2478ca50c224c9ac5e Mon Sep 17 00:00:00 2001 From: Nicholas Hulston Date: Thu, 6 Feb 2025 15:47:34 -0500 Subject: [PATCH 12/16] add build.gradle comment --- dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle | 1 + 1 file changed, 1 insertion(+) diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle index 90796d3ff32..29d759d1026 100644 --- a/dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle @@ -2,6 +2,7 @@ muzzle { pass { group = "software.amazon.awssdk" module = "sfn" + // 2.15.35 is the minimum version with step functions versions = "[2.15.35,)" assertInverse = true } From 96f7f39e3db12c37c09f75cbb10828482146d130 Mon Sep 17 00:00:00 2001 From: Nicholas Hulston Date: Thu, 6 Feb 2025 15:48:50 -0500 Subject: [PATCH 13/16] fix SfnClientInstrumentation class header --- .../instrumentation/aws/v2/sfn/SfnClientInstrumentation.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnClientInstrumentation.java b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnClientInstrumentation.java index b89b8aa9dbe..97ccde2b40e 100644 --- a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnClientInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnClientInstrumentation.java @@ -13,7 +13,7 @@ /** AWS SDK v2 Step Function instrumentation */ @AutoService(InstrumenterModule.class) public final class SfnClientInstrumentation extends InstrumenterModule.Tracing - implements Instrumenter.ForSingleType { + implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { public SfnClientInstrumentation() { super("sfn", "aws-sdk"); From fdb01a2c35dd6879b67e0066857193cd3301c43f Mon Sep 17 00:00:00 2001 From: Nicholas Hulston Date: Thu, 6 Feb 2025 15:50:09 -0500 Subject: [PATCH 14/16] update InputAttributeInjector to use `datadog.json` component https://github.com/DataDog/dd-trace-java/pull/7973/ --- .../aws/v2/sfn/InputAttributeInjector.java | 47 +++++++++---------- 1 file changed, 23 insertions(+), 24 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java index 9df6549229f..20bbae9ae4e 100644 --- a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java @@ -1,33 +1,32 @@ package datadog.trace.instrumentation.aws.v2.sfn; -import datadog.trace.bootstrap.JsonBuffer; +import datadog.json.JsonMapper; +import datadog.json.JsonWriter; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; public class InputAttributeInjector { public static String buildTraceContext(AgentSpan span) { - // Extract span tags - JsonBuffer spanTagsJSON = new JsonBuffer(); - spanTagsJSON.beginObject(); - span.getTags() - .forEach((tagKey, tagValue) -> spanTagsJSON.name(tagKey).value(tagValue.toString())); - spanTagsJSON.endObject(); + String tagsJson = JsonMapper.toJson(span.getTags()); - // Build DD trace context object - JsonBuffer ddTraceContextJSON = new JsonBuffer(); - ddTraceContextJSON - .beginObject() - .name("_datadog") - .beginObject() - .name("x-datadog-trace-id") - .value(span.getTraceId().toString()) - .name("x-datadog-parent-id") - .value(String.valueOf(span.getSpanId())) - .name("x-datadog-tags") - .value(spanTagsJSON) - .endObject() - .endObject(); + try { + JsonWriter ddTraceContextJSON = new JsonWriter(); + ddTraceContextJSON + .beginObject() + .name("_datadog") + .beginObject() + .name("x-datadog-trace-id") + .value(span.getTraceId().toString()) + .name("x-datadog-parent-id") + .value(String.valueOf(span.getSpanId())) + .name("x-datadog-tags") + .jsonValue(tagsJson) + .endObject() + .endObject(); - return ddTraceContextJSON.toString(); + return ddTraceContextJSON.toString(); + } catch (Exception e) { + return "{}"; + } } public static String getModifiedInput(String request, String ddTraceContextJSON) { @@ -38,8 +37,8 @@ public static String getModifiedInput(String request, String ddTraceContextJSON) if (inputContent.isEmpty()) { modifiedInput.insert(endPos, ddTraceContextJSON); } else { - modifiedInput.insert( - endPos, ",".concat(ddTraceContextJSON)); // prepend comma to existing input + // Prepend comma to separate from existing content + modifiedInput.insert(endPos, "," + ddTraceContextJSON); } return modifiedInput.toString(); } From 8417f6e60b9b029a5bb94e475f070cd91ca466d5 Mon Sep 17 00:00:00 2001 From: Nicholas Hulston Date: Thu, 6 Feb 2025 16:18:39 -0500 Subject: [PATCH 15/16] fix injection --- .../aws/v2/sfn/InputAttributeInjector.java | 33 ++++++++----------- 1 file changed, 13 insertions(+), 20 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java index 20bbae9ae4e..3cefbbb72b9 100644 --- a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java @@ -5,25 +5,17 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan; public class InputAttributeInjector { + private static final String DATADOG_KEY = "_datadog"; + public static String buildTraceContext(AgentSpan span) { String tagsJson = JsonMapper.toJson(span.getTags()); - - try { - JsonWriter ddTraceContextJSON = new JsonWriter(); - ddTraceContextJSON - .beginObject() - .name("_datadog") - .beginObject() - .name("x-datadog-trace-id") - .value(span.getTraceId().toString()) - .name("x-datadog-parent-id") - .value(String.valueOf(span.getSpanId())) - .name("x-datadog-tags") - .jsonValue(tagsJson) - .endObject() - .endObject(); - - return ddTraceContextJSON.toString(); + try (JsonWriter writer = new JsonWriter()) { + writer.beginObject(); + writer.name("x-datadog-trace-id").value(span.getTraceId().toString()); + writer.name("x-datadog-parent-id").value(String.valueOf(span.getSpanId())); + writer.name("x-datadog-tags").jsonValue(tagsJson); + writer.endObject(); + return writer.toString(); } catch (Exception e) { return "{}"; } @@ -33,12 +25,13 @@ public static String getModifiedInput(String request, String ddTraceContextJSON) StringBuilder modifiedInput = new StringBuilder(request); int startPos = modifiedInput.indexOf("{"); int endPos = modifiedInput.lastIndexOf("}"); - String inputContent = modifiedInput.substring(startPos + 1, endPos); + + String inputContent = modifiedInput.substring(startPos + 1, endPos).trim(); if (inputContent.isEmpty()) { - modifiedInput.insert(endPos, ddTraceContextJSON); + modifiedInput.insert(endPos, String.format("\"%s\":%s", DATADOG_KEY, ddTraceContextJSON)); } else { // Prepend comma to separate from existing content - modifiedInput.insert(endPos, "," + ddTraceContextJSON); + modifiedInput.insert(endPos, String.format(",\"%s\":%s", DATADOG_KEY, ddTraceContextJSON)); } return modifiedInput.toString(); } From 0aac23904342f686880a3ef75faf68a3b2b3db5d Mon Sep 17 00:00:00 2001 From: Nicholas Hulston Date: Thu, 6 Feb 2025 16:32:11 -0500 Subject: [PATCH 16/16] test edge cases --- .../aws/v2/sfn/InputAttributeInjector.java | 2 +- .../src/test/groovy/SfnClientTest.groovy | 26 +++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java index 3cefbbb72b9..0fba2d97bd7 100644 --- a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java @@ -22,7 +22,7 @@ public static String buildTraceContext(AgentSpan span) { } public static String getModifiedInput(String request, String ddTraceContextJSON) { - StringBuilder modifiedInput = new StringBuilder(request); + StringBuilder modifiedInput = new StringBuilder(request.trim()); int startPos = modifiedInput.indexOf("{"); int endPos = modifiedInput.lastIndexOf("}"); diff --git a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy index c70268988ee..fd0d6a6966e 100644 --- a/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy @@ -115,6 +115,32 @@ abstract class SfnClientTest extends VersionedNamingTestBase { input["_datadog"]["x-datadog-parent-id"] != null input["_datadog"]["x-datadog-tags"] != null } + + def "Doesn't cause error for Step Functions input edge cases"() { + def inputs = [ + '''{}''', + '''{ }''', + ''' { } ''', + '''{"foo": "bar"}''', + ''' { "foo" : "bar" } ''', + '''{"key1": "val1", "key2": "val2"}''', + ''' { "key1" : "val1" , "key2" : "val2" } ''' + ] + + when: + inputs.forEach { input -> + TraceUtils.runUnderTrace('parent', { + sfnClient.startExecution { builder -> + builder.stateMachineArn(testStateMachineARN) + .input(input) + .build() + } + }) + } + + then: + noExceptionThrown() + } } class SfnClientV0Test extends SfnClientTest {