From 0620e0c180903af50b642f1b75a8e8d7003b1f4d Mon Sep 17 00:00:00 2001 From: Olivier John Ndjike Nzia Date: Wed, 26 Aug 2026 17:51:57 -0400 Subject: [PATCH] fix sfn x-datadog-tags input --- .../aws/v2/sfn/InputAttributeInjector.java | 16 +++++++++------- .../aws/v2/sfn/SfnClientInstrumentation.java | 6 +++++- .../aws/v2/sfn/SfnInterceptor.java | 16 ++++++---------- .../aws/v2/sfn/TextMapInjectAdapter.java | 14 ++++++++++++++ .../src/test/groovy/SfnClientTest.groovy | 4 +++- 5 files changed, 37 insertions(+), 19 deletions(-) create mode 100644 dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/TextMapInjectAdapter.java 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..ec27e4217fb 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,21 @@ 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()); + 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 5ab99ecae96..53e1ece95bb 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..37037dedaba 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 @@ -112,7 +112,9 @@ abstract class SfnClientTest extends VersionedNamingTestBase { 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-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"() {