Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -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) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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
Expand All @@ -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);
Expand Down
Original file line number Diff line number Diff line change
@@ -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);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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"() {
Expand Down
Loading