Flamegraph generation using DTrace

generate_flamegraph.sh

#!/bin/bash

# This script generates flame graphs by using DTrace to sample the active stack frames of the
# specified process and then https://github.com/brendangregg/FlameGraph to convert the
# stack traces into the flame graph.
# Note: On all current MacOS versions System Integrity Protection (SIP) is enabled by default and
#       prevents most uses of dtrace. The usual way to make dtrace work on MacOS is to boot into
#       recovery mode and disable some of the SIP protections with `csrutil enable --without dtrace`

# Usage:
# `./generate_flamegraph <duration(seconds)> <process-name> <(optional)flamegraph-title>`

# this is based off some work by fanquake: https://github.com/fanquake/core-review/tree/master/flamegraph
# NOTE: for this to work, you need to clone the Flame Graph repo and set the appropriate variable
ABSOLUTE_PATH_TO_FLAME_GRAPH_REPO=""
RUNNING_PROCESS=true        # is this a running process? or we need to start it before profiling
OUTPUT_FOLDER="flamegraphs" # relative path to flamegraph output

DURATION=${1:-30} # profiling duration in seconds
PROCESS=$2        # process for which the flame graph will be generated

PROCESS_CMD="./$PROCESS" # the process for which profiling will run

NOW=$(date '+%F_%H:%M:%S') # timestamp
TITLE=${3:-$NOW} # default is timestamp
OUTPUT="$(pwd)/${OUTPUT_FOLDER}/${TITLE}_flamegraph"

# Verify that required files exists
[ ! -f "${ABSOLUTE_PATH_TO_FLAME_GRAPH_REPO}/stackcollapse.pl" ] || [ ! -f "${ABSOLUTE_PATH_TO_FLAME_GRAPH_REPO}/flamegraph.pl" ] && {
    echo "The specified path for the FlameGraph repo does not contain the neccessary files"
    echo "Make sure that 'ABSOLUTE_PATH_TO_FLAME_GRAPH_REPO=${ABSOLUTE_PATH_TO_FLAME_GRAPH_REPO}' is the correct path"
    exit 1
}

generate_flamegraph () {
    # Converting stack traces into the flame graph
    "${ABSOLUTE_PATH_TO_FLAME_GRAPH_REPO}"/stackcollapse.pl "${OUTPUT}.stacks" >"${OUTPUT}.folded"
    "${ABSOLUTE_PATH_TO_FLAME_GRAPH_REPO}"/flamegraph.pl "${OUTPUT}.folded" --color bitcoin --title "${PROCESS}" --width 1600 >"${OUTPUT}.svg"
    echo "Generated flame graph can be found at $OUTPUT"
}

# Profiling
mkdir -p $OUTPUT_FOLDER # create output directory if not exists
sudo true               # make sure that we clear sudo's password prompt before we start running anything time-sensitive

if [[ "$RUNNING_PROCESS" == false ]]; then
    $PROCESS_CMD & # run the specified process in the background
    PID=$!
else
    # check that specified process is running
    PID=$(pgrep "$PROCESS")
    [ "$PID" == "" ] && {
        echo "Process $PROCESS is not running!"
        exit 1
    }
fi

trap generate_flamegraph EXIT # on exit generate flamegraph 
echo "Profiling of $PROCESS($PID) has started"
# capture $DURATION worth of stackframes using DTrace at 99 Hertz for the specified process
if [[ ${DURATION} == 0 ]]; then
    echo "Capturing stackframes until interruption(Ctrl+C)..."
    sudo dtrace -x ustackframes=100 -o "${OUTPUT}.stacks" -n 'profile-99 /pid == $1 && arg1/ { @[ustack()] = count(); }' $PID
else
    echo "Capturing stackframes for ${DURATION}s..."
    sudo dtrace -x ustackframes=100 -o "${OUTPUT}.stacks" -n 'profile-99 /pid == $1 && arg1/ { @[ustack()] = count(); } $2 { exit(0); }' $PID "tick-${DURATION}s"
fi
[ "$RUNNING_PROCESS" == false ] && {
    kill $! # kill running process otherwise it might run for ever
}