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