Migrate instrumentations to Context API

dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java

@@ -1,13 +1,16 @@
 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;
 import java.util.Map;
 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<AgentSpan> SPAN_ATTRIBUTE =
+  public static final ExecutionAttribute<Context> 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;

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<AgentSpan> SPAN_ATTRIBUTE =
+  public static final ExecutionAttribute<Context> 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);

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<AgentSpan> SPAN_ATTRIBUTE =
+  public static final ExecutionAttribute<Context> 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;

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()

dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java

@@ -1,13 +1,15 @@
 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;
 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.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<AgentSpan> SPAN_ATTRIBUTE =
+  public static final ExecutionAttribute<Context> 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<Object, String> 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<SdkResponse> 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());

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<AgentSpan> SPAN_ATTRIBUTE =
+  public static final ExecutionAttribute<Context> 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();