Profiling utils (#595)

* Profiling utils

* Fix CI trigger

Author: levkk

.github/workflows/ci.yml

@@ -1,6 +1,8 @@
 name: ci
 on:
   push:
+    branches:
+      - main
   pull_request:
     types: [opened, synchronize, reopened]
 

cli.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+#
+# Dev CLI.
+#
+
+# Connect to the admin database.
+function admin() {
+    PGPASSWORD=pgdog psql -h 127.0.0.1 -p 6432 -U admin admin
+}
+
+# Run pgbench.
+#
+# Arguments:
+#
+# - protocol: simple|extended|prepared
+#
+function bench() {
+    PGPASSWORD=pgdog pgbench -h 127.0.0.1 -p 6432 -U pgdog pgdog --protocol ${1:-simple}
+}
+
+# Parse command
+case "$1" in
+    admin)
+        admin
+        ;;
+    bench)
+        bench $1
+        ;;
+    *)
+        echo "Usage: $0 {admin} {bench}"
+        exit 1
+        ;;
+esac

integration/pgbench/run.sh

@@ -7,5 +7,6 @@ run_pgdog
 wait_for_pgdog
 
 bash ${SCRIPT_DIR}/dev.sh
+bash ${SCRIPT_DIR}/stress.sh
 
 stop_pgdog

integration/pgbench/stress.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+#
+# This will create a decent amount of churn on the network & message
+# buffers, by generating different size responses.
+#
+set -e
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+export PGPASSWORD=pgdog
+
+pushd ${SCRIPT_DIR}
+pgbench -h 127.0.0.1 -U pgdog -p 6432 pgdog -t 100 -c 10 --protocol simple -f stress.sql -P 1
+popd

integration/pgbench/stress.sql

@@ -0,0 +1,7 @@
+\set response_size (random(2, 1000000))
+
+-- Generate large response from server.
+-- Range: 2 bytes to 1M
+SELECT
+    string_agg(substr('abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789', (random()*61+1)::int, 1), '')
+FROM generate_series(1, :response_size);

pgdog/tests/flame/base/pgdog.toml

@@ -0,0 +1,6 @@
+[[databases]]
+name = "pgdog"
+host = "127.0.0.1"
+
+[admin]
+password = "pgdog"

pgdog/tests/flame/base/users.toml

@@ -0,0 +1,4 @@
+[[users]]
+name = "pgdog"
+password = "pgdog"
+database = "pgdog"

pgdog/tests/flame/profile.sh

@@ -0,0 +1,69 @@
+#!/bin/bash
+#
+# Utils for performance testing.
+#
+set -e
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+export PGDOG_BIN=${SCRIPT_DIR}/../../../target/release/pgdog
+export PGPASSWORD=pgdog
+export PGUSER=pgdog
+export PGDATABASE=pgdog
+export PGHOST=127.0.0.1
+export PGPORT=6432
+
+if [ ! -f ${PGDOG_BIN} ]; then
+    echo "PgDog is not compiled in release mode (target/release/pgdog is missing)"
+    echo "Please compile PgDog with:"
+    echo
+    printf "\tcargo build --release"
+    echo
+    echo
+    exit 1
+fi
+
+if ! which samply > /dev/null; then
+    echo "Samply profiler is not installed on this system"
+    echo "Please install Samply with:"
+    echo
+    printf "\tcargo install samply"
+    echo
+    echo
+    exit 1
+fi
+
+paranoid=$(cat /proc/sys/kernel/perf_event_paranoid)
+
+if [ ! "$paranoid" -eq "-1" ]; then
+    echo "\"/proc/sys/kernel/perf_event_paranoid\" need to be set to -1 for sampling profiler to work"
+    echo "Please set it manually by running:"
+    echo
+    printf "\techo '-1' | sudo tee /proc/sys/kernel/perf_event_paranoid"
+    echo
+    echo
+    exit 1
+fi
+
+function profile_pgdog() {
+    pushd $1
+    samply record ${PGDOG_BIN} &
+    pid=$!
+
+    while ! pg_isready > /dev/null; do
+        sleep 1
+    done
+
+    pgbench -i
+    pgbench -c 10 -t 1000000 -S -P 1 --protocol extended
+
+    kill -TERM $pid
+    popd
+}
+
+if [ ! -d "$1" ]; then
+    echo "Benchmark $1 doesn't exist."
+    echo "Available benchmarks:"
+    echo "$(ls -d */)"
+    exit 1
+fi
+
+profile_pgdog $1

pgdog/tests/flame/read_write/pgdog.toml

@@ -0,0 +1,11 @@
+[[databases]]
+name = "pgdog"
+host = "127.0.0.1"
+
+[[databases]]
+name = "pgdog"
+host = "127.0.0.1"
+role = "replica"
+
+[admin]
+password = "pgdog"

pgdog/tests/flame/read_write/users.toml

@@ -0,0 +1,4 @@
+[[users]]
+name = "pgdog"
+password = "pgdog"
+database = "pgdog"