Enable Configurable Profiler Output Formats

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

@@ -58,6 +58,10 @@ jobs:
         working-directory: ./experimental/aws-lambda-java-profiler
         run: ./integration_tests/invoke_function.sh
 
+      - name: Invoke Java Custom Options function
+        working-directory: ./experimental/aws-lambda-java-profiler
+        run: ./integration_tests/invoke_function_custom_options.sh
+
       - name: Download from s3
         working-directory: ./experimental/aws-lambda-java-profiler
         run: ./integration_tests/download_from_s3.sh

experimental/aws-lambda-java-profiler/README.md

@@ -83,6 +83,25 @@ When the agent is constructed, it starts the profiler and registers itself as a
 A new thread is created to handle calling `/next` and uploading the results of the profiler to S3. The bucket to upload
 the result to is configurable using an environment variable.
 
+### Custom Parameters for the Profiler
+
+Users can configure the profiler output by setting environment variables.
+
+```
+# Example: Output as JFR format instead of HTML
+AWS_LAMBDA_PROFILER_START_COMMAND="start,event=wall,interval=1us,file=/tmp/profile.jfr"
+AWS_LAMBDA_PROFILER_STOP_COMMAND="stop,file=%s"
+```
+
+Defaults are the following:
+
+```
+AWS_LAMBDA_PROFILER_START_COMMAND="start,event=wall,interval=1us"
+AWS_LAMBDA_PROFILER_STOP_COMMAND="stop,file=%s,include=*AWSLambda.main,include=start_thread"
+```
+
+See [async-profiler's ProfilerOptions](https://github.com/async-profiler/async-profiler/blob/master/docs/ProfilerOptions.md) for all available profiler parameters.
+
 ### Troubleshooting
 
 - Ensure the Lambda function execution role has the necessary permissions to write to the S3 bucket.

experimental/aws-lambda-java-profiler/extension/src/main/java/com/amazonaws/services/lambda/extension/Constants.java

@@ -0,0 +1,29 @@
+package com.amazonaws.services.lambda.extension;
+
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+public class Constants {
+
+    private static final String DEFAULT_AWS_LAMBDA_PROFILER_START_COMMAND =
+            "start,event=wall,interval=1us";
+    private static final String DEFAULT_AWS_LAMBDA_PROFILER_STOP_COMMAND =
+            "stop,file=%s,include=*AWSLambda.main,include=start_thread";
+    public static final String PROFILER_START_COMMAND =
+            System.getenv().getOrDefault(
+                    "AWS_LAMBDA_PROFILER_START_COMMAND",
+                    DEFAULT_AWS_LAMBDA_PROFILER_START_COMMAND
+            );
+    public static final String PROFILER_STOP_COMMAND =
+            System.getenv().getOrDefault(
+                    "AWS_LAMBDA_PROFILER_STOP_COMMAND",
+                    DEFAULT_AWS_LAMBDA_PROFILER_STOP_COMMAND
+            );
+
+    public static String getFilePathFromEnv(){
+        Pattern pattern = Pattern.compile("file=([^,]+)");
+        Matcher matcher = pattern.matcher(PROFILER_START_COMMAND);
+
+        return matcher.find() ? matcher.group(1) : "/tmp/profiling-data-%s.html";
+    }
+}

experimental/aws-lambda-java-profiler/extension/src/main/java/com/amazonaws/services/lambda/extension/PreMain.java

@@ -2,49 +2,57 @@
 // SPDX-License-Identifier: MIT-0
 package com.amazonaws.services.lambda.extension;
 
+import com.sun.net.httpserver.HttpExchange;
+import com.sun.net.httpserver.HttpHandler;
+import com.sun.net.httpserver.HttpServer;
 import java.io.File;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.lang.instrument.Instrumentation;
 import java.net.InetSocketAddress;
 import java.nio.charset.StandardCharsets;
+import one.profiler.AsyncProfiler;
 
-import com.sun.net.httpserver.HttpExchange;
-import com.sun.net.httpserver.HttpHandler;
-import com.sun.net.httpserver.HttpServer;
+import static com.amazonaws.services.lambda.extension.Constants.PROFILER_START_COMMAND;
+import static com.amazonaws.services.lambda.extension.Constants.PROFILER_STOP_COMMAND;
+
+public class PreMain {
 
-import one.profiler.AsyncProfiler;
 
-public class PreMain {    
+    private static final String INTERNAL_COMMUNICATION_PORT =
+        System.getenv().getOrDefault(
+            "AWS_LAMBDA_PROFILER_COMMUNICATION_PORT",
+            "1234"
+        );
 
-    private static final String DEFAULT_AWS_LAMBDA_PROFILER_START_COMMAND = "start,event=wall,interval=1us";
-    private static final String DEFAULT_AWS_LAMBDA_PROFILER_STOP_COMMAND = "stop,file=%s,include=*AWSLambda.main,include=start_thread";
-    private static final String PROFILER_START_COMMAND = System.getenv().getOrDefault("AWS_LAMBDA_PROFILER_START_COMMAND", DEFAULT_AWS_LAMBDA_PROFILER_START_COMMAND);
-    private static final String PROFILER_STOP_COMMAND = System.getenv().getOrDefault("AWS_LAMBDA_PROFILER_STOP_COMMAND", DEFAULT_AWS_LAMBDA_PROFILER_STOP_COMMAND);
-    private static final String INTERNAL_COMMUNICATION_PORT = System.getenv().getOrDefault("AWS_LAMBDA_PROFILER_COMMUNICATION_PORT", "1234");
+
+    private String filepath;
 
     public static void premain(String agentArgs, Instrumentation inst) {
         Logger.debug("premain is starting");
-        if(!createFileIfNotExist("/tmp/aws-lambda-java-profiler")) {
+        if (!createFileIfNotExist("/tmp/aws-lambda-java-profiler")) {
             Logger.debug("starting the profiler for coldstart");
             startProfiler();
             registerShutdownHook();
             try {
                 Integer port = Integer.parseInt(INTERNAL_COMMUNICATION_PORT);
                 Logger.debug("using profile communication port = " + port);
-                HttpServer server = HttpServer.create(new InetSocketAddress(port), 0);
+                HttpServer server = HttpServer.create(
+                    new InetSocketAddress(port),
+                    0
+                );
                 server.createContext("/profiler/start", new StartProfiler());
                 server.createContext("/profiler/stop", new StopProfiler());
                 server.setExecutor(null); // Use the default executor
                 server.start();
-            } catch(Exception e) {
+            } catch (Exception e) {
                 e.printStackTrace();
             }
         }
     }
 
     private static boolean createFileIfNotExist(String filePath) {
-        File file = new File(filePath);   
+        File file = new File(filePath);
         try {
             return file.createNewFile();
         } catch (IOException e) {
@@ -54,10 +62,13 @@ private static boolean createFileIfNotExist(String filePath) {
     }
 
     public static class StopProfiler implements HttpHandler {
+
         @Override
         public void handle(HttpExchange exchange) throws IOException {
             Logger.debug("hit /profiler/stop");
-            final String fileName = exchange.getRequestHeaders().getFirst(ExtensionMain.HEADER_NAME);
+            final String fileName = exchange
+                .getRequestHeaders()
+                .getFirst(ExtensionMain.HEADER_NAME);
             stopProfiler(fileName);
             String response = "ok";
             exchange.sendResponseHeaders(200, response.length());
@@ -68,6 +79,7 @@ public void handle(HttpExchange exchange) throws IOException {
     }
 
     public static class StartProfiler implements HttpHandler {
+
         @Override
         public void handle(HttpExchange exchange) throws IOException {
             Logger.debug("hit /profiler/start");
@@ -80,31 +92,40 @@ public void handle(HttpExchange exchange) throws IOException {
         }
     }
 
-
     public static void stopProfiler(String fileNameSuffix) {
         try {
-            final String fileName = String.format("/tmp/profiling-data-%s.html", fileNameSuffix);
-            Logger.debug("stopping the profiler with filename  = " + fileName + " with command = " + PROFILER_STOP_COMMAND);
-            AsyncProfiler.getInstance().execute(String.format(PROFILER_STOP_COMMAND, fileName));
-        } catch(Exception e) {
+            final String fileName = String.format(
+                Constants.getFilePathFromEnv(),
+                fileNameSuffix
+            );
+            Logger.debug(
+                "stopping the profiler with filename  = " + fileName
+            );
+            AsyncProfiler.getInstance().execute(
+                String.format(PROFILER_STOP_COMMAND, fileName)
+            );
+        } catch (Exception e) {
             Logger.error("could not stop the profiler");
             e.printStackTrace();
         }
     }
 
     public static void startProfiler() {
         try {
-            Logger.debug("staring the profiler with command = " + PROFILER_START_COMMAND);
+            Logger.debug(
+                "starting the profiler with command = " + PROFILER_START_COMMAND
+            );
             AsyncProfiler.getInstance().execute(PROFILER_START_COMMAND);
         } catch (IOException e) {
             throw new RuntimeException(e);
         }
     }
 
     public static void registerShutdownHook() {
-        Logger.debug("registering shutdown hook");
-        Thread shutdownHook = new Thread(new ShutdownHook(PROFILER_STOP_COMMAND));
+        Logger.debug("registering shutdown hook wit command = " + PROFILER_STOP_COMMAND);
+        Thread shutdownHook = new Thread(
+            new ShutdownHook(PROFILER_STOP_COMMAND)
+        );
         Runtime.getRuntime().addShutdownHook(shutdownHook);
     }
-
-}
+}

experimental/aws-lambda-java-profiler/extension/src/main/java/com/amazonaws/services/lambda/extension/S3Manager.java

@@ -4,7 +4,6 @@
 
 import java.io.File;
 import java.time.format.DateTimeFormatter;
-import java.time.Instant;
 import java.time.LocalDate;
 
 import software.amazon.awssdk.core.sync.RequestBody;
@@ -39,7 +38,7 @@ public void upload(String fileName, boolean isShutDownEvent) {
                     .bucket(bucketName)
                     .key(key)
                     .build();
-            File file = new File(String.format("/tmp/profiling-data-%s.html", suffix));
+            File file = new File(String.format(Constants.getFilePathFromEnv(), suffix));
             if (file.exists()) {
                 Logger.debug("file size is " + file.length());
                 RequestBody requestBody = RequestBody.fromFile(file);

experimental/aws-lambda-java-profiler/integration_tests/create_function.sh

@@ -2,13 +2,16 @@
 
 # Set variables
 FUNCTION_NAME="aws-lambda-java-profiler-function-${GITHUB_RUN_ID}"
+FUNCTION_NAME_CUSTOM_PROFILER_OPTIONS="aws-lambda-java-profiler-function-custom-${GITHUB_RUN_ID}"
 ROLE_NAME="aws-lambda-java-profiler-role-${GITHUB_RUN_ID}"
 HANDLER="helloworld.Handler::handleRequest"
 RUNTIME="java21"
 LAYER_ARN=$(cat /tmp/layer_arn)
 
 JAVA_TOOL_OPTIONS="-XX:+UnlockDiagnosticVMOptions -XX:+DebugNonSafepoints -javaagent:/opt/profiler-extension.jar"
 AWS_LAMBDA_PROFILER_RESULTS_BUCKET_NAME="aws-lambda-java-profiler-bucket-${GITHUB_RUN_ID}"
+AWS_LAMBDA_PROFILER_START_COMMAND="start,event=wall,interval=1us,file=/tmp/profile.jfr"
+AWS_LAMBDA_PROFILER_STOP_COMMAND="stop,file=%s"
 
 # Compile the Hello World project
 cd integration_tests/helloworld
@@ -63,6 +66,19 @@ aws lambda create-function \
     --environment "Variables={JAVA_TOOL_OPTIONS='$JAVA_TOOL_OPTIONS',AWS_LAMBDA_PROFILER_RESULTS_BUCKET_NAME='$AWS_LAMBDA_PROFILER_RESULTS_BUCKET_NAME',AWS_LAMBDA_PROFILER_DEBUG='true'}" \
     --layers "$LAYER_ARN"
 
+
+# Create Lambda function custom profiler options
+aws lambda create-function \
+    --function-name "$FUNCTION_NAME_CUSTOM_PROFILER_OPTIONS" \
+    --runtime "$RUNTIME" \
+    --role "$ROLE_ARN" \
+    --handler "$HANDLER" \
+    --timeout 30 \
+    --memory-size 512 \
+    --zip-file fileb://integration_tests/helloworld/build/distributions/code.zip \
+    --environment "Variables={JAVA_TOOL_OPTIONS='$JAVA_TOOL_OPTIONS',AWS_LAMBDA_PROFILER_RESULTS_BUCKET_NAME='$AWS_LAMBDA_PROFILER_RESULTS_BUCKET_NAME',AWS_LAMBDA_PROFILER_DEBUG='true',AWS_LAMBDA_PROFILER_START_COMMAND='$AWS_LAMBDA_PROFILER_START_COMMAND',AWS_LAMBDA_PROFILER_STOP_COMMAND='$AWS_LAMBDA_PROFILER_STOP_COMMAND'}" \
+    --layers "$LAYER_ARN"
+
 echo "Lambda function '$FUNCTION_NAME' created successfully with Java 21 runtime"
 
 echo "Waiting the function to be ready so we can invoke it..."

experimental/aws-lambda-java-profiler/integration_tests/helloworld/build.gradle

@@ -1,11 +1,15 @@
-apply plugin: 'java'
+plugins {
+    id 'java'
+}
 
 repositories {
     mavenCentral()
 }
 
-sourceCompatibility = 21
-targetCompatibility = 21
+java {
+    sourceCompatibility = JavaVersion.VERSION_21
+    targetCompatibility = JavaVersion.VERSION_21
+}
 
 dependencies {
     implementation (
@@ -24,4 +28,5 @@ task buildZip(type: Zip) {
     }
 }
 
+
 build.dependsOn buildZip

experimental/aws-lambda-java-profiler/integration_tests/invoke_function.sh

@@ -32,8 +32,8 @@ fi
 echo "Function output:"
 cat output.json
 
-echo "$LOG_RESULT" | base64 --decode | grep "starting the profiler for coldstart" || exit 1
-echo "$LOG_RESULT" | base64 --decode | grep -v "uploading" || exit 1
+echo "$LOG_RESULT" | base64 --decode | grep "starting the profiler for coldstart" || { echo "ERROR: Profiler did not start for coldstart"; exit 1; }
+echo "$LOG_RESULT" | base64 --decode | grep -v "uploading" || { echo "ERROR: Unexpected upload detected on cold start"; exit 1; }
 
 # Clean up the output file
 rm output.json
@@ -68,7 +68,7 @@ fi
 echo "Function output:"
 cat output.json
 
-echo "$LOG_RESULT" | base64 --decode | grep "uploading" || exit 1
+echo "$LOG_RESULT" | base64 --decode | grep "uploading" || { echo "ERROR: Upload not detected on warm start"; exit 1; }
 
 # Clean up the output file
 rm output.json