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..29d759d1026 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/build.gradle @@ -0,0 +1,30 @@ +muzzle { + pass { + group = "software.amazon.awssdk" + module = "sfn" + // 2.15.35 is the minimum version with step functions + versions = "[2.15.35,)" + 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.15.35' + + // 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.15.35' + testImplementation libs.testcontainers + + 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..0fba2d97bd7 --- /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,38 @@ +package datadog.trace.instrumentation.aws.v2.sfn; + +import datadog.json.JsonMapper; +import datadog.json.JsonWriter; +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 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 "{}"; + } + } + + public static String getModifiedInput(String request, String ddTraceContextJSON) { + StringBuilder modifiedInput = new StringBuilder(request.trim()); + int startPos = modifiedInput.indexOf("{"); + int endPos = modifiedInput.lastIndexOf("}"); + + String inputContent = modifiedInput.substring(startPos + 1, endPos).trim(); + if (inputContent.isEmpty()) { + modifiedInput.insert(endPos, String.format("\"%s\":%s", DATADOG_KEY, ddTraceContextJSON)); + } else { + // Prepend comma to separate from existing content + modifiedInput.insert(endPos, String.format(",\"%s\":%s", DATADOG_KEY, ddTraceContextJSON)); + } + return modifiedInput.toString(); + } +} 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..97ccde2b40e --- /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, Instrumenter.HasMethodAdvice { + + 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..e0acf28ca3c --- /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,72 @@ +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) { + 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) { + StartExecutionRequest request = (StartExecutionRequest) context.request(); + if (request.input() == null) { + return request; + } + return injectTraceContext(span, request); + } + + // StartSyncExecutionRequest + if (context.request() instanceof StartSyncExecutionRequest) { + StartSyncExecutionRequest request = (StartSyncExecutionRequest) context.request(); + if (request.input() == null) { + return request; + } + 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 StartExecutionRequest 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 StartSyncExecutionRequest input + String modifiedInput = + InputAttributeInjector.getModifiedInput(request.input(), ddTraceContextJSON); + + return request.toBuilder().input(modifiedInput).build(); + } +} 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..fd0d6a6966e --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy @@ -0,0 +1,161 @@ +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 { + @Shared GenericContainer localStack + @Shared SfnClient sfnClient + @Shared String testStateMachineARN + @Shared Object endPoint + + def setupSpec() { + 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) + .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() { + sfnClient.close() + 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() + } + }) + + then: + assertTraces(1) { + trace(2) { + basicSpan(it, "parent") + span { + serviceName service() + operationName operation() + 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 + } + + 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 { + @Override + int version() { + 0 + } + + @Override + String service() { + return "java-aws-sdk" + } + + @Override + String operation() { + return "aws.http" + } +} diff --git a/settings.gradle b/settings.gradle index f6082cc6b63..f9100b07213 100644 --- a/settings.gradle +++ b/settings.gradle @@ -196,6 +196,7 @@ include ':dd-java-agent:instrumentation:aws-common' include ':dd-java-agent:instrumentation:aws-java-eventbridge-2.0' include ':dd-java-agent:instrumentation:aws-java-sdk-1.11.0' include ':dd-java-agent:instrumentation:aws-java-sdk-2.2' +include ':dd-java-agent:instrumentation:aws-java-sfn-2.0' 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'