Merge branch 'main' into develop

Author: peniakoff

.github/workflows/aws-lambda-java-profiler.yml

@@ -0,0 +1,74 @@
+name: Run integration tests for aws-lambda-java-profiler
+
+on:
+  pull_request:
+    branches: [ '*' ]
+    paths:
+      - 'experimental/aws-lambda-java-profiler/**'
+      - '.github/workflows/aws-lambda-java-profiler.yml'
+  push:
+    branches: ['*']
+    paths:
+      - 'experimental/aws-lambda-java-profiler/**'
+      - '.github/workflows/aws-lambda-java-profiler.yml'
+
+jobs:
+
+  build:
+    runs-on: ubuntu-latest
+
+    permissions:
+      id-token: write
+      contents: read
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up JDK
+        uses: actions/setup-java@v4
+        with:
+          java-version: 21
+          distribution: corretto
+
+      - name: Issue AWS credentials
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          aws-region: ${{ secrets.AWS_REGION_PROFILER_EXTENSION_INTEGRATION_TEST }}
+          role-to-assume: ${{ secrets.AWS_ROLE_PROFILER_EXTENSION_INTEGRATION_TEST }}
+          role-session-name: GitHubActionsRunIntegrationTests
+          role-duration-seconds: 900
+
+      - name: Build layer
+        working-directory: ./experimental/aws-lambda-java-profiler/extension
+        run: ./build_layer.sh
+      
+      - name: Publish layer
+        working-directory: ./experimental/aws-lambda-java-profiler
+        run: ./integration_tests/publish_layer.sh
+
+      - name: Create the bucket layer
+        working-directory: ./experimental/aws-lambda-java-profiler
+        run: ./integration_tests/create_bucket.sh
+
+      - name: Create Java function
+        working-directory: ./experimental/aws-lambda-java-profiler
+        run: ./integration_tests/create_function.sh
+
+      - name: Invoke Java function
+        working-directory: ./experimental/aws-lambda-java-profiler
+        run: ./integration_tests/invoke_function.sh
+
+      - name: Download from s3
+        working-directory: ./experimental/aws-lambda-java-profiler
+        run: ./integration_tests/download_from_s3.sh
+
+      - name: Upload profiles
+        uses: actions/upload-artifact@v4
+        with:
+          name: profiles
+          path: /tmp/s3-artifacts 
+        
+      - name: cleanup
+        if: always()
+        working-directory: ./experimental/aws-lambda-java-profiler
+        run: ./integration_tests/cleanup.sh

.github/workflows/samples.yml

@@ -15,7 +15,7 @@ on:
       - '.github/workflows/samples.yml'
 
 jobs:
-  build-kinesis-sample:
+  build:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
@@ -28,6 +28,9 @@ jobs:
       # Install events module
       - name: Install events with Maven
         run: mvn -B install --file aws-lambda-java-events/pom.xml
+      # Install serialization module
+      - name: Install serialization with Maven
+        run: mvn -B install --file aws-lambda-java-serialization/pom.xml
       # Install tests module
       - name: Install tests with Maven
         run: mvn -B install --file aws-lambda-java-tests/pom.xml

.gitignore

@@ -28,3 +28,10 @@ dependency-reduced-pom.xml
 
 # snapshot process
 aws-lambda-java-runtime-interface-client/pom.xml.versionsBackup
