feat: pgbench with postgres, pgbouncer and proxy

.github/actions/setup-test/action.yml

@@ -0,0 +1,22 @@
+name: Setup Test
+description: |
+  This action sets up the environment for running tests in a Rust project.
+  It installs Rust, sets up caching, and installs mise.
+
+runs:
+  using: composite
+  steps:
+    - uses: actions/checkout@v4
+    - name: Install rust
+      shell: /bin/bash -l {0}
+      run: rustup toolchain install stable --profile minimal --no-self-update
+    - name: Setup Rust cache
+      uses: Swatinem/rust-cache@v2
+      with:
+        cache-provider: buildjet
+        cache-all-crates: true
+    - uses: jdx/mise-action@v2
+      with:
+        version: 2025.1.0 # [default: latest] mise version to install
+        install: true # [default: true] run `mise install`
+        cache: true # [default: true] cache mise using GitHub's cache

.github/workflows/benchmark.yml

@@ -0,0 +1,59 @@
+name: Benchmark Proxy
+on:
+  pull_request:
+    branches:
+      - main
+  push:
+    branches:
+      - main
+  workflow_dispatch:
+
+env:
+  CARGO_TERM_COLOR: always
+
+permissions:
+  # deployments permission to deploy GitHub pages website
+  deployments: write
+  # contents permission to update benchmark contents in gh-pages branch
+  contents: write
+
+jobs:
+  benchmark:
+    name: Performance regression check
+    runs-on: buildjet-16vcpu-ubuntu-2204
+    steps:
+      - uses: actions/checkout@v4
+      - uses: ./.github/actions/setup-test
+      - run: |
+          mise run postgres:up --extra-args "--detach --wait"
+      - name: Run benchmark
+        working-directory: tests/benchmark
+        env:
+          CS_WORKSPACE_ID: ${{ secrets.CS_WORKSPACE_ID }}
+          CS_CLIENT_ACCESS_KEY: ${{ secrets.CS_CLIENT_ACCESS_KEY }}
+          CS_DEFAULT_KEYSET_ID: ${{ secrets.CS_DEFAULT_KEYSET_ID }}
+          CS_CLIENT_ID: ${{ secrets.CS_CLIENT_ID }}
+          CS_CLIENT_KEY: ${{ secrets.CS_CLIENT_KEY }}
+          RUST_BACKTRACE: "1"
+        run: mise run benchmark:continuous
+      # Download previous benchmark result from cache (if exists)
+      - name: Download previous benchmark data
+        uses: actions/cache@v4
+        with:
+          path: ./cache
+          key: ${{ runner.os }}-benchmark
+      # Run `github-action-benchmark` action
+      - name: Store benchmark result
+        uses: benchmark-action/github-action-benchmark@v1
+        with:
+          # What benchmark tool the output.txt came from
+          tool: 'customSmallerIsBetter'
+          # Where the output from the benchmark tool is stored
+          output-file-path: tests/benchmark/results/output.json
+
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          fail-on-alert: true
+          comment-on-alert: true
+          summary-always: true
+          auto-push: true
+          benchmark-data-dir-path: docs

.github/workflows/test.yml

@@ -1,4 +1,4 @@
-name: test
+name: Test
 on:
   pull_request:
     branches:
@@ -10,25 +10,14 @@ on:
 
 env:
   CARGO_TERM_COLOR: always
+
 jobs:
   test:
     name: Test
     runs-on: buildjet-16vcpu-ubuntu-2204
     steps:
       - uses: actions/checkout@v4
-      - name: Install rust
-        shell: /bin/bash -l {0}
-        run: rustup toolchain install stable --profile minimal --no-self-update
-      - name: Setup Rust cache
-        uses: Swatinem/rust-cache@v2
-        with:
-          cache-provider: buildjet
-          cache-all-crates: true
-      - uses: jdx/mise-action@v2
-        with:
-          version: 2025.1.0 # [default: latest] mise version to install
-          install: true # [default: true] run `mise install`
-          cache: true # [default: true] cache mise using GitHub's cache
+      - uses: ./.github/actions/setup-test
       - run: |
           mise run postgres:up --extra-args "--detach --wait"
       - env:

.gitignore

@@ -18,3 +18,7 @@ rust-toolchain.toml
 
 # credentials for local dev
 .env.proxy.docker
