From bbe1cbe2b27bddadd3f3cfcfa9bbb5b4e8f4f535 Mon Sep 17 00:00:00 2001
From: daijro <daijro.dev@gmail.com>
Date: Sat, 30 Nov 2024 21:15:43 -0600
Subject: [PATCH] Memory benchmark scripts via podman #87

---
 .gitignore                     |  1 -
 scripts/benchmark/Dockerfile   | 25 +++++++++
 scripts/benchmark/Makefile     |  7 +++
 scripts/benchmark/benchmark.py | 97 ++++++++++++++++++++++++++++++++++
 4 files changed, 129 insertions(+), 1 deletion(-)
 create mode 100644 scripts/benchmark/Dockerfile
 create mode 100644 scripts/benchmark/Makefile
 create mode 100644 scripts/benchmark/benchmark.py

diff --git a/.gitignore b/.gitignore
index ad5a06c..8cefd1e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -18,7 +18,6 @@ scripts/*.png
 .vscode
 /tests/*.disabled
 k8s/
-benchmark/
 
 # Old data
 _old/
diff --git a/scripts/benchmark/Dockerfile b/scripts/benchmark/Dockerfile
new file mode 100644
index 0000000..7039c3d
--- /dev/null
+++ b/scripts/benchmark/Dockerfile
@@ -0,0 +1,25 @@
+FROM debian:latest
+
+# Set working directory
+WORKDIR /app
+
+# Install Python and essential tools
+RUN apt-get update && apt-get install -y \
+    python3 \
+    python3-pip \
+    curl \
+    wget \
+    xvfb \
+    && apt-get clean
+
+# Install Playwright dependencies
+RUN pip3 install --no-cache-dir playwright tabulate camoufox[geoip] --break-system-packages && \
+    playwright install-deps && \
+    playwright install firefox && \
+    python3 -m camoufox fetch
+
+# Copy the benchmark script
+COPY benchmark.py /app/
+
+# Default entrypoint to run the benchmark
+ENTRYPOINT ["python3", "/app/benchmark.py"]
diff --git a/scripts/benchmark/Makefile b/scripts/benchmark/Makefile
new file mode 100644
index 0000000..8a6e902
--- /dev/null
+++ b/scripts/benchmark/Makefile
@@ -0,0 +1,7 @@
+build:
+	podman build -t camoufox-benchmark .
+
+run:
+	podman run camoufox-benchmark --mode headless --browser firefox
+	podman run camoufox-benchmark --mode headless --browser camoufox
+	podman run camoufox-benchmark --mode headless --browser camoufox-ubo
diff --git a/scripts/benchmark/benchmark.py b/scripts/benchmark/benchmark.py
new file mode 100644
index 0000000..52895af
--- /dev/null
+++ b/scripts/benchmark/benchmark.py
@@ -0,0 +1,97 @@
+import argparse
+import subprocess
+import time
+
+from camoufox.pkgman import launch_path
+from camoufox.sync_api import Camoufox
+from camoufox.virtdisplay import VirtualDisplay
+from playwright.sync_api import sync_playwright
+from tabulate import tabulate
+
+# URLs to benchmark
+urls = ["about:blank", "https://google.com", "https://yahoo.com"]
+
+
+def get_firefox_memory(name):
+    """Get the total memory usage of all processes named 'firefox'."""
+    try:
+        result = subprocess.run(["ps", "-C", name, "-o", "rss="], capture_output=True, text=True)
+        memory_kb = sum(int(line.strip()) for line in result.stdout.splitlines())
+        memory_mb = memory_kb / 1024  # Convert KB to MB
+        return memory_mb
+    except Exception as e:
+        print(f"Error getting Firefox memory: {e}")
+        return 0
+
+
+def get_average_memory(name, duration):
+    """Monitor memory usage for Firefox over a duration (seconds) and return the average."""
+    memory_samples = []
+    for n in range(duration):
+        memory_samples.append(get_firefox_memory(name))
+        # print(f"> Mem ({n}sec): {memory_samples[-1]} MB")
+        time.sleep(1)
+    return sum(memory_samples) / len(memory_samples) if memory_samples else 0
+
+
+def run_playwright(mode, browser_name):
+    headless = mode == "headless"
+    memory_usage = []
+    # Set up virtual display
+    virt = VirtualDisplay()
+    env = {"DISPLAY": virt.get()}
+
+    if browser_name == "camoufox-ubo":
+        camoufox = Camoufox(headless=headless, env=env)
+        browser = camoufox.start()
+    elif browser_name == "firefox":
+        playwright = sync_playwright().start()
+        browser = playwright.firefox.launch(headless=headless, env=env)
+    elif browser_name == "camoufox":
+        playwright = sync_playwright().start()
+        browser = playwright.firefox.launch(
+            headless=headless, env=env, executable_path=launch_path()
+        )
+
+    for url in urls:
+        page = browser.new_page()
+        page.goto(url)
+        time.sleep(5)  # Allow the page to load
+        memory = get_average_memory(
+            name="camoufox-bin" if browser_name.startswith('camoufox') else 'firefox', duration=10
+        )
+        memory_usage.append((url, memory))
+        page.close()
+
+    browser.close()
+
+    return memory_usage
+
+
+if __name__ == "__main__":
+    # Set up argument parsing
+    parser = argparse.ArgumentParser(description="Run a browser memory benchmark.")
+    parser.add_argument(
+        "--mode",
+        type=str,
+        choices=["headless", "headful"],
+        required=True,
+        help="Mode to run the browser in (headless or headful).",
+    )
+    parser.add_argument(
+        "--browser",
+        type=str,
+        choices=["firefox", "camoufox", "camoufox-ubo"],
+        required=True,
+        help="Browser to use for the benchmark.",
+    )
+
+    args = parser.parse_args()
+
+    # Run the benchmark
+    results = run_playwright(args.mode, args.browser)
+
+    # Format results as a table
+    print(f"\n=== MEMORY RESULTS FOR {args.browser.upper()} ===")
+    table = [["URL", "Memory Usage (MB)"]] + results
+    print(tabulate(table, headers="firstrow", tablefmt="grid"))