From 717223bd9be0f01ed6922869c103e439a30f1ede Mon Sep 17 00:00:00 2001
From: Wei Su <wsu@meta.com>
Date: Fri, 23 Aug 2024 15:16:18 -0700
Subject: [PATCH] Infer a JAVA_HOME path in DjangoBench

Summary: Try finding a JAVA_HOME path when running DjangoBench's DB role to reduce the possibility of Cassandra not being able to find JVM.

Reviewed By: ahmadelyoussef

Differential Revision: D61638351

fbshipit-source-id: 8539b77adb8d1341682d8250b21b2163ae0048c1
---
 packages/common/find_java_home.py             | 42 +++++++++++++++++++
 packages/django_workload/templates/run.sh     | 16 ++++++-
 .../templates/proj_root/scripts/utils.py      |  4 +-
 3 files changed, 59 insertions(+), 3 deletions(-)
 create mode 100755 packages/common/find_java_home.py

diff --git a/packages/common/find_java_home.py b/packages/common/find_java_home.py
new file mode 100755
index 0000000..62cb7a3
--- /dev/null
+++ b/packages/common/find_java_home.py
@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+
+import os
+import pathlib
+import platform
+import subprocess
+
+
+def find_java_home() -> str:
+    # Try finding a home path for java 8
+    candidates = [
+        "/usr/lib/jvm/java-1.8.0-openjdk",
+        "/usr/lib/jvm/java-1.8.0-jre",
+        "/usr/lib/jvm/java-8-openjdk",
+        "/usr/lib/jvm/java-8-jre",
+        "/usr/lib/jvm/openjdk-8",
+        "/usr/lib/jvm/jre-1.8.0",
+        "/usr/lib/jvm/jre-1.8.0-openjdk",
+    ]
+    archname = platform.machine()
+    if archname == "x86_64":
+        archname = "amd64"
+    elif archname == "aarch64":
+        archname = "arm64"
+    for path in candidates:
+        if os.path.exists(f"{path}/bin/java"):
+            return path
+        path_with_arch = f"{path}-{archname}"
+        if os.path.exists(f"{path_with_arch}/bin/java"):
+            return path_with_arch
+    # If none of the candidate exists, try find through `java` command
+    try:
+        java_path = subprocess.check_output(["which", "java"], text=True).strip()
+        java_home = str(pathlib.Path(os.path.realpath(java_path)).parents[1])
+    except subprocess.CalledProcessError:
+        java_home = ""
+
+    return java_home
+
+
+if __name__ == "__main__":
+    print(find_java_home())
diff --git a/packages/django_workload/templates/run.sh b/packages/django_workload/templates/run.sh
index 1516a00..457be7f 100755
--- a/packages/django_workload/templates/run.sh
+++ b/packages/django_workload/templates/run.sh
@@ -4,13 +4,20 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+# Abs path of benchmarks/django_workload/bin
 SCRIPT_ROOT="$(dirname "$(readlink -f "$0")")"
+# Abs path to DCPerf root
+BENCHPRESS_ROOT="$(readlink -f "${SCRIPT_ROOT}/../../..")"
 MEMCACHED_PID=
+CLEANUP_REQS=0
 
 if [ -z "$JAVA_HOME" ]; then
-    export JAVA_HOME="/etc/alternatives/jre_1.8.0_openjdk"
+  _JAVA_HOME="$("${BENCHPRESS_ROOT}"/packages/common/find_java_home.py)"
+  export JAVA_HOME="${_JAVA_HOME}"
+  echo "JAVA_HOME is not set, so setting it to ${JAVA_HOME}."
 fi
 
+# shellcheck disable=SC2317
 cleanup() {
   echo "Stopping services ..."
   cd "${SCRIPT_ROOT}/.." || exit 1
@@ -21,7 +28,14 @@ cleanup() {
   [ -n "$MEMCACHED_PID" ] && { echo "Stopping memcached"; kill "$MEMCACHED_PID" || true; }
   # Stop Cassandra
   [ -f cassandra.pid ] && { echo "Stopping cassandra"; kill "$(cat cassandra.pid)" || true; }
+  # Kill Siege
+  SIEGE_PID="$(pgrep siege)"
+  [ -n "$SIEGE_PID" ] && { echo "Killing siege"; kill -9 "$SIEGE_PID" || true; }
   echo "Done"
+  if [ "$CLEANUP_REQS" -gt 0 ]; then
+    exit
+  fi
+  CLEANUP_REQS=$((CLEANUP_REQS + 1))
 }
 
 trap 'cleanup' ERR EXIT SIGINT SIGTERM
diff --git a/packages/spark_standalone/templates/proj_root/scripts/utils.py b/packages/spark_standalone/templates/proj_root/scripts/utils.py
index 3fd61a4..a6febbb 100755
--- a/packages/spark_standalone/templates/proj_root/scripts/utils.py
+++ b/packages/spark_standalone/templates/proj_root/scripts/utils.py
@@ -37,7 +37,7 @@ def launch_proc(cmd, cwd, stdout, stderr, env):
 def run_cmd(
     cmd: List[str],
     cwd: str,
-    outfile: str,
+    outfile: Optional[str],
     env: Dict[str, str],
     for_real: bool,
     print_cmd: bool = True,
@@ -112,7 +112,7 @@ def find_java_home() -> str:
     # If none of the candidate exists, try find through `java` command
     try:
         java_path = subprocess.check_output(["which", "java"], text=True).strip()
-        java_home = str(pathlib.Path(os.path.realpath(java_path)).parents[2])
+        java_home = str(pathlib.Path(os.path.realpath(java_path)).parents[1])
     except subprocess.CalledProcessError:
         java_home = ""