Skip to content

Commit

Permalink
Modifying instrumentation to send other attrs
Browse files Browse the repository at this point in the history
  • Loading branch information
meiao committed Jun 21, 2024
1 parent 4d4fac2 commit 084bb2a
Show file tree
Hide file tree
Showing 24 changed files with 735 additions and 227 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -20,28 +20,28 @@
import com.newrelic.api.agent.weaver.MatchType;
import com.newrelic.api.agent.weaver.Weave;
import com.newrelic.api.agent.weaver.Weaver;
import com.newrelic.utils.Util;
import com.newrelic.utils.MetricUtil;

@Weave(type = MatchType.Interface, originalName = "com.amazonaws.services.sqs.AmazonSQS")
public class AmazonSQS_Instrumentation {

@Trace
public SendMessageBatchResult sendMessageBatch(SendMessageBatchRequest sendMessageBatchRequest) {
MessageProduceParameters messageProduceParameters = Util.generateExternalProduceMetrics(sendMessageBatchRequest.getQueueUrl());
MessageProduceParameters messageProduceParameters = MetricUtil.generateExternalProduceMetrics(sendMessageBatchRequest.getQueueUrl());
NewRelic.getAgent().getTracedMethod().reportAsExternal(messageProduceParameters);
return Weaver.callOriginal();
}

@Trace
public SendMessageResult sendMessage(SendMessageRequest sendMessageRequest) {
MessageProduceParameters messageProduceParameters = Util.generateExternalProduceMetrics(sendMessageRequest.getQueueUrl());
MessageProduceParameters messageProduceParameters = MetricUtil.generateExternalProduceMetrics(sendMessageRequest.getQueueUrl());
NewRelic.getAgent().getTracedMethod().reportAsExternal(messageProduceParameters);
return Weaver.callOriginal();
}

@Trace
public ReceiveMessageResult receiveMessage(ReceiveMessageRequest receiveMessageRequest) {
MessageConsumeParameters messageConsumeParameters = Util.generateExternalConsumeMetrics(receiveMessageRequest.getQueueUrl());
MessageConsumeParameters messageConsumeParameters = MetricUtil.generateExternalConsumeMetrics(receiveMessageRequest.getQueueUrl());
NewRelic.getAgent().getTracedMethod().reportAsExternal(messageConsumeParameters);
return Weaver.callOriginal();
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
/*
*
* * Copyright 2024 New Relic Corporation. All rights reserved.
* * SPDX-License-Identifier: Apache-2.0
*
*/

package com.newrelic.utils;

import com.newrelic.agent.bridge.AgentBridge;

import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

class DestinationData {
private final String region;
private final String accountId;
private final String queueName;

private static final DestinationData UNKNOWN = new DestinationData(null, null, "unknown");
private static final Pattern QUEUE_URL_PATTERN = Pattern.compile(
"https://sqs.(?<region>[^.]+).amazonaws.com/(?<accountId>[^/]+)/(?<queueName>.*)");
private static final Map<String, DestinationData> CACHE = AgentBridge.collectionFactory
.createConcurrentTimeBasedEvictionMap(3600 * 8); // 8 hours

private DestinationData(String region, String accountId, String queueName) {
this.region = region;
this.accountId = accountId;
this.queueName = queueName;
}

public String getRegion() {
return region;
}

public String getAccountId() {
return accountId;
}

public String getQueueName() {
return queueName;
}

public static DestinationData parse(String queueUrl) {
return CACHE.computeIfAbsent(queueUrl, DestinationData::doParse);
}

private static DestinationData doParse(String queueUrl) {
Matcher matcher = QUEUE_URL_PATTERN.matcher(queueUrl);
if (matcher.matches()) {
return new DestinationData(
matcher.group("region"),
matcher.group("accountId"),
matcher.group("queueName")
);
} else {
int index = queueUrl.lastIndexOf('/');
if (index <= 0) {
return UNKNOWN;
}
String queueName = queueUrl.substring(index + 1);
return new DestinationData(null, null, queueName);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -14,37 +14,35 @@
import com.newrelic.api.agent.Segment;
import com.newrelic.api.agent.weaver.Weaver;

public class Util {
public class MetricUtil {

public static final String LIBRARY = "SQS";
public static final String UNKOWN = "unknown";
public static final String OTEL_LIBRARY = "aws_sqs";

public static MessageProduceParameters generateExternalProduceMetrics(String queueUrl) {
String queueName = UNKOWN;
int index = queueUrl.lastIndexOf('/');
if (index > 0) {
queueName = queueUrl.substring(index + 1);
}
DestinationData destinationData = DestinationData.parse(queueUrl);

MessageProduceParameters params = MessageProduceParameters
.library(LIBRARY)
.library(LIBRARY, OTEL_LIBRARY)
.destinationType(DestinationType.NAMED_QUEUE)
.destinationName(queueName)
.destinationName(destinationData.getQueueName())
.outboundHeaders(null)
.cloudRegion(destinationData.getRegion())
.cloudAccountId(destinationData.getAccountId())
.build();
return params;
}

public static MessageConsumeParameters generateExternalConsumeMetrics(String queueUrl) {
String queueName = UNKOWN;
int index = queueUrl.lastIndexOf('/');
if (index > 0) {
queueName = queueUrl.substring(index + 1);
}
DestinationData destinationData = DestinationData.parse(queueUrl);

MessageConsumeParameters params = MessageConsumeParameters
.library(LIBRARY)
.library(LIBRARY, OTEL_LIBRARY)
.destinationType(DestinationType.NAMED_QUEUE)
.destinationName(queueName)
.destinationName(destinationData.getQueueName())
.inboundHeaders(null)
.cloudRegion(destinationData.getRegion())
.cloudAccountId(destinationData.getAccountId())
.build();
return params;
}
Expand All @@ -56,4 +54,5 @@ public static void finishSegment(Segment segment) {
AgentBridge.instrumentation.noticeInstrumentationError(t, Weaver.getImplementationTitle());
}
}

}

This file was deleted.

This file was deleted.

Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,10 @@
*
*/

package com.amazonaws.services.sqs;

import com.amazonaws.auth.BasicAWSCredentials;
import com.amazonaws.internal.StaticCredentialsProvider;
import com.amazonaws.services.sqs.AmazonSQSClient;
import com.amazonaws.services.sqs.model.CreateQueueRequest;
import com.amazonaws.services.sqs.model.CreateQueueResult;
import com.amazonaws.services.sqs.model.ReceiveMessageRequest;
Expand All @@ -23,13 +24,15 @@
import org.elasticmq.rest.sqs.SQSRestServerBuilder;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Ignore;
import org.junit.Test;
import org.junit.runner.RunWith;

import static org.junit.Assert.assertEquals;

@RunWith(InstrumentationTestRunner.class)
@InstrumentationTestConfig(includePrefixes = { "com.amazonaws.services.sqs" })
@InstrumentationTestConfig(includePrefixes = { "com.amazonaws.services.sqs" }, configName = "dt_enabled.yml")
@Ignore("This tests uses a new API that is obscured by the security API. When the security API is updated, this test can run again.")
public class SqsClientTest {

private static AmazonSQSClient sqsClient;
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
/*
*
* * Copyright 2024 New Relic Corporation. All rights reserved.
* * SPDX-License-Identifier: Apache-2.0
*
*/

package com.newrelic.utils;

import org.junit.Test;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;

public class DestinationDataTest {

@Test
public void parse() {
DestinationData data = DestinationData.parse("https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue");

assertNotNull(data);
assertEquals("us-east-2", data.getRegion());
assertEquals("123456789012", data.getAccountId());
assertEquals("MyQueue", data.getQueueName());
}

@Test
public void parse_bogus() {
DestinationData data = DestinationData.parse("bogus");

assertNotNull(data);
assertNull(data.getRegion());
assertNull(data.getAccountId());
assertEquals("unknown", data.getQueueName());
}

@Test
public void parse_nonMatchingUrl() {
DestinationData data = DestinationData.parse("https://not-sqs.us-east-2.amazonaws.com/123456789012/MyQueue");

assertNotNull(data);
assertNull(data.getRegion());
assertNull(data.getAccountId());
assertEquals("MyQueue", data.getQueueName());
}

}
Loading

0 comments on commit 084bb2a

Please sign in to comment.