+
+# profiler
+experimental/aws-lambda-java-profiler/integration_tests/helloworld/build
+experimental/aws-lambda-java-profiler/extension/build/
+experimental/aws-lambda-java-profiler/integration_tests/helloworld/bin
+!experimental/aws-lambda-java-profiler/extension/gradle/wrapper/*.jar
+/scratch/

README.md

@@ -139,6 +139,18 @@ See the [README](aws-lambda-java-log4j2/README.md) or the [official documentatio
 </dependency>
 ```
 
+## Lambda Profiler Extension for Java - aws-lambda-java-profiler
+
+<p align="center">
+    <img src="experimental/aws-lambda-java-profiler/docs/example-cold-start-flame-graph-small.png" alt="A flame graph of a Java Lambda function">
+</p>
+
+This project allows you to profile your Java functions invoke by invoke, with high fidelity, and no code changes. It 
+uses the [async-profiler](https://github.com/async-profiler/async-profiler) project to produce profiling data and 
+automatically uploads the data as flame graphs to S3.
+
+Follow our [Quick Start](experimental/aws-lambda-java-profiler#quick-start) to profile your functions.
+
 ## Java implementation of the Runtime Interface Client API - aws-lambda-java-runtime-interface-client
 [![Maven](https://img.shields.io/maven-central/v/com.amazonaws/aws-lambda-java-runtime-interface-client.svg?label=Maven)](https://central.sonatype.com/artifact/com.amazonaws/aws-lambda-java-runtime-interface-client)
 

aws-lambda-java-runtime-interface-client/pom.xml

@@ -87,6 +87,12 @@
             <version>4.11.0</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-junit-jupiter</artifactId>
+            <version>4.11.0</version>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>com.squareup.okhttp3</groupId>
             <artifactId>mockwebserver</artifactId>

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/AWSLambda.java

@@ -2,6 +2,7 @@
 Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 SPDX-License-Identifier: Apache-2.0
 */
+
 package com.amazonaws.services.lambda.runtime.api.client;
 
 import com.amazonaws.services.lambda.crac.Core;
@@ -35,7 +36,6 @@
 import java.security.Security;
 import java.util.Properties;
 
