Migrate instrumentations to Context API

dd-java-agent/instrumentation/armeria/armeria-grpc-0.84/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/GrpcClientDecorator.java

@@ -1,6 +1,7 @@
 package datadog.trace.instrumentation.armeria.grpc.client;
 
 import static datadog.context.propagation.Propagators.defaultPropagator;
+import static datadog.trace.api.datastreams.DataStreamsTags.Direction.OUTBOUND;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
 
@@ -32,8 +33,7 @@ public class GrpcClientDecorator extends ClientDecorator {
   public static final CharSequence GRPC_MESSAGE = UTF8BytesString.create("grpc.message");
 
   private static DataStreamsContext createDsmContext() {
-    return DataStreamsContext.fromTags(
-        DataStreamsTags.create("grpc", DataStreamsTags.Direction.Outbound));
+    return DataStreamsContext.fromTags(DataStreamsTags.create("grpc", OUTBOUND));
   }
 
   public static final GrpcClientDecorator DECORATE = new GrpcClientDecorator();

dd-java-agent/instrumentation/armeria/armeria-grpc-0.84/src/main/java/datadog/trace/instrumentation/armeria/grpc/server/GrpcServerDecorator.java

@@ -1,5 +1,7 @@
 package datadog.trace.instrumentation.armeria.grpc.server;
 
+import static datadog.trace.api.datastreams.DataStreamsTags.Direction.INBOUND;
+
 import datadog.trace.api.Config;
 import datadog.trace.api.cache.DDCache;
 import datadog.trace.api.cache.DDCaches;
@@ -30,7 +32,7 @@ public class GrpcServerDecorator extends ServerDecorator {
   public static final CharSequence GRPC_MESSAGE = UTF8BytesString.create("grpc.message");
 
   private static DataStreamsTags createServerPathwaySortedTags() {
-    return DataStreamsTags.create("grpc", DataStreamsTags.Direction.Inbound);
+    return DataStreamsTags.create("grpc", INBOUND);
   }
 
   public static final DataStreamsTags SERVER_PATHWAY_EDGE_TAGS = createServerPathwaySortedTags();

dd-java-agent/instrumentation/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java

@@ -1,6 +1,7 @@
 package datadog.trace.instrumentation.aws.v2.eventbridge;
 
 import static datadog.context.propagation.Propagators.defaultPropagator;
+import static datadog.trace.api.datastreams.DataStreamsTags.Direction.OUTBOUND;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
 import static datadog.trace.instrumentation.aws.v2.eventbridge.TextMapInjectAdapter.SETTER;
 
@@ -85,8 +86,7 @@ private String getTraceContextToInject(
     // Inject context
     datadog.context.Context context = span;
     if (traceConfig().isDataStreamsEnabled()) {
-      DataStreamsTags tags =
-          DataStreamsTags.createWithBus(DataStreamsTags.Direction.Outbound, eventBusName);
+      DataStreamsTags tags = DataStreamsTags.createWithBus(OUTBOUND, eventBusName);
       DataStreamsContext dsmContext = DataStreamsContext.fromTags(tags);
       context = context.with(dsmContext);
     }

dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AWSHttpClientInstrumentation.java

@@ -3,13 +3,15 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closeActive;
+import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
+import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.CONTEXT_CONTEXT_KEY;
 import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.DECORATE;
-import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.SPAN_CONTEXT_KEY;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 
 import com.amazonaws.AmazonClientException;
 import com.amazonaws.Request;
 import com.amazonaws.handlers.RequestHandler2;
+import datadog.context.Context;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import net.bytebuddy.asm.Advice;
@@ -54,12 +56,15 @@ public static void methodExit(
       }
 
       if (throwable != null && request != null) {
-        final AgentSpan span = request.getHandlerContext(SPAN_CONTEXT_KEY);
-        if (span != null) {
-          request.addHandlerContext(SPAN_CONTEXT_KEY, null);
-          DECORATE.onError(span, throwable);
-          DECORATE.beforeFinish(span);
-          span.finish();
+        final Context context = request.getHandlerContext(CONTEXT_CONTEXT_KEY);
+        if (context != null) {
+          request.addHandlerContext(CONTEXT_CONTEXT_KEY, null);
+          final AgentSpan span = spanFromContext(context);
+          if (span != null) {
+            DECORATE.onError(span, throwable);
+            DECORATE.beforeFinish(span);
+            span.finish();
+          }
         }
       }
     }

dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AwsSdkClientDecorator.java

@@ -1,6 +1,8 @@
 package datadog.trace.instrumentation.aws.v0;
 
 import static datadog.trace.api.datastreams.DataStreamsContext.create;
+import static datadog.trace.api.datastreams.DataStreamsTags.Direction.INBOUND;
+import static datadog.trace.api.datastreams.DataStreamsTags.Direction.OUTBOUND;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
 import static datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities.RPC_COMMAND_NAME;
 
@@ -264,8 +266,7 @@ && traceConfig().isDataStreamsEnabled()) {
             && ("GetObjectMetadataRequest".equalsIgnoreCase(awsOperation)
                 || "GetObjectRequest".equalsIgnoreCase(awsOperation))) {
           DataStreamsTags tags =
-              DataStreamsTags.createWithDataset(
-                  "s3", DataStreamsTags.Direction.Inbound, bucket, key, bucket);
+              DataStreamsTags.createWithDataset("s3", INBOUND, bucket, key, bucket);
           AgentTracer.get()
               .getDataStreamsMonitoring()
               .setCheckpoint(span, create(tags, 0, responseSize));
@@ -279,8 +280,7 @@ && traceConfig().isDataStreamsEnabled()) {
             payloadSize = (long) requestSize;
           }
           DataStreamsTags tags =
-              DataStreamsTags.createWithDataset(
-                  "s3", DataStreamsTags.Direction.Outbound, bucket, key, bucket);
+              DataStreamsTags.createWithDataset("s3", OUTBOUND, bucket, key, bucket);
           AgentTracer.get()
               .getDataStreamsMonitoring()
               .setCheckpoint(span, create(tags, 0, payloadSize));

dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AwsSdkModule.java

@@ -37,9 +37,7 @@ public String[] helperClassNames() {
   public Map<String, String> contextStore() {
     Map<String, String> map = new java.util.HashMap<>();
     map.put(namespace + ".services.sqs.model.ReceiveMessageResult", "java.lang.String");
-    map.put(
-        namespace + ".AmazonWebServiceRequest",
-        "datadog.trace.bootstrap.instrumentation.api.AgentSpan");
+    map.put(namespace + ".AmazonWebServiceRequest", "datadog.context.Context");
     return map;
   }
 

dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/HandlerChainFactoryInstrumentation.java

@@ -46,8 +46,7 @@ public static void addHandler(@Advice.Return final List<RequestHandler2> handler
               InstrumentationContext.get(
                   "com.amazonaws.services.sqs.model.ReceiveMessageResult", "java.lang.String"),
               InstrumentationContext.get(
-                  "com.amazonaws.AmazonWebServiceRequest",
-                  "datadog.trace.bootstrap.instrumentation.api.AgentSpan")));
+                  "com.amazonaws.AmazonWebServiceRequest", "datadog.context.Context")));
     }
   }
 }

dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/OnErrorDecorator.java

@@ -1,14 +1,15 @@
 package datadog.trace.instrumentation.aws.v0;
 
 import com.amazonaws.handlers.HandlerContextKey;
-import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.context.Context;
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
 import datadog.trace.bootstrap.instrumentation.decorator.BaseDecorator;
 
 public class OnErrorDecorator extends BaseDecorator {
 
-  public static final HandlerContextKey<AgentSpan> SPAN_CONTEXT_KEY =
-      new HandlerContextKey<>("DatadogSpan"); // same as TracingRequestHandler.SPAN_CONTEXT_KEY
+  public static final HandlerContextKey<Context> CONTEXT_CONTEXT_KEY =
+      new HandlerContextKey<>(
+          "DatadogContext"); // same as TracingRequestHandler.CONTEXT_CONTEXT_KEY
 
   public static final OnErrorDecorator DECORATE = new OnErrorDecorator();
 

dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/RequestExecutorInstrumentation.java

@@ -3,11 +3,13 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closeActive;
+import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
+import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.CONTEXT_CONTEXT_KEY;
 import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.DECORATE;
-import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.SPAN_CONTEXT_KEY;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 
 import com.amazonaws.Request;
+import datadog.context.Context;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import net.bytebuddy.asm.Advice;
@@ -51,12 +53,15 @@ public static void methodExit(
       }
 
       if (throwable != null) {
-        final AgentSpan span = request.getHandlerContext(SPAN_CONTEXT_KEY);
-        if (span != null) {
-          request.addHandlerContext(SPAN_CONTEXT_KEY, null);
-          DECORATE.onError(span, throwable);
-          DECORATE.beforeFinish(span);
-          span.finish();
+        final Context context = request.getHandlerContext(CONTEXT_CONTEXT_KEY);
+        if (context != null) {
+          request.addHandlerContext(CONTEXT_CONTEXT_KEY, null);
+          final AgentSpan span = spanFromContext(context);
+          if (span != null) {
+            DECORATE.onError(span, throwable);
+            DECORATE.beforeFinish(span);
+            span.finish();
+          }
         }
       }
     }

dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java

@@ -1,11 +1,14 @@
 package datadog.trace.instrumentation.aws.v0;
 
 import static datadog.trace.api.datastreams.DataStreamsContext.create;
+import static datadog.trace.api.datastreams.DataStreamsTags.Direction.INBOUND;
+import static datadog.trace.api.datastreams.DataStreamsTags.create;
 import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN;
 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.bootstrap.instrumentation.api.AgentTracer.traceConfig;
+import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
 import static datadog.trace.instrumentation.aws.v0.AwsSdkClientDecorator.AWS_LEGACY_TRACING;
 import static datadog.trace.instrumentation.aws.v0.AwsSdkClientDecorator.DECORATE;
 
@@ -14,9 +17,13 @@
 import com.amazonaws.Response;
 import com.amazonaws.handlers.HandlerContextKey;
 import com.amazonaws.handlers.RequestHandler2;
+import datadog.context.Context;
 import datadog.context.propagation.Propagators;
 import datadog.trace.api.Config;
-import datadog.trace.api.datastreams.*;
+import datadog.trace.api.datastreams.AgentDataStreamsMonitoring;
+import datadog.trace.api.datastreams.DataStreamsContext;
+import datadog.trace.api.datastreams.DataStreamsTags;
+import datadog.trace.api.datastreams.PathwayContext;
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -28,41 +35,42 @@
 /** Tracing Request Handler */
 public class TracingRequestHandler extends RequestHandler2 {
 
-  public static final HandlerContextKey<AgentSpan> SPAN_CONTEXT_KEY =
-      new HandlerContextKey<>("DatadogSpan"); // same as OnErrorDecorator.SPAN_CONTEXT_KEY
+  public static final HandlerContextKey<Context> CONTEXT_CONTEXT_KEY =
+      new HandlerContextKey<>("DatadogContext"); // same as OnErrorDecorator.CONTEXT_CONTEXT_KEY
 
   private static final Logger log = LoggerFactory.getLogger(TracingRequestHandler.class);
 
   private final ContextStore<Object, String> responseQueueStore;
-  private final ContextStore<AmazonWebServiceRequest, AgentSpan> requestSpanStore;
+  private final ContextStore<AmazonWebServiceRequest, Context> requestContextStore;
 
   public TracingRequestHandler(
       ContextStore<Object, String> responseQueueStore,
-      ContextStore<AmazonWebServiceRequest, AgentSpan> requestSpanStore) {
+      ContextStore<AmazonWebServiceRequest, Context> requestContextStore) {
     this.responseQueueStore = responseQueueStore;
-    this.requestSpanStore = requestSpanStore;
+    this.requestContextStore = requestContextStore;
   }
 
   @Override
   public void beforeRequest(final Request<?> request) {
-    AgentSpan span;
     if (!AWS_LEGACY_TRACING && isPollingRequest(request.getOriginalRequest())) {
       // SQS messages spans are created by aws-java-sqs-1.0 - replace client scope with no-op,
       // so we can tell when receive call is complete without affecting the rest of the trace
       activateSpanWithoutScope(blackholeSpan());
     } else {
-      span = requestSpanStore.remove(request.getOriginalRequest());
+      Context context = requestContextStore.remove(request.getOriginalRequest());
+      AgentSpan span = spanFromContext(context);
       if (span != null) {
         // we'll land here for SQS send requests when DSM is enabled. In that case, we create the
         // span in SqsInterceptor to inject DSM tags.
         span.setOperationName(AwsNameCache.spanName(request));
       } else {
         // this is the most common code path
-        span = startSpan(AwsNameCache.spanName(request));
+        span = startSpan("aws-sdk", AwsNameCache.spanName(request));
+        context = span; // TODO If DSM is enabled, add DSM context here too
       }
       DECORATE.afterStart(span);
       DECORATE.onRequest(span, request);
-      request.addHandlerContext(SPAN_CONTEXT_KEY, span);
+      request.addHandlerContext(CONTEXT_CONTEXT_KEY, context);
       if (Config.get().isAwsPropagationEnabled()) {
         try {
           Propagators.forConcern(XRAY_TRACING_CONCERN).inject(span, request, DECORATE);
@@ -81,13 +89,18 @@ public void beforeRequest(final Request<?> request) {
 
   @Override
   public void afterResponse(final Request<?> request, final Response<?> response) {
-    final AgentSpan span = request.getHandlerContext(SPAN_CONTEXT_KEY);
-    if (span != null) {
-      request.addHandlerContext(SPAN_CONTEXT_KEY, null);
-      DECORATE.onResponse(span, response);
-      DECORATE.onServiceResponse(span, request.getServiceName(), response);
-      DECORATE.beforeFinish(span);
-      span.finish();
+    final Context context = request.getHandlerContext(CONTEXT_CONTEXT_KEY);
+    log.warn("context {}", context);
+    AgentSpan span = null;
+    if (context != null) {
+      request.addHandlerContext(CONTEXT_CONTEXT_KEY, null);
+      span = spanFromContext(context);
+      if (span != null) {
+        DECORATE.onResponse(span, response);
+        DECORATE.onServiceResponse(span, request.getServiceName(), response);
+        DECORATE.beforeFinish(span);
+        span.finish();
+      }
     }
     AmazonWebServiceRequest originalRequest = request.getOriginalRequest();
     GetterAccess requestAccess = GetterAccess.of(originalRequest);
@@ -105,22 +118,25 @@ && traceConfig().isDataStreamsEnabled()
         && "AmazonKinesis".equals(request.getServiceName())
         && "GetRecords".equals(requestAccess.getOperationNameFromType())) {
       String streamArn = requestAccess.getStreamARN(originalRequest);
-      if (null != streamArn) {
-        List<?> records =
-            GetterAccess.of(response.getAwsResponse()).getRecords(response.getAwsResponse());
-        if (null != records) {
-          DataStreamsTags tags =
-              DataStreamsTags.create("kinesis", DataStreamsTags.Direction.Inbound, streamArn);
-          for (Object record : records) {
-            Date arrivalTime = GetterAccess.of(record).getApproximateArrivalTimestamp(record);
-            AgentDataStreamsMonitoring dataStreamsMonitoring =
-                AgentTracer.get().getDataStreamsMonitoring();
-            PathwayContext pathwayContext = dataStreamsMonitoring.newPathwayContext();
-            DataStreamsContext context = create(tags, arrivalTime.getTime(), 0);
-            pathwayContext.setCheckpoint(context, dataStreamsMonitoring::add);
-            if (!span.context().getPathwayContext().isStarted()) {
-              span.context().mergePathwayContext(pathwayContext);
-            }
+      dsmCheckpoint(span, streamArn, response);
+    }
+  }
+
+  private void dsmCheckpoint(AgentSpan span, String streamArn, Response<?> response) {
+    if (null != streamArn) {
+      List<?> records =
+          GetterAccess.of(response.getAwsResponse()).getRecords(response.getAwsResponse());
+      if (null != records) {
+        DataStreamsTags tags = create("kinesis", INBOUND, streamArn);
+        for (Object record : records) {
+          Date arrivalTime = GetterAccess.of(record).getApproximateArrivalTimestamp(record);
+          AgentDataStreamsMonitoring dataStreamsMonitoring =
+              AgentTracer.get().getDataStreamsMonitoring();
+          PathwayContext pathwayContext = dataStreamsMonitoring.newPathwayContext();
+          DataStreamsContext dataStreamsContext = create(tags, arrivalTime.getTime(), 0);
+          pathwayContext.setCheckpoint(dataStreamsContext, dataStreamsMonitoring::add);
+          if (!span.context().getPathwayContext().isStarted()) {
+            span.context().mergePathwayContext(pathwayContext);
           }
         }
       }
@@ -129,24 +145,27 @@ && traceConfig().isDataStreamsEnabled()
 
   @Override
   public void afterError(final Request<?> request, final Response<?> response, final Exception e) {
-    AgentSpan span = request.getHandlerContext(SPAN_CONTEXT_KEY);
-    if (span == null) {
-      // also try getting the span from the context store, if the error happened early
-      span = requestSpanStore.remove(request.getOriginalRequest());
+    Context context = request.getHandlerContext(CONTEXT_CONTEXT_KEY);
+    if (context == null) {
+      // also try getting the context from the context store, if the error happened early
+      context = requestContextStore.remove(request.getOriginalRequest());
     }
 
-    if (span != null) {
-      request.addHandlerContext(SPAN_CONTEXT_KEY, null);
-      if (response != null) {
-        DECORATE.onResponse(span, response);
-        if (span.isError()) {
+    if (context != null) {
+      request.addHandlerContext(CONTEXT_CONTEXT_KEY, null);
+      final AgentSpan span = spanFromContext(context);
+      if (span != null) {
+        if (response != null) {
+          DECORATE.onResponse(span, response);
+          if (span.isError()) {
+            DECORATE.onError(span, e);
+          }
+        } else {
           DECORATE.onError(span, e);
         }
-      } else {
-        DECORATE.onError(span, e);
+        DECORATE.beforeFinish(span);
+        span.finish();
       }
-      DECORATE.beforeFinish(span);
-      span.finish();
     }
   }