Fix sfn x-datadog-tags input

dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java

@@ -1,19 +1,23 @@
 package datadog.trace.instrumentation.aws.v2.sfn;
 
-import datadog.json.JsonMapper;
+import static datadog.context.propagation.Propagators.defaultPropagator;
+import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
+
+import datadog.context.Context;
 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) {
+  public static String buildTraceContext(Context context) {
+    if (fromContext(context) == null) {
+      return null;
+    }
     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");
-      JsonMapper.writeAsJsonValue(writer, span.getTags());
+      // note: injection allows non-datadog style propogation (W3C, B3)
+      // which the extension does not yet extract
+      defaultPropagator().inject(context, writer, TextMapInjectAdapter.SETTER);
       writer.endObject();
       return writer.toString();
     } catch (Exception e) {

dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnClientInstrumentation.java

@@ -33,7 +33,11 @@ public void methodAdvice(MethodTransformer transformer) {
 
   @Override
   public String[] helperClassNames() {
-    return new String[] {packageName + ".SfnInterceptor", packageName + ".InputAttributeInjector"};
+    return new String[] {
+      packageName + ".SfnInterceptor",
+      packageName + ".InputAttributeInjector",
+      packageName + ".TextMapInjectAdapter"
+    };
   }
 
   public static class AwsSfnBuilderAdvice {

dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java

@@ -1,11 +1,8 @@
 package datadog.trace.instrumentation.aws.v2.sfn;
 
-import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
-
 import datadog.context.Context;
 import datadog.trace.api.Config;
 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.ModifyRequest;
 import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
@@ -38,14 +35,13 @@ public SdkRequest modifyRequest(ModifyRequest context, ExecutionAttributes execu
   public SdkRequest modifyRequestImpl(
       ModifyRequest context, ExecutionAttributes executionAttributes) {
     final Context ddContext = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE);
-    final AgentSpan span = fromContext(ddContext);
     // StartExecutionRequest
     if (context.request() instanceof StartExecutionRequest) {
       StartExecutionRequest request = (StartExecutionRequest) context.request();
       if (request.input() == null) {
         return request;
       }
-      return injectTraceContext(span, request);
+      return injectTraceContext(ddContext, request);
     }
 
     // StartSyncExecutionRequest
@@ -54,23 +50,23 @@ public SdkRequest modifyRequestImpl(
       if (request.input() == null) {
         return request;
       }
-      return injectTraceContext(span, request);
+      return injectTraceContext(ddContext, request);
     }
 
     return context.request();
   }
 
-  private SdkRequest injectTraceContext(AgentSpan span, StartExecutionRequest request) {
-    String ddTraceContextJSON = InputAttributeInjector.buildTraceContext(span);
+  private SdkRequest injectTraceContext(Context ddContext, StartExecutionRequest request) {
+    String ddTraceContextJSON = InputAttributeInjector.buildTraceContext(ddContext);
     // 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);
+  private SdkRequest injectTraceContext(Context ddContext, StartSyncExecutionRequest request) {
+    String ddTraceContextJSON = InputAttributeInjector.buildTraceContext(ddContext);
     // Inject the trace context into the StartSyncExecutionRequest input
     String modifiedInput =
         InputAttributeInjector.getModifiedInput(request.input(), ddTraceContextJSON);

dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/TextMapInjectAdapter.java

@@ -0,0 +1,14 @@
+package datadog.trace.instrumentation.aws.v2.sfn;
+
+import datadog.context.propagation.CarrierSetter;
+import datadog.json.JsonWriter;
+
+public class TextMapInjectAdapter implements CarrierSetter<JsonWriter> {
+
+  public static final TextMapInjectAdapter SETTER = new TextMapInjectAdapter();
+
+  @Override
+  public void set(final JsonWriter writer, final String key, final String value) {
+    writer.name(key).value(value);
+  }
+}

dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy

@@ -2,6 +2,7 @@ import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
 
 import datadog.trace.agent.test.naming.VersionedNamingTestBase
 import datadog.trace.agent.test.utils.TraceUtils
+import datadog.trace.api.DDSpanId
 import datadog.trace.api.DDSpanTypes
 import datadog.trace.bootstrap.instrumentation.api.Tags
 import groovy.json.JsonSlurper
@@ -102,6 +103,8 @@ abstract class SfnClientTest extends VersionedNamingTestBase {
           .build()
       }
     })
+    TEST_WRITER.waitForTraces(1)
+    def sfnSpan = TEST_WRITER.flatten().find { it.resourceName.toString() == "Sfn.StartExecution" }
 
     then:
     def execution = sfnClient.describeExecution { builder ->
@@ -110,9 +113,11 @@ abstract class SfnClientTest extends VersionedNamingTestBase {
     }
     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
+    input["_datadog"]["x-datadog-trace-id"] == sfnSpan.traceId.toString()
+    input["_datadog"]["x-datadog-parent-id"] == DDSpanId.toString(sfnSpan.spanId)
+    input["_datadog"]["x-datadog-sampling-priority"] != null
+    input["_datadog"]["x-datadog-tags"] instanceof String
+    input["_datadog"]["x-datadog-tags"].contains("_dd.p.")
   }
 
   def "datadog context is not injected when SfnInjectDatadogAttribute is disabled"() {