perf: build Linux releases with profile-guided optimization

.github/workflows/package.yml

@@ -65,6 +65,14 @@ jobs:
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
 
+      - name: Generate PGO profile
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y build-essential cmake clang pkg-config libssl-dev mold
+          bash integration/ci/install-deps.sh
+          bash integration/ci/setup.sh --with-toxi
+          bash scripts/pgo-profile.sh
+
       - name: Build and push Docker image
         id: build
         uses: docker/build-push-action@v6

.github/workflows/release-binaries.yml

@@ -45,8 +45,20 @@ jobs:
         with:
           targets: ${{ matrix.target }}
 
+      - name: Generate PGO profile
+        if: runner.os == 'Linux'
+        run: |
+          bash integration/ci/install-deps.sh
+          bash integration/ci/setup.sh --with-toxi
+          bash scripts/pgo-profile.sh ${{ matrix.target }}
+
       - name: Build
-        run: cargo build --release --target ${{ matrix.target }} -p pgdog
+        run: |
+          pgo_args=()
+          if [ -f pgo.profdata ]; then
+            pgo_args=(--config "target.'cfg(all())'.rustflags=['-Cprofile-use=$PWD/pgo.profdata']")
+          fi
+          cargo build --release --target ${{ matrix.target }} -p pgdog "${pgo_args[@]}"
 
       - name: Package tarball
         run: |

.gitignore

@@ -74,3 +74,5 @@ enterprise/web2/.vite/*
 # Generated by integration/tls/backend/run.sh
 integration/tls/backend/rotate/
 integration/tls/backend/docker/certs/
+
+pgo.profdata

Dockerfile

@@ -14,8 +14,12 @@ RUN source ~/.cargo/env && \
     if [ -n "${FEATURES}" ]; then \
         cargo_features=(--no-default-features --features "${FEATURES}"); \
     fi && \
+    pgo_args=(); \
+    if [ -f pgo.profdata ]; then \
+        pgo_args=(--config "target.'cfg(all())'.rustflags=['-Cprofile-use=/build/pgo.profdata']"); \
+    fi && \
     cd pgdog && \
-    cargo build --release "${cargo_features[@]}" && \
+    cargo build --release "${cargo_features[@]}" "${pgo_args[@]}" && \
     cd .. && \
     cargo build --release -p pgdog-primary-only-tables "${cargo_features[@]}"
 

integration/setup.sh

@@ -2,11 +2,10 @@
 set -e
 SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 OS=$(uname -s | tr '[:upper:]' '[:lower:]')
-if [[ "$OS" == "darwin" ]]; then
-    ARCH=arm64
-else
-    ARCH=amd64
-fi
+case "$(uname -m)" in
+    arm64|aarch64) ARCH=arm64 ;;
+    *) ARCH=amd64 ;;
+esac
 
 # Verify and apply required PostgreSQL system settings.
 # ALTER SYSTEM cannot run inside a function/DO block, so checks are done in bash.

scripts/pgo-profile.sh

@@ -0,0 +1,65 @@
+#!/bin/bash
+#
+# Profile-guided optimization (PGO) compiles PgDog with a profile of how often
+# each branch within our code runs under realistic traffic.
+#
+# The compiler optimizes frequently run code for speed. HOWEVER: Code the training
+# workload doesn't cover can be slightly slower than in a regular build.
+# This is why this uses our integration suites as opposed to custom scripts.
+#
+# On x86, PGO builds handle roughly 10% more transactions per second while using
+# about 10% less CPU per transaction. Code outside the hot path, e.g., COPY
+# and connection setup, were benchmarked in a test and found to be
+# unchanged on x86 and actually slightly faster on Graviton (than a bare release build).
+#
+# TODO(jkaczman): the integration suites weight code by how **often** tests hit them, not by
+# production traffic. A replay of real traffic or a TPC-C run would be more representative.
+#
+# TODO(jkaczman): COPY is measured, but still need to test some things to ensure no regressions,
+# e.g., schema sync, data sync and logical replication.
+#
+# TODO(jkaczman): compare PGO and regular release builds via Bencher to check for regressions & speedup
+
+set -euo pipefail
+
+REPO_ROOT=$(cd "$(dirname "$0")/.." && pwd)
+RAW_PROFILES_DIR=$REPO_ROOT/target/pgo-raw
+INSTRUMENTED_TARGET_DIR=$REPO_ROOT/target/pgo-instrumented
+TARGET_TRIPLE=${1:-}
+
+TRAINING_SUITES=(sql python pgbench complex schema_sync tls two_pc copy_data/data_sync)
+
+build_instrumented_pgdog() {
+    local target_args=()
+    if [ -n "$TARGET_TRIPLE" ]; then
+        target_args=(--target "$TARGET_TRIPLE")
+    fi
+
+    cargo build --release -p pgdog "${target_args[@]}" \
+        --target-dir "$INSTRUMENTED_TARGET_DIR" \
+        --config "target.'cfg(all())'.rustflags=['-Cprofile-generate=$RAW_PROFILES_DIR']"
+}
+
+record_training_profiles() {
+    export PGDOG_BIN=$INSTRUMENTED_TARGET_DIR/$TARGET_TRIPLE/release/pgdog
+
+    local suite
+    for suite in "${TRAINING_SUITES[@]}"; do
+        bash "integration/$suite/run.sh"
+    done
+}
+
+merge_profiles() {
+    local llvm_profdata
+    llvm_profdata=$(find "$(rustc --print sysroot)" -name llvm-profdata -type f | head -1)
+
+    "$llvm_profdata" merge -o pgo.profdata "$RAW_PROFILES_DIR"/*.profraw
+    echo "Profile written to $REPO_ROOT/pgo.profdata"
+}
+
+cd "$REPO_ROOT"
+rm -rf "$RAW_PROFILES_DIR"
+
+build_instrumented_pgdog
+record_training_profiles
+merge_profiles