+
+## benchmark result data
+tests/benchmark/results/*.csv
+tests/benchmark/benchmark-*.png

Cargo.toml

@@ -10,17 +10,29 @@ edition = "2021"
 incremental = true
 debug = true
 
+# [profile.dev.package]# aws-lc-sys.opt-level = 3
+# proc-macro2.opt-level = 3
+# quote.opt-level = 3
+# serde_derive.opt-level = 3
+# sqlparser.opt-level = 3
+# syn.opt-level = 3
+
 [profile.dev.build-override]
 opt-level = 3
 
 [profile.test]
 incremental = true
 debug = true
 
+#  Default release profile (https://doc.rust-lang.org/cargo/reference/profiles.html)
 [profile.release]
 codegen-units = 1
 strip = "symbols"
-lto = true
+
+[profile.profiling]
+inherits = "release"
+strip = "none"
+debug = true
 
 [workspace.dependencies]
 sqlparser = { version = "^0.52", features = ["bigdecimal", "serde"] }

README.md

@@ -786,6 +786,41 @@ If the proxy is running on a host other than localhost, access on that host.
 | `-h`, `--help`          | Displays this help message                                     | -               |
 
 
+## Benchmarks
+
+Benchmarking is integrated into CI, and can be run locally.
+
+The benchmark uses `pgbench` to compare direct access to PostgreSQL against `pgbouncer` and Proxy.
+Benchmarks are executed in Docker containers, as the focus is on providing a repeatable baseline for performance.
+Your mileage may vary when comparing against a "real-world" production configuration.
+
+The benchmarks use the extended protocol option and cover:
+- the default `pgbench` transaction
+- insert, update & select of plaintext data
+- insert, update & select of encrypted data (CipherStash Proxy only)
+
+The benchmark setup includes the database configuration, but does requires access to a CipherStash account environment.
+
+### Require environment variables
+```
+CS_WORKSPACE_ID
+CS_CLIENT_ACCESS_KEY
+CS_DEFAULT_KEYSET_ID
+CS_CLIENT_ID
+CS_CLIENT_KEY
+```
+
+
+### Running the benchmark
+
+```bash
+cd tests/benchmark
+mise run benchmark
+```
+
+Results are graphed in a file called `benchmark-{YmdHM}.png` where `YmdHM` is a generated timestamp.
+Detailed results are generated in `csv` format and in the `results` directory.
+
 
 ## More info
 

mise.toml

@@ -111,6 +111,15 @@ mise run test:unit
 mise run test:integration
 """
 
+[tasks."reset"]
+alias = 'r'
+description = "Reset database"
+run = """
+mise run postgres:down
+mise run postgres:up --extra-args "--detach --wait"
+mise run postgres:setup
+"""
+
 [tasks."test:local:unit"]
 alias = 'lu'
 description = "Runs test/s"

packages/cipherstash-proxy/Cargo.toml

@@ -56,7 +56,6 @@ uuid = { version = "1.11.0", features = ["serde", "v4"] }
 x509-parser = "0.17.0"
 vitaminc-protected = "0.1.0-pre2"
 
-
 [dev-dependencies]
 recipher = "0.1.3"
 temp-env = "0.3.6"

packages/cipherstash-proxy/src/postgresql/messages/terminate.rs

@@ -1,6 +1,5 @@
-use bytes::{BufMut, BytesMut};
-
 use super::FrontendCode;
+use bytes::{BufMut, BytesMut};
 
 pub struct Terminate;
 

tests/benchmark/docker/pgbouncer/Dockerfile

@@ -0,0 +1,11 @@
+FROM alpine:3.19
+
+RUN apk add pgbouncer
+RUN adduser -D -S pgbouncer
+
+ADD pgbouncer.ini /etc/pgbouncer/pgbouncer.ini
+ADD userlist.txt /etc/pgbouncer/userlist.txt
+
+EXPOSE 6433
+USER pgbouncer
+CMD /usr/bin/pgbouncer /etc/pgbouncer/pgbouncer.ini

tests/benchmark/docker/pgbouncer/pgbouncer.ini

@@ -0,0 +1,24 @@
+[databases]
+cipherstash = host=host.docker.internal port=5532 dbname=cipherstash user=cipherstash password=password
+
+[users]
+
+[pgbouncer]
+listen_addr = *
+listen_port = 6433
+listen_backlog = 1000
+so_reuseport = 1
+
+default_pool_size = 100
+min_pool_size = 20
+
+max_db_connections = 100
+max_client_conn = 1000
+
+syslog = 0
+log_connections = 0
+log_disconnections = 0
+verbose = 0
+
+auth_type = scram-sha-256
+auth_file = /etc/pgbouncer/userlist.txt

tests/benchmark/docker/pgbouncer/userlist.txt

@@ -0,0 +1 @@
+"cipherstash" "password"

tests/benchmark/docker/pgcat/pgcat.toml

@@ -0,0 +1,111 @@
+#
+# PgCat config example.
+#
+
+#
+# General pooler settings
+[general]
+# What IP to run on, 0.0.0.0 means accessible from everywhere.
+host = "0.0.0.0"
+
+# Port to run on, same as PgBouncer used in this example.
+port = 6434
+
+# Whether to enable prometheus exporter or not.
+enable_prometheus_exporter = true
+
+# Port at which prometheus exporter listens on.
+prometheus_exporter_port = 9930
+
+# How long to wait before aborting a server connection (ms).
+connect_timeout = 5000
+
+# How much time to give `SELECT 1` health check query to return with a result (ms).
+healthcheck_timeout = 1000
+
+# How long to keep connection available for immediate re-use, without running a healthcheck query on it
+healthcheck_delay = 30000
+
+# How much time to give clients during shutdown before forcibly killing client connections (ms).
+shutdown_timeout = 60000
+
+# For how long to ban a server if it fails a health check (seconds).
+ban_time = 60 # seconds
+
+# If we should log client connections
+log_client_connections = false
+
+# If we should log client disconnections
+log_client_disconnections = false
+
+# TLS
+# tls_certificate = "server.cert"
+# tls_private_key = "server.key"
+
+# Credentials to access the virtual administrative database (pgbouncer or pgcat)
+# Connecting to that database allows running commands like `SHOW POOLS`, `SHOW DATABASES`, etc..
+admin_username = "cipherstash"
+admin_password = "password"
+
+# pool
+# configs are structured as pool.<pool_name>
+# the pool_name is what clients use as database name when connecting
+# For the example below a client can connect using "postgres://sharding_user:sharding_user@pgcat_host:pgcat_port/sharded"
+[pools.cipherstash]
+# Pool mode (see PgBouncer docs for more).
+# session: one server connection per connected client
+# transaction: one server connection per client transaction
+pool_mode = "session"
+
+# If the client doesn't specify, route traffic to
+# this role by default.
+#
+# any: round-robin between primary and replicas,
+# replica: round-robin between replicas only without touching the primary,
+# primary: all queries go to the primary unless otherwise specified.
+default_role = "any"
+
+# Query parser. If enabled, we'll attempt to parse
+# every incoming query to determine if it's a read or a write.
+# If it's a read query, we'll direct it to a replica. Otherwise, if it's a write,
+# we'll direct it to the primary.
+query_parser_enabled = false
+
+# If the query parser is enabled and this setting is enabled, we'll attempt to
+# infer the role from the query itself.
+query_parser_read_write_splitting = false
+
+# If the query parser is enabled and this setting is enabled, the primary will be part of the pool of databases used for
+# load balancing of read queries. Otherwise, the primary will only be used for write
+# queries. The primary can always be explicitly selected with our custom protocol.
+primary_reads_enabled = true
+
+# So what if you wanted to implement a different hashing function,
+# or you've already built one and you want this pooler to use it?
+#
+# Current options:
+#
+# pg_bigint_hash: PARTITION BY HASH (Postgres hashing function)
+# sha1: A hashing function based on SHA1
+#
+sharding_function = "pg_bigint_hash"
+
+# Credentials for users that may connect to this cluster
+[pools.cipherstash.users.0]
+username = "cipherstash"
+password = "password"
+# Maximum number of server connections that can be established for this user
+# The maximum number of connection from a single Pgcat process to any database in the cluster
+# is the sum of pool_size across all users.
+pool_size = 100
+min_pool_size = 20
+
+# Maximum query duration. Dangerous, but protects against DBs that died in a non-obvious way.
+statement_timeout = 0
+
+# Shard 0
+[pools.cipherstash.shards.0]
+
+servers = [["host.docker.internal", 5532, "primary"]]
+
+database = "cipherstash"