diff --git a/experimental/aws-lambda-java-profiler/extension/src/main/java/com/amazonaws/services/lambda/extension/Constants.java b/experimental/aws-lambda-java-profiler/extension/src/main/java/com/amazonaws/services/lambda/extension/Constants.java new file mode 100644 index 00000000..f9ca3010 --- /dev/null +++ b/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"; + } +} diff --git a/experimental/aws-lambda-java-profiler/extension/src/main/java/com/amazonaws/services/lambda/extension/PreMain.java b/experimental/aws-lambda-java-profiler/extension/src/main/java/com/amazonaws/services/lambda/extension/PreMain.java index c0522641..1a51684e 100644 --- a/experimental/aws-lambda-java-profiler/extension/src/main/java/com/amazonaws/services/lambda/extension/PreMain.java +++ b/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,13 +92,22 @@ 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(); } @@ -94,7 +115,9 @@ public static void stopProfiler(String fileNameSuffix) { 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); } - -} \ No newline at end of file +} diff --git a/experimental/aws-lambda-java-profiler/extension/src/main/java/com/amazonaws/services/lambda/extension/S3Manager.java b/experimental/aws-lambda-java-profiler/extension/src/main/java/com/amazonaws/services/lambda/extension/S3Manager.java index 3b55984c..0e31a242 100644 --- a/experimental/aws-lambda-java-profiler/extension/src/main/java/com/amazonaws/services/lambda/extension/S3Manager.java +++ b/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);