Add cassandra profile command with async-profiler pass-through to Pyroscope

bin/end-to-end-test

@@ -577,6 +577,19 @@ step_cassandra_start_stop() {
     easy-db-lab cassandra restart
 }
 
+step_cassandra_profile() {
+    if [ "$ENABLE_CASSANDRA" != true ]; then
+        echo "=== Skipping Cassandra profile test (use --cassandra to enable) ==="
+        return 0
+    fi
+    echo "=== Testing cassandra profile start/stop ==="
+    # Use a 5-second interval with alloc profiling; start runs in background via systemd
+    easy-db-lab cassandra profile start --loop-interval 5 -e alloc
+    echo "Profiling running, waiting 10 seconds for at least one upload cycle..."
+    sleep 10
+    easy-db-lab cassandra profile stop
+}
+
 step_test_ssh_nodetool() {
     if [ "$ENABLE_CASSANDRA" != true ]; then
         echo "=== Skipping SSH nodetool test (use --cassandra to enable) ==="
@@ -1243,6 +1256,7 @@ STEPS_WORKDIR=(
     "step_bulk_writer_sidecar:Test bulk writer (sidecar)"
     "step_connector_writer:Test connector writer"
     "step_cassandra_start_stop:Cassandra start/stop cycle"
+    "step_cassandra_profile:Test cassandra profile command"
     "step_clickhouse_start:Start ClickHouse"
     "step_clickhouse_test:Test ClickHouse"
     "step_clickhouse_stop:Stop ClickHouse"

packer/cassandra/bin/flamegraph-to-pyroscope

@@ -0,0 +1,122 @@
+#!/bin/bash
+
+# flamegraph-to-pyroscope: Continuous profiling using async-profiler --loop, uploading to Pyroscope.
+#
+# Uses asprof's native --loop mode: profiles continuously, writing one JFR file per interval.
+# Each completed file is uploaded to Pyroscope and removed. Runs until stopped (SIGTERM).
+#
+# Usage: flamegraph-to-pyroscope [--interval N] [-- asprof-args...]
+#   --interval N   Loop duration in seconds per upload cycle (default: 60)
+#   --             Everything after is passed to asprof
+#
+# Examples:
+#   flamegraph-to-pyroscope
+#   flamegraph-to-pyroscope --interval 30
+#   flamegraph-to-pyroscope --interval 5 -- -e alloc -t
+#
+# The Cassandra PID is auto-detected via cassandra-pid.
+# PYROSCOPE_SERVER_ADDRESS is sourced from /etc/default/cassandra.
+
+set -e
+
+ASPROF=/usr/local/async-profiler/bin/asprof
+APP_NAME="cassandra.$(hostname)"
+OUTDIR=$(mktemp -d /tmp/flamegraph-XXXXXX)
+ASPROF_PID=""
+INTERVAL=60
+EXTRA_ARGS=()
+
+# Parse CLI args
+while [ $# -gt 0 ]; do
+    case "$1" in
+        --interval)
+            INTERVAL="$2"
+            shift 2
+            ;;
+        --)
+            shift
+            EXTRA_ARGS=("$@")
+            break
+            ;;
+        *)
+            echo "Unknown option: $1" >&2
+            echo "Usage: flamegraph-to-pyroscope [--interval N] [-- asprof-args...]" >&2
+            exit 1
+            ;;
+    esac
+done
+
+# Source environment for Pyroscope address
+# shellcheck source=/dev/null
+[ -f /etc/default/cassandra ] && source /etc/default/cassandra
+
+PYROSCOPE_ADDR="${PYROSCOPE_SERVER_ADDRESS:-http://localhost:4040}"
+
+# Returns 0 if the file hasn't been modified in the last 2 seconds (considered complete).
+# Prevents uploading a file that asprof is still writing.
+file_is_stable() {
+    local f="$1"
+    local now mtime
+    now=$(date +%s)
+    mtime=$(stat -c %Y "$f")
+    [ $((now - mtime)) -ge 2 ]
+}
+
+upload_file() {
+    local f="$1"
+    local until from
+    # Use file mtime as the end of the profiling interval.
+    until=$(stat -c %Y "$f")
+    from=$((until - INTERVAL))
+    curl -sf -X POST \
+        "${PYROSCOPE_ADDR}/ingest?name=${APP_NAME}&spyName=asyncprofiler&format=jfr&from=${from}&until=${until}" \
+        -H "Content-Type: application/octet-stream" \
+        --data-binary @"$f" \
+        || echo "Upload failed at $(date -Iseconds), continuing..."
+    rm -f "$f"
+    echo "Profile sent to Pyroscope ($(date -d @"${from}" -Iseconds) to $(date -d @"${until}" -Iseconds))"
+}
+
+cleanup() {
+    echo "Stopping profiler..."
+    if [ -n "$ASPROF_PID" ]; then
+        kill "$ASPROF_PID" 2>/dev/null || true
+        wait "$ASPROF_PID" 2>/dev/null || true
+    fi
+    # Upload any remaining completed files before exiting.
+    for f in "${OUTDIR}"/profile-*.jfr; do
+        [ -f "$f" ] && upload_file "$f" || true
+    done
+    rm -rf "$OUTDIR"
+}
+trap cleanup EXIT
+
+# Get the Cassandra PID
+PID=$(cassandra-pid)
+
+if [ -z "$PID" ] || [ "$PID" = "0" ]; then
+    echo "Error: Could not determine Cassandra PID" >&2
+    exit 1
+fi
+
+echo "Starting continuous profiling (interval: ${INTERVAL}s, PID: $PID, app: $APP_NAME)..."
+
+# Start async-profiler with --loop for native continuous profiling.
+# --loop <N>s: profile in a loop, writing one output file per interval, then reset.
+# %t in the filename is replaced with a timestamp; .jfr extension sets the output format.
+sudo "$ASPROF" --loop "${INTERVAL}s" \
+    -f "${OUTDIR}/profile-%t.jfr" \
+    "${EXTRA_ARGS[@]}" "$PID" &
+ASPROF_PID=$!
+
+# Poll for completed profile files and upload them. Skip files modified in the last
+# 2 seconds so we never upload a JFR that asprof is still writing.
+# Files are removed after upload so only unprocessed files remain at any time.
+while kill -0 "$ASPROF_PID" 2>/dev/null; do
+    for f in "${OUTDIR}"/profile-*.jfr; do
+        [ -f "$f" ] && file_is_stable "$f" && upload_file "$f" || true
+    done
+    sleep 5
+done
+
+wait "$ASPROF_PID" 2>/dev/null || true

