From 6d746a31f99ce79d4841616e351e0c0544281064 Mon Sep 17 00:00:00 2001 From: Bruce Bujon Date: Tue, 19 Aug 2025 11:19:09 +0200 Subject: [PATCH] feat(aws): Refactor AWS instrumentation around Context --- .../aws/v2/dynamodb/DynamoDbInterceptor.java | 15 +++-- .../eventbridge/EventBridgeInterceptor.java | 14 ++--- .../aws/v2/s3/S3Interceptor.java | 13 ++-- .../aws/v2/AwsSdkClientDecorator.java | 14 +++-- .../aws/v2/TracingExecutionInterceptor.java | 62 ++++++++++++------- .../aws/v2/sfn/SfnInterceptor.java | 17 ++--- .../aws/v2/sns/SnsInterceptor.java | 14 ++--- .../aws/v2/sqs/SqsInterceptor.java | 22 +++---- 8 files changed, 95 insertions(+), 76 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java index 86cf378bc45..3468915c7f1 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java @@ -1,5 +1,8 @@ package datadog.trace.instrumentation.aws.v2.dynamodb; +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; @@ -7,7 +10,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.awssdk.core.SdkRequest; -import software.amazon.awssdk.core.interceptor.Context; +import software.amazon.awssdk.core.interceptor.Context.AfterExecution; import software.amazon.awssdk.core.interceptor.ExecutionAttribute; import software.amazon.awssdk.core.interceptor.ExecutionAttributes; import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; @@ -18,20 +21,20 @@ public class DynamoDbInterceptor implements ExecutionInterceptor { private static final Logger log = LoggerFactory.getLogger(DynamoDbInterceptor.class); - public static final ExecutionAttribute SPAN_ATTRIBUTE = + public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan")); + .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); private static final boolean CAN_ADD_SPAN_POINTERS = Config.get().isAddSpanPointers("aws"); @Override - public void afterExecution( - Context.AfterExecution context, ExecutionAttributes executionAttributes) { + public void afterExecution(AfterExecution context, ExecutionAttributes executionAttributes) { if (!CAN_ADD_SPAN_POINTERS) { return; } - AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + Context ddContext = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE); + AgentSpan span = fromContext(ddContext); if (span == null) { log.debug("Unable to find DynamoDb request span. Not creating span pointer."); return; diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java index 8b01b573464..7e7cd81b053 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java @@ -5,17 +5,17 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig; import static datadog.trace.instrumentation.aws.v2.eventbridge.TextMapInjectAdapter.SETTER; +import datadog.context.Context; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.api.datastreams.DataStreamsTags; import datadog.trace.api.datastreams.PathwayContext; import datadog.trace.bootstrap.InstanceStore; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.util.ArrayList; import java.util.List; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.awssdk.core.SdkRequest; -import software.amazon.awssdk.core.interceptor.Context; +import software.amazon.awssdk.core.interceptor.Context.ModifyRequest; import software.amazon.awssdk.core.interceptor.ExecutionAttribute; import software.amazon.awssdk.core.interceptor.ExecutionAttributes; import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; @@ -25,16 +25,15 @@ public class EventBridgeInterceptor implements ExecutionInterceptor { private static final Logger log = LoggerFactory.getLogger(EventBridgeInterceptor.class); - public static final ExecutionAttribute SPAN_ATTRIBUTE = + public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan")); + .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); private static final String START_TIME_KEY = "x-datadog-start-time"; private static final String RESOURCE_NAME_KEY = "x-datadog-resource-name"; @Override - public SdkRequest modifyRequest( - Context.ModifyRequest context, ExecutionAttributes executionAttributes) { + public SdkRequest modifyRequest(ModifyRequest context, ExecutionAttributes executionAttributes) { if (!(context.request() instanceof PutEventsRequest)) { return context.request(); } @@ -79,12 +78,11 @@ public SdkRequest modifyRequest( private String getTraceContextToInject( ExecutionAttributes executionAttributes, String eventBusName, long startTime) { - final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + Context context = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE); StringBuilder jsonBuilder = new StringBuilder(); jsonBuilder.append('{'); // Inject context - datadog.context.Context context = span; if (traceConfig().isDataStreamsEnabled()) { DataStreamsTags tags = DataStreamsTags.createWithBus(OUTBOUND, eventBusName); DataStreamsContext dsmContext = DataStreamsContext.fromTags(tags); diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3Interceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3Interceptor.java index bdaf1ba6559..e03b6bb2427 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3Interceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3Interceptor.java @@ -2,12 +2,13 @@ import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.S3_ETAG; +import datadog.context.Context; import datadog.trace.api.Config; import datadog.trace.bootstrap.InstanceStore; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import software.amazon.awssdk.core.interceptor.Context; +import software.amazon.awssdk.core.interceptor.Context.AfterExecution; import software.amazon.awssdk.core.interceptor.ExecutionAttribute; import software.amazon.awssdk.core.interceptor.ExecutionAttributes; import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; @@ -18,20 +19,20 @@ public class S3Interceptor implements ExecutionInterceptor { private static final Logger log = LoggerFactory.getLogger(S3Interceptor.class); - public static final ExecutionAttribute SPAN_ATTRIBUTE = + public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan")); + .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); private static final boolean CAN_ADD_SPAN_POINTERS = Config.get().isAddSpanPointers("aws"); @Override - public void afterExecution( - Context.AfterExecution context, ExecutionAttributes executionAttributes) { + public void afterExecution(AfterExecution context, ExecutionAttributes executionAttributes) { if (!CAN_ADD_SPAN_POINTERS) { return; } - AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + Context ddContext = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE); + AgentSpan span = AgentSpan.fromContext(ddContext); if (span == null) { log.debug("Unable to find S3 request span. Not creating span pointer."); return; diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java b/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java index abb0a85cd0e..a5aef08372f 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java @@ -5,8 +5,10 @@ import static datadog.trace.api.datastreams.DataStreamsTags.Direction.OUTBOUND; import static datadog.trace.api.datastreams.DataStreamsTags.create; import static datadog.trace.api.datastreams.DataStreamsTags.createWithDataset; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig; +import datadog.context.Context; import datadog.context.propagation.CarrierSetter; import datadog.trace.api.Config; import datadog.trace.api.ConfigDefaults; @@ -113,11 +115,12 @@ public CharSequence spanName(final ExecutionAttributes attributes) { "aws", attributes.getAttribute(SdkExecutionAttribute.SERVICE_NAME), s)); } - public AgentSpan onSdkRequest( - final AgentSpan span, + public Context onSdkRequest( + final Context context, final SdkRequest request, final SdkHttpRequest httpRequest, final ExecutionAttributes attributes) { + final AgentSpan span = fromContext(context); final String awsServiceName = attributes.getAttribute(SdkExecutionAttribute.SERVICE_NAME); final String awsOperationName = attributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME); onOperation(span, awsServiceName, awsOperationName); @@ -224,7 +227,7 @@ public AgentSpan onSdkRequest( span.setTag(DDTags.PEER_SERVICE_SOURCE, "peer.service"); } - return span; + return context; } private static AgentSpan onOperation( @@ -315,12 +318,13 @@ private static void setTableName(AgentSpan span, String name) { setPeerService(span, InstrumentationTags.AWS_TABLE_NAME, name); } - public AgentSpan onSdkResponse( - final AgentSpan span, + public Context onSdkResponse( + final Context context, final SdkResponse response, final SdkHttpResponse httpResponse, final ExecutionAttributes attributes) { + final AgentSpan span = fromContext(context); Config config = Config.get(); String serviceName = attributes.getAttribute(SdkExecutionAttribute.SERVICE_NAME); if (config.isCloudResponsePayloadTaggingEnabled() diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java index a739a7bb5de..57bfceab407 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.aws.v2; import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpanWithoutScope; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.blackholeSpan; @@ -8,6 +9,7 @@ import static datadog.trace.instrumentation.aws.v2.AwsSdkClientDecorator.AWS_LEGACY_TRACING; import static datadog.trace.instrumentation.aws.v2.AwsSdkClientDecorator.DECORATE; +import datadog.context.Context; import datadog.context.propagation.Propagators; import datadog.trace.api.Config; import datadog.trace.bootstrap.ContextStore; @@ -19,7 +21,12 @@ import org.slf4j.LoggerFactory; import software.amazon.awssdk.core.SdkRequest; import software.amazon.awssdk.core.SdkResponse; -import software.amazon.awssdk.core.interceptor.Context; +import software.amazon.awssdk.core.interceptor.Context.AfterExecution; +import software.amazon.awssdk.core.interceptor.Context.AfterMarshalling; +import software.amazon.awssdk.core.interceptor.Context.BeforeExecution; +import software.amazon.awssdk.core.interceptor.Context.BeforeTransmission; +import software.amazon.awssdk.core.interceptor.Context.FailedExecution; +import software.amazon.awssdk.core.interceptor.Context.ModifyHttpRequest; import software.amazon.awssdk.core.interceptor.ExecutionAttribute; import software.amazon.awssdk.core.interceptor.ExecutionAttributes; import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; @@ -28,9 +35,9 @@ /** AWS request execution interceptor */ public class TracingExecutionInterceptor implements ExecutionInterceptor { - public static final ExecutionAttribute SPAN_ATTRIBUTE = + public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan")); + .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); private static final Logger log = LoggerFactory.getLogger(TracingExecutionInterceptor.class); @@ -42,37 +49,40 @@ public TracingExecutionInterceptor(ContextStore responseQueueSto @Override public void beforeExecution( - final Context.BeforeExecution context, final ExecutionAttributes executionAttributes) { + final BeforeExecution context, final ExecutionAttributes executionAttributes) { if (!AWS_LEGACY_TRACING && isPollingRequest(context.request())) { return; // SQS messages spans are created by aws-java-sqs-2.0 } - final AgentSpan span = startSpan(DECORATE.spanName(executionAttributes)); + final AgentSpan span = startSpan("aws-sdk", DECORATE.spanName(executionAttributes)); + // TODO If DSM is enabled, add DSM context here too DECORATE.afterStart(span); - executionAttributes.putAttribute(SPAN_ATTRIBUTE, span); + executionAttributes.putAttribute(CONTEXT_ATTRIBUTE, span); } @Override public void afterMarshalling( - final Context.AfterMarshalling context, final ExecutionAttributes executionAttributes) { - final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); - if (span != null) { + final AfterMarshalling context, final ExecutionAttributes executionAttributes) { + final Context ddContext = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE); + final AgentSpan span = fromContext(ddContext); + if (context != null && span != null) { try (AgentScope ignored = activateSpan(span)) { DECORATE.onRequest(span, context.httpRequest()); - DECORATE.onSdkRequest(span, context.request(), context.httpRequest(), executionAttributes); + DECORATE.onSdkRequest( + ddContext, context.request(), context.httpRequest(), executionAttributes); } } } @Override public SdkHttpRequest modifyHttpRequest( - Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) { + ModifyHttpRequest context, ExecutionAttributes executionAttributes) { if (Config.get().isAwsPropagationEnabled()) { try { - final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); - if (span != null) { + final Context ddContext = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE); + if (ddContext != null) { SdkHttpRequest.Builder requestBuilder = context.httpRequest().toBuilder(); - Propagators.forConcern(XRAY_TRACING_CONCERN).inject(span, requestBuilder, DECORATE); + Propagators.forConcern(XRAY_TRACING_CONCERN).inject(ddContext, requestBuilder, DECORATE); return requestBuilder.build(); } } catch (Throwable e) { @@ -84,12 +94,13 @@ public SdkHttpRequest modifyHttpRequest( @Override public void beforeTransmission( - final Context.BeforeTransmission context, final ExecutionAttributes executionAttributes) { + final BeforeTransmission context, final ExecutionAttributes executionAttributes) { final AgentSpan span; if (!AWS_LEGACY_TRACING) { span = blackholeSpan(); } else { - span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + final Context ddContext = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE); + span = fromContext(ddContext); } if (span != null) { // This scope will be closed by AwsHttpClientInstrumentation since ExecutionInterceptor API @@ -100,10 +111,11 @@ public void beforeTransmission( @Override public void afterExecution( - final Context.AfterExecution context, final ExecutionAttributes executionAttributes) { - final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + final AfterExecution context, final ExecutionAttributes executionAttributes) { + final Context ddContext = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE); + final AgentSpan span = fromContext(ddContext); if (span != null) { - executionAttributes.putAttribute(SPAN_ATTRIBUTE, null); + executionAttributes.putAttribute(CONTEXT_ATTRIBUTE, null); // Call onResponse on both types of responses: DECORATE.onSdkResponse(span, context.response(), context.httpResponse(), executionAttributes); DECORATE.onResponse(span, context.httpResponse()); @@ -121,14 +133,16 @@ public void afterExecution( @Override public void onExecutionFailure( - final Context.FailedExecution context, final ExecutionAttributes executionAttributes) { - final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); - if (span != null) { - executionAttributes.putAttribute(SPAN_ATTRIBUTE, null); + final FailedExecution context, final ExecutionAttributes executionAttributes) { + final Context ddContext = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE); + final AgentSpan span = fromContext(ddContext); + if (ddContext != null && span != null) { + executionAttributes.putAttribute(CONTEXT_ATTRIBUTE, null); Optional responseOpt = context.response(); if (responseOpt.isPresent()) { SdkResponse response = responseOpt.get(); - DECORATE.onSdkResponse(span, response, response.sdkHttpResponse(), executionAttributes); + DECORATE.onSdkResponse( + ddContext, response, response.sdkHttpResponse(), executionAttributes); DECORATE.onResponse(span, response.sdkHttpResponse()); if (span.isError()) { DECORATE.onError(span, context.exception()); 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 e0acf28ca3c..abc6b72e15b 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,9 +1,12 @@ package datadog.trace.instrumentation.aws.v2.sfn; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; + +import datadog.context.Context; 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; +import software.amazon.awssdk.core.interceptor.Context.ModifyRequest; import software.amazon.awssdk.core.interceptor.ExecutionAttribute; import software.amazon.awssdk.core.interceptor.ExecutionAttributes; import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; @@ -12,15 +15,14 @@ public class SfnInterceptor implements ExecutionInterceptor { - public static final ExecutionAttribute SPAN_ATTRIBUTE = + public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan")); + .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); public SfnInterceptor() {} @Override - public SdkRequest modifyRequest( - Context.ModifyRequest context, ExecutionAttributes executionAttributes) { + public SdkRequest modifyRequest(ModifyRequest context, ExecutionAttributes executionAttributes) { try { return modifyRequestImpl(context, executionAttributes); } catch (Exception e) { @@ -29,8 +31,9 @@ public SdkRequest modifyRequest( } public SdkRequest modifyRequestImpl( - Context.ModifyRequest context, ExecutionAttributes executionAttributes) { - final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + 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(); diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java index a81be0759f8..4090a46c0c6 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java @@ -6,17 +6,17 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig; import static datadog.trace.instrumentation.aws.v2.sns.TextMapInjectAdapter.SETTER; +import datadog.context.Context; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.api.datastreams.DataStreamsTags; import datadog.trace.bootstrap.InstanceStore; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.HashMap; import java.util.Map; import software.amazon.awssdk.core.SdkBytes; import software.amazon.awssdk.core.SdkRequest; -import software.amazon.awssdk.core.interceptor.Context; +import software.amazon.awssdk.core.interceptor.Context.ModifyRequest; import software.amazon.awssdk.core.interceptor.ExecutionAttribute; import software.amazon.awssdk.core.interceptor.ExecutionAttributes; import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; @@ -27,16 +27,15 @@ public class SnsInterceptor implements ExecutionInterceptor { - public static final ExecutionAttribute SPAN_ATTRIBUTE = + public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan")); + .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); private SdkBytes getMessageAttributeValueToInject( ExecutionAttributes executionAttributes, String snsTopicName) { - final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + Context context = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE); StringBuilder jsonBuilder = new StringBuilder(); jsonBuilder.append('{'); - datadog.context.Context context = span; if (traceConfig().isDataStreamsEnabled()) { DataStreamsContext dsmContext = DataStreamsContext.fromTags(getTags(snsTopicName)); context = context.with(dsmContext); @@ -50,8 +49,7 @@ private SdkBytes getMessageAttributeValueToInject( public SnsInterceptor() {} @Override - public SdkRequest modifyRequest( - Context.ModifyRequest context, ExecutionAttributes executionAttributes) { + public SdkRequest modifyRequest(ModifyRequest context, ExecutionAttributes executionAttributes) { // Injecting the trace context into SNS messageAttributes. if (context.request() instanceof PublishRequest) { PublishRequest request = (PublishRequest) context.request(); diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java index cc09d8686c5..c78e041fa7e 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java @@ -7,19 +7,19 @@ import static datadog.trace.bootstrap.instrumentation.api.URIUtils.urlFileName; import static datadog.trace.instrumentation.aws.v2.sqs.MessageAttributeInjector.SETTER; +import datadog.context.Context; import datadog.context.propagation.Propagator; import datadog.context.propagation.Propagators; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.api.datastreams.DataStreamsTags; import datadog.trace.bootstrap.InstanceStore; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; import software.amazon.awssdk.core.SdkRequest; -import software.amazon.awssdk.core.interceptor.Context; +import software.amazon.awssdk.core.interceptor.Context.ModifyRequest; import software.amazon.awssdk.core.interceptor.ExecutionAttribute; import software.amazon.awssdk.core.interceptor.ExecutionAttributes; import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; @@ -31,15 +31,14 @@ public class SqsInterceptor implements ExecutionInterceptor { - public static final ExecutionAttribute SPAN_ATTRIBUTE = + public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan")); + .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); public SqsInterceptor() {} @Override - public SdkRequest modifyRequest( - Context.ModifyRequest context, ExecutionAttributes executionAttributes) { + public SdkRequest modifyRequest(ModifyRequest context, ExecutionAttributes executionAttributes) { if (context.request() instanceof SendMessageRequest) { SendMessageRequest request = (SendMessageRequest) context.request(); Optional optionalQueueUrl = request.getValueForField("QueueUrl", String.class); @@ -48,7 +47,7 @@ public SdkRequest modifyRequest( } Propagator dsmPropagator = Propagators.forConcern(DSM_CONCERN); - datadog.context.Context ctx = getContext(executionAttributes, optionalQueueUrl.get()); + Context ctx = getContext(executionAttributes, optionalQueueUrl.get()); Map messageAttributes = new HashMap<>(request.messageAttributes()); dsmPropagator.inject(ctx, messageAttributes, SETTER); @@ -63,7 +62,7 @@ public SdkRequest modifyRequest( } Propagator dsmPropagator = Propagators.forConcern(DSM_CONCERN); - datadog.context.Context ctx = getContext(executionAttributes, optionalQueueUrl.get()); + Context ctx = getContext(executionAttributes, optionalQueueUrl.get()); List entries = new ArrayList<>(); for (SendMessageBatchRequestEntry entry : request.entries()) { @@ -90,12 +89,11 @@ public SdkRequest modifyRequest( } } - private datadog.context.Context getContext( - ExecutionAttributes executionAttributes, String queueUrl) { - AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + private Context getContext(ExecutionAttributes executionAttributes, String queueUrl) { + Context context = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE); DataStreamsTags tags = create("sqs", OUTBOUND, urlFileName(queueUrl)); DataStreamsContext dsmContext = DataStreamsContext.fromTags(tags); - return span.with(dsmContext); + return context.with(dsmContext); } }