From b889f577e0c01509b61bad0118188ba8f753deac Mon Sep 17 00:00:00 2001
From: Laszlo Bodor <bodorlaszlo0202@gmail.com>
Date: Wed, 5 Jun 2024 17:44:52 +0200
Subject: [PATCH] TEZ-4568: ProfileServlet: add html to output formats and
 prepare for profiler 3.0

---
 .../apache/tez/common/web/ProfileServlet.java | 32 +++++++++++++------
 1 file changed, 22 insertions(+), 10 deletions(-)

diff --git a/tez-common/src/main/java/org/apache/tez/common/web/ProfileServlet.java b/tez-common/src/main/java/org/apache/tez/common/web/ProfileServlet.java
index 8d55e6f8cf..2884bbb4d7 100644
--- a/tez-common/src/main/java/org/apache/tez/common/web/ProfileServlet.java
+++ b/tez-common/src/main/java/org/apache/tez/common/web/ProfileServlet.java
@@ -31,6 +31,9 @@
 import java.io.File;
 import java.io.IOException;
 import java.io.PrintStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
@@ -50,19 +53,19 @@
  * //  -m method         fully qualified method name: 'ClassName.methodName'
  * //  -t                profile different threads separately
  * //  -s                simple class names instead of FQN
- * //  -o fmt[,fmt...]   output format: summary|traces|flat|collapsed|svg|tree|jfr
+ * //  -o fmt[,fmt...]   output format: summary|traces|flat|collapsed|svg|tree|jfr|html
  * //  --width px        SVG width pixels (integer)
  * //  --height px       SVG frame height pixels (integer)
  * //  --minwidth px     skip frames smaller than px (double)
  * //  --reverse         generate stack-reversed FlameGraph / Call tree
  * Example:
- * - To collect 30 second CPU profile of current process (returns FlameGraph svg)
+ * - To collect 10 second CPU profile of current process (returns FlameGraph html)
  * {@literal curl "http://localhost:10002/prof"}
  * - To collect 1 minute CPU profile of current process and output in tree format (html)
  * {@literal curl  "http://localhost:10002/prof?output=tree&duration=60"}
- * - To collect 30 second heap allocation profile of current process (returns FlameGraph svg)
+ * - To collect 10 second heap allocation profile of current process (returns FlameGraph html)
  * {@literal curl "http://localhost:10002/prof?event=alloc"}
- * - To collect lock contention profile of current process (returns FlameGraph svg)
+ * - To collect lock contention profile of current process (returns FlameGraph html)
  * {@literal curl "http://localhost:10002/prof?event=lock"}
  * Following event types are supported (default is 'cpu') (NOTE: not all OS'es support all events)
  * // Perf events:
@@ -95,7 +98,6 @@ public class ProfileServlet extends HttpServlet {
   private static final String CONTENT_TYPE_TEXT = "text/plain; charset=utf-8";
   private static final String ASYNC_PROFILER_HOME_ENV = "ASYNC_PROFILER_HOME";
   private static final String ASYNC_PROFILER_HOME_SYSTEM_PROPERTY = "async.profiler.home";
-  private static final String PROFILER_SCRIPT = "/profiler.sh";
   private static final int DEFAULT_DURATION_SECONDS = 10;
   private static final AtomicInteger ID_GEN = new AtomicInteger(0);
   public static final String OUTPUT_DIR = System.getProperty("java.io.tmpdir") + "/prof-output";
@@ -140,7 +142,7 @@ public static Event fromInternalName(final String name) {
   }
 
   enum Output {
-    SUMMARY, TRACES, FLAT, COLLAPSED, SVG, TREE, JFR
+    SUMMARY, TRACES, FLAT, COLLAPSED, SVG, TREE, JFR, HTML
   }
 
   private final Lock profilerLock = new ReentrantLock();
@@ -216,7 +218,7 @@ public void doGet(HttpServletRequest request, HttpServletResponse response) thro
                     + (method == null ? event.name().toLowerCase() : method) + "-" + ID_GEN.incrementAndGet() + "."
                     + output.name().toLowerCase());
             List<String> cmd = new ArrayList<>();
-            cmd.add(asyncProfilerHome + PROFILER_SCRIPT);
+            cmd.add(getProfilerScriptPath());
             cmd.add("-e");
             cmd.add(method == null ? event.getInternalName() : method);
             cmd.add("-d");
@@ -298,6 +300,16 @@ public void doGet(HttpServletRequest request, HttpServletResponse response) thro
     out.close();
   }
 
+  /**
+   * Get the path of the profiler script to be executed.
+   * Before async-profiler 3.0, the script was named profiler.sh, and after 3.0 it's bin/asprof
+   * @return
+   */
+  private String getProfilerScriptPath() {
+    Path defaultPath = Paths.get(asyncProfilerHome + "/bin/asprof");
+    return Files.exists(defaultPath)? defaultPath.toString() : asyncProfilerHome + "/profiler.sh";
+  }
+
   private Integer getInteger(final HttpServletRequest req, final String param, final Integer defaultValue) {
     final String value = req.getParameter(param);
     if (value != null) {
@@ -349,11 +361,11 @@ private Output getOutput(final HttpServletRequest req) {
       try {
         return Output.valueOf(outputArg.trim().toUpperCase());
       } catch (IllegalArgumentException e) {
-        LOG.warn("Output format value is invalid, returning with default SVG");
-        return Output.SVG;
+        LOG.warn("Output format value is invalid, returning with default HTML");
+        return Output.HTML;
       }
     }
-    return Output.SVG;
+    return Output.HTML;
   }
 
   private void setResponseHeader(final HttpServletResponse response) {