From 2d815fb20b4f7569d422a4ec8b6d95ae2ea1fc8e Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Fri, 19 Apr 2024 11:02:13 -0400 Subject: [PATCH 01/18] add SNS integration injection part v1 --- .../aws-java-sns-1.0/build.gradle | 18 +++++ .../v1/sns/AbstractSnsInstrumentation.java | 9 +++ .../aws/v1/sns/MessageAttributeInjector.java | 24 +++++++ .../aws/v1/sns/SnsClientInstrumentation.java | 69 +++++++++++++++++++ .../aws/v1/sns/SnsInterceptor.java | 67 ++++++++++++++++++ .../src/test/groovy/SnsClientTest.groovy | 46 +++++++++++++ settings.gradle | 1 + 7 files changed, 234 insertions(+) create mode 100644 dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle create mode 100644 dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/AbstractSnsInstrumentation.java create mode 100644 dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java create mode 100644 dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsClientInstrumentation.java create mode 100644 dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java create mode 100644 dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle new file mode 100644 index 00000000000..17bf0ec0ccd --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle @@ -0,0 +1,18 @@ +muzzle { + pass { + group = "com.amazonaws" + module = "aws-java-sdk-sns" + versions = "[1.0.0,)" + assertInverse = true + } +} + +apply from: "$rootDir/gradle/java.gradle" + +addTestSuiteForDir('latestDepTest', 'test') + +dependencies { + compileOnly group: 'com.amazonaws', name: 'aws-java-sdk-sns', version: '1.12.676' + + latestDepTestImplementation group: 'com.amazonaws', name: 'aws-java-sdk-sns', version: '+' +} diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/AbstractSnsInstrumentation.java b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/AbstractSnsInstrumentation.java new file mode 100644 index 00000000000..5336266d982 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/AbstractSnsInstrumentation.java @@ -0,0 +1,9 @@ +package datadog.trace.instrumentation.aws.v1.sns; + +import datadog.trace.agent.tooling.InstrumenterModule; + +public abstract class AbstractSnsInstrumentation extends InstrumenterModule.Tracing { + public AbstractSnsInstrumentation() { + super("sns", "aws-sdk"); + } +} diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java new file mode 100644 index 00000000000..fd8edd996d0 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java @@ -0,0 +1,24 @@ +package datadog.trace.instrumentation.aws.v1.sns; + +import com.amazonaws.services.sns.model.MessageAttributeValue; +import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; +import java.util.Map; + +public class MessageAttributeInjector + implements AgentPropagation.Setter> { + + public static final MessageAttributeInjector SETTER = new MessageAttributeInjector(); + + @Override + public void set( + final Map carrier, final String key, final String value) { + // The injector would use "X-Amzn-Trace-Id" key, but to keep the behavior same as SQS, use "AWSTraceHeader" + // Also checks if the key already exists because AWS could in the future automatically injects AWSTraceHeader + // (as it does today in SQS) + if (key == "X-Amzn-Trace-Id" && carrier.size() < 10 && !carrier.containsKey("AWSTraceHeader")) { + carrier.put( + "AWSTraceHeader", + new MessageAttributeValue().withDataType("String").withStringValue(value)); + } + } +} 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 new file mode 100644 index 00000000000..38f6941647f --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsClientInstrumentation.java @@ -0,0 +1,69 @@ +package datadog.trace.instrumentation.aws.v1.sns; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.isMethod; + +import com.amazonaws.handlers.RequestHandler2; +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.bootstrap.InstrumentationContext; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import net.bytebuddy.asm.Advice; + +/** AWS SDK v1 instrumentation */ +@AutoService(InstrumenterModule.class) +public final class SnsClientInstrumentation extends InstrumenterModule.Tracing + implements Instrumenter.ForSingleType { + private static final String INSTRUMENTATION_NAME = "aws-sdk"; + + public SnsClientInstrumentation() { + super(INSTRUMENTATION_NAME); + } + + @Override + public String instrumentedType() { + return "com.amazonaws.handlers.HandlerChainFactory"; + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isMethod().and(named("newRequestHandler2Chain")), + SnsClientInstrumentation.class.getName() + "$HandlerChainAdvice"); + } + + @Override + public String[] helperClassNames() { + return new String[] { + packageName + ".SnsInterceptor", packageName + ".MessageAttributeInjector" + }; + } + + @Override + public Map contextStore() { + return Collections.singletonMap( + "com.amazonaws.AmazonWebServiceRequest", + "datadog.trace.bootstrap.instrumentation.api.AgentSpan"); + } + + public static class HandlerChainAdvice { + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void addHandler(@Advice.Return final List handlers) { + // if (Config.get().isDataStreamsEnabled()) { + for (RequestHandler2 interceptor : handlers) { + if (interceptor instanceof SnsInterceptor) { + return; // list already has our interceptor, return to builder + } + } + handlers.add( + new SnsInterceptor( + InstrumentationContext.get( + "com.amazonaws.AmazonWebServiceRequest", + "datadog.trace.bootstrap.instrumentation.api.AgentSpan"))); + } + // } + } +} 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 new file mode 100644 index 00000000000..ef660111133 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java @@ -0,0 +1,67 @@ +package datadog.trace.instrumentation.aws.v1.sns; + +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.URIUtils.urlFileName; +import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_OUT; +import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_TAG; +import static datadog.trace.core.datastreams.TagsProcessor.TOPIC_TAG; +import static datadog.trace.core.datastreams.TagsProcessor.TYPE_TAG; +import static datadog.trace.instrumentation.aws.v1.sns.MessageAttributeInjector.SETTER; + +import com.amazonaws.AmazonWebServiceRequest; +import com.amazonaws.handlers.RequestHandler2; +import com.amazonaws.services.sns.model.PublishBatchRequest; +import com.amazonaws.services.sns.model.PublishBatchRequestEntry; +import com.amazonaws.services.sns.model.PublishRequest; +import datadog.trace.api.TracePropagationStyle; +import datadog.trace.bootstrap.ContextStore; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import java.util.LinkedHashMap; + + +public class SnsInterceptor extends RequestHandler2 { + + private final ContextStore contextStore; + + public SnsInterceptor(ContextStore contextStore) { + this.contextStore = contextStore; + } + + @Override + public AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request) { + // Injecting the AWSTraceHeader into SNS messageAttributes. This uses fewer keys and is consistent with SQS cases. + if (request instanceof PublishRequest) { + PublishRequest pRequest = (PublishRequest) request; + final AgentSpan span = newSpan(request); + // note: modifying message attributes has to be done before marshalling, otherwise the changes + // are not reflected in the actual request (and the MD5 check on send will fail). + propagate().inject(span, pRequest.getMessageAttributes(), SETTER, TracePropagationStyle.XRAY); + + } else if (request instanceof PublishBatchRequest) { + PublishBatchRequest pmbRequest = (PublishBatchRequest) request; + + final AgentSpan span = newSpan(request); + for (PublishBatchRequestEntry entry : pmbRequest.getPublishBatchRequestEntries()) { + propagate().inject(span, entry.getMessageAttributes(), SETTER, TracePropagationStyle.XRAY); + } + } + return request; + } + + private AgentSpan newSpan(AmazonWebServiceRequest request) { + final AgentSpan span = startSpan("aws.sns.send"); + // pass the span to TracingRequestHandler in the sdk instrumentation where it'll be enriched & + // activated + contextStore.put(request, span); + return span; + } + + private static LinkedHashMap getTags(String topicArn) { + LinkedHashMap sortedTags = new LinkedHashMap<>(); + sortedTags.put(DIRECTION_TAG, DIRECTION_OUT); + sortedTags.put(TOPIC_TAG, urlFileName(topicArn)); + sortedTags.put(TYPE_TAG, "sns"); + return sortedTags; + } +} diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy new file mode 100644 index 00000000000..2e995666765 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy @@ -0,0 +1,46 @@ +// +//import com.amazonaws.SDKGlobalConfiguration +////import com.amazonaws.auth.AWSStaticCredentialsProvider +////import com.amazonaws.auth.AnonymousAWSCredentials +//import datadog.trace.agent.test.naming.VersionedNamingTestBase +////import datadog.trace.api.config.GeneralConfig +////import spock.lang.Shared +// +// +//abstract class SnsClientTest extends VersionedNamingTestBase { +// +//// def setup() { +//// System.setProperty(SDKGlobalConfiguration.ACCESS_KEY_SYSTEM_PROPERTY, "my-access-key") +//// System.setProperty(SDKGlobalConfiguration.SECRET_KEY_SYSTEM_PROPERTY, "my-secret-key") +//// } +// +//// @Override +//// protected void configurePreAgent() { +////// super.configurePreAgent() +////// // Set a service name that gets sorted early with SORT_BY_NAMES +////// injectSysConfig(GeneralConfig.SERVICE_NAME, "A-service") +////// injectSysConfig(GeneralConfig.DATA_STREAMS_ENABLED, isDataStreamsEnabled().toString()) +//// } +// +//// @Shared +//// def credentialsProvider = new AWSStaticCredentialsProvider(new AnonymousAWSCredentials()) +// +// @Override +// String operation() { +// null +// } +// +// @Override +// String service() { +// null +// } +// +//// boolean hasTimeInQueueSpan() { +//// false +//// } +// +//// abstract String expectedOperation(String awsService, String awsOperation) +// +//// abstract String expectedService(String awsService, String awsOperation) +// +//} diff --git a/settings.gradle b/settings.gradle index 83f67f043d4..4e1d282d44b 100644 --- a/settings.gradle +++ b/settings.gradle @@ -171,6 +171,7 @@ include ':dd-java-agent:instrumentation:armeria-grpc' include ':dd-java-agent:instrumentation:armeria-jetty' include ':dd-java-agent:instrumentation:aws-java-sdk-1.11.0' include ':dd-java-agent:instrumentation:aws-java-sdk-2.2' +include ':dd-java-agent:instrumentation:aws-java-sns-1.0' include ':dd-java-agent:instrumentation:aws-java-sqs-1.0' include ':dd-java-agent:instrumentation:aws-java-sqs-2.0' include ':dd-java-agent:instrumentation:aws-lambda-handler' From 91e1dac912f61358eabc855b5c1e2709072ba498 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Tue, 23 Apr 2024 14:59:41 -0400 Subject: [PATCH 02/18] add SNS integration injection part v2 --- .../aws/v1/sns/MessageAttributeInjector.java | 11 +++- .../aws/v1/sns/SnsClientInstrumentation.java | 4 +- .../aws/v1/sns/SnsInterceptor.java | 21 +----- .../aws-java-sns-2.0/build.gradle | 18 +++++ .../v2/sns/AbstractSnsInstrumentation.java | 9 +++ .../aws/v2/sns/MessageAttributeInjector.java | 33 ++++++++++ .../aws/v2/sns/SnsClientInstrumentation.java | 65 +++++++++++++++++++ .../aws/v2/sns/SnsInterceptor.java | 56 ++++++++++++++++ .../src/test/groovy/SnsClientTest.groovy | 46 +++++++++++++ settings.gradle | 1 + 10 files changed, 240 insertions(+), 24 deletions(-) create mode 100644 dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle create mode 100644 dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/AbstractSnsInstrumentation.java create mode 100644 dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java create mode 100644 dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsClientInstrumentation.java create mode 100644 dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java create mode 100644 dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java index fd8edd996d0..5ae7e2db772 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java @@ -3,6 +3,7 @@ import com.amazonaws.services.sns.model.MessageAttributeValue; import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; import java.util.Map; +import java.util.Objects; public class MessageAttributeInjector implements AgentPropagation.Setter> { @@ -12,10 +13,14 @@ public class MessageAttributeInjector @Override public void set( final Map carrier, final String key, final String value) { - // The injector would use "X-Amzn-Trace-Id" key, but to keep the behavior same as SQS, use "AWSTraceHeader" - // Also checks if the key already exists because AWS could in the future automatically injects AWSTraceHeader + // The injector would use "X-Amzn-Trace-Id" key, but to keep the behavior same as SQS, use + // "AWSTraceHeader" + // Also checks if the key already exists because AWS could in the future automatically injects + // AWSTraceHeader // (as it does today in SQS) - if (key == "X-Amzn-Trace-Id" && carrier.size() < 10 && !carrier.containsKey("AWSTraceHeader")) { + if (Objects.equals(key, "X-Amzn-Trace-Id") + && carrier.size() < 10 + && !carrier.containsKey("AWSTraceHeader")) { carrier.put( "AWSTraceHeader", new MessageAttributeValue().withDataType("String").withStringValue(value)); 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 index 38f6941647f..440986b4508 100644 --- 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 @@ -13,7 +13,7 @@ import java.util.Map; import net.bytebuddy.asm.Advice; -/** AWS SDK v1 instrumentation */ +/** AWS SDK v1 SNS instrumentation */ @AutoService(InstrumenterModule.class) public final class SnsClientInstrumentation extends InstrumenterModule.Tracing implements Instrumenter.ForSingleType { @@ -52,7 +52,6 @@ public Map contextStore() { public static class HandlerChainAdvice { @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void addHandler(@Advice.Return final List handlers) { - // if (Config.get().isDataStreamsEnabled()) { for (RequestHandler2 interceptor : handlers) { if (interceptor instanceof SnsInterceptor) { return; // list already has our interceptor, return to builder @@ -64,6 +63,5 @@ public static void addHandler(@Advice.Return final List handler "com.amazonaws.AmazonWebServiceRequest", "datadog.trace.bootstrap.instrumentation.api.AgentSpan"))); } - // } } } 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 index ef660111133..a63288e4c83 100644 --- 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 @@ -1,12 +1,6 @@ package datadog.trace.instrumentation.aws.v1.sns; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; -import static datadog.trace.bootstrap.instrumentation.api.URIUtils.urlFileName; -import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_OUT; -import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_TAG; -import static datadog.trace.core.datastreams.TagsProcessor.TOPIC_TAG; -import static datadog.trace.core.datastreams.TagsProcessor.TYPE_TAG; import static datadog.trace.instrumentation.aws.v1.sns.MessageAttributeInjector.SETTER; import com.amazonaws.AmazonWebServiceRequest; @@ -17,8 +11,7 @@ import datadog.trace.api.TracePropagationStyle; import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; -import java.util.LinkedHashMap; - +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; public class SnsInterceptor extends RequestHandler2 { @@ -30,7 +23,7 @@ public SnsInterceptor(ContextStore contextSt @Override public AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request) { - // Injecting the AWSTraceHeader into SNS messageAttributes. This uses fewer keys and is consistent with SQS cases. + // Injecting the AWSTraceHeader into SNS messageAttributes. This is consistent with SQS cases. if (request instanceof PublishRequest) { PublishRequest pRequest = (PublishRequest) request; final AgentSpan span = newSpan(request); @@ -50,18 +43,10 @@ public AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request } private AgentSpan newSpan(AmazonWebServiceRequest request) { - final AgentSpan span = startSpan("aws.sns.send"); + final AgentSpan span = AgentTracer.startSpan("aws.sns.send"); // pass the span to TracingRequestHandler in the sdk instrumentation where it'll be enriched & // activated contextStore.put(request, span); return span; } - - private static LinkedHashMap getTags(String topicArn) { - LinkedHashMap sortedTags = new LinkedHashMap<>(); - sortedTags.put(DIRECTION_TAG, DIRECTION_OUT); - sortedTags.put(TOPIC_TAG, urlFileName(topicArn)); - sortedTags.put(TYPE_TAG, "sns"); - return sortedTags; - } } diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle new file mode 100644 index 00000000000..5789665cc07 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle @@ -0,0 +1,18 @@ +muzzle { + pass { + group = "software.amazon.awssdk" + module = "sns" + versions = "[2.0.0,)" + assertInverse = true + } +} + +apply from: "$rootDir/gradle/java.gradle" + +addTestSuiteForDir('latestDepTest', 'test') + +dependencies { + compileOnly group: 'software.amazon.awssdk', name: 'sns', version: '2.24.4' + + latestDepTestImplementation group: 'software.amazon.awssdk', name: 'sns', version: '+' +} diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/AbstractSnsInstrumentation.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/AbstractSnsInstrumentation.java new file mode 100644 index 00000000000..56d059e3a63 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/AbstractSnsInstrumentation.java @@ -0,0 +1,9 @@ +package datadog.trace.instrumentation.aws.v2.sns; + +import datadog.trace.agent.tooling.InstrumenterModule; + +public abstract class AbstractSnsInstrumentation extends InstrumenterModule.Tracing { + public AbstractSnsInstrumentation() { + super("sns", "aws-sdk"); + } +} diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java new file mode 100644 index 00000000000..12cffd47ad8 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java @@ -0,0 +1,33 @@ +package datadog.trace.instrumentation.aws.v2.sns; + +import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; +import java.util.Map; +import java.util.Objects; +import software.amazon.awssdk.services.sns.model.MessageAttributeValue; + +public class MessageAttributeInjector + implements AgentPropagation.Setter> { + + public static final MessageAttributeInjector SETTER = new MessageAttributeInjector(); + + @Override + public void set( + final Map carrier, final String key, final String value) { + // The injector would use "X-Amzn-Trace-Id" key, but to keep the behavior same as SQS, use + // "AWSTraceHeader" + // Also checks if the key already exists because AWS could in the future automatically injects + // AWSTraceHeader + // (as it does today in SQS) + System.out.println("[JOEY]2"); + + if (Objects.equals(key, "X-Amzn-Trace-Id") + && carrier.size() < 10 + && !carrier.containsKey("AWSTraceHeader")) { + System.out.println("[JOEY]3"); + System.out.println(value); + carrier.put( + "AWSTraceHeader", + MessageAttributeValue.builder().dataType("String").stringValue(value).build()); + } + } +} diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsClientInstrumentation.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsClientInstrumentation.java new file mode 100644 index 00000000000..0e551be9f52 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsClientInstrumentation.java @@ -0,0 +1,65 @@ +package datadog.trace.instrumentation.aws.v2.sns; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.isMethod; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import net.bytebuddy.asm.Advice; +import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; + +/** AWS SDK v2 SNS instrumentation */ +@AutoService(InstrumenterModule.class) +public final class SnsClientInstrumentation extends InstrumenterModule.Tracing + implements Instrumenter.ForSingleType { + private static final String INSTRUMENTATION_NAME = "aws-sdk"; + + public SnsClientInstrumentation() { + super(INSTRUMENTATION_NAME); + } + + @Override + public String instrumentedType() { + return "software.amazon.awssdk.core.client.builder.SdkDefaultClientBuilder"; + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isMethod().and(named("resolveExecutionInterceptors")), + SnsClientInstrumentation.class.getName() + "$AwsSnsBuilderAdvice"); + } + + @Override + public String[] helperClassNames() { + return new String[] { + packageName + ".SnsInterceptor", packageName + ".MessageAttributeInjector" + }; + } + + @Override + public Map contextStore() { + return Collections.singletonMap( + "com.amazonaws.AmazonWebServiceRequest", + "datadog.trace.bootstrap.instrumentation.api.AgentSpan"); + } + + public static class AwsSnsBuilderAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + public static void addHandler(@Advice.Return final List interceptors) { + for (ExecutionInterceptor interceptor : interceptors) { + if (interceptor instanceof SnsInterceptor) { + return; // list already has our interceptor, return to builder + } + } + interceptors.add(new SnsInterceptor()); + // InstrumentationContext.get( + // "software.amazon.awssdk.core.SdkRequest", + // "datadog.trace.bootstrap.instrumentation.api.AgentSpan"))); + } + } +} diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java new file mode 100644 index 00000000000..2ea01ed9cdb --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java @@ -0,0 +1,56 @@ +package datadog.trace.instrumentation.aws.v2.sns; + +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate; +import static datadog.trace.instrumentation.aws.v2.sns.MessageAttributeInjector.SETTER; + +import datadog.trace.api.TracePropagationStyle; +import datadog.trace.bootstrap.InstanceStore; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Map; +import software.amazon.awssdk.core.SdkRequest; +import software.amazon.awssdk.core.interceptor.Context; +import software.amazon.awssdk.core.interceptor.ExecutionAttribute; +import software.amazon.awssdk.core.interceptor.ExecutionAttributes; +import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; +import software.amazon.awssdk.services.sns.model.MessageAttributeValue; +import software.amazon.awssdk.services.sns.model.PublishBatchRequest; +import software.amazon.awssdk.services.sns.model.PublishBatchRequestEntry; +import software.amazon.awssdk.services.sns.model.PublishRequest; + +public class SnsInterceptor implements ExecutionInterceptor { + + public static final ExecutionAttribute SPAN_ATTRIBUTE = + InstanceStore.of(ExecutionAttribute.class) + .putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan")); + + public SnsInterceptor() {} + + @Override + public SdkRequest modifyRequest( + Context.ModifyRequest context, ExecutionAttributes executionAttributes) { + if (context.request() instanceof PublishRequest) { + PublishRequest request = (PublishRequest) context.request(); + Map messageAttributes = + new HashMap<>(request.messageAttributes()); + final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + propagate().inject(span, messageAttributes, SETTER, TracePropagationStyle.XRAY); + System.out.println("[JOEY]"); + System.out.println(messageAttributes); + return request.toBuilder().messageAttributes(messageAttributes).build(); + } else if (context.request() instanceof PublishBatchRequest) { + PublishBatchRequest request = (PublishBatchRequest) context.request(); + final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + ArrayList entries = new ArrayList<>(); + for (PublishBatchRequestEntry entry : request.publishBatchRequestEntries()) { + Map messageAttributes = + new HashMap<>(entry.messageAttributes()); + propagate().inject(span, messageAttributes, SETTER, TracePropagationStyle.XRAY); + entries.add(entry.toBuilder().messageAttributes(messageAttributes).build()); + } + return request.toBuilder().publishBatchRequestEntries(entries).build(); + } + return context.request(); + } +} diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy new file mode 100644 index 00000000000..2e995666765 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy @@ -0,0 +1,46 @@ +// +//import com.amazonaws.SDKGlobalConfiguration +////import com.amazonaws.auth.AWSStaticCredentialsProvider +////import com.amazonaws.auth.AnonymousAWSCredentials +//import datadog.trace.agent.test.naming.VersionedNamingTestBase +////import datadog.trace.api.config.GeneralConfig +////import spock.lang.Shared +// +// +//abstract class SnsClientTest extends VersionedNamingTestBase { +// +//// def setup() { +//// System.setProperty(SDKGlobalConfiguration.ACCESS_KEY_SYSTEM_PROPERTY, "my-access-key") +//// System.setProperty(SDKGlobalConfiguration.SECRET_KEY_SYSTEM_PROPERTY, "my-secret-key") +//// } +// +//// @Override +//// protected void configurePreAgent() { +////// super.configurePreAgent() +////// // Set a service name that gets sorted early with SORT_BY_NAMES +////// injectSysConfig(GeneralConfig.SERVICE_NAME, "A-service") +////// injectSysConfig(GeneralConfig.DATA_STREAMS_ENABLED, isDataStreamsEnabled().toString()) +//// } +// +//// @Shared +//// def credentialsProvider = new AWSStaticCredentialsProvider(new AnonymousAWSCredentials()) +// +// @Override +// String operation() { +// null +// } +// +// @Override +// String service() { +// null +// } +// +//// boolean hasTimeInQueueSpan() { +//// false +//// } +// +//// abstract String expectedOperation(String awsService, String awsOperation) +// +//// abstract String expectedService(String awsService, String awsOperation) +// +//} diff --git a/settings.gradle b/settings.gradle index 4e1d282d44b..8632adc8d4c 100644 --- a/settings.gradle +++ b/settings.gradle @@ -172,6 +172,7 @@ include ':dd-java-agent:instrumentation:armeria-jetty' include ':dd-java-agent:instrumentation:aws-java-sdk-1.11.0' include ':dd-java-agent:instrumentation:aws-java-sdk-2.2' include ':dd-java-agent:instrumentation:aws-java-sns-1.0' +include ':dd-java-agent:instrumentation:aws-java-sns-2.0' include ':dd-java-agent:instrumentation:aws-java-sqs-1.0' include ':dd-java-agent:instrumentation:aws-java-sqs-2.0' include ':dd-java-agent:instrumentation:aws-lambda-handler' From 7d21fcb589b546c931f32fd857d51ee851957a2a Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Mon, 29 Apr 2024 23:31:33 -0400 Subject: [PATCH 03/18] add SNS integration tests v1 --- .../aws-java-sns-1.0/build.gradle | 14 +- .../v1/sns/AbstractSnsInstrumentation.java | 9 - .../src/test/groovy/SnsClientTest.groovy | 245 ++++++++++++++---- .../v2/sns/AbstractSnsInstrumentation.java | 9 - .../aws/v2/sns/MessageAttributeInjector.java | 3 - .../aws/v2/sns/SnsInterceptor.java | 2 - 6 files changed, 211 insertions(+), 71 deletions(-) delete mode 100644 dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/AbstractSnsInstrumentation.java delete mode 100644 dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/AbstractSnsInstrumentation.java diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle index 17bf0ec0ccd..afee603328b 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle @@ -2,7 +2,7 @@ muzzle { pass { group = "com.amazonaws" module = "aws-java-sdk-sns" - versions = "[1.0.0,)" + versions = "[1.12.0,)" assertInverse = true } } @@ -10,9 +10,19 @@ muzzle { apply from: "$rootDir/gradle/java.gradle" addTestSuiteForDir('latestDepTest', 'test') +addTestSuiteExtendingForDir('latestDepForkedTest', 'latestDepTest', 'test') dependencies { - compileOnly group: 'com.amazonaws', name: 'aws-java-sdk-sns', version: '1.12.676' + compileOnly group: 'com.amazonaws', name: 'aws-java-sdk-sns', version: '1.12.710' + + // Include httpclient instrumentation for testing because it is a dependency for aws-sdk. + testImplementation project(':dd-java-agent:instrumentation:apache-httpclient-4') + testImplementation project(':dd-java-agent:instrumentation:aws-java-sdk-1.11.0') + testImplementation group: 'com.amazonaws', name: 'aws-java-sdk-sns', version: '1.12.710' + // There's a problem with sqs sdk v1 with localstack+testcontainers testing. so use sdk v2 for sqs + testImplementation 'software.amazon.awssdk:sqs:2.25.40' + testImplementation 'org.testcontainers:testcontainers:1.19.7' + testImplementation 'org.testcontainers:localstack:1.19.7' latestDepTestImplementation group: 'com.amazonaws', name: 'aws-java-sdk-sns', version: '+' } diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/AbstractSnsInstrumentation.java b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/AbstractSnsInstrumentation.java deleted file mode 100644 index 5336266d982..00000000000 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/AbstractSnsInstrumentation.java +++ /dev/null @@ -1,9 +0,0 @@ -package datadog.trace.instrumentation.aws.v1.sns; - -import datadog.trace.agent.tooling.InstrumenterModule; - -public abstract class AbstractSnsInstrumentation extends InstrumenterModule.Tracing { - public AbstractSnsInstrumentation() { - super("sns", "aws-sdk"); - } -} diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy index 2e995666765..3c22c7b819a 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy @@ -1,46 +1,199 @@ -// -//import com.amazonaws.SDKGlobalConfiguration -////import com.amazonaws.auth.AWSStaticCredentialsProvider -////import com.amazonaws.auth.AnonymousAWSCredentials -//import datadog.trace.agent.test.naming.VersionedNamingTestBase -////import datadog.trace.api.config.GeneralConfig -////import spock.lang.Shared -// -// -//abstract class SnsClientTest extends VersionedNamingTestBase { -// -//// def setup() { -//// System.setProperty(SDKGlobalConfiguration.ACCESS_KEY_SYSTEM_PROPERTY, "my-access-key") -//// System.setProperty(SDKGlobalConfiguration.SECRET_KEY_SYSTEM_PROPERTY, "my-secret-key") -//// } -// -//// @Override -//// protected void configurePreAgent() { -////// super.configurePreAgent() -////// // Set a service name that gets sorted early with SORT_BY_NAMES -////// injectSysConfig(GeneralConfig.SERVICE_NAME, "A-service") -////// injectSysConfig(GeneralConfig.DATA_STREAMS_ENABLED, isDataStreamsEnabled().toString()) -//// } -// -//// @Shared -//// def credentialsProvider = new AWSStaticCredentialsProvider(new AnonymousAWSCredentials()) -// -// @Override -// String operation() { -// null -// } -// -// @Override -// String service() { -// null -// } -// -//// boolean hasTimeInQueueSpan() { -//// false -//// } -// -//// abstract String expectedOperation(String awsService, String awsOperation) -// -//// abstract String expectedService(String awsService, String awsOperation) -// -//} +import com.amazonaws.SDKGlobalConfiguration +import com.amazonaws.auth.AWSStaticCredentialsProvider +import com.amazonaws.auth.BasicAWSCredentials +import com.amazonaws.client.builder.AwsClientBuilder +import com.amazonaws.services.sns.AmazonSNSClient +import com.amazonaws.services.sns.AmazonSNSClientBuilder +import datadog.trace.agent.test.naming.VersionedNamingTestBase +import datadog.trace.agent.test.utils.TraceUtils +import datadog.trace.api.DDSpanId +import datadog.trace.api.DDSpanTypes +import datadog.trace.api.DDTags +import datadog.trace.api.config.GeneralConfig +import datadog.trace.bootstrap.instrumentation.api.Tags +import org.testcontainers.containers.localstack.LocalStackContainer +import org.testcontainers.utility.DockerImageName +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider +import software.amazon.awssdk.regions.Region +import software.amazon.awssdk.services.sqs.SqsClient +import software.amazon.awssdk.services.sqs.model.QueueAttributeName +import spock.lang.Shared +import groovy.json.JsonSlurper + +import static datadog.trace.agent.test.utils.TraceUtils.basicSpan +import static org.testcontainers.containers.localstack.LocalStackContainer.Service.SNS +import static org.testcontainers.containers.localstack.LocalStackContainer.Service.SQS + +abstract class SnsClientTest extends VersionedNamingTestBase { + static final localstack = new LocalStackContainer(DockerImageName.parse("localstack/localstack:latest")) + .withServices( SQS, SNS) + @Shared AmazonSNSClient snsClient + @Shared SqsClient sqsClient + + @Shared String testQueueURL + @Shared String testQueueARN + @Shared String testTopicARN + + def setupSpec() { + localstack.start() + snsClient = AmazonSNSClientBuilder.standard() + .withEndpointConfiguration(new AwsClientBuilder.EndpointConfiguration( localstack.getEndpointOverride(SNS).toString(), localstack.getRegion())) + .withCredentials( new AWSStaticCredentialsProvider(new BasicAWSCredentials("test", "test"))) + .build() + sqsClient = SqsClient.builder() + .endpointOverride(localstack.getEndpointOverride(SQS)) + .region(Region.of(localstack.getRegion())) + .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test", "test"))) + .build() + testQueueURL = sqsClient.createQueue { it.queueName("testqueue") }.queueUrl() + testQueueARN = sqsClient.getQueueAttributes {it.queueUrl(testQueueURL).attributeNames(QueueAttributeName.QUEUE_ARN)}.attributes().get(QueueAttributeName.QUEUE_ARN) + testTopicARN = snsClient.createTopic("testtopic").topicArn + snsClient.subscribe(testTopicARN, "sqs", testQueueARN) + } + + def cleanupSpec() { + localstack.stop() + } + + @Override + protected void configurePreAgent() { + super.configurePreAgent() + // Set a service name that gets sorted early with SORT_BY_NAMES + injectSysConfig(GeneralConfig.SERVICE_NAME, "A-service") + } + + @Override + String operation() { + null + } + + @Override + String service() { + null + } + + abstract String expectedOperation(String awsService, String awsOperation) + abstract String expectedService(String awsService, String awsOperation) + + def "trace details propagated via SNS system message attributes"() { + when: + TEST_WRITER.clear() + TraceUtils.runUnderTrace('parent', { + snsClient.publish(testTopicARN, 'sometext') + }) + + def message = sqsClient.receiveMessage {it.queueUrl(testQueueURL).waitTimeSeconds(3)}.messages().get(0) + def jsonSlurper = new JsonSlurper() + def messageBody = jsonSlurper.parseText(message.body()) + + then: + def sendSpan + assertTraces(2) { + trace(2) { + basicSpan(it, "parent") + span { + serviceName expectedService("SNS", "Publish") + operationName expectedOperation("SNS", "Publish") + resourceName "SNS.Publish" + spanType DDSpanTypes.HTTP_CLIENT + errored false + measured true + childOf(span(0)) + tags { + "$Tags.COMPONENT" "java-aws-sdk" + "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT + "$Tags.HTTP_URL" localstack.getEndpointOverride(SNS).toString()+'/' + "$Tags.HTTP_METHOD" "POST" + "$Tags.HTTP_STATUS" 200 + "$Tags.PEER_PORT" localstack.getEndpointOverride(SNS).port + "$Tags.PEER_HOSTNAME" localstack.getEndpointOverride(SNS).host + "aws.service" "AmazonSNS" + "aws_service" "sns" + "aws.endpoint" localstack.getEndpointOverride(SNS).toString() + "aws.operation" "PublishRequest" + "aws.agent" "java-aws-sdk" + "aws.topic.name" "testtopic" + "topicname" "testtopic" + defaultTags() + } + } + sendSpan = span(1) + } + trace(1) { + span { + serviceName expectedService("None", "http.post") + operationName expectedOperation("None", "http.post") + resourceName "POST /" + spanType DDSpanTypes.HTTP_CLIENT + errored false + measured true + childOf(sendSpan) + tags { + "$Tags.COMPONENT" "apache-httpclient" + "$Tags.SPAN_KIND" Tags.SPAN_KIND_CONSUMER + "$Tags.HTTP_URL" localstack.getEndpointOverride(SNS).toString()+'/' + "$Tags.HTTP_METHOD" "POST" + "$Tags.HTTP_STATUS" 200 + "$Tags.PEER_PORT" localstack.getEndpointOverride(SNS).port + "$Tags.PEER_HOSTNAME" localstack.getEndpointOverride(SNS).host + defaultTags(true) + } + } + } + } + + and: + messageBody["Message"] == "sometext" + messageBody["MessageAttributes"]["AWSTraceHeader"]["Value"] =~ + /Root=1-[0-9a-f]{8}-00000000${sendSpan.traceId.toHexStringPadded(16)};Parent=${DDSpanId.toHexStringPadded(sendSpan.spanId)};Sampled=1/ + + } +} + +class SnsClientV0Test extends SnsClientTest { + + @Override + String expectedOperation(String awsService, String awsOperation) { + if ("SNS" == awsService) { + return "aws.http" + } + return "http.request" + + } + + @Override + String expectedService(String awsService, String awsOperation) { + if ("SNS" == awsService) { + return "sns" + } + return "A-service" + } + + @Override + int version() { + 0 + } +} + +class SnsClientV1ForkedTest extends SnsClientTest { + + @Override + String expectedOperation(String awsService, String awsOperation) { + if (awsService == "SNS"&& awsOperation == "Publish") { + return "aws.sns.send" + } + return "aws.${awsService.toLowerCase()}.request" + } + + @Override + String expectedService(String awsService, String awsOperation) { + "A-service" + } + + @Override + int version() { + 1 + } +} + diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/AbstractSnsInstrumentation.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/AbstractSnsInstrumentation.java deleted file mode 100644 index 56d059e3a63..00000000000 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/AbstractSnsInstrumentation.java +++ /dev/null @@ -1,9 +0,0 @@ -package datadog.trace.instrumentation.aws.v2.sns; - -import datadog.trace.agent.tooling.InstrumenterModule; - -public abstract class AbstractSnsInstrumentation extends InstrumenterModule.Tracing { - public AbstractSnsInstrumentation() { - super("sns", "aws-sdk"); - } -} diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java index 12cffd47ad8..e97d80186ec 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java @@ -18,13 +18,10 @@ public void set( // Also checks if the key already exists because AWS could in the future automatically injects // AWSTraceHeader // (as it does today in SQS) - System.out.println("[JOEY]2"); if (Objects.equals(key, "X-Amzn-Trace-Id") && carrier.size() < 10 && !carrier.containsKey("AWSTraceHeader")) { - System.out.println("[JOEY]3"); - System.out.println(value); carrier.put( "AWSTraceHeader", MessageAttributeValue.builder().dataType("String").stringValue(value).build()); diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java index 2ea01ed9cdb..e32de673414 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java @@ -36,8 +36,6 @@ public SdkRequest modifyRequest( new HashMap<>(request.messageAttributes()); final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); propagate().inject(span, messageAttributes, SETTER, TracePropagationStyle.XRAY); - System.out.println("[JOEY]"); - System.out.println(messageAttributes); return request.toBuilder().messageAttributes(messageAttributes).build(); } else if (context.request() instanceof PublishBatchRequest) { PublishBatchRequest request = (PublishBatchRequest) context.request(); From d191ee131b5135509447314d18e833a37bae2e60 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Mon, 29 Apr 2024 23:36:04 -0400 Subject: [PATCH 04/18] small fix --- .../aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy index 3c22c7b819a..48cf2e5b987 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy @@ -128,10 +128,9 @@ abstract class SnsClientTest extends VersionedNamingTestBase { spanType DDSpanTypes.HTTP_CLIENT errored false measured true - childOf(sendSpan) tags { "$Tags.COMPONENT" "apache-httpclient" - "$Tags.SPAN_KIND" Tags.SPAN_KIND_CONSUMER + "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT "$Tags.HTTP_URL" localstack.getEndpointOverride(SNS).toString()+'/' "$Tags.HTTP_METHOD" "POST" "$Tags.HTTP_STATUS" 200 From 4826bbcce90969bc7af0c1e055657f80fb218a6a Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 1 May 2024 22:17:34 -0400 Subject: [PATCH 05/18] add tests --- .../aws-java-sns-1.0/build.gradle | 3 +- .../src/test/groovy/SnsClientTest.groovy | 8 +- .../aws-java-sns-2.0/build.gradle | 12 +- .../src/test/groovy/SnsClientTest.groovy | 221 ++++++++++++++---- 4 files changed, 190 insertions(+), 54 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle index afee603328b..b03ad70fc20 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle @@ -2,7 +2,7 @@ muzzle { pass { group = "com.amazonaws" module = "aws-java-sdk-sns" - versions = "[1.12.0,)" + versions = "[1.0.0,)" assertInverse = true } } @@ -19,6 +19,7 @@ dependencies { testImplementation project(':dd-java-agent:instrumentation:apache-httpclient-4') testImplementation project(':dd-java-agent:instrumentation:aws-java-sdk-1.11.0') testImplementation group: 'com.amazonaws', name: 'aws-java-sdk-sns', version: '1.12.710' + // SQS is used to act as the "Subscriber" of the SNS topic. // There's a problem with sqs sdk v1 with localstack+testcontainers testing. so use sdk v2 for sqs testImplementation 'software.amazon.awssdk:sqs:2.25.40' testImplementation 'org.testcontainers:testcontainers:1.19.7' diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy index 48cf2e5b987..2ee55aa520e 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy @@ -1,4 +1,3 @@ -import com.amazonaws.SDKGlobalConfiguration import com.amazonaws.auth.AWSStaticCredentialsProvider import com.amazonaws.auth.BasicAWSCredentials import com.amazonaws.client.builder.AwsClientBuilder @@ -8,7 +7,6 @@ import datadog.trace.agent.test.naming.VersionedNamingTestBase import datadog.trace.agent.test.utils.TraceUtils import datadog.trace.api.DDSpanId import datadog.trace.api.DDSpanTypes -import datadog.trace.api.DDTags import datadog.trace.api.config.GeneralConfig import datadog.trace.bootstrap.instrumentation.api.Tags import org.testcontainers.containers.localstack.LocalStackContainer @@ -27,7 +25,7 @@ import static org.testcontainers.containers.localstack.LocalStackContainer.Servi abstract class SnsClientTest extends VersionedNamingTestBase { static final localstack = new LocalStackContainer(DockerImageName.parse("localstack/localstack:latest")) - .withServices( SQS, SNS) + .withServices( SQS, SNS) @Shared AmazonSNSClient snsClient @Shared SqsClient sqsClient @@ -146,7 +144,6 @@ abstract class SnsClientTest extends VersionedNamingTestBase { messageBody["Message"] == "sometext" messageBody["MessageAttributes"]["AWSTraceHeader"]["Value"] =~ /Root=1-[0-9a-f]{8}-00000000${sendSpan.traceId.toHexStringPadded(16)};Parent=${DDSpanId.toHexStringPadded(sendSpan.spanId)};Sampled=1/ - } } @@ -158,7 +155,6 @@ class SnsClientV0Test extends SnsClientTest { return "aws.http" } return "http.request" - } @Override @@ -182,7 +178,7 @@ class SnsClientV1ForkedTest extends SnsClientTest { if (awsService == "SNS"&& awsOperation == "Publish") { return "aws.sns.send" } - return "aws.${awsService.toLowerCase()}.request" + return "http.client.request" } @Override diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle index 5789665cc07..b82a6320d01 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle @@ -10,9 +10,19 @@ muzzle { apply from: "$rootDir/gradle/java.gradle" addTestSuiteForDir('latestDepTest', 'test') +addTestSuiteExtendingForDir('latestDepForkedTest', 'latestDepTest', 'test') dependencies { - compileOnly group: 'software.amazon.awssdk', name: 'sns', version: '2.24.4' + compileOnly group: 'software.amazon.awssdk', name: 'sns', version: '2.25.40' + + // Include httpclient instrumentation for testing because it is a dependency for aws-sdk. + testImplementation project(':dd-java-agent:instrumentation:apache-httpclient-4') + testImplementation project(':dd-java-agent:instrumentation:aws-java-sdk-2.2') + testImplementation 'software.amazon.awssdk:sns:2.25.40' + // SQS is used to act as the "Subscriber" of the SNS topic. + testImplementation 'software.amazon.awssdk:sqs:2.25.40' + testImplementation 'org.testcontainers:testcontainers:1.19.7' + testImplementation 'org.testcontainers:localstack:1.19.7' latestDepTestImplementation group: 'software.amazon.awssdk', name: 'sns', version: '+' } diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy index 2e995666765..c9805bb9c86 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy @@ -1,46 +1,175 @@ -// -//import com.amazonaws.SDKGlobalConfiguration -////import com.amazonaws.auth.AWSStaticCredentialsProvider -////import com.amazonaws.auth.AnonymousAWSCredentials -//import datadog.trace.agent.test.naming.VersionedNamingTestBase -////import datadog.trace.api.config.GeneralConfig -////import spock.lang.Shared -// -// -//abstract class SnsClientTest extends VersionedNamingTestBase { -// -//// def setup() { -//// System.setProperty(SDKGlobalConfiguration.ACCESS_KEY_SYSTEM_PROPERTY, "my-access-key") -//// System.setProperty(SDKGlobalConfiguration.SECRET_KEY_SYSTEM_PROPERTY, "my-secret-key") -//// } -// -//// @Override -//// protected void configurePreAgent() { -////// super.configurePreAgent() -////// // Set a service name that gets sorted early with SORT_BY_NAMES -////// injectSysConfig(GeneralConfig.SERVICE_NAME, "A-service") -////// injectSysConfig(GeneralConfig.DATA_STREAMS_ENABLED, isDataStreamsEnabled().toString()) -//// } -// -//// @Shared -//// def credentialsProvider = new AWSStaticCredentialsProvider(new AnonymousAWSCredentials()) -// -// @Override -// String operation() { -// null -// } -// -// @Override -// String service() { -// null -// } -// -//// boolean hasTimeInQueueSpan() { -//// false -//// } -// -//// abstract String expectedOperation(String awsService, String awsOperation) -// -//// abstract String expectedService(String awsService, String awsOperation) -// -//} + +import datadog.trace.agent.test.naming.VersionedNamingTestBase +import datadog.trace.agent.test.utils.TraceUtils +import datadog.trace.api.DDSpanId +import datadog.trace.api.DDSpanTypes +import datadog.trace.api.config.GeneralConfig +import datadog.trace.bootstrap.instrumentation.api.Tags +import org.testcontainers.containers.localstack.LocalStackContainer +import org.testcontainers.utility.DockerImageName +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider +import software.amazon.awssdk.regions.Region +import software.amazon.awssdk.services.sns.SnsClient +import software.amazon.awssdk.services.sns.model.PublishRequest +import software.amazon.awssdk.services.sns.model.PublishResponse +import software.amazon.awssdk.services.sqs.SqsClient +import software.amazon.awssdk.services.sqs.model.QueueAttributeName +import spock.lang.Shared +import groovy.json.JsonSlurper + +import static datadog.trace.agent.test.utils.TraceUtils.basicSpan +import static org.testcontainers.containers.localstack.LocalStackContainer.Service.SNS +import static org.testcontainers.containers.localstack.LocalStackContainer.Service.SQS + +abstract class SnsClientTest extends VersionedNamingTestBase { + static final localstack = new LocalStackContainer(DockerImageName.parse("localstack/localstack:latest")) + .withServices( SQS, SNS) + @Shared SnsClient snsClient + @Shared SqsClient sqsClient + + @Shared String testQueueURL + @Shared String testQueueARN + @Shared String testTopicARN + + def setupSpec() { + localstack.start() + snsClient = SnsClient.builder() + .endpointOverride(localstack.getEndpointOverride(LocalStackContainer.Service.SNS)) + .region(Region.of(localstack.getRegion())) + .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test", "test"))) + .build() + sqsClient = SqsClient.builder() + .endpointOverride(localstack.getEndpointOverride(SQS)) + .region(Region.of(localstack.getRegion())) + .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test", "test"))) + .build() + testQueueURL = sqsClient.createQueue { it.queueName("testqueue") }.queueUrl() + testQueueARN = sqsClient.getQueueAttributes {it.queueUrl(testQueueURL).attributeNames(QueueAttributeName.QUEUE_ARN)}.attributes().get(QueueAttributeName.QUEUE_ARN) + testTopicARN = snsClient.createTopic { it.name("testtopic") }.topicArn() + snsClient.subscribe {it.topicArn(testTopicARN).protocol("sqs").endpoint(testQueueARN)} + } + + def cleanupSpec() { + localstack.stop() + } + + @Override + protected void configurePreAgent() { + super.configurePreAgent() + // Set a service name that gets sorted early with SORT_BY_NAMES + injectSysConfig(GeneralConfig.SERVICE_NAME, "A-service") + } + + @Override + String operation() { + null + } + + @Override + String service() { + null + } + + abstract String expectedOperation(String awsService, String awsOperation) + abstract String expectedService(String awsService, String awsOperation) + + def "trace details propagated via SNS system message attributes"() { + when: + TEST_WRITER.clear() + PublishResponse response + TraceUtils.runUnderTrace('parent', { + response = snsClient.publish { it.message("sometext").topicArn(testTopicARN)} + }) + + def message = sqsClient.receiveMessage {it.queueUrl(testQueueURL).waitTimeSeconds(3)}.messages().get(0) + def jsonSlurper = new JsonSlurper() + def messageBody = jsonSlurper.parseText(message.body()) + + then: + def sendSpan + assertTraces(1) { + trace(2) { + basicSpan(it, "parent") + span { + serviceName expectedService("SNS", "Publish") + operationName expectedOperation("SNS", "Publish") + resourceName "Sns.Publish" + spanType DDSpanTypes.HTTP_CLIENT + errored false + measured true + childOf(span(0)) + tags { + "$Tags.COMPONENT" "java-aws-sdk" + "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT + "$Tags.HTTP_URL" localstack.getEndpointOverride(SNS).toString()+'/' + "$Tags.HTTP_METHOD" "POST" + "$Tags.HTTP_STATUS" 200 + "$Tags.PEER_PORT" localstack.getEndpointOverride(SNS).port + "$Tags.PEER_HOSTNAME" localstack.getEndpointOverride(SNS).host + "aws.service" "Sns" + "aws_service" "Sns" + "aws.operation" "Publish" + "aws.agent" "java-aws-sdk" + "aws.topic.name" "testtopic" + "topicname" "testtopic" + "aws.requestId" response.responseMetadata().requestId() + defaultTags() + } + } + sendSpan = span(1) + } + } + + and: + messageBody["Message"] == "sometext" + messageBody["MessageAttributes"]["AWSTraceHeader"]["Value"] =~ + /Root=1-[0-9a-f]{8}-00000000${sendSpan.traceId.toHexStringPadded(16)};Parent=${DDSpanId.toHexStringPadded(sendSpan.spanId)};Sampled=1/ + } +} + +class SnsClientV0Test extends SnsClientTest { + + @Override + String expectedOperation(String awsService, String awsOperation) { + if ("SNS" == awsService) { + return "aws.http" + } + return "http.request" + } + + @Override + String expectedService(String awsService, String awsOperation) { + if ("SNS" == awsService) { + return "sns" + } + return "A-service" + } + + @Override + int version() { + 0 + } +} + +class SnsClientV1ForkedTest extends SnsClientTest { + + @Override + String expectedOperation(String awsService, String awsOperation) { + if (awsService == "SNS"&& awsOperation == "Publish") { + return "aws.sns.send" + } + return "http.client.request" + } + + @Override + String expectedService(String awsService, String awsOperation) { + "A-service" + } + + @Override + int version() { + 1 + } +} + From ea237e118009c65cc7d6a988fb6e488472c739b0 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 1 May 2024 23:17:07 -0400 Subject: [PATCH 06/18] use X-Amzn-Trace-Id directly --- .../aws-java-sns-1.0/build.gradle | 2 +- .../aws/v1/sns/MessageAttributeInjector.java | 16 ++++------------ .../src/test/groovy/SnsClientTest.groovy | 2 +- .../aws-java-sns-2.0/build.gradle | 2 +- .../aws/v2/sns/MessageAttributeInjector.java | 16 ++++------------ .../src/test/groovy/SnsClientTest.groovy | 2 +- 6 files changed, 12 insertions(+), 28 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle index b03ad70fc20..05a8dbca4c8 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle @@ -2,7 +2,7 @@ muzzle { pass { group = "com.amazonaws" module = "aws-java-sdk-sns" - versions = "[1.0.0,)" + versions = "[1.12.0,)" assertInverse = true } } diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java index 5ae7e2db772..e56450d1ddd 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java @@ -3,7 +3,6 @@ import com.amazonaws.services.sns.model.MessageAttributeValue; import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; import java.util.Map; -import java.util.Objects; public class MessageAttributeInjector implements AgentPropagation.Setter> { @@ -13,17 +12,10 @@ public class MessageAttributeInjector @Override public void set( final Map carrier, final String key, final String value) { - // The injector would use "X-Amzn-Trace-Id" key, but to keep the behavior same as SQS, use - // "AWSTraceHeader" - // Also checks if the key already exists because AWS could in the future automatically injects - // AWSTraceHeader - // (as it does today in SQS) - if (Objects.equals(key, "X-Amzn-Trace-Id") - && carrier.size() < 10 - && !carrier.containsKey("AWSTraceHeader")) { - carrier.put( - "AWSTraceHeader", - new MessageAttributeValue().withDataType("String").withStringValue(value)); + // 10 messageAttributes is a limit from SQS, which is often used as a subscriber and therefore + // still apply here + if (carrier.size() < 10 && !carrier.containsKey(key)) { + carrier.put(key, new MessageAttributeValue().withDataType("String").withStringValue(value)); } } } diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy index 2ee55aa520e..05a8281f359 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy @@ -142,7 +142,7 @@ abstract class SnsClientTest extends VersionedNamingTestBase { and: messageBody["Message"] == "sometext" - messageBody["MessageAttributes"]["AWSTraceHeader"]["Value"] =~ + messageBody["MessageAttributes"]["X-Amzn-Trace-Id"]["Value"] =~ /Root=1-[0-9a-f]{8}-00000000${sendSpan.traceId.toHexStringPadded(16)};Parent=${DDSpanId.toHexStringPadded(sendSpan.spanId)};Sampled=1/ } } diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle index b82a6320d01..d9e25837058 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle @@ -2,7 +2,7 @@ muzzle { pass { group = "software.amazon.awssdk" module = "sns" - versions = "[2.0.0,)" + versions = "[2.2.0,)" assertInverse = true } } diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java index e97d80186ec..ee60d19301b 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java @@ -2,7 +2,6 @@ import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; import java.util.Map; -import java.util.Objects; import software.amazon.awssdk.services.sns.model.MessageAttributeValue; public class MessageAttributeInjector @@ -13,18 +12,11 @@ public class MessageAttributeInjector @Override public void set( final Map carrier, final String key, final String value) { - // The injector would use "X-Amzn-Trace-Id" key, but to keep the behavior same as SQS, use - // "AWSTraceHeader" - // Also checks if the key already exists because AWS could in the future automatically injects - // AWSTraceHeader - // (as it does today in SQS) - - if (Objects.equals(key, "X-Amzn-Trace-Id") - && carrier.size() < 10 - && !carrier.containsKey("AWSTraceHeader")) { + // 10 messageAttributes is a limit from SQS, which is often used as a subscriber and therefore + // still apply here + if (carrier.size() < 10 && !carrier.containsKey(key)) { carrier.put( - "AWSTraceHeader", - MessageAttributeValue.builder().dataType("String").stringValue(value).build()); + key, MessageAttributeValue.builder().dataType("String").stringValue(value).build()); } } } diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy index c9805bb9c86..48e9921296f 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy @@ -123,7 +123,7 @@ abstract class SnsClientTest extends VersionedNamingTestBase { and: messageBody["Message"] == "sometext" - messageBody["MessageAttributes"]["AWSTraceHeader"]["Value"] =~ + messageBody["MessageAttributes"]["X-Amzn-Trace-Id"]["Value"] =~ /Root=1-[0-9a-f]{8}-00000000${sendSpan.traceId.toHexStringPadded(16)};Parent=${DDSpanId.toHexStringPadded(sendSpan.spanId)};Sampled=1/ } } From 93c87509087170a861beb0d6f0424261342397a0 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Thu, 2 May 2024 00:35:03 -0400 Subject: [PATCH 07/18] remove comments --- .../instrumentation/aws/v2/sns/SnsClientInstrumentation.java | 3 --- 1 file changed, 3 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsClientInstrumentation.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsClientInstrumentation.java index 0e551be9f52..aad95b7b235 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsClientInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsClientInstrumentation.java @@ -57,9 +57,6 @@ public static void addHandler(@Advice.Return final List in } } interceptors.add(new SnsInterceptor()); - // InstrumentationContext.get( - // "software.amazon.awssdk.core.SdkRequest", - // "datadog.trace.bootstrap.instrumentation.api.AgentSpan"))); } } } From 74673c692b1d7ba0e2ed43cee8bff4374d02817a Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Mon, 20 May 2024 14:34:38 -0400 Subject: [PATCH 08/18] change to the same instrumentation as other dd-trace libraries --- .../aws-java-sns-1.0/build.gradle | 2 +- .../aws/v1/sns/MessageAttributeInjector.java | 21 ---------- .../aws/v1/sns/SnsClientInstrumentation.java | 7 +--- .../aws/v1/sns/SnsInterceptor.java | 42 +++++++++++++++++-- .../aws/v1/sns/TextMapInjectAdapter.java | 13 ++++++ .../src/test/groovy/SnsClientTest.groovy | 13 ++++-- .../aws-java-sns-2.0/build.gradle | 2 +- .../aws/v2/sns/MessageAttributeInjector.java | 22 ---------- .../aws/v2/sns/SnsClientInstrumentation.java | 8 +--- .../aws/v2/sns/SnsInterceptor.java | 32 ++++++++++++-- .../aws/v2/sns/TextMapInjectAdapter.java | 13 ++++++ .../src/test/groovy/SnsClientTest.groovy | 14 ++++--- 12 files changed, 117 insertions(+), 72 deletions(-) delete mode 100644 dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java create mode 100644 dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/TextMapInjectAdapter.java delete mode 100644 dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java create mode 100644 dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/TextMapInjectAdapter.java diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle index 05a8dbca4c8..6c8b9068364 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle @@ -2,7 +2,7 @@ muzzle { pass { group = "com.amazonaws" module = "aws-java-sdk-sns" - versions = "[1.12.0,)" + versions = "[1.12.0,2)" assertInverse = true } } diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java deleted file mode 100644 index e56450d1ddd..00000000000 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/MessageAttributeInjector.java +++ /dev/null @@ -1,21 +0,0 @@ -package datadog.trace.instrumentation.aws.v1.sns; - -import com.amazonaws.services.sns.model.MessageAttributeValue; -import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; -import java.util.Map; - -public class MessageAttributeInjector - implements AgentPropagation.Setter> { - - public static final MessageAttributeInjector SETTER = new MessageAttributeInjector(); - - @Override - public void set( - final Map carrier, final String key, final String value) { - // 10 messageAttributes is a limit from SQS, which is often used as a subscriber and therefore - // still apply here - if (carrier.size() < 10 && !carrier.containsKey(key)) { - carrier.put(key, new MessageAttributeValue().withDataType("String").withStringValue(value)); - } - } -} 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 index 440986b4508..b6565469b19 100644 --- 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 @@ -17,10 +17,9 @@ @AutoService(InstrumenterModule.class) public final class SnsClientInstrumentation extends InstrumenterModule.Tracing implements Instrumenter.ForSingleType { - private static final String INSTRUMENTATION_NAME = "aws-sdk"; public SnsClientInstrumentation() { - super(INSTRUMENTATION_NAME); + super("sns", "aws-sdk"); } @Override @@ -37,9 +36,7 @@ public void methodAdvice(MethodTransformer transformer) { @Override public String[] helperClassNames() { - return new String[] { - packageName + ".SnsInterceptor", packageName + ".MessageAttributeInjector" - }; + return new String[] {packageName + ".SnsInterceptor", packageName + ".TextMapInjectAdapter"}; } @Override 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 index a63288e4c83..270a02dda85 100644 --- 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 @@ -1,10 +1,11 @@ package datadog.trace.instrumentation.aws.v1.sns; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate; -import static datadog.trace.instrumentation.aws.v1.sns.MessageAttributeInjector.SETTER; +import static datadog.trace.instrumentation.aws.v1.sns.TextMapInjectAdapter.SETTER; import com.amazonaws.AmazonWebServiceRequest; import com.amazonaws.handlers.RequestHandler2; +import com.amazonaws.services.sns.model.MessageAttributeValue; import com.amazonaws.services.sns.model.PublishBatchRequest; import com.amazonaws.services.sns.model.PublishBatchRequestEntry; import com.amazonaws.services.sns.model.PublishRequest; @@ -12,6 +13,9 @@ import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.Map; public class SnsInterceptor extends RequestHandler2 { @@ -29,14 +33,44 @@ public AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request final AgentSpan span = newSpan(request); // note: modifying message attributes has to be done before marshalling, otherwise the changes // are not reflected in the actual request (and the MD5 check on send will fail). - propagate().inject(span, pRequest.getMessageAttributes(), SETTER, TracePropagationStyle.XRAY); - + Map messageAttributes = pRequest.getMessageAttributes(); + // 10 messageAttributes is a limit from SQS, which is often used as a subscriber, therefore + // the limit still applies here + if (messageAttributes.size() < 10) { + StringBuilder jsonBuilder = new StringBuilder(); + jsonBuilder.append("{"); + propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); + jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma + jsonBuilder.append("}"); + messageAttributes.put( + "_datadog", + new MessageAttributeValue() + .withDataType( + "Binary") // Use Binary since SNS subscription filter policies fail silently + // with JSON strings + // https://github.com/DataDog/datadog-lambda-js/pull/269 + .withBinaryValue( + ByteBuffer.wrap(jsonBuilder.toString().getBytes(StandardCharsets.UTF_8)))); + } } else if (request instanceof PublishBatchRequest) { PublishBatchRequest pmbRequest = (PublishBatchRequest) request; final AgentSpan span = newSpan(request); + StringBuilder jsonBuilder = new StringBuilder(); + jsonBuilder.append("{"); + propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); + jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma + jsonBuilder.append("}"); + ByteBuffer binaryValue = + ByteBuffer.wrap(jsonBuilder.toString().getBytes(StandardCharsets.UTF_8)); for (PublishBatchRequestEntry entry : pmbRequest.getPublishBatchRequestEntries()) { - propagate().inject(span, entry.getMessageAttributes(), SETTER, TracePropagationStyle.XRAY); + Map messageAttributes = entry.getMessageAttributes(); + if (messageAttributes.size() < 10) { + binaryValue.rewind(); + messageAttributes.put( + "_datadog", + new MessageAttributeValue().withDataType("Binary").withBinaryValue(binaryValue)); + } } } return request; diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/TextMapInjectAdapter.java b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/TextMapInjectAdapter.java new file mode 100644 index 00000000000..6f2bb8f888b --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/TextMapInjectAdapter.java @@ -0,0 +1,13 @@ +package datadog.trace.instrumentation.aws.v1.sns; + +import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; + +public class TextMapInjectAdapter implements AgentPropagation.Setter { + + public static final TextMapInjectAdapter SETTER = new TextMapInjectAdapter(); + + @Override + public void set(final StringBuilder builder, final String key, final String value) { + builder.append("\"").append(key).append("\":\"").append(value).append("\","); + } +} diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy index 05a8281f359..c54ee6cba62 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy @@ -5,7 +5,6 @@ import com.amazonaws.services.sns.AmazonSNSClient import com.amazonaws.services.sns.AmazonSNSClientBuilder import datadog.trace.agent.test.naming.VersionedNamingTestBase import datadog.trace.agent.test.utils.TraceUtils -import datadog.trace.api.DDSpanId import datadog.trace.api.DDSpanTypes import datadog.trace.api.config.GeneralConfig import datadog.trace.bootstrap.instrumentation.api.Tags @@ -142,8 +141,14 @@ abstract class SnsClientTest extends VersionedNamingTestBase { and: messageBody["Message"] == "sometext" - messageBody["MessageAttributes"]["X-Amzn-Trace-Id"]["Value"] =~ - /Root=1-[0-9a-f]{8}-00000000${sendSpan.traceId.toHexStringPadded(16)};Parent=${DDSpanId.toHexStringPadded(sendSpan.spanId)};Sampled=1/ + String base64EncodedString = messageBody["MessageAttributes"]["_datadog"]["Value"] + byte[] decodedBytes = base64EncodedString.decodeBase64() + String decodedString = new String(decodedBytes, "UTF-8") + JsonSlurper slurper = new JsonSlurper() + Map traceContextInJson = slurper.parseText(decodedString) + traceContextInJson['x-datadog-trace-id'] == sendSpan.traceId.toString() + traceContextInJson['x-datadog-parent-id'] == sendSpan.spanId.toString() + traceContextInJson['x-datadog-sampling-priority'] == "1" } } @@ -175,7 +180,7 @@ class SnsClientV1ForkedTest extends SnsClientTest { @Override String expectedOperation(String awsService, String awsOperation) { - if (awsService == "SNS"&& awsOperation == "Publish") { + if (awsService == "SNS" && awsOperation == "Publish") { return "aws.sns.send" } return "http.client.request" diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle index d9e25837058..945e4f30c2a 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle @@ -2,7 +2,7 @@ muzzle { pass { group = "software.amazon.awssdk" module = "sns" - versions = "[2.2.0,)" + versions = "[2.2.0,3)" assertInverse = true } } diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java deleted file mode 100644 index ee60d19301b..00000000000 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/MessageAttributeInjector.java +++ /dev/null @@ -1,22 +0,0 @@ -package datadog.trace.instrumentation.aws.v2.sns; - -import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; -import java.util.Map; -import software.amazon.awssdk.services.sns.model.MessageAttributeValue; - -public class MessageAttributeInjector - implements AgentPropagation.Setter> { - - public static final MessageAttributeInjector SETTER = new MessageAttributeInjector(); - - @Override - public void set( - final Map carrier, final String key, final String value) { - // 10 messageAttributes is a limit from SQS, which is often used as a subscriber and therefore - // still apply here - if (carrier.size() < 10 && !carrier.containsKey(key)) { - carrier.put( - key, MessageAttributeValue.builder().dataType("String").stringValue(value).build()); - } - } -} diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsClientInstrumentation.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsClientInstrumentation.java index aad95b7b235..0e24d267baf 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsClientInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsClientInstrumentation.java @@ -16,10 +16,8 @@ @AutoService(InstrumenterModule.class) public final class SnsClientInstrumentation extends InstrumenterModule.Tracing implements Instrumenter.ForSingleType { - private static final String INSTRUMENTATION_NAME = "aws-sdk"; - public SnsClientInstrumentation() { - super(INSTRUMENTATION_NAME); + super("sns", "aws-sdk"); } @Override @@ -36,9 +34,7 @@ public void methodAdvice(MethodTransformer transformer) { @Override public String[] helperClassNames() { - return new String[] { - packageName + ".SnsInterceptor", packageName + ".MessageAttributeInjector" - }; + return new String[] {packageName + ".SnsInterceptor", packageName + ".TextMapInjectAdapter"}; } @Override diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java index e32de673414..9c995f52c39 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java @@ -1,14 +1,16 @@ package datadog.trace.instrumentation.aws.v2.sns; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate; -import static datadog.trace.instrumentation.aws.v2.sns.MessageAttributeInjector.SETTER; +import static datadog.trace.instrumentation.aws.v2.sns.TextMapInjectAdapter.SETTER; import datadog.trace.api.TracePropagationStyle; 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.ExecutionAttribute; @@ -35,16 +37,40 @@ public SdkRequest modifyRequest( Map messageAttributes = new HashMap<>(request.messageAttributes()); final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); - propagate().inject(span, messageAttributes, SETTER, TracePropagationStyle.XRAY); + // 10 messageAttributes is a limit from SQS, which is often used as a subscriber, therefore + // the limit still applies here + if (messageAttributes.size() < 10) { + StringBuilder jsonBuilder = new StringBuilder(); + jsonBuilder.append("{"); + propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); + jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma + jsonBuilder.append("}"); + + messageAttributes.put( + "_datadog", // Use Binary since SNS subscription filter policies fail silently with JSON + // strings https://github.com/DataDog/datadog-lambda-js/pull/269 + MessageAttributeValue.builder() + .dataType("Binary") + .binaryValue(SdkBytes.fromString(jsonBuilder.toString(), StandardCharsets.UTF_8)) + .build()); + } return request.toBuilder().messageAttributes(messageAttributes).build(); } else if (context.request() instanceof PublishBatchRequest) { PublishBatchRequest request = (PublishBatchRequest) context.request(); final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); ArrayList entries = new ArrayList<>(); + StringBuilder jsonBuilder = new StringBuilder(); + jsonBuilder.append("{"); + propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); + jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma + jsonBuilder.append("}"); + SdkBytes binaryValue = SdkBytes.fromString(jsonBuilder.toString(), StandardCharsets.UTF_8); for (PublishBatchRequestEntry entry : request.publishBatchRequestEntries()) { Map messageAttributes = new HashMap<>(entry.messageAttributes()); - propagate().inject(span, messageAttributes, SETTER, TracePropagationStyle.XRAY); + messageAttributes.put( + "_datadog", + MessageAttributeValue.builder().dataType("Binary").binaryValue(binaryValue).build()); entries.add(entry.toBuilder().messageAttributes(messageAttributes).build()); } return request.toBuilder().publishBatchRequestEntries(entries).build(); diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/TextMapInjectAdapter.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/TextMapInjectAdapter.java new file mode 100644 index 00000000000..cfe0368e298 --- /dev/null +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/TextMapInjectAdapter.java @@ -0,0 +1,13 @@ +package datadog.trace.instrumentation.aws.v2.sns; + +import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; + +public class TextMapInjectAdapter implements AgentPropagation.Setter { + + public static final TextMapInjectAdapter SETTER = new TextMapInjectAdapter(); + + @Override + public void set(final StringBuilder builder, final String key, final String value) { + builder.append("\"").append(key).append("\":\"").append(value).append("\","); + } +} diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy index 48e9921296f..f18dc377dcb 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy @@ -1,7 +1,6 @@ import datadog.trace.agent.test.naming.VersionedNamingTestBase import datadog.trace.agent.test.utils.TraceUtils -import datadog.trace.api.DDSpanId import datadog.trace.api.DDSpanTypes import datadog.trace.api.config.GeneralConfig import datadog.trace.bootstrap.instrumentation.api.Tags @@ -11,7 +10,6 @@ import software.amazon.awssdk.auth.credentials.AwsBasicCredentials import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider import software.amazon.awssdk.regions.Region import software.amazon.awssdk.services.sns.SnsClient -import software.amazon.awssdk.services.sns.model.PublishRequest import software.amazon.awssdk.services.sns.model.PublishResponse import software.amazon.awssdk.services.sqs.SqsClient import software.amazon.awssdk.services.sqs.model.QueueAttributeName @@ -123,8 +121,14 @@ abstract class SnsClientTest extends VersionedNamingTestBase { and: messageBody["Message"] == "sometext" - messageBody["MessageAttributes"]["X-Amzn-Trace-Id"]["Value"] =~ - /Root=1-[0-9a-f]{8}-00000000${sendSpan.traceId.toHexStringPadded(16)};Parent=${DDSpanId.toHexStringPadded(sendSpan.spanId)};Sampled=1/ + String base64EncodedString = messageBody["MessageAttributes"]["_datadog"]["Value"] + byte[] decodedBytes = base64EncodedString.decodeBase64() + String decodedString = new String(decodedBytes, "UTF-8") + JsonSlurper slurper = new JsonSlurper() + Map traceContextInJson = slurper.parseText(decodedString) + traceContextInJson['x-datadog-trace-id'] == sendSpan.traceId.toString() + traceContextInJson['x-datadog-parent-id'] == sendSpan.spanId.toString() + traceContextInJson['x-datadog-sampling-priority'] == "1" } } @@ -156,7 +160,7 @@ class SnsClientV1ForkedTest extends SnsClientTest { @Override String expectedOperation(String awsService, String awsOperation) { - if (awsService == "SNS"&& awsOperation == "Publish") { + if (awsService == "Sns" && awsOperation == "Publish") { return "aws.sns.send" } return "http.client.request" From 519206e94719e0d8a295b2933092a6dcf13c7fa3 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Mon, 20 May 2024 14:39:07 -0400 Subject: [PATCH 09/18] fix comments --- .../trace/instrumentation/aws/v1/sns/SnsInterceptor.java | 2 +- .../trace/instrumentation/aws/v2/sns/SnsInterceptor.java | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) 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 index 270a02dda85..d96077184d9 100644 --- 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 @@ -27,7 +27,7 @@ public SnsInterceptor(ContextStore contextSt @Override public AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request) { - // Injecting the AWSTraceHeader into SNS messageAttributes. This is consistent with SQS cases. + // Injecting the trace context into SNS messageAttributes. if (request instanceof PublishRequest) { PublishRequest pRequest = (PublishRequest) request; final AgentSpan span = newSpan(request); diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java index 9c995f52c39..1103f8a84e7 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java @@ -32,6 +32,7 @@ public SnsInterceptor() {} @Override public SdkRequest modifyRequest( Context.ModifyRequest context, ExecutionAttributes executionAttributes) { + // Injecting the trace context into SNS messageAttributes. if (context.request() instanceof PublishRequest) { PublishRequest request = (PublishRequest) context.request(); Map messageAttributes = From b91ad2b725045f4142e8f107af66545827ee8ebe Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Mon, 20 May 2024 15:31:41 -0400 Subject: [PATCH 10/18] no need to rewind --- .../datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java | 1 - 1 file changed, 1 deletion(-) 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 index d96077184d9..14b8c1ea75b 100644 --- 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 @@ -66,7 +66,6 @@ public AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request for (PublishBatchRequestEntry entry : pmbRequest.getPublishBatchRequestEntries()) { Map messageAttributes = entry.getMessageAttributes(); if (messageAttributes.size() < 10) { - binaryValue.rewind(); messageAttributes.put( "_datadog", new MessageAttributeValue().withDataType("Binary").withBinaryValue(binaryValue)); From 5b4a7afaf6df8e0836871fdf8fe0837b6ebee103 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 22 May 2024 11:59:20 -0400 Subject: [PATCH 11/18] fix for muzzle tests --- dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle | 2 +- dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle index 6c8b9068364..caf85ecafe7 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle @@ -2,7 +2,7 @@ muzzle { pass { group = "com.amazonaws" module = "aws-java-sdk-sns" - versions = "[1.12.0,2)" + versions = "[1.12.113,2)" assertInverse = true } } diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle index 945e4f30c2a..9ad3a49aad7 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle @@ -2,7 +2,7 @@ muzzle { pass { group = "software.amazon.awssdk" module = "sns" - versions = "[2.2.0,3)" + versions = "[2.17.84,3)" assertInverse = true } } From 503dcf2c019bacecc2f453633786e356af9f6593 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 22 May 2024 12:37:01 -0400 Subject: [PATCH 12/18] fix for codenarcLatestDepForkedTest --- .../src/test/groovy/SnsClientTest.groovy | 26 +++++++++---------- .../src/test/groovy/SnsClientTest.groovy | 20 +++++++------- 2 files changed, 23 insertions(+), 23 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy index c54ee6cba62..4b716986e81 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy @@ -23,7 +23,7 @@ import static org.testcontainers.containers.localstack.LocalStackContainer.Servi import static org.testcontainers.containers.localstack.LocalStackContainer.Service.SQS abstract class SnsClientTest extends VersionedNamingTestBase { - static final localstack = new LocalStackContainer(DockerImageName.parse("localstack/localstack:latest")) + static final LOCALSTACK = new LocalStackContainer(DockerImageName.parse("localstack/localstack:latest")) .withServices( SQS, SNS) @Shared AmazonSNSClient snsClient @Shared SqsClient sqsClient @@ -33,14 +33,14 @@ abstract class SnsClientTest extends VersionedNamingTestBase { @Shared String testTopicARN def setupSpec() { - localstack.start() + LOCALSTACK.start() snsClient = AmazonSNSClientBuilder.standard() - .withEndpointConfiguration(new AwsClientBuilder.EndpointConfiguration( localstack.getEndpointOverride(SNS).toString(), localstack.getRegion())) + .withEndpointConfiguration(new AwsClientBuilder.EndpointConfiguration( LOCALSTACK.getEndpointOverride(SNS).toString(), LOCALSTACK.getRegion())) .withCredentials( new AWSStaticCredentialsProvider(new BasicAWSCredentials("test", "test"))) .build() sqsClient = SqsClient.builder() - .endpointOverride(localstack.getEndpointOverride(SQS)) - .region(Region.of(localstack.getRegion())) + .endpointOverride(LOCALSTACK.getEndpointOverride(SQS)) + .region(Region.of(LOCALSTACK.getRegion())) .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test", "test"))) .build() testQueueURL = sqsClient.createQueue { it.queueName("testqueue") }.queueUrl() @@ -50,7 +50,7 @@ abstract class SnsClientTest extends VersionedNamingTestBase { } def cleanupSpec() { - localstack.stop() + LOCALSTACK.stop() } @Override @@ -100,14 +100,14 @@ abstract class SnsClientTest extends VersionedNamingTestBase { tags { "$Tags.COMPONENT" "java-aws-sdk" "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT - "$Tags.HTTP_URL" localstack.getEndpointOverride(SNS).toString()+'/' + "$Tags.HTTP_URL" LOCALSTACK.getEndpointOverride(SNS).toString()+'/' "$Tags.HTTP_METHOD" "POST" "$Tags.HTTP_STATUS" 200 - "$Tags.PEER_PORT" localstack.getEndpointOverride(SNS).port - "$Tags.PEER_HOSTNAME" localstack.getEndpointOverride(SNS).host + "$Tags.PEER_PORT" LOCALSTACK.getEndpointOverride(SNS).port + "$Tags.PEER_HOSTNAME" LOCALSTACK.getEndpointOverride(SNS).host "aws.service" "AmazonSNS" "aws_service" "sns" - "aws.endpoint" localstack.getEndpointOverride(SNS).toString() + "aws.endpoint" LOCALSTACK.getEndpointOverride(SNS).toString() "aws.operation" "PublishRequest" "aws.agent" "java-aws-sdk" "aws.topic.name" "testtopic" @@ -128,11 +128,11 @@ abstract class SnsClientTest extends VersionedNamingTestBase { tags { "$Tags.COMPONENT" "apache-httpclient" "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT - "$Tags.HTTP_URL" localstack.getEndpointOverride(SNS).toString()+'/' + "$Tags.HTTP_URL" LOCALSTACK.getEndpointOverride(SNS).toString()+'/' "$Tags.HTTP_METHOD" "POST" "$Tags.HTTP_STATUS" 200 - "$Tags.PEER_PORT" localstack.getEndpointOverride(SNS).port - "$Tags.PEER_HOSTNAME" localstack.getEndpointOverride(SNS).host + "$Tags.PEER_PORT" LOCALSTACK.getEndpointOverride(SNS).port + "$Tags.PEER_HOSTNAME" LOCALSTACK.getEndpointOverride(SNS).host defaultTags(true) } } diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy index f18dc377dcb..4a5ad7fc889 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy @@ -21,7 +21,7 @@ import static org.testcontainers.containers.localstack.LocalStackContainer.Servi import static org.testcontainers.containers.localstack.LocalStackContainer.Service.SQS abstract class SnsClientTest extends VersionedNamingTestBase { - static final localstack = new LocalStackContainer(DockerImageName.parse("localstack/localstack:latest")) + static final LOCALSTACK = new LocalStackContainer(DockerImageName.parse("localstack/localstack:latest")) .withServices( SQS, SNS) @Shared SnsClient snsClient @Shared SqsClient sqsClient @@ -31,15 +31,15 @@ abstract class SnsClientTest extends VersionedNamingTestBase { @Shared String testTopicARN def setupSpec() { - localstack.start() + LOCALSTACK.start() snsClient = SnsClient.builder() - .endpointOverride(localstack.getEndpointOverride(LocalStackContainer.Service.SNS)) - .region(Region.of(localstack.getRegion())) + .endpointOverride(LOCALSTACK.getEndpointOverride(LocalStackContainer.Service.SNS)) + .region(Region.of(LOCALSTACK.getRegion())) .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test", "test"))) .build() sqsClient = SqsClient.builder() - .endpointOverride(localstack.getEndpointOverride(SQS)) - .region(Region.of(localstack.getRegion())) + .endpointOverride(LOCALSTACK.getEndpointOverride(SQS)) + .region(Region.of(LOCALSTACK.getRegion())) .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test", "test"))) .build() testQueueURL = sqsClient.createQueue { it.queueName("testqueue") }.queueUrl() @@ -49,7 +49,7 @@ abstract class SnsClientTest extends VersionedNamingTestBase { } def cleanupSpec() { - localstack.stop() + LOCALSTACK.stop() } @Override @@ -100,11 +100,11 @@ abstract class SnsClientTest extends VersionedNamingTestBase { tags { "$Tags.COMPONENT" "java-aws-sdk" "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT - "$Tags.HTTP_URL" localstack.getEndpointOverride(SNS).toString()+'/' + "$Tags.HTTP_URL" LOCALSTACK.getEndpointOverride(SNS).toString()+'/' "$Tags.HTTP_METHOD" "POST" "$Tags.HTTP_STATUS" 200 - "$Tags.PEER_PORT" localstack.getEndpointOverride(SNS).port - "$Tags.PEER_HOSTNAME" localstack.getEndpointOverride(SNS).host + "$Tags.PEER_PORT" LOCALSTACK.getEndpointOverride(SNS).port + "$Tags.PEER_HOSTNAME" LOCALSTACK.getEndpointOverride(SNS).host "aws.service" "Sns" "aws_service" "Sns" "aws.operation" "Publish" From 876d505d88a2e09d1fbaca9bcf404cb6f2870cab Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Thu, 23 May 2024 09:39:29 -0400 Subject: [PATCH 13/18] use service testcontainersLimit --- dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle | 4 ++++ dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle index caf85ecafe7..fb5880444a2 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle @@ -27,3 +27,7 @@ dependencies { latestDepTestImplementation group: 'com.amazonaws', name: 'aws-java-sdk-sns', version: '+' } + +tasks.withType(Test).configureEach { + usesService(testcontainersLimit) +} diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle index 9ad3a49aad7..cb2c6dbd87f 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle @@ -26,3 +26,7 @@ dependencies { latestDepTestImplementation group: 'software.amazon.awssdk', name: 'sns', version: '+' } + +tasks.withType(Test).configureEach { + usesService(testcontainersLimit) +} From 3d68f7d9b1fcdf8a2022d96755d8f41b3643f8e6 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Thu, 23 May 2024 10:34:10 -0400 Subject: [PATCH 14/18] remove redundant testcontainers dep --- dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle | 1 - dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle | 1 - 2 files changed, 2 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle index fb5880444a2..480ea864e22 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/build.gradle @@ -22,7 +22,6 @@ dependencies { // SQS is used to act as the "Subscriber" of the SNS topic. // There's a problem with sqs sdk v1 with localstack+testcontainers testing. so use sdk v2 for sqs testImplementation 'software.amazon.awssdk:sqs:2.25.40' - testImplementation 'org.testcontainers:testcontainers:1.19.7' testImplementation 'org.testcontainers:localstack:1.19.7' latestDepTestImplementation group: 'com.amazonaws', name: 'aws-java-sdk-sns', version: '+' diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle b/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle index cb2c6dbd87f..41099604996 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/build.gradle @@ -21,7 +21,6 @@ dependencies { testImplementation 'software.amazon.awssdk:sns:2.25.40' // SQS is used to act as the "Subscriber" of the SNS topic. testImplementation 'software.amazon.awssdk:sqs:2.25.40' - testImplementation 'org.testcontainers:testcontainers:1.19.7' testImplementation 'org.testcontainers:localstack:1.19.7' latestDepTestImplementation group: 'software.amazon.awssdk', name: 'sns', version: '+' From c3fd396b2efa627008e1e494d393a9081f7206af Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Fri, 24 May 2024 10:30:52 -0400 Subject: [PATCH 15/18] use GenericContainer (fix for CircleCI) --- .../src/test/groovy/SnsClientTest.groovy | 37 +++++++++++-------- .../src/test/groovy/SnsClientTest.groovy | 30 +++++++++------ 2 files changed, 40 insertions(+), 27 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy index 4b716986e81..20c1fa2d171 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/test/groovy/SnsClientTest.groovy @@ -8,7 +8,6 @@ import datadog.trace.agent.test.utils.TraceUtils import datadog.trace.api.DDSpanTypes import datadog.trace.api.config.GeneralConfig import datadog.trace.bootstrap.instrumentation.api.Tags -import org.testcontainers.containers.localstack.LocalStackContainer import org.testcontainers.utility.DockerImageName import software.amazon.awssdk.auth.credentials.AwsBasicCredentials import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider @@ -18,13 +17,18 @@ import software.amazon.awssdk.services.sqs.model.QueueAttributeName import spock.lang.Shared import groovy.json.JsonSlurper +import java.time.Duration +import org.testcontainers.containers.GenericContainer import static datadog.trace.agent.test.utils.TraceUtils.basicSpan -import static org.testcontainers.containers.localstack.LocalStackContainer.Service.SNS -import static org.testcontainers.containers.localstack.LocalStackContainer.Service.SQS abstract class SnsClientTest extends VersionedNamingTestBase { - static final LOCALSTACK = new LocalStackContainer(DockerImageName.parse("localstack/localstack:latest")) - .withServices( SQS, SNS) + + static final LOCALSTACK = new GenericContainer(DockerImageName.parse("localstack/localstack")) + .withExposedPorts(4566) // Default LocalStack port + .withEnv("SERVICES", "sns,sqs") // Enable SNS and SQS service + .withReuse(true) + .withStartupTimeout(Duration.ofSeconds(120)) + @Shared AmazonSNSClient snsClient @Shared SqsClient sqsClient @@ -32,15 +36,17 @@ abstract class SnsClientTest extends VersionedNamingTestBase { @Shared String testQueueARN @Shared String testTopicARN + def setupSpec() { LOCALSTACK.start() + def endPoint = "http://" + LOCALSTACK.getHost() + ":" + LOCALSTACK.getMappedPort(4566) snsClient = AmazonSNSClientBuilder.standard() - .withEndpointConfiguration(new AwsClientBuilder.EndpointConfiguration( LOCALSTACK.getEndpointOverride(SNS).toString(), LOCALSTACK.getRegion())) + .withEndpointConfiguration(new AwsClientBuilder.EndpointConfiguration(endPoint, "us-east-1")) .withCredentials( new AWSStaticCredentialsProvider(new BasicAWSCredentials("test", "test"))) .build() sqsClient = SqsClient.builder() - .endpointOverride(LOCALSTACK.getEndpointOverride(SQS)) - .region(Region.of(LOCALSTACK.getRegion())) + .endpointOverride(URI.create(endPoint)) + .region(Region.of("us-east-1")) .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test", "test"))) .build() testQueueURL = sqsClient.createQueue { it.queueName("testqueue") }.queueUrl() @@ -83,6 +89,7 @@ abstract class SnsClientTest extends VersionedNamingTestBase { def message = sqsClient.receiveMessage {it.queueUrl(testQueueURL).waitTimeSeconds(3)}.messages().get(0) def jsonSlurper = new JsonSlurper() def messageBody = jsonSlurper.parseText(message.body()) + def endPoint = "http://" + LOCALSTACK.getHost() + ":" + LOCALSTACK.getMappedPort(4566) then: def sendSpan @@ -100,14 +107,14 @@ abstract class SnsClientTest extends VersionedNamingTestBase { tags { "$Tags.COMPONENT" "java-aws-sdk" "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT - "$Tags.HTTP_URL" LOCALSTACK.getEndpointOverride(SNS).toString()+'/' + "$Tags.HTTP_URL" endPoint+'/' "$Tags.HTTP_METHOD" "POST" "$Tags.HTTP_STATUS" 200 - "$Tags.PEER_PORT" LOCALSTACK.getEndpointOverride(SNS).port - "$Tags.PEER_HOSTNAME" LOCALSTACK.getEndpointOverride(SNS).host + "$Tags.PEER_PORT" LOCALSTACK.getMappedPort(4566) + "$Tags.PEER_HOSTNAME" LOCALSTACK.getHost() "aws.service" "AmazonSNS" "aws_service" "sns" - "aws.endpoint" LOCALSTACK.getEndpointOverride(SNS).toString() + "aws.endpoint" endPoint "aws.operation" "PublishRequest" "aws.agent" "java-aws-sdk" "aws.topic.name" "testtopic" @@ -128,11 +135,11 @@ abstract class SnsClientTest extends VersionedNamingTestBase { tags { "$Tags.COMPONENT" "apache-httpclient" "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT - "$Tags.HTTP_URL" LOCALSTACK.getEndpointOverride(SNS).toString()+'/' + "$Tags.HTTP_URL" endPoint+'/' "$Tags.HTTP_METHOD" "POST" "$Tags.HTTP_STATUS" 200 - "$Tags.PEER_PORT" LOCALSTACK.getEndpointOverride(SNS).port - "$Tags.PEER_HOSTNAME" LOCALSTACK.getEndpointOverride(SNS).host + "$Tags.PEER_PORT" LOCALSTACK.getMappedPort(4566) + "$Tags.PEER_HOSTNAME" LOCALSTACK.getHost() defaultTags(true) } } diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy index 4a5ad7fc889..e65668ba873 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy @@ -4,7 +4,7 @@ import datadog.trace.agent.test.utils.TraceUtils import datadog.trace.api.DDSpanTypes import datadog.trace.api.config.GeneralConfig import datadog.trace.bootstrap.instrumentation.api.Tags -import org.testcontainers.containers.localstack.LocalStackContainer +import org.testcontainers.containers.GenericContainer import org.testcontainers.utility.DockerImageName import software.amazon.awssdk.auth.credentials.AwsBasicCredentials import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider @@ -16,13 +16,17 @@ import software.amazon.awssdk.services.sqs.model.QueueAttributeName import spock.lang.Shared import groovy.json.JsonSlurper +import java.time.Duration + import static datadog.trace.agent.test.utils.TraceUtils.basicSpan -import static org.testcontainers.containers.localstack.LocalStackContainer.Service.SNS -import static org.testcontainers.containers.localstack.LocalStackContainer.Service.SQS abstract class SnsClientTest extends VersionedNamingTestBase { - static final LOCALSTACK = new LocalStackContainer(DockerImageName.parse("localstack/localstack:latest")) - .withServices( SQS, SNS) + static final LOCALSTACK = new GenericContainer(DockerImageName.parse("localstack/localstack")) + .withExposedPorts(4566) // Default LocalStack port + .withEnv("SERVICES", "sns,sqs") // Enable SNS and SQS service + .withReuse(true) + .withStartupTimeout(Duration.ofSeconds(120)) + @Shared SnsClient snsClient @Shared SqsClient sqsClient @@ -32,14 +36,15 @@ abstract class SnsClientTest extends VersionedNamingTestBase { def setupSpec() { LOCALSTACK.start() + def endPoint = "http://" + LOCALSTACK.getHost() + ":" + LOCALSTACK.getMappedPort(4566) snsClient = SnsClient.builder() - .endpointOverride(LOCALSTACK.getEndpointOverride(LocalStackContainer.Service.SNS)) - .region(Region.of(LOCALSTACK.getRegion())) + .endpointOverride(URI.create(endPoint)) + .region(Region.of("us-east-1")) .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test", "test"))) .build() sqsClient = SqsClient.builder() - .endpointOverride(LOCALSTACK.getEndpointOverride(SQS)) - .region(Region.of(LOCALSTACK.getRegion())) + .endpointOverride(URI.create(endPoint)) + .region(Region.of("us-east-1")) .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create("test", "test"))) .build() testQueueURL = sqsClient.createQueue { it.queueName("testqueue") }.queueUrl() @@ -83,6 +88,7 @@ abstract class SnsClientTest extends VersionedNamingTestBase { def message = sqsClient.receiveMessage {it.queueUrl(testQueueURL).waitTimeSeconds(3)}.messages().get(0) def jsonSlurper = new JsonSlurper() def messageBody = jsonSlurper.parseText(message.body()) + def endPoint = "http://" + LOCALSTACK.getHost() + ":" + LOCALSTACK.getMappedPort(4566) then: def sendSpan @@ -100,11 +106,11 @@ abstract class SnsClientTest extends VersionedNamingTestBase { tags { "$Tags.COMPONENT" "java-aws-sdk" "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT - "$Tags.HTTP_URL" LOCALSTACK.getEndpointOverride(SNS).toString()+'/' + "$Tags.HTTP_URL" endPoint+'/' "$Tags.HTTP_METHOD" "POST" "$Tags.HTTP_STATUS" 200 - "$Tags.PEER_PORT" LOCALSTACK.getEndpointOverride(SNS).port - "$Tags.PEER_HOSTNAME" LOCALSTACK.getEndpointOverride(SNS).host + "$Tags.PEER_PORT" LOCALSTACK.getMappedPort(4566) + "$Tags.PEER_HOSTNAME" LOCALSTACK.getHost() "aws.service" "Sns" "aws_service" "Sns" "aws.operation" "Publish" From 42a0ee3208b55a38658fd032c78a69a6e50c70ec Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Fri, 24 May 2024 11:30:11 -0400 Subject: [PATCH 16/18] fix --- .../aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy index e65668ba873..5db1c312636 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/test/groovy/SnsClientTest.groovy @@ -166,7 +166,7 @@ class SnsClientV1ForkedTest extends SnsClientTest { @Override String expectedOperation(String awsService, String awsOperation) { - if (awsService == "Sns" && awsOperation == "Publish") { + if (awsService == "SNS" && awsOperation == "Publish") { return "aws.sns.send" } return "http.client.request" From 345f066bfff5fabbf5159877ed43da51843d5178 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Thu, 30 May 2024 23:58:05 -0400 Subject: [PATCH 17/18] refactoring --- .../aws/v1/sns/SnsInterceptor.java | 37 ++++++++++--------- .../aws/v2/sns/SnsInterceptor.java | 36 ++++++++++-------- 2 files changed, 40 insertions(+), 33 deletions(-) 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 index 14b8c1ea75b..f6d04e0cebe 100644 --- 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 @@ -20,28 +20,38 @@ public class SnsInterceptor extends RequestHandler2 { private final ContextStore contextStore; + private ByteBuffer messageAttributeValueToInject; public SnsInterceptor(ContextStore contextStore) { this.contextStore = contextStore; } + private ByteBuffer getMessageAttributeValueToInject(AmazonWebServiceRequest request) { + if (this.messageAttributeValueToInject == null) { + final AgentSpan span = newSpan(request); + StringBuilder jsonBuilder = new StringBuilder(); + jsonBuilder.append("{"); + propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); + jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma + jsonBuilder.append("}"); + this.messageAttributeValueToInject = + ByteBuffer.wrap(jsonBuilder.toString().getBytes(StandardCharsets.UTF_8)); + } + + return this.messageAttributeValueToInject; + } + @Override public AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request) { // Injecting the trace context into SNS messageAttributes. if (request instanceof PublishRequest) { PublishRequest pRequest = (PublishRequest) request; - final AgentSpan span = newSpan(request); // note: modifying message attributes has to be done before marshalling, otherwise the changes // are not reflected in the actual request (and the MD5 check on send will fail). Map messageAttributes = pRequest.getMessageAttributes(); // 10 messageAttributes is a limit from SQS, which is often used as a subscriber, therefore // the limit still applies here if (messageAttributes.size() < 10) { - StringBuilder jsonBuilder = new StringBuilder(); - jsonBuilder.append("{"); - propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); - jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma - jsonBuilder.append("}"); messageAttributes.put( "_datadog", new MessageAttributeValue() @@ -49,26 +59,19 @@ public AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request "Binary") // Use Binary since SNS subscription filter policies fail silently // with JSON strings // https://github.com/DataDog/datadog-lambda-js/pull/269 - .withBinaryValue( - ByteBuffer.wrap(jsonBuilder.toString().getBytes(StandardCharsets.UTF_8)))); + .withBinaryValue(this.getMessageAttributeValueToInject(request))); } } else if (request instanceof PublishBatchRequest) { PublishBatchRequest pmbRequest = (PublishBatchRequest) request; - final AgentSpan span = newSpan(request); - StringBuilder jsonBuilder = new StringBuilder(); - jsonBuilder.append("{"); - propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); - jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma - jsonBuilder.append("}"); - ByteBuffer binaryValue = - ByteBuffer.wrap(jsonBuilder.toString().getBytes(StandardCharsets.UTF_8)); for (PublishBatchRequestEntry entry : pmbRequest.getPublishBatchRequestEntries()) { Map messageAttributes = entry.getMessageAttributes(); if (messageAttributes.size() < 10) { messageAttributes.put( "_datadog", - new MessageAttributeValue().withDataType("Binary").withBinaryValue(binaryValue)); + new MessageAttributeValue() + .withDataType("Binary") + .withBinaryValue(this.getMessageAttributeValueToInject(request))); } } } diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java index 1103f8a84e7..bc41348a87c 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java @@ -26,6 +26,21 @@ public class SnsInterceptor implements ExecutionInterceptor { public static final ExecutionAttribute SPAN_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) .putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan")); + private SdkBytes messageAttributeValueToInject; + + private SdkBytes getMessageAttributeValueToInject(ExecutionAttributes executionAttributes) { + if (this.messageAttributeValueToInject == null) { + final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + StringBuilder jsonBuilder = new StringBuilder(); + jsonBuilder.append("{"); + propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); + jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma + jsonBuilder.append("}"); + this.messageAttributeValueToInject = + SdkBytes.fromString(jsonBuilder.toString(), StandardCharsets.UTF_8); + } + return this.messageAttributeValueToInject; + } public SnsInterceptor() {} @@ -37,41 +52,30 @@ public SdkRequest modifyRequest( PublishRequest request = (PublishRequest) context.request(); Map messageAttributes = new HashMap<>(request.messageAttributes()); - final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); // 10 messageAttributes is a limit from SQS, which is often used as a subscriber, therefore // the limit still applies here if (messageAttributes.size() < 10) { - StringBuilder jsonBuilder = new StringBuilder(); - jsonBuilder.append("{"); - propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); - jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma - jsonBuilder.append("}"); - messageAttributes.put( "_datadog", // Use Binary since SNS subscription filter policies fail silently with JSON // strings https://github.com/DataDog/datadog-lambda-js/pull/269 MessageAttributeValue.builder() .dataType("Binary") - .binaryValue(SdkBytes.fromString(jsonBuilder.toString(), StandardCharsets.UTF_8)) + .binaryValue(this.getMessageAttributeValueToInject(executionAttributes)) .build()); } return request.toBuilder().messageAttributes(messageAttributes).build(); } else if (context.request() instanceof PublishBatchRequest) { PublishBatchRequest request = (PublishBatchRequest) context.request(); - final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); ArrayList entries = new ArrayList<>(); - StringBuilder jsonBuilder = new StringBuilder(); - jsonBuilder.append("{"); - propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); - jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma - jsonBuilder.append("}"); - SdkBytes binaryValue = SdkBytes.fromString(jsonBuilder.toString(), StandardCharsets.UTF_8); for (PublishBatchRequestEntry entry : request.publishBatchRequestEntries()) { Map messageAttributes = new HashMap<>(entry.messageAttributes()); messageAttributes.put( "_datadog", - MessageAttributeValue.builder().dataType("Binary").binaryValue(binaryValue).build()); + MessageAttributeValue.builder() + .dataType("Binary") + .binaryValue(this.getMessageAttributeValueToInject(executionAttributes)) + .build()); entries.add(entry.toBuilder().messageAttributes(messageAttributes).build()); } return request.toBuilder().publishBatchRequestEntries(entries).build(); From d32837dc5fcbd0c370f907fd52c7364ef64cc9d1 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Fri, 31 May 2024 00:09:39 -0400 Subject: [PATCH 18/18] refactoring (not caching the byte values to be sure) --- .../aws/v1/sns/SnsInterceptor.java | 26 +++++++------------ .../aws/v2/sns/SnsInterceptor.java | 25 +++++++----------- 2 files changed, 18 insertions(+), 33 deletions(-) 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 index f6d04e0cebe..4912c89f15b 100644 --- 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 @@ -20,25 +20,19 @@ public class SnsInterceptor extends RequestHandler2 { private final ContextStore contextStore; - private ByteBuffer messageAttributeValueToInject; public SnsInterceptor(ContextStore contextStore) { this.contextStore = contextStore; } private ByteBuffer getMessageAttributeValueToInject(AmazonWebServiceRequest request) { - if (this.messageAttributeValueToInject == null) { - final AgentSpan span = newSpan(request); - StringBuilder jsonBuilder = new StringBuilder(); - jsonBuilder.append("{"); - propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); - jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma - jsonBuilder.append("}"); - this.messageAttributeValueToInject = - ByteBuffer.wrap(jsonBuilder.toString().getBytes(StandardCharsets.UTF_8)); - } - - return this.messageAttributeValueToInject; + final AgentSpan span = newSpan(request); + StringBuilder jsonBuilder = new StringBuilder(); + jsonBuilder.append("{"); + propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); + jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma + jsonBuilder.append("}"); + return ByteBuffer.wrap(jsonBuilder.toString().getBytes(StandardCharsets.UTF_8)); } @Override @@ -63,15 +57,13 @@ public AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request } } else if (request instanceof PublishBatchRequest) { PublishBatchRequest pmbRequest = (PublishBatchRequest) request; - + final ByteBuffer bytebuffer = this.getMessageAttributeValueToInject(request); for (PublishBatchRequestEntry entry : pmbRequest.getPublishBatchRequestEntries()) { Map messageAttributes = entry.getMessageAttributes(); if (messageAttributes.size() < 10) { messageAttributes.put( "_datadog", - new MessageAttributeValue() - .withDataType("Binary") - .withBinaryValue(this.getMessageAttributeValueToInject(request))); + new MessageAttributeValue().withDataType("Binary").withBinaryValue(bytebuffer)); } } } diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java index bc41348a87c..8bba0b41e48 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java @@ -26,20 +26,15 @@ public class SnsInterceptor implements ExecutionInterceptor { public static final ExecutionAttribute SPAN_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) .putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan")); - private SdkBytes messageAttributeValueToInject; private SdkBytes getMessageAttributeValueToInject(ExecutionAttributes executionAttributes) { - if (this.messageAttributeValueToInject == null) { - final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); - StringBuilder jsonBuilder = new StringBuilder(); - jsonBuilder.append("{"); - propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); - jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma - jsonBuilder.append("}"); - this.messageAttributeValueToInject = - SdkBytes.fromString(jsonBuilder.toString(), StandardCharsets.UTF_8); - } - return this.messageAttributeValueToInject; + final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); + StringBuilder jsonBuilder = new StringBuilder(); + jsonBuilder.append("{"); + propagate().inject(span, jsonBuilder, SETTER, TracePropagationStyle.DATADOG); + jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma + jsonBuilder.append("}"); + return SdkBytes.fromString(jsonBuilder.toString(), StandardCharsets.UTF_8); } public SnsInterceptor() {} @@ -67,15 +62,13 @@ public SdkRequest modifyRequest( } else if (context.request() instanceof PublishBatchRequest) { PublishBatchRequest request = (PublishBatchRequest) context.request(); ArrayList entries = new ArrayList<>(); + final SdkBytes sdkBytes = this.getMessageAttributeValueToInject(executionAttributes); for (PublishBatchRequestEntry entry : request.publishBatchRequestEntries()) { Map messageAttributes = new HashMap<>(entry.messageAttributes()); messageAttributes.put( "_datadog", - MessageAttributeValue.builder() - .dataType("Binary") - .binaryValue(this.getMessageAttributeValueToInject(executionAttributes)) - .build()); + MessageAttributeValue.builder().dataType("Binary").binaryValue(sdkBytes).build()); entries.add(entry.toBuilder().messageAttributes(messageAttributes).build()); } return request.toBuilder().publishBatchRequestEntries(entries).build();