Inject trace context into AWS Step Functions input

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)
+}

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();
+  }
+}

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<ExecutionInterceptor> interceptors) {
+      for (ExecutionInterceptor interceptor : interceptors) {
+        if (interceptor instanceof SfnInterceptor) {
+          return;
+        }
+      }
+      interceptors.add(new SfnInterceptor());
+    }
+  }
+}

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<AgentSpan> 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();
+  }
+}

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"
+  }
+}

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'