From 017e5a450bc5f3e2b2500cdccd3cad5fa820f0c9 Mon Sep 17 00:00:00 2001 From: Bruce Bujon Date: Wed, 22 Jan 2025 15:12:20 +0100 Subject: [PATCH] feat(core): Migrate XRax inject calls to new API --- .../aws/v0/AwsSdkClientDecorator.java | 1 + .../aws/v0/TracingRequestHandler.java | 5 +-- .../aws/v2/AwsSdkClientDecorator.java | 2 + .../aws/v2/TracingExecutionInterceptor.java | 5 +-- .../core/propagation/TracingPropagator.java | 20 +++++---- .../core/propagation/XRayPropagator.java | 45 +++++++++++++++++++ 6 files changed, 64 insertions(+), 14 deletions(-) create mode 100644 dd-trace-core/src/main/java/datadog/trace/core/propagation/XRayPropagator.java diff --git a/dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AwsSdkClientDecorator.java b/dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AwsSdkClientDecorator.java index 6ac29370dd6f..eba5b1a9cbb5 100644 --- a/dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AwsSdkClientDecorator.java +++ b/dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AwsSdkClientDecorator.java @@ -332,6 +332,7 @@ protected int status(final Response response) { return response.getHttpResponse().getStatusCode(); } + @ParametersAreNonnullByDefault @Override public void set(Request carrier, String key, String value) { carrier.addHeader(key, value); diff --git a/dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java b/dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java index 4773c2090bd2..2ee003eb10f0 100644 --- a/dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java +++ b/dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java @@ -2,7 +2,6 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.blackholeSpan; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_IN; import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_TAG; @@ -17,12 +16,12 @@ import com.amazonaws.handlers.HandlerContextKey; import com.amazonaws.handlers.RequestHandler2; import datadog.trace.api.Config; -import datadog.trace.api.TracePropagationStyle; import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.instrumentation.api.AgentDataStreamsMonitoring; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.PathwayContext; +import datadog.trace.core.propagation.XRayPropagator; import java.util.Date; import java.util.LinkedHashMap; import java.util.List; @@ -69,7 +68,7 @@ public void beforeRequest(final Request request) { request.addHandlerContext(SPAN_CONTEXT_KEY, span); if (Config.get().isAwsPropagationEnabled()) { try { - propagate().inject(span, request, DECORATE, TracePropagationStyle.XRAY); + XRayPropagator.INSTANCE.inject(span, request, DECORATE); } catch (Throwable e) { log.warn("Unable to inject trace header", e); } diff --git a/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java b/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java index 94227b681243..ebd52e3e35f1 100644 --- a/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java +++ b/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java @@ -37,6 +37,7 @@ import java.util.Optional; import java.util.Set; import javax.annotation.Nonnull; +import javax.annotation.ParametersAreNonnullByDefault; import software.amazon.awssdk.awscore.AwsResponse; import software.amazon.awssdk.core.SdkBytes; import software.amazon.awssdk.core.SdkField; @@ -443,6 +444,7 @@ protected int status(final SdkHttpResponse response) { return response.statusCode(); } + @ParametersAreNonnullByDefault @Override public void set(SdkHttpRequest.Builder carrier, String key, String value) { carrier.putHeader(key, value); diff --git a/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java b/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java index 6234af389cc9..58a16685a497 100644 --- a/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java @@ -2,17 +2,16 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.blackholeSpan; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.instrumentation.aws.v2.AwsSdkClientDecorator.AWS_LEGACY_TRACING; import static datadog.trace.instrumentation.aws.v2.AwsSdkClientDecorator.DECORATE; import datadog.trace.api.Config; -import datadog.trace.api.TracePropagationStyle; import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.InstanceStore; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.core.propagation.XRayPropagator; import java.util.Optional; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -71,7 +70,7 @@ public SdkHttpRequest modifyHttpRequest( final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); if (span != null) { SdkHttpRequest.Builder requestBuilder = context.httpRequest().toBuilder(); - propagate().inject(span, requestBuilder, DECORATE, TracePropagationStyle.XRAY); + XRayPropagator.INSTANCE.inject(span, requestBuilder, DECORATE); return requestBuilder.build(); } } catch (Throwable e) { diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/TracingPropagator.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/TracingPropagator.java index 632e3b6dbede..0b26d98ada5f 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/TracingPropagator.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/TracingPropagator.java @@ -1,5 +1,6 @@ package datadog.trace.core.propagation; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromSpanContext; import datadog.context.Context; @@ -21,6 +22,8 @@ public class TracingPropagator implements Propagator { private final HttpCodec.Injector injector; private final HttpCodec.Extractor extractor; + // Create HttpCodecPropagator abstract class + // TODO Javadoc public TracingPropagator(HttpCodec.Injector injector, HttpCodec.Extractor extractor) { this.injector = injector; @@ -34,17 +37,18 @@ public static Concern concern() { @Override public void inject(Context context, C carrier, CarrierSetter setter) { + AgentSpan span; //noinspection ConstantValue - if (context == null || carrier == null || setter == null) { + if (context == null + || carrier == null + || setter == null + || (span = fromContext(context)) == null) { return; } - AgentSpan span = AgentSpan.fromContext(context); - if (span != null) { - AgentSpanContext spanContext = span.context(); - if (spanContext instanceof DDSpanContext) { - DDSpanContext ddSpanContext = (DDSpanContext) context; - injector.inject(ddSpanContext, carrier, setter::set); - } + AgentSpanContext spanContext = span.context(); + if (spanContext instanceof DDSpanContext) { + DDSpanContext ddSpanContext = (DDSpanContext) context; + this.injector.inject(ddSpanContext, carrier, setter::set); } } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/XRayPropagator.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/XRayPropagator.java new file mode 100644 index 000000000000..cae5eaf353b2 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/XRayPropagator.java @@ -0,0 +1,45 @@ +package datadog.trace.core.propagation; + +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; + +import datadog.context.Context; +import datadog.context.propagation.CarrierSetter; +import datadog.context.propagation.CarrierVisitor; +import datadog.context.propagation.Propagator; +import datadog.trace.api.Config; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; +import datadog.trace.core.DDSpanContext; +import javax.annotation.ParametersAreNonnullByDefault; + +@ParametersAreNonnullByDefault +public class XRayPropagator implements Propagator { + public static final Propagator INSTANCE = new XRayPropagator(); + private final HttpCodec.Injector injector; + + public XRayPropagator() { + this.injector = XRayHttpCodec.newInjector(Config.get().getBaggageMapping()); + } + + @Override + public void inject(Context context, C carrier, CarrierSetter setter) { + AgentSpan span; + //noinspection ConstantValue + if (context == null + || carrier == null + || setter == null + || (span = fromContext(context)) == null) { + return; + } + AgentSpanContext spanContext = span.context(); + if (spanContext instanceof DDSpanContext) { + DDSpanContext ddSpanContext = (DDSpanContext) context; + this.injector.inject(ddSpanContext, carrier, setter::set); + } + } + + @Override + public Context extract(Context context, C carrier, CarrierVisitor visitor) { + return context; + } +}