src/main/kotlin/com/rustyrazorblade/easydblab/commands/cassandra/Cassandra.kt

@@ -32,6 +32,7 @@ import picocli.CommandLine.Spec
         UseCassandra::class,
         WriteConfig::class,
         Stress::class,
+        Profile::class,
     ],
 )
 class Cassandra : Runnable {

src/main/kotlin/com/rustyrazorblade/easydblab/commands/cassandra/Profile.kt

@@ -0,0 +1,33 @@
+package com.rustyrazorblade.easydblab.commands.cassandra
+
+import picocli.CommandLine.Command
+import picocli.CommandLine.Model.CommandSpec
+import picocli.CommandLine.Spec
+
+/**
+ * Parent command for async-profiler flamegraph operations on Cassandra nodes.
+ *
+ * Profiling runs as a long-lived systemd service on each node, uploading flamegraph
+ * data to Pyroscope once per cycle (default: every 60 seconds).
+ *
+ * Sub-commands:
+ * - start: Start continuous profiling on all Cassandra nodes
+ * - stop:  Stop continuous profiling on all Cassandra nodes
+ */
+@Command(
+    name = "profile",
+    description = ["Manage async-profiler flamegraph collection on Cassandra nodes."],
+    mixinStandardHelpOptions = true,
+    subcommands = [
+        ProfileStart::class,
+        ProfileStop::class,
+    ],
+)
+class Profile : Runnable {
+    @Spec
+    lateinit var spec: CommandSpec
+
+    override fun run() {
+        spec.commandLine().usage(System.out)
+    }
+}

src/main/kotlin/com/rustyrazorblade/easydblab/commands/cassandra/ProfileStart.kt

@@ -0,0 +1,77 @@
+package com.rustyrazorblade.easydblab.commands.cassandra
+
+import com.rustyrazorblade.easydblab.annotations.RequireProfileSetup
+import com.rustyrazorblade.easydblab.annotations.RequireSSHKey
+import com.rustyrazorblade.easydblab.commands.PicoBaseCommand
+import com.rustyrazorblade.easydblab.configuration.ServerType
+import com.rustyrazorblade.easydblab.configuration.toHost
+import com.rustyrazorblade.easydblab.events.Event
+import com.rustyrazorblade.easydblab.services.HostOperationsService
+import com.rustyrazorblade.easydblab.services.ProfilingService
+import org.koin.core.component.inject
+import picocli.CommandLine.Command
+import picocli.CommandLine.Option
+import picocli.CommandLine.Unmatched
+
+/**
+ * Start continuous async-profiler flamegraph collection on all Cassandra nodes.
+ *
+ * Starts the `flamegraph-cassandra` systemd service on each node. The service loops
+ * indefinitely: profile for --loop-interval seconds (default: 60), upload to
+ * Pyroscope, repeat — until stopped with `cassandra profile stop`.
+ *
+ * Unrecognized options are passed directly to asprof. Common options:
+ *   -e <event>          Profiling event: cpu, alloc, lock, wall, nativemem (default: cpu)
+ *   -e cpu,alloc,lock   Multiple events (JFR format required, used automatically)
+ *   --alloc <size>      Allocation sampling interval (e.g. 500k)
+ *   --lock <duration>   Lock profiling threshold (e.g. 10ms)
+ *   --nativemem <size>  Native memory leak profiling
+ *   --all               Enable cpu, wall, alloc, live, lock, nativemem simultaneously
+ *   -i <interval>       Sampling interval (e.g. 50ms)
+ *   -t                  Split stack traces by thread
+ *
+ * Examples:
+ *   cassandra profile start
+ *   cassandra profile start -e alloc
+ *   cassandra profile start --loop-interval 30 -e cpu,alloc,lock
+ *   cassandra profile start --all --alloc 2m
+ */
+@RequireProfileSetup
+@RequireSSHKey
+@Command(
+    name = "start",
+    description = [
+        "Start continuous profiling on all Cassandra nodes.",
+        "Data is uploaded to Pyroscope every --loop-interval seconds (default: 60).",
+        "Unrecognized options are passed directly to asprof.",
+    ],
+    mixinStandardHelpOptions = true,
+)
+class ProfileStart : PicoBaseCommand() {
+    private val profilingService: ProfilingService by inject()
+    private val hostOperationsService: HostOperationsService by inject()
+
+    @Option(
+        names = ["--loop-interval"],
+        description = ["Loop duration in seconds: profile this long, upload, repeat (default: 60)."],
+    )
+    var loopInterval: Int? = null
+
+    @Unmatched
+    var profilerArgs: MutableList<String> = mutableListOf()
+
+    override fun execute() {
+        val cassandraHosts = clusterState.hosts[ServerType.Cassandra]
+        if (cassandraHosts.isNullOrEmpty()) {
+            eventBus.emit(Event.Profiling.Error("cluster", "No Cassandra nodes found. Is the cluster running?"))
+            return
+        }
+
+        // Best-effort: failures on individual nodes emit Event.Profiling.Error
+        // but don't abort profiling on remaining nodes.
+        val args = profilerArgs.toList()
+        hostOperationsService.withHosts(clusterState.hosts, ServerType.Cassandra, parallel = true) { host ->
+            profilingService.startProfiling(host.toHost(), args, loopInterval)
+        }
+    }
+}

src/main/kotlin/com/rustyrazorblade/easydblab/commands/cassandra/ProfileStop.kt

@@ -0,0 +1,43 @@
+package com.rustyrazorblade.easydblab.commands.cassandra
+
+import com.rustyrazorblade.easydblab.annotations.RequireProfileSetup
+import com.rustyrazorblade.easydblab.annotations.RequireSSHKey
+import com.rustyrazorblade.easydblab.commands.PicoBaseCommand
+import com.rustyrazorblade.easydblab.configuration.ServerType
+import com.rustyrazorblade.easydblab.configuration.toHost
+import com.rustyrazorblade.easydblab.events.Event
+import com.rustyrazorblade.easydblab.services.HostOperationsService
+import com.rustyrazorblade.easydblab.services.ProfilingService
+import org.koin.core.component.inject
+import picocli.CommandLine.Command
+
+/**
+ * Stop continuous async-profiler flamegraph collection on all Cassandra nodes.
+ *
+ * Stops the `flamegraph-cassandra` systemd service on each node in parallel.
+ */
+@RequireProfileSetup
+@RequireSSHKey
+@Command(
+    name = "stop",
+    description = ["Stop continuous profiling on all Cassandra nodes."],
+    mixinStandardHelpOptions = true,
+)
+class ProfileStop : PicoBaseCommand() {
+    private val profilingService: ProfilingService by inject()
+    private val hostOperationsService: HostOperationsService by inject()
+
+    override fun execute() {
+        val cassandraHosts = clusterState.hosts[ServerType.Cassandra]
+        if (cassandraHosts.isNullOrEmpty()) {
+            eventBus.emit(Event.Profiling.Error("cluster", "No Cassandra nodes found. Is the cluster running?"))
+            return
+        }
+
+        // Best-effort: failures on individual nodes emit Event.Profiling.Error
+        // but don't abort stopping on remaining nodes.
+        hostOperationsService.withHosts(clusterState.hosts, ServerType.Cassandra, parallel = true) { host ->
+            profilingService.stop(host.toHost())
+        }
+    }
+}

src/main/kotlin/com/rustyrazorblade/easydblab/commands/exec/ExecRun.kt

@@ -12,6 +12,7 @@ import com.rustyrazorblade.easydblab.configuration.getHosts
 import com.rustyrazorblade.easydblab.configuration.toHost
 import com.rustyrazorblade.easydblab.events.Event
 import com.rustyrazorblade.easydblab.services.HostOperationsService
+import com.rustyrazorblade.easydblab.shell.shellQuote
 import org.koin.core.component.inject
 import picocli.CommandLine.Command
 import picocli.CommandLine.Mixin
@@ -125,18 +126,6 @@ class ExecRun : PicoBaseCommand() {
         return "$toolName-$epoch"
     }
 
-    /**
-     * Shell-quotes a string if it contains characters that need escaping.
-     * Wraps in single quotes with proper escaping of embedded single quotes.
-     */
-    internal fun String.shellQuote(): String {
-        if (isEmpty()) return "''"
-        // If it only contains safe characters, no quoting needed
-        if (matches(Regex("[a-zA-Z0-9_./:=@%+,-]+"))) return this
-        // Wrap in single quotes, escaping any embedded single quotes
-        return "'" + replace("'", "'\\''") + "'"
-    }
-
     internal fun buildSystemdRunCommand(
         unitName: String,
         commandString: String,