test(perf): improve replication benchmark

benches/compare.py

@@ -113,10 +113,10 @@ def peak_mem(result: dict) -> float | None:
         if delta is None or abs(delta) < 1.0:
             pct_colored = pct_str
         elif delta < 0:
-            # lower is better for time; higher is worse for memory
-            pct_colored = color(pct_str, RED if label == "max_mem" else GREEN)
+            # lower is better for both time and memory
+            pct_colored = color(pct_str, GREEN)
         else:
-            pct_colored = color(pct_str, GREEN if label == "max_mem" else RED)
+            pct_colored = color(pct_str, RED)
 
         fmt = fmt_bytes if label == "max_mem" else fmt_time
         print(

benches/resharding/README.md

@@ -69,14 +69,18 @@ Two proxies are created:
 | `resharding_source` | 15400 | pgdog1, pgdog2, pgdog3 | downstream (PG → pgdog) |
 | `resharding_destination` | 15401 | shard_0–shard_3 | upstream (pgdog → PG) |
 
-Toxic settings are configurable via environment variables:
-
-| Variable | Effect |
-|---|---|
-| `SOURCE_LATENCY_MS` | Added latency on source reads |
-| `DEST_LATENCY_MS` | Added latency on destination writes |
-| `SOURCE_BW_KBPS` | Bandwidth cap on source reads (KB/s) |
-| `DEST_BW_KBPS` | Bandwidth cap on destination writes (KB/s) |
+The default toxics model a cross-region setup: pgdog next to the source, and the
+destination in another region (for example us-east-1 → us-east-2). Latency is
+one-way, so it adds the same amount to each round trip.
+
+These environment variables change the toxics:
+
+| Variable | Default | Effect |
+|---|---|---|
+| `SOURCE_LATENCY_MS` | `1` | Added latency on source reads |
+| `DEST_LATENCY_MS` | `15` | Added latency on destination writes |
+| `SOURCE_BW_KBPS` | `125000` | Bandwidth cap on source reads (KB/s per connection, `0` = no cap) |
+| `DEST_BW_KBPS` | `125000` | Bandwidth cap on destination writes (KB/s per connection, `0` = no cap) |
 
 ```sh
 # run with default toxic settings

benches/resharding/fill.sql

@@ -9,6 +9,7 @@
 --   documents  : scale/100 rows — body ≈ 32 KB (STORAGE EXTERNAL).
 --   files      : scale/100 rows — content ≈ 4 KB bytea (TOAST).
 --   ledger     : scale/100 rows — notes ≈ 16 KB (STORAGE EXTERNAL).
+--   small txns : scale/25 single-row sessions inserts, each its own transaction.
 --
 -- The TOAST tables share their row count so a single BENCH_SCALE knob controls
 -- everything: inline rows = scale, TOAST rows = scale/100.
@@ -19,6 +20,7 @@
 \getenv scale BENCH_SCALE
 \set QUIET on
 SET client_min_messages TO warning;
+SET synchronous_commit TO off;
 \if :{?scale}
 \else
 \set scale 100000
@@ -119,3 +121,16 @@ SELECT format(
 )
 FROM unnest(ARRAY['sessions', 'documents', 'files', 'ledger']) AS t,
      generate_series(2, :copies) AS g \gexec
+
+-- ── small transactions (scale/25 single-row inserts, autocommit) ──────────────
+SELECT format(
+    'INSERT INTO bench_copy.sessions (tenant_id, user_id, device_id, session_id, correlation_id, '
+    'request_id, trace_id, started_at, ended_at, last_seen_at, expires_at) '
+    'VALUES (%s, %L, %L, gen_random_uuid(), gen_random_uuid(), gen_random_uuid(), %L, '
+    'now(), now() + interval ''1 hour'', now(), now() + interval ''30 days'')',
+    gs.i,
+    md5((gs.i % 100000)::text)::uuid,
+    md5((gs.i % 10000)::text)::uuid,
+    md5(gs.i::text)::uuid
+)
+FROM generate_series(:shard_index + 1, :scale / 25, :num_shards) AS gs(i) \gexec

benches/resharding/replication/prepare.sh

@@ -1,10 +1,14 @@
 #!/usr/bin/env bash
-# hyperfine --prepare: reset schema, COPY empty tables (creates/reuses slot), fill sources.
+# hyperfine --prepare: reset schemas, sync the destination schema, create slots, fill sources.
 #
-# The COPY phase runs against empty tables — trivial, just establishes the slot
-# at the current LSN. fill.sql then queues the WAL backlog the timed command drains.
-# Slot lifecycle (bench_copy_slot_replication) is owned by this script:
-# dropped at the start of each prepare so every run starts from a fresh LSN.
+# The destination schema comes from pgdog schema-sync (pre-data, then post-data).
+# The slots are created with SQL, one per source shard, named "bench_copy_slot_replication_<shard>"
+# as pgdog expects. No data-sync runs here, so nothing else writes WAL between the
+# slot creation and the fill. fill.sql then queues the WAL backlog the timed command
+# drains: few large INSERT ... SELECT transactions, then BENCH_SCALE/25 single-row
+# transactions. Every bench_copy slot is dropped at the start of each prepare so every
+# run starts from a fresh LSN. A CHECKPOINT after the drop lets Postgres recycle the
+# old WAL, and a CHECKPOINT at the end keeps checkpoint I/O out of the timed command.
 
 set -euo pipefail
 
@@ -21,12 +25,29 @@ echo "============================================================"
 
 echo ""
 echo "============================================================"
-echo ">>>>> [1/6] dropping stale replication slots"
-echo "============================================================"
-for db in "${SOURCE_DBS[@]}"; do
-    psql -d "${db}" -qX -c \
-        "SELECT pg_drop_replication_slot(slot_name) FROM pg_replication_slots WHERE slot_name LIKE 'bench_copy%' AND NOT active;"
-done
+echo ">>>>> [1/6] dropping replication slots and recycling WAL"
+echo "============================================================"
+psql -d "${SOURCE_DBS[0]}" -qX -v ON_ERROR_STOP=1 <<'SQL'
+SELECT pg_terminate_backend(active_pid)
+FROM pg_replication_slots
+WHERE slot_name LIKE 'bench_copy%' AND active_pid IS NOT NULL;
+DO $$
+BEGIN
+    FOR _ IN 1..100 LOOP
+        EXIT WHEN NOT EXISTS (
+            SELECT 1 FROM pg_replication_slots WHERE slot_name LIKE 'bench_copy%' AND active
+        );
+        PERFORM pg_sleep(0.1);
+    END LOOP;
+END
+$$;
+SELECT pg_drop_replication_slot(slot_name)
+FROM pg_replication_slots
+WHERE slot_name LIKE 'bench_copy%';
+CHECKPOINT;
+SELECT pg_switch_wal();
+CHECKPOINT;
+SQL
 
 echo ""
 echo "============================================================"
@@ -41,35 +62,45 @@ echo "============================================================"
 echo ">>>>> [3/6] resetting destination schema"
 echo "============================================================"
 for db in "${DEST_DBS[@]}"; do
-    psql -d "${db}" -qX -c "DROP SCHEMA IF EXISTS bench_copy CASCADE" > /dev/null
+    psql -d "${db}" -qX -c "SET client_min_messages TO warning" -c "DROP SCHEMA IF EXISTS bench_copy CASCADE" > /dev/null
 done
 
 echo ""
 echo "============================================================"
-echo ">>>>> [4/6] establishing replication slot via COPY (empty tables)"
-echo "============================================================"
-# COPY empty tables; creates the slot at the current LSN.  Exit immediately after COPY (--sync-only).
-"${PGDOG_BIN}" \
-    --config "${PGDOG_CONFIG}" \
-    --users  "${SETUP_DIR}/users.toml" \
-    data-sync \
-    --from-database source \
-    --to-database   destination \
-    --publication   bench_copy \
-    --replication-slot bench_copy_slot_replication \
-    --sync-only
+echo ">>>>> [4/6] syncing destination schema and creating replication slots"
+echo "============================================================"
+# --data-sync-complete selects the post-data phase; older pgdog versions have no --phase.
+for phase_args in "" "--data-sync-complete"; do
+    # shellcheck disable=SC2086
+    "${PGDOG_BIN}" \
+        --config "${PGDOG_CONFIG}" \
+        --users  "${SETUP_DIR}/users.toml" \
+        schema-sync \
+        --from-database source \
+        --to-database   destination \
+        --publication   bench_copy \
+        ${phase_args}
+done
+
+for shard in "${!SOURCE_DBS[@]}"; do
+    psql -d "${SOURCE_DBS[${shard}]}" -qX -v ON_ERROR_STOP=1 -c \
+        "SELECT pg_create_logical_replication_slot('bench_copy_slot_replication_${shard}', 'pgoutput')" > /dev/null
+    echo "  [${SOURCE_DBS[${shard}]}] slot bench_copy_slot_replication_${shard} created"
+done
 
 echo ""
 echo "============================================================"
-echo ">>>>> [5/6] filling source shards in parallel (WAL backlog)"
+echo ">>>>> [5/6] filling source shards in parallel (WAL backlog: big + small transactions)"
 echo "============================================================"
 FILL_PIDS=()
 for db in "${SOURCE_DBS[@]}"; do
-    psql -d "${db}" -qX -f "${SETUP_DIR}/fill.sql" &
+    psql -d "${db}" -qX -v ON_ERROR_STOP=1 -f "${SETUP_DIR}/fill.sql" &
     FILL_PIDS+=("$!")
 done
 echo "  waiting for all fills to complete..."
-wait "${FILL_PIDS[@]}"
+for pid in "${FILL_PIDS[@]}"; do
+    wait "${pid}"
+done
 
 echo ""
 echo "============================================================"
@@ -85,6 +116,8 @@ for db in "${SOURCE_DBS[@]}"; do
     echo "  [${db}] sentinel inserted"
 done
 
+psql -d "${SOURCE_DBS[0]}" -qX -v ON_ERROR_STOP=1 -c "CHECKPOINT"
+
 echo ""
 echo "============================================================"
 echo ">>>>> WAL backlog + sentinels ready"

benches/resharding/replication/run.sh

@@ -2,18 +2,22 @@
 # Benchmark: WAL replication throughput (3 sources → 4 destination shards).
 #
 # Measures time for pgdog to drain a pre-seeded WAL backlog via --replicate-only.
-# COPY happens in prepare (untimed); prepare also writes per-source sentinel rows.
+# Prepare (untimed) runs schema-sync, creates the slots and writes per-source sentinel rows.
 # The timed command streams WAL and exits once all sentinels appear on all shards.
 #
+# The backlog comes from fill.sql on every source:
+#   ~160 large INSERT ... SELECT transactions, then scale/25 single-row transactions
+# Result name: resharding.streaming[.toxi]
+#
 # Usage:
 #   bash benches/resharding/replication/run.sh [bench.sh flags]
 #   bash benches/resharding/replication/run.sh --save-baseline main
 #   bash benches/resharding/replication/run.sh --baseline main
 #
 # Tuning:
-#   BENCH_SCALE=N    sessions per source shard; TOAST rows = N/100
-#   RUNS=N WARMUP=N  forwarded to hyperfine via environment
-#   USE_TOXI=1       route pgdog through toxiproxy (started/stopped automatically)
+#   BENCH_SCALE=N             sessions per source shard; TOAST rows = N/100; small transactions = N/25
+#   RUNS=N WARMUP=N           forwarded to hyperfine via environment
+#   USE_TOXI=1                route pgdog through toxiproxy (started/stopped automatically)
 
 set -euo pipefail
 
@@ -30,11 +34,13 @@ export PGPASSWORD=pgdog
 export SETUP_DIR  # inherited by prepare.sh
 
 DEST_DBS=(shard_0 shard_1 shard_2 shard_3)
-export BENCH_SCALE=${BENCH_SCALE:-50000}
+export BENCH_SCALE=${BENCH_SCALE:-1000000}
 
+NAME_SUFFIX=""
 PGDOG_CONFIG="${SETUP_DIR}/pgdog.toml"
 if [[ "${USE_TOXI:-0}" == "1" ]]; then
     PGDOG_CONFIG="${SETUP_DIR}/pgdog.toxi.toml"
+    NAME_SUFFIX=".toxi"
     # Register teardown before setup so a partial-setup failure still tears down.
     trap 'bash "${SETUP_DIR}/toxi/teardown.sh"' EXIT
     bash "${SETUP_DIR}/toxi/setup.sh"
@@ -44,7 +50,7 @@ export PGDOG_CONFIG
 echo "============================================================"
 echo ">>>>> benchmark"
 echo "============================================================"
-bench_run "resharding.streaming" \
+bench_run "resharding.streaming${NAME_SUFFIX}" \
     "${SCRIPT_DIR}/stream.sh" \
     --prepare "${SCRIPT_DIR}/prepare.sh" \
     "$@"

benches/resharding/replication/stream.sh

@@ -15,10 +15,10 @@ SENTINEL_ID=-1
 
 # Inherit PGDOG_CONFIG from run.sh (defaults to pgdog.toml for standalone use).
 PGDOG_CONFIG="${PGDOG_CONFIG:-${SETUP_DIR}/pgdog.toml}"
-POLL_INTERVAL=0.5  # seconds between destination polls
+POLL_INTERVAL=0.1  # seconds between destination polls
 
 # Start replication in the background.  Schema sync is skipped because prepare.sh
-# already ran it via --sync-only.
+# already ran schema-sync and created the slots.
 "${PGDOG_BIN}" \
     --config "${PGDOG_CONFIG}" \
     --users  "${SETUP_DIR}/users.toml" \
@@ -31,8 +31,17 @@ POLL_INTERVAL=0.5  # seconds between destination polls
     --replicate-only &
 PGDOG_PID=$!
 
-# Ensure pgdog is killed on any exit path
-trap 'kill "${PGDOG_PID}" 2>/dev/null || true; sleep 5; kill -9 "${PGDOG_PID}" 2>/dev/null || true' EXIT
+stop_pgdog() {
+    kill "${PGDOG_PID}" 2>/dev/null || return 0
+    for _ in $(seq 100); do
+        kill -0 "${PGDOG_PID}" 2>/dev/null || return 0
+        sleep 0.05
+    done
+    kill -9 "${PGDOG_PID}" 2>/dev/null || true
+}
+
+# Ensure pgdog is stopped on any exit path; returns as soon as it exits (max 5s before SIGKILL).
+trap stop_pgdog EXIT
 
 echo "  pgdog replication started (pid=${PGDOG_PID}), polling for sentinels..."
 

benches/resharding/toxi/setup.sh

@@ -10,16 +10,20 @@
 #   resharding_replica_2     15411         copy worker entries (pgdog.toxi.toml)
 #   resharding_destination   15401         shard_0, shard_1, shard_2, shard_3
 #
+# Defaults model a cross-region setup: pgdog next to the source, destination in
+# another region (for example us-east-1 -> us-east-2). Latency is one-way, so it
+# adds the same amount to each request/response round trip.
+#
 # Toxics applied to each proxy (all configurable via env):
 #   SOURCE_LATENCY_MS   -- latency on source reads, fast source proxies,       --downstream
 #   SLOW_LATENCY_MS     -- latency on source reads, resharding_replica_2,      --downstream
 #   SOURCE_BW_KBPS      -- bandwidth cap per connection, all source proxies,   --downstream
 #   DEST_LATENCY_MS     -- latency on destination writes, --upstream
-#   DEST_BW_KBPS        -- bandwidth cap on destination writes, --upstream
+#   DEST_BW_KBPS        -- bandwidth cap on destination writes, --upstream (0 = no cap)
 #
 # Usage:
 #   bash benches/resharding/toxi/setup.sh
-#   SOURCE_LATENCY_MS=20 DEST_BW_KBPS=10000 bash benches/resharding/toxi/setup.sh
+#   DEST_LATENCY_MS=30 bash benches/resharding/toxi/setup.sh
 #
 set -euo pipefail
 
@@ -29,11 +33,11 @@ CLI="${REPO_ROOT}/integration/toxiproxy-cli"
 SERVER="${REPO_ROOT}/integration/toxiproxy-server"
 UPSTREAM="127.0.0.1:5432"
 
-SOURCE_LATENCY_MS="${SOURCE_LATENCY_MS:-0}"
-DEST_LATENCY_MS="${DEST_LATENCY_MS:-0}"
-SOURCE_BW_KBPS="${SOURCE_BW_KBPS:-10000}"
+SOURCE_LATENCY_MS="${SOURCE_LATENCY_MS:-1}"
+DEST_LATENCY_MS="${DEST_LATENCY_MS:-15}"
+SOURCE_BW_KBPS="${SOURCE_BW_KBPS:-125000}"
 SLOW_LATENCY_MS="${SLOW_LATENCY_MS:-500}"
-DEST_BW_KBPS="${DEST_BW_KBPS:-10000}"
+DEST_BW_KBPS="${DEST_BW_KBPS:-125000}"
 
 # ── Start server ─────────────────────────────────────────────────────────────
 echo "Starting toxiproxy-server..."
@@ -76,14 +80,16 @@ done
     --attribute  latency="${SLOW_LATENCY_MS}" \
     resharding_replica_2
 
-for proxy in resharding_source resharding_replica_1 resharding_replica_2; do
-    "${CLI}" toxic add \
-        --toxicName bandwidth \
-        --type      bandwidth \
-        --downstream \
-        --attribute  rate="${SOURCE_BW_KBPS}" \
-        "${proxy}"
-done
+if [[ "${SOURCE_BW_KBPS}" -gt 0 ]]; then
+    for proxy in resharding_source resharding_replica_1 resharding_replica_2; do
+        "${CLI}" toxic add \
+            --toxicName bandwidth \
+            --type      bandwidth \
+            --downstream \
+            --attribute  rate="${SOURCE_BW_KBPS}" \
+            "${proxy}"
+    done
+fi
 
 "${CLI}" toxic add \
     --toxicName latency \
@@ -92,19 +98,21 @@ done
     --attribute  latency="${DEST_LATENCY_MS}" \
     resharding_destination
 
-"${CLI}" toxic add \
-    --toxicName bandwidth \
-    --type      bandwidth \
-    --upstream \
-    --attribute  rate="${DEST_BW_KBPS}" \
-    resharding_destination
+if [[ "${DEST_BW_KBPS}" -gt 0 ]]; then
+    "${CLI}" toxic add \
+        --toxicName bandwidth \
+        --type      bandwidth \
+        --upstream \
+        --attribute  rate="${DEST_BW_KBPS}" \
+        resharding_destination
+fi
 
 echo ""
 "${CLI}" list
 echo ""
 echo "Toxiproxy ready (API port=8474)."
 echo "  fast source pipes latency=${SOURCE_LATENCY_MS}ms  bandwidth=${SOURCE_BW_KBPS} KB/s per connection"
 echo "  slow replica      latency=${SLOW_LATENCY_MS}ms per round trip (resharding_replica_2)"
-echo "  destination  latency=${DEST_LATENCY_MS}ms  bandwidth=${DEST_BW_KBPS} KB/s"
+echo "  destination       latency=${DEST_LATENCY_MS}ms  bandwidth=${DEST_BW_KBPS} KB/s per connection"
 echo ""
 echo "Run teardown.sh to stop, or just re-run setup.sh to restart."