feat(aws): Refactor AWS instrumentation around Context

PerfectSlayer · PerfectSlayer · commit 11d43a3e967c · 2025-08-18T17:38:00.000+02:00
diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsClientInstrumentation.java b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsClientInstrumentation.java
@@ -42,8 +42,7 @@ public String[] helperClassNames() {
   @Override
   public Map<String, String> contextStore() {
     return Collections.singletonMap(
-        "com.amazonaws.AmazonWebServiceRequest",
-        "datadog.trace.bootstrap.instrumentation.api.AgentSpan");
+        "com.amazonaws.AmazonWebServiceRequest", "datadog.context.Context");
   }
 
   public static class HandlerChainAdvice {
@@ -57,8 +56,7 @@ public static void addHandler(@Advice.Return final List<RequestHandler2> handler
       handlers.add(
           new SnsInterceptor(
               InstrumentationContext.get(
-                  "com.amazonaws.AmazonWebServiceRequest",
-                  "datadog.trace.bootstrap.instrumentation.api.AgentSpan")));
+                  "com.amazonaws.AmazonWebServiceRequest", "datadog.context.Context")));
     }
   }
 }
diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java
@@ -25,9 +25,9 @@
 
 public class SnsInterceptor extends RequestHandler2 {
 
-  private final ContextStore<AmazonWebServiceRequest, AgentSpan> contextStore;
+  private final ContextStore<AmazonWebServiceRequest, Context> contextStore;
 
-  public SnsInterceptor(ContextStore<AmazonWebServiceRequest, AgentSpan> contextStore) {
+  public SnsInterceptor(ContextStore<AmazonWebServiceRequest, Context> contextStore) {
     this.contextStore = contextStore;
   }
 
@@ -108,6 +108,7 @@ private AgentSpan newSpan(AmazonWebServiceRequest request) {
     final AgentSpan span = AgentTracer.startSpan("aws.sns.send");
     // pass the span to TracingRequestHandler in the sdk instrumentation where it'll be enriched &
     // activated
+    // TODO If DSM is enabled, add DSM context here too
     contextStore.put(request, span);
     return span;
   }

Original file line number	Diff line number	Diff line change
`@@ -42,8 +42,7 @@ public String[] helperClassNames() {`
`42`	`42`	`@Override`
`43`	`43`	`public Map<String, String> contextStore() {`
`44`	`44`	`return Collections.singletonMap(`
`45`		`- "com.amazonaws.AmazonWebServiceRequest",`
`46`		`- "datadog.trace.bootstrap.instrumentation.api.AgentSpan");`
	`45`	`+ "com.amazonaws.AmazonWebServiceRequest", "datadog.context.Context");`
`47`	`46`	`}`
`48`	`47`
`49`	`48`	`public static class HandlerChainAdvice {`
`@@ -57,8 +56,7 @@ public static void addHandler(@Advice.Return final List<RequestHandler2> handler`
`57`	`56`	`handlers.add(`
`58`	`57`	`new SnsInterceptor(`
`59`	`58`	`InstrumentationContext.get(`
`60`		`- "com.amazonaws.AmazonWebServiceRequest",`
`61`		`- "datadog.trace.bootstrap.instrumentation.api.AgentSpan")));`
	`59`	`+ "com.amazonaws.AmazonWebServiceRequest", "datadog.context.Context")));`
`62`	`60`	`}`
`63`	`61`	`}`
`64`	`62`	`}`