chore: enabling different format outputs

Author: Davide Melfi

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,21 +92,32 @@ 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 +
+                    " with command = " +
+                    PROFILER_STOP_COMMAND
+            );
+            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(
+                "staring the profiler with command = " + PROFILER_START_COMMAND
+            );
             AsyncProfiler.getInstance().execute(PROFILER_START_COMMAND);
         } catch (IOException e) {
             throw new RuntimeException(e);
@@ -103,8 +126,9 @@ public static void startProfiler() {
 
     public static void registerShutdownHook() {
         Logger.debug("registering shutdown hook");
-        Thread shutdownHook = new Thread(new ShutdownHook(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);