AWS Lambda instrumentation (#1998)

agent-bridge/src/main/java/com/newrelic/agent/bridge/CollectionFactory.java

@@ -49,4 +49,16 @@ public interface CollectionFactory {
      * @param <V> the type of value stored/returned
      */
     <K, V> Function<K, V> memorize(Function<K, V> loader, int maxSize);
+
+    /**
+     * Create a time based eviction cache in which an entry's age is determined on a last-access basis.
+     *
+     * @param <K>             key type
+     * @param <V>             cached type
+     * @param ageInSeconds    how old, in seconds, a cache entry must be to be evicted after last access
+     * @param initialCapacity the initial capacity of the cache
+     * @param loader          the function to calculate the value for a key, used if the key is not cached
+     * @return a time based concurrent cache
+     */
+    <K, V> Function<K, V> createAccessTimeBasedCache(long ageInSeconds, int initialCapacity, Function<K, V> loader);
 }

agent-bridge/src/main/java/com/newrelic/agent/bridge/DefaultCollectionFactory.java

@@ -43,4 +43,12 @@ public <K, V> Function<K, V> memorize(Function<K, V> loader, int maxSize) {
             return loader.apply(k1);
         });
     }
+
+    /**
+     * Note: In this implementation, this method will return the loader function as is.
+     */
+    @Override
+    public <K, V> Function<K, V> createAccessTimeBasedCache(long ageInSeconds, int initialCapacity, Function<K, V> loader) {
+        return loader;
+    }
 }

instrumentation/aws-java-sdk-lambda-1.11.280/build.gradle

@@ -0,0 +1,20 @@
+jar {
+    manifest { attributes 'Implementation-Title': 'com.newrelic.instrumentation.aws-java-sdk-lambda-1.11.280' }
+}
+
+dependencies {
+    implementation(project(":agent-bridge"))
+    implementation(project(":agent-bridge-datastore"))
+
+    implementation("com.amazonaws:aws-java-sdk-lambda:1.12.763")
+}
+
+verifyInstrumentation {
+    // not using passesOnly to decrease the number of artifacts this is tested against
+    passes 'com.amazonaws:aws-java-sdk-lambda:[1.11.280,)'
+}
+
+site {
+    title 'AWS Lambda'
+    type 'Framework'
+}

instrumentation/aws-java-sdk-lambda-1.11.280/src/main/java/com/agent/instrumentation/awsjavasdk1/services/lambda/FunctionProcessedData.java

@@ -0,0 +1,29 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.agent.instrumentation.awsjavasdk1.services.lambda;
+
+/**
+ * Function data extracted from the request and config.
+ */
+class FunctionProcessedData {
+    private final String functionName;
+    private final String arn;
+
+    public FunctionProcessedData(String functionName, String arn) {
+        this.functionName = functionName;
+        this.arn = arn;
+    }
+
+    public String getFunctionName() {
+        return functionName;
+    }
+
+    public String getArn() {
+        return arn;
+    }
+}

instrumentation/aws-java-sdk-lambda-1.11.280/src/main/java/com/agent/instrumentation/awsjavasdk1/services/lambda/FunctionRawData.java

@@ -0,0 +1,55 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.agent.instrumentation.awsjavasdk1.services.lambda;
+
+import java.util.Objects;
+
+/**
+ * Data necessary to calculate the ARN. This class is used as the key for the ARN cache.
+ */
+public class FunctionRawData {
+    private final String functionRef;
+    private final String qualifier;
+    private final String region;
+
+    public FunctionRawData(String functionRef, String qualifier, String region) {
+        this.functionRef = functionRef;
+        this.qualifier = qualifier;
+        this.region = region;
+    }
+
+    public String getFunctionRef() {
+        return functionRef;
+    }
+
+    public String getQualifier() {
+        return qualifier;
+    }
+
+    public String getRegion() {
+        return region;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        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);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(functionRef, qualifier, region);
+    }
+}

instrumentation/aws-java-sdk-lambda-1.11.280/src/main/java/com/agent/instrumentation/awsjavasdk1/services/lambda/LambdaUtil.java

