diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java b/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java index 75b6db833f7..e78fbcb797c 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/InputAttributeInjector.java +++ b/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) { diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnClientInstrumentation.java b/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnClientInstrumentation.java index 97ccde2b40e..5c88a55278b 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnClientInstrumentation.java +++ b/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 { diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java index e78d3b2d188..0b6b36b945d 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java +++ b/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,14 +50,14 @@ 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); @@ -69,8 +65,8 @@ private SdkRequest injectTraceContext(AgentSpan span, StartExecutionRequest requ 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); diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/TextMapInjectAdapter.java b/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/TextMapInjectAdapter.java new file mode 100644 index 00000000000..0c7311129fc --- /dev/null +++ b/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 { + + 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); + } +} diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy b/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy index d921d09e17a..18b2c36ee06 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy +++ b/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"() {