Skip to content

Commit

Permalink
aws-lambda: use Cloud API
Browse files Browse the repository at this point in the history
  • Loading branch information
@meiao
meiao committed Oct 30, 2024
1 parent ec234dd commit fc8cf8f
Show file tree
Hide file tree
Showing 10 changed files with 265 additions and 68 deletions.
21 changes: 17 additions & 4 deletions .../src/main/java/com/agent/instrumentation/awsjavasdk1/services/lambda/FunctionRawData.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

package com.agent.instrumentation.awsjavasdk1.services.lambda;

import java.lang.ref.WeakReference;
import java.util.Objects;

/**
Expand All @@ -16,11 +17,13 @@ public class FunctionRawData {
private final String functionRef;
private final String qualifier;
private final String region;
private final WeakReference<Object> sdkClient;

public FunctionRawData(String functionRef, String qualifier, String region) {
public FunctionRawData(String functionRef, String qualifier, String region, Object sdkClient) {
this.functionRef = functionRef;
this.qualifier = qualifier;
this.region = region;
this.sdkClient = new WeakReference<>(sdkClient);
}

public String getFunctionRef() {
Expand All @@ -35,6 +38,10 @@ public String getRegion() {
return region;
}

public Object getSdkClient() {
return sdkClient.get();
}

@Override
public boolean equals(Object o) {
if (this == o) {
Expand All @@ -43,13 +50,19 @@ public boolean equals(Object o) {
if (!(o instanceof FunctionRawData)) {
return false;
}

FunctionRawData that = (FunctionRawData) o;
return Objects.equals(functionRef, that.functionRef) && Objects.equals(qualifier, that.qualifier) &&
Objects.equals(region, that.region);
if (this.sdkClient.get() == null || that.sdkClient.get() == null) {
return false;
}
return Objects.equals(functionRef, that.functionRef) &&
Objects.equals(qualifier, that.qualifier) &&
Objects.equals(region, that.region) &&
Objects.equals(sdkClient.get(), that.sdkClient.get());
}

@Override
public int hashCode() {
return Objects.hash(functionRef, qualifier, region);
return Objects.hash(functionRef, qualifier, region, sdkClient.get());
}
}
28 changes: 24 additions & 4 deletions ...1.280/src/main/java/com/agent/instrumentation/awsjavasdk1/services/lambda/LambdaUtil.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
package com.agent.instrumentation.awsjavasdk1.services.lambda;

import com.newrelic.agent.bridge.AgentBridge;
import com.newrelic.api.agent.CloudAccountInfo;
import com.newrelic.api.agent.CloudParameters;

import java.util.function.Function;
Expand Down Expand Up @@ -60,13 +61,26 @@ static FunctionProcessedData processData(FunctionRawData data) {
String arn = NULL_ARN;

if (parts.length == 1) {
// function ref is only function name
// does not have the account id, so cannot assemble the ARN.
// function name: {function-name}
String accountId = getAccountId(data.getSdkClient());
if (accountId != null) {
String qualifier = data.getQualifier();
if (qualifier == null) {
arn = PREFIX + data.getRegion() + ":" + accountId + ":function:" + functionRef;
} else {
arn = PREFIX + data.getRegion() + ":" + accountId + ":function:" + functionRef + ":" + qualifier;
}
}
functionName = functionRef;

} else if (parts.length == 2) {
// function ref is only function name with alias/version
// does not have the account id, so cannot assemble the ARN.
// function name + qualifier: {function-name}:{qualifier}
String accountId = getAccountId(data.getSdkClient());
if (accountId != null) {
arn = PREFIX + data.getRegion() + ":" + accountId + ":function:" + functionRef;
}
functionName = parts[0];

} else if (parts.length == 3) {
// partial ARN: {account-id}:function:{function-name}
functionName = parts[2];
Expand All @@ -76,10 +90,12 @@ static FunctionProcessedData processData(FunctionRawData data) {
} else {
arn = PREFIX + data.getRegion() + ":" + functionRef + ":" + qualifier;
}

} else if (parts.length == 4) {
// partial ARN with qualifier: {account-id}:function:{function-name}:{qualifier}
functionName = parts[2];
arn = PREFIX + data.getRegion() + ":" + functionRef;

} else if (parts.length == 7) {
// full ARN: arn:aws:lambda:{region}:{account-id}:function:{function-name}
functionName = parts[6];
Expand All @@ -89,6 +105,7 @@ static FunctionProcessedData processData(FunctionRawData data) {
} else {
arn = functionRef + ":" + qualifier;
}

} else if (parts.length == 8) {
// full ARN with qualifier: arn:aws:lambda:{region}:{account-id}:function:{function-name}:{qualifier}
functionName = parts[6];
Expand All @@ -99,6 +116,9 @@ static FunctionProcessedData processData(FunctionRawData data) {
return new FunctionProcessedData(functionName, arn);
}

private static String getAccountId(Object sdkClient) {
return AgentBridge.cloud.getAccountInfo(sdkClient, CloudAccountInfo.AWS_ACCOUNT_ID);
}

public static String getSimpleFunctionName(FunctionRawData functionRawData) {
return CACHE.apply(functionRawData).getFunctionName();
Expand Down
2 changes: 1 addition & 1 deletion ...280/src/main/java/com/amazonaws/services/lambda/AWSLambdaAsyncClient_Instrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ public abstract class AWSLambdaAsyncClient_Instrumentation {
protected abstract String getSigningRegion();

public Future<InvokeResult> invokeAsync(final InvokeRequest request, AsyncHandler<InvokeRequest, InvokeResult> asyncHandler) {
FunctionRawData functionRawData = new FunctionRawData(request.getFunctionName(), request.getQualifier(), getSigningRegion());
FunctionRawData functionRawData = new FunctionRawData(request.getFunctionName(), request.getQualifier(), getSigningRegion(), this);
CloudParameters cloudParameters = LambdaUtil.getCloudParameters(functionRawData);
String functionName = LambdaUtil.getSimpleFunctionName(functionRawData);
Segment segment = NewRelic.getAgent().getTransaction().startSegment("Lambda", "invoke/" + functionName);
Expand Down
3 changes: 1 addition & 2 deletions ...1.11.280/src/main/java/com/amazonaws/services/lambda/AWSLambdaClient_Instrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
import com.newrelic.api.agent.NewRelic;
import com.newrelic.api.agent.Trace;
import com.newrelic.api.agent.TracedMethod;
import com.newrelic.api.agent.weaver.CatchAndLog;
import com.newrelic.api.agent.weaver.MatchType;
import com.newrelic.api.agent.weaver.Weave;
import com.newrelic.api.agent.weaver.Weaver;
Expand All @@ -27,7 +26,7 @@ public abstract class AWSLambdaClient_Instrumentation {

@Trace(leaf = true)
public InvokeResult invoke(InvokeRequest invokeRequest) {
FunctionRawData functionRawData = new FunctionRawData(invokeRequest.getFunctionName(), invokeRequest.getQualifier(), getSigningRegion());
FunctionRawData functionRawData = new FunctionRawData(invokeRequest.getFunctionName(), invokeRequest.getQualifier(), getSigningRegion(), this);
CloudParameters cloudParameters = LambdaUtil.getCloudParameters(functionRawData);
TracedMethod tracedMethod = NewRelic.getAgent().getTracedMethod();
tracedMethod.reportAsExternal(cloudParameters);
Expand Down
107 changes: 85 additions & 22 deletions ...0/src/test/java/com/agent/instrumentation/awsjavasdk1/services/lambda/LambdaUtilTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,19 +8,38 @@
package com.agent.instrumentation.awsjavasdk1.services.lambda;

import com.amazonaws.regions.Regions;
import com.amazonaws.services.lambda.model.InvokeRequest;
import com.newrelic.agent.bridge.AgentBridge;
import com.newrelic.agent.bridge.CloudApi;
import com.newrelic.agent.bridge.NoOpCloud;
import com.newrelic.api.agent.CloudAccountInfo;
import com.newrelic.api.agent.CloudParameters;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

public class LambdaUtilTest {

@Before
public void before() {
AgentBridge.cloud = mock(CloudApi.class);
}

@After
public void after() {
AgentBridge.cloud = NoOpCloud.INSTANCE;
}

@Test
public void testGetCloudParamFunctionName() {
FunctionRawData functionRawData = new FunctionRawData("my-function", null, getRegion());
FunctionRawData functionRawData = data("my-function", null);
CloudParameters cloudParameters = LambdaUtil.getCloudParameters(functionRawData);
assertNotNull(cloudParameters);
assertEquals("aws_lambda", cloudParameters.getPlatform());
Expand All @@ -29,7 +48,7 @@ public void testGetCloudParamFunctionName() {

@Test
public void testGetCloudParamPartialArn() {
FunctionRawData functionRawData = new FunctionRawData("123456789012:function:my-function", null, getRegion());
FunctionRawData functionRawData = data("123456789012:function:my-function", null);
CloudParameters cloudParameters = LambdaUtil.getCloudParameters(functionRawData);
assertNotNull(cloudParameters);
assertEquals("aws_lambda", cloudParameters.getPlatform());
Expand All @@ -38,7 +57,7 @@ public void testGetCloudParamPartialArn() {

@Test
public void testGetCloudParamArnQualifier() {
FunctionRawData functionRawData = new FunctionRawData("arn:aws:lambda:us-east-1:123456789012:function:my-function", "alias", getRegion());
FunctionRawData functionRawData = data("arn:aws:lambda:us-east-1:123456789012:function:my-function", "alias");
CloudParameters cloudParameters = LambdaUtil.getCloudParameters(functionRawData);
assertNotNull(cloudParameters);
assertEquals("aws_lambda", cloudParameters.getPlatform());
Expand All @@ -47,118 +66,162 @@ public void testGetCloudParamArnQualifier() {

@Test
public void testGetArnFunctionName() {
FunctionProcessedData data = LambdaUtil.processData(new FunctionRawData("my-function", null, getRegion()));
FunctionProcessedData data = LambdaUtil.processData(data("my-function", null));
assertEquals("", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameClient() {
mockCloudApiClient();
FunctionProcessedData data = LambdaUtil.processData(data("my-function", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameWithAlias() {
FunctionProcessedData data = LambdaUtil.processData(new FunctionRawData("my-function:alias", null, getRegion()));
FunctionProcessedData data = LambdaUtil.processData(data("my-function:alias", null));
assertEquals("", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameWithAliasClient() {
mockCloudApiClient();
FunctionProcessedData data = LambdaUtil.processData(data("my-function:alias", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:alias", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameWithVersion() {
FunctionProcessedData data = LambdaUtil.processData(new FunctionRawData("my-function:123", null, getRegion()));
FunctionProcessedData data = LambdaUtil.processData(data("my-function:123", null));
assertEquals("", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameWithVersionClient() {
mockCloudApiClient();
FunctionProcessedData data = LambdaUtil.processData(data("my-function:123", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:123", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameAndAliasQualifier() {
FunctionProcessedData data = LambdaUtil.processData(new FunctionRawData("my-function", "alias", getRegion()));
FunctionProcessedData data = LambdaUtil.processData(data("my-function", "alias"));
assertEquals("", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameAndAliasQualifierClient() {
mockCloudApiClient();
FunctionProcessedData data = LambdaUtil.processData(data("my-function", "alias"));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:alias", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameAndVersionQualifier() {
FunctionProcessedData data = LambdaUtil.processData(new FunctionRawData("my-function", "123", getRegion()));
FunctionProcessedData data = LambdaUtil.processData(data("my-function", "123"));
assertEquals("", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameAndVersionQualifierClient() {
mockCloudApiClient();
FunctionProcessedData data = LambdaUtil.processData(data("my-function", "123"));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:123", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnPartialArn() {
FunctionProcessedData data = LambdaUtil.processData(new FunctionRawData("123456789012:function:my-function", null, getRegion()));
FunctionProcessedData data = LambdaUtil.processData(data("123456789012:function:my-function", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnPartialArnWithAlias() {
FunctionProcessedData data = LambdaUtil.processData(new FunctionRawData("123456789012:function:my-function:alias", null, getRegion()));
FunctionProcessedData data = LambdaUtil.processData(data("123456789012:function:my-function:alias", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:alias", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnPartialArnWithVersion() {
FunctionProcessedData data = LambdaUtil.processData(new FunctionRawData("123456789012:function:my-function:123", null, getRegion()));
FunctionProcessedData data = LambdaUtil.processData(data("123456789012:function:my-function:123", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:123", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnPartialArnAndAliasQualifier() {
FunctionProcessedData data = LambdaUtil.processData(new FunctionRawData("123456789012:function:my-function", "alias", getRegion()));
FunctionProcessedData data = LambdaUtil.processData(data("123456789012:function:my-function", "alias"));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:alias", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnPartialArnAndVersionQualifier() {
FunctionProcessedData data = LambdaUtil.processData(new FunctionRawData("123456789012:function:my-function", "123", getRegion()));
FunctionProcessedData data = LambdaUtil.processData(data("123456789012:function:my-function", "123"));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:123", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFullArn() {
FunctionProcessedData data = LambdaUtil.processData(new FunctionRawData("arn:aws:lambda:us-east-1:123456789012:function:my-function", null, getRegion()));
FunctionProcessedData data = LambdaUtil.processData(data("arn:aws:lambda:us-east-1:123456789012:function:my-function", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFullArnWithAlias() {
FunctionProcessedData data = LambdaUtil.processData(new FunctionRawData("arn:aws:lambda:us-east-1:123456789012:function:my-function:alias", null, getRegion()));
FunctionProcessedData data = LambdaUtil.processData(data("arn:aws:lambda:us-east-1:123456789012:function:my-function:alias", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:alias", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFullArnWithVersion() {
FunctionProcessedData data = LambdaUtil.processData(new FunctionRawData("arn:aws:lambda:us-east-1:123456789012:function:my-function:123", null, getRegion()));
FunctionProcessedData data = LambdaUtil.processData(data("arn:aws:lambda:us-east-1:123456789012:function:my-function:123", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:123", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFullArnAndAliasQualifier() {
FunctionProcessedData data = LambdaUtil.processData(new FunctionRawData("arn:aws:lambda:us-east-1:123456789012:function:my-function", "alias", getRegion()));
FunctionProcessedData data = LambdaUtil.processData(data("arn:aws:lambda:us-east-1:123456789012:function:my-function", "alias"));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:alias", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFullArnAndVersionQualifier() {
FunctionProcessedData data = LambdaUtil.processData(new FunctionRawData("arn:aws:lambda:us-east-1:123456789012:function:my-function", "123", getRegion()));
FunctionProcessedData data = LambdaUtil.processData(data("arn:aws:lambda:us-east-1:123456789012:function:my-function", "123"));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:123", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnDifferentRegion() {
FunctionProcessedData data = LambdaUtil.processData(new FunctionRawData("arn:aws:lambda:us-west-2:123456789012:function:my-function", null, getRegion()));
FunctionProcessedData data = LambdaUtil.processData(data("arn:aws:lambda:us-west-2:123456789012:function:my-function", null));
assertEquals("arn:aws:lambda:us-west-2:123456789012:function:my-function", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

public String getRegion() {
return Regions.US_EAST_1.getName();
private static void mockCloudApiClient() {
when(AgentBridge.cloud.getAccountInfo(any(), eq(CloudAccountInfo.AWS_ACCOUNT_ID)))
.thenReturn("123456789012");
}

private FunctionRawData data(String functionRef, String qualifier) {
return new FunctionRawData(functionRef, qualifier, Regions.US_EAST_1.getName(), this);
}
}
Loading

0 comments on commit fc8cf8f

Please sign in to comment.