@@ -0,0 +1,106 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.agent.instrumentation.awsjavasdk1.services.lambda;
+
+import com.newrelic.agent.bridge.AgentBridge;
+import com.newrelic.api.agent.CloudParameters;
+
+import java.util.function.Function;
+
+public class LambdaUtil {
+
+    private static final String PLATFORM = "aws_lambda";
+    private static final String NULL_ARN = "";
+    private static final String PREFIX = "arn:aws:lambda:";
+    private static final Function<FunctionRawData, FunctionProcessedData> CACHE =
+            AgentBridge.collectionFactory.createAccessTimeBasedCache(3600, 8, LambdaUtil::processData);
+
+    public static CloudParameters getCloudParameters(FunctionRawData functionRawData) {
+        FunctionProcessedData data = CACHE.apply(functionRawData);
+        String arn = data.getArn();
+        CloudParameters.ResourceIdParameter cloudParameters = CloudParameters.provider(PLATFORM);
+        // the cache will always return the NULL_ARN when it is not possible to calculate the ARN
+        // so saving a few cycles by using != instead of equals.
+        if (arn != NULL_ARN) {
+            cloudParameters.resourceId(arn);
+        }
+
+        return cloudParameters.build();
+    }
+
+    /**
+     * <p>
+     *     Calculates the simple function name and ARN given
+     *     the function name, qualifier, and possibly region (provided by config).
+     * </p>
+     * <p>
+     *     Aliases are returned as part of the ARN, but versions are removed
+     *     because they would make it harder to link to Lambdas/Alias entities.
+     * </p>
+     * <p>
+     *     If qualifiers are provided both in the function ref, and as a qualifier, the one in function ref "wins".
+     *     If they differ, the LambdaClient will throw an exception.
+     * </p>
+     *
+     * @return a FunctionProcessedData object with the function name and ARN.
+     *         If any of its values cannot be calculated, it will be the NULL_ARN.
+     */
+    // Visible for testing
+    static FunctionProcessedData processData(FunctionRawData data) {
+        String functionRef = data.getFunctionRef();
+
+        String[] parts = functionRef.split(":");
+
+        String functionName = NULL_ARN;
+        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.
+            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.
+            functionName = parts[0];
+        } else if (parts.length == 3) {
+            // partial ARN: {account-id}:function:{function-name}
+            functionName = parts[2];
+            String qualifier = data.getQualifier();
+            if (qualifier == null) {
+                arn = PREFIX + data.getRegion() + ":" + functionRef;
+            } 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];
+            String qualifier = data.getQualifier();
+            if (qualifier == null) {
+                arn = functionRef;
+            } 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];
+            arn = functionRef;
+        }
+        // reference should be invalid if the number of parts do not match any of the expected cases
+
+        return new FunctionProcessedData(functionName, arn);
+    }
+
+
+    public static String getSimpleFunctionName(FunctionRawData functionRawData) {
+        return CACHE.apply(functionRawData).getFunctionName();
+    }
+}

instrumentation/aws-java-sdk-lambda-1.11.280/src/main/java/com/amazonaws/services/lambda/AWSLambdaAsyncClient_Instrumentation.java

@@ -0,0 +1,71 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.amazonaws.services.lambda;
+
+import com.agent.instrumentation.awsjavasdk1.services.lambda.FunctionRawData;
+import com.agent.instrumentation.awsjavasdk1.services.lambda.LambdaUtil;
+import com.amazonaws.handlers.AsyncHandler;
+import com.amazonaws.services.lambda.model.InvokeRequest;
+import com.amazonaws.services.lambda.model.InvokeResult;
+import com.newrelic.api.agent.CloudParameters;
+import com.newrelic.api.agent.NewRelic;
+import com.newrelic.api.agent.Segment;
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+
+import java.util.concurrent.Future;
+
+@Weave(type = MatchType.ExactClass, originalName = "com.amazonaws.services.lambda.AWSLambdaAsyncClient")
+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());
+        CloudParameters cloudParameters = LambdaUtil.getCloudParameters(functionRawData);
+        String functionName = LambdaUtil.getSimpleFunctionName(functionRawData);
+        Segment segment = NewRelic.getAgent().getTransaction().startSegment("Lambda", "invoke/" + functionName);
+
+        try {
+            segment.reportAsExternal(cloudParameters);
+            asyncHandler = new SegmentEndingAsyncHandler(asyncHandler, segment);
+            return Weaver.callOriginal();
+        } catch (Throwable t) {
+            segment.end();
+            throw t;
+        }
+    }
+
+    private static class SegmentEndingAsyncHandler implements AsyncHandler<InvokeRequest, InvokeResult> {
+        private final AsyncHandler<InvokeRequest, InvokeResult> originalHandler;
+        private final Segment segment;
+
+        public SegmentEndingAsyncHandler(
+                AsyncHandler<InvokeRequest, InvokeResult> asyncHandler, Segment segment) {
+            this.segment = segment;
+            this.originalHandler = asyncHandler;
+        }
+
+        @Override
+        public void onError(Exception exception) {
+            segment.end();
+            if (originalHandler != null) {
+                originalHandler.onError(exception);
+            }
+        }
+
+        @Override
+        public void onSuccess(InvokeRequest request, InvokeResult invokeResult) {
+            segment.end();
+            if (originalHandler != null) {
+                originalHandler.onSuccess(request, invokeResult);
+            }
+        }
+    }
+}

instrumentation/aws-java-sdk-lambda-1.11.280/src/main/java/com/amazonaws/services/lambda/AWSLambdaClient_Instrumentation.java

@@ -0,0 +1,39 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.amazonaws.services.lambda;
+
+import com.agent.instrumentation.awsjavasdk1.services.lambda.FunctionRawData;
+import com.agent.instrumentation.awsjavasdk1.services.lambda.LambdaUtil;
+import com.amazonaws.services.lambda.model.InvokeRequest;
+import com.amazonaws.services.lambda.model.InvokeResult;
+import com.newrelic.api.agent.CloudParameters;
+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;
+
+@Weave(type = MatchType.ExactClass, originalName = "com.amazonaws.services.lambda.AWSLambdaClient")
+public abstract class AWSLambdaClient_Instrumentation {
+
+    abstract protected String getSigningRegion();
+
+    @Trace(leaf = true)
+    public InvokeResult invoke(InvokeRequest invokeRequest) {
+        FunctionRawData functionRawData = new FunctionRawData(invokeRequest.getFunctionName(), invokeRequest.getQualifier(), getSigningRegion());
+        CloudParameters cloudParameters = LambdaUtil.getCloudParameters(functionRawData);
+        TracedMethod tracedMethod = NewRelic.getAgent().getTracedMethod();
+        tracedMethod.reportAsExternal(cloudParameters);
+        tracedMethod.setMetricName("Lambda", "invoke", LambdaUtil.getSimpleFunctionName(functionRawData));
+        return Weaver.callOriginal();
+    }
+
+}
+