-
 /**
  * The entrypoint of this class is {@link AWSLambda#startRuntime}. It performs two main tasks:
  *
@@ -137,6 +137,42 @@ private static LambdaRequestHandler findRequestHandler(final String handlerStrin
         return requestHandler;
     }
 
+    private static LambdaRequestHandler getLambdaRequestHandlerObject(String handler, LambdaContextLogger lambdaLogger) throws ClassNotFoundException, IOException {
+        UnsafeUtil.disableIllegalAccessWarning();
+
+        System.setOut(new PrintStream(new LambdaOutputStream(System.out), false, "UTF-8"));
+        System.setErr(new PrintStream(new LambdaOutputStream(System.err), false, "UTF-8"));
+        setupRuntimeLogger(lambdaLogger);
+
+        runtimeClient = new LambdaRuntimeApiClientImpl(LambdaEnvironment.RUNTIME_API);
+
+        String taskRoot = System.getProperty("user.dir");
+        String libRoot = "/opt/java";
+        // Make system classloader the customer classloader's parent to ensure any aws-lambda-java-core classes
+        // are loaded from the system classloader.
+        customerClassLoader = new CustomerClassLoader(taskRoot, libRoot, ClassLoader.getSystemClassLoader());
+        Thread.currentThread().setContextClassLoader(customerClassLoader);
+
+        // Load the user's handler
+        LambdaRequestHandler requestHandler = null;
+        try {
+            requestHandler = findRequestHandler(handler, customerClassLoader);
+        } catch (UserFault userFault) {
+            lambdaLogger.log(userFault.reportableError(), lambdaLogger.getLogFormat() == LogFormat.JSON ? LogLevel.ERROR : LogLevel.UNDEFINED);
+            LambdaError error = new LambdaError(
+                    LambdaErrorConverter.fromUserFault(userFault),
+                    RapidErrorType.BadFunctionCode);
+            runtimeClient.reportInitError(error);
+            System.exit(1);
+        }
+
+        if (INIT_TYPE_SNAP_START.equals(AWS_LAMBDA_INITIALIZATION_TYPE)) {
+            onInitComplete(lambdaLogger);
+        }
+
+        return requestHandler;
+    }
+
     public static void setupRuntimeLogger(LambdaLogger lambdaLogger)
             throws ClassNotFoundException {
         ReflectUtil.setStaticField(
@@ -176,55 +212,27 @@ private static LogSink createLogSink() {
         }
     }
 
-    public static void main(String[] args) {
-        startRuntime(args[0]);
-    }
+    public static void main(String[] args) throws Throwable {
+        try (LambdaContextLogger logger = initLogger()) {
+            LambdaRequestHandler lambdaRequestHandler = getLambdaRequestHandlerObject(args[0], logger);
+            startRuntimeLoop(lambdaRequestHandler, logger);
 
-    private static void startRuntime(String handler) {
-        try (LogSink logSink = createLogSink()) {
-            LambdaContextLogger logger = new LambdaContextLogger(
-                    logSink,
-                    LogLevel.fromString(LambdaEnvironment.LAMBDA_LOG_LEVEL),
-                    LogFormat.fromString(LambdaEnvironment.LAMBDA_LOG_FORMAT)
-            );
-            startRuntime(handler, logger);
-        } catch (Throwable t) {
+        } catch (IOException | ClassNotFoundException t) {
             throw new Error(t);
         }
     }
 
-    private static void startRuntime(String handler, LambdaContextLogger lambdaLogger) throws Throwable {
-        UnsafeUtil.disableIllegalAccessWarning();
-
-        System.setOut(new PrintStream(new LambdaOutputStream(System.out), false, "UTF-8"));
-        System.setErr(new PrintStream(new LambdaOutputStream(System.err), false, "UTF-8"));
-        setupRuntimeLogger(lambdaLogger);
+    private static LambdaContextLogger initLogger() {
+        LogSink logSink = createLogSink();
+        LambdaContextLogger logger = new LambdaContextLogger(
+                logSink,
+                LogLevel.fromString(LambdaEnvironment.LAMBDA_LOG_LEVEL),
+                LogFormat.fromString(LambdaEnvironment.LAMBDA_LOG_FORMAT));
 
-        runtimeClient = new LambdaRuntimeApiClientImpl(LambdaEnvironment.RUNTIME_API);
-
-        String taskRoot = System.getProperty("user.dir");
-        String libRoot = "/opt/java";
-        // Make system classloader the customer classloader's parent to ensure any aws-lambda-java-core classes
-        // are loaded from the system classloader.
-        customerClassLoader = new CustomerClassLoader(taskRoot, libRoot, ClassLoader.getSystemClassLoader());
-        Thread.currentThread().setContextClassLoader(customerClassLoader);
+        return logger;
+    }
 
-        // Load the user's handler
-        LambdaRequestHandler requestHandler;
-        try {
-            requestHandler = findRequestHandler(handler, customerClassLoader);
-        } catch (UserFault userFault) {
-            lambdaLogger.log(userFault.reportableError(), lambdaLogger.getLogFormat() == LogFormat.JSON ? LogLevel.ERROR : LogLevel.UNDEFINED);
-            LambdaError error = new LambdaError(
-                    LambdaErrorConverter.fromUserFault(userFault),
-                    RapidErrorType.BadFunctionCode);
-            runtimeClient.reportInitError(error);
-            System.exit(1);
-            return;
-        }
-        if (INIT_TYPE_SNAP_START.equals(AWS_LAMBDA_INITIALIZATION_TYPE)) {
-            onInitComplete(lambdaLogger);
-        }
+    private static void startRuntimeLoop(LambdaRequestHandler requestHandler, LambdaContextLogger lambdaLogger) throws Throwable {
         boolean shouldExit = false;
         while (!shouldExit) {
             UserFault userFault = null;
@@ -240,7 +248,7 @@ private static void startRuntime(String handler, LambdaContextLogger lambdaLogge
                 payload = requestHandler.call(request);
                 runtimeClient.reportInvocationSuccess(request.getId(), payload.toByteArray());
                 // clear interrupted flag in case if it was set by user's code
-                boolean ignored = Thread.interrupted();
+                Thread.interrupted();
             } catch (UserFault f) {
                 shouldExit = f.fatal;
                 userFault = f;
@@ -278,6 +286,7 @@ static void onInitComplete(final LambdaContextLogger lambdaLogger) throws IOExce
                     RapidErrorType.BeforeCheckpointError));
             System.exit(64);
         }
+
         try {
             Core.getGlobalContext().afterRestore(null);
         } catch (Exception restoreExc) {

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/EventHandlerLoader.java

@@ -581,6 +581,7 @@ public ByteArrayOutputStream call(InvocationRequest request) throws Error, Excep
                         cognitoIdentity,
                         LambdaEnvironment.FUNCTION_VERSION,
                         request.getInvokedFunctionArn(),
+                        request.getTenantId(),
                         clientContext
                 );
 

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaContext.java

@@ -22,6 +22,7 @@ public class LambdaContext implements Context {
     private final long deadlineTimeInMs;
     private final CognitoIdentity cognitoIdentity;
     private final ClientContext clientContext;
+    private final String tenantId;
     private final LambdaLogger logger;
 
     public LambdaContext(
@@ -34,6 +35,7 @@ public LambdaContext(
             CognitoIdentity identity,
             String functionVersion,
             String invokedFunctionArn,
+            String tenantId,
             ClientContext clientContext
     ) {
         this.memoryLimit = memoryLimit;
@@ -46,6 +48,7 @@ public LambdaContext(
         this.clientContext = clientContext;
         this.functionVersion = functionVersion;
         this.invokedFunctionArn = invokedFunctionArn;
+        this.tenantId = tenantId;
         this.logger = com.amazonaws.services.lambda.runtime.LambdaRuntime.getLogger();
     }
 
@@ -91,6 +94,10 @@ public int getRemainingTimeInMillis() {
         return delta > 0 ? delta : 0;
     }
 
+    public String getTenantId() {
+        return tenantId;
+    }
+
     public LambdaLogger getLogger() {
         return logger;
     }

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/JsonLogFormatter.java

@@ -41,6 +41,7 @@ private StructuredLogMessage createLogMessage(String message, LogLevel logLevel)
 
         if (lambdaContext != null) {
             msg.AWSRequestId = lambdaContext.getAwsRequestId();
+            msg.tenantId = lambdaContext.getTenantId();
         }
         return msg;
     }

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/LambdaContextLogger.java

@@ -7,9 +7,11 @@
 
 import com.amazonaws.services.lambda.runtime.logging.LogFormat;
 import com.amazonaws.services.lambda.runtime.logging.LogLevel;
+import java.io.Closeable;
+import java.io.IOException;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
-public class LambdaContextLogger extends AbstractLambdaLogger {
+public class LambdaContextLogger extends AbstractLambdaLogger implements Closeable {
     // If a null string is passed in, replace it with "null",
     // replicating the behavior of System.out.println(null);
     private static final byte[] NULL_BYTES_VALUE = "null".getBytes(UTF_8);
@@ -29,4 +31,10 @@ protected void logMessage(byte[] message, LogLevel logLevel) {
             sink.log(logLevel, this.logFormat, message);
         }
     }
+
+    @Override
+    public void close() throws IOException {
+        sink.close();
+        
+    }
 }

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/StructuredLogMessage.java

@@ -12,4 +12,5 @@ class StructuredLogMessage {
     public String message;
     public LogLevel level;
     public String AWSRequestId;
+    public String tenantId;
 }

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/runtimeapi/dto/InvocationRequest.java

@@ -40,6 +40,11 @@ public class InvocationRequest {
      */
     private String cognitoIdentity;
 
+    /**
+     * The tenant ID associated with the request.
+     */
+    private String tenantId;
+
     private byte[] content;
 
     public String getId() {
@@ -94,6 +99,14 @@ public void setCognitoIdentity(String cognitoIdentity) {
         this.cognitoIdentity = cognitoIdentity;
     }
 
+    public String getTenantId() {
+        return tenantId;
+    }
+
+    public void setTenantId(String tenantId) {
+        this.tenantId = tenantId;
+    }
+
     public byte[] getContent() {
         return content;
     }