#!/bin/bash
#
# ssi_lock_reclamation_repro.sh
#
# Reproduces unbounded predicate lock table growth under SSI, which shows
# up as
#
#   ERROR:  out of shared memory
#   HINT:  You might need to increase "max_pred_locks_per_transaction".
#
# in backends that are not themselves doing anything unusual.
#
# Background
# ----------
# When the SERIALIZABLEXACT pool is exhausted, SummarizeOldestCommittedSxact()
# folds a committed transaction's SIREAD locks onto the dummy transaction
# OldCommittedSxact.  Without the patch, such a folded lock is only released
# once PredXact->WritableSxactCount drops to 0, i.e. once no read-write
# serializable transaction is running anywhere.  Under a sustained
# serializable read-write workload that never happens, so folded locks
# accumulate until the predicate lock table is full.
#
# The patch tracks finishedBefore per predicate lock, so ClearOldPredicateLocks()
# can release a folded lock as soon as SxactGlobalXmin has advanced past it,
# without waiting for the workload to go quiet.
#
# Workload
# --------
#   - one "anchor" session running short serializable transactions that each
#     hold a snapshot for 5s.  This repeatedly pins and then releases
#     SxactGlobalXmin, which is what gives the patched code its opportunity
#     to reclaim.
#   - 32 clients running a serializable read-one-row / write-another-row
#     transaction, which is what drives the pool to exhaustion and keeps
#     WritableSxactCount permanently above 0.  Each transaction includes a
#     small pg_sleep(0.05): on fast hardware, 32 clients running this
#     workload flat out can exhaust the predicate lock table in well under
#     a second, before the anchor's first 5s snapshot has even ended,
#     with or without the patch.
#     The small sleep keeps the same concurrency and the same continuous
#     WritableSxactCount > 0, just paced slowly enough that whether
#     reclaiming happens during that first window actually matters.
#
# Expected result
# ---------------
#   unpatched: fails partway through, "out of shared memory" (exit status 1)
#   patched:   completes the full run cleanly     (exit status 0)
#
# Everything runs with initdb defaults; no postgresql.conf changes needed.
# Point it at a disposable cluster: it creates a database and drops and
# recreates a table named "sibench" inside it.
#
# Usage:
#   ./ssi_lock_reclamation_repro.sh [dbname] [duration_seconds]
#
# Connection is taken from the usual PG* environment variables
# (PGHOST/PGPORT/PGUSER/...), so point those at the cluster under test.

set -u

DBNAME="${1:-sibench}"
DURATION="${2:-60}"
ROWS=10000
CLIENTS=32
THREADS=8
ANCHOR_SLEEP=5
TX_PACE=0.05

WORKDIR=$(mktemp -d "${TMPDIR:-/tmp}/ssi_repro.XXXXXX")
ANCHOR_PID=""

cleanup()
{
    [ -n "$ANCHOR_PID" ] && kill "$ANCHOR_PID" 2>/dev/null
    wait "$ANCHOR_PID" 2>/dev/null
    rm -rf "$WORKDIR"
}
trap cleanup EXIT

# ---------------------------------------------------------------- schema
cat > "$WORKDIR/schema.sql" <<EOF
DROP TABLE IF EXISTS sibench;
CREATE TABLE sibench (id int PRIMARY KEY, val int NOT NULL);
INSERT INTO sibench (id, val) SELECT g, 0 FROM generate_series(1, $ROWS) g;
VACUUM ANALYZE sibench;
EOF

# The anchor holds a serializable snapshot for a few seconds at a time.
# Each time one of these commits, SxactGlobalXmin can advance, which is
# when the patched ClearOldPredicateLocks() gets to release folded locks.
cat > "$WORKDIR/anchor.sql" <<EOF
BEGIN ISOLATION LEVEL SERIALIZABLE;
\\set id random(1, $ROWS)
SELECT val FROM sibench WHERE id = :id;
SELECT pg_sleep($ANCHOR_SLEEP);
COMMIT;
EOF

# Read one row, write a different one, at SERIALIZABLE.  Keeps a steady
# stream of committed serializable transactions to be summarized, and
# keeps WritableSxactCount above 0 for the whole run.  See TX_PACE above
# for why each transaction includes a small sleep.
cat > "$WORKDIR/write_skew_pair.sql" <<EOF
BEGIN ISOLATION LEVEL SERIALIZABLE;
\\set k1 random(1, $ROWS)
\\set k2 random(1, $ROWS)
SELECT val FROM sibench WHERE id = :k1;
SELECT pg_sleep($TX_PACE);
UPDATE sibench SET val = val + 1 WHERE id = :k2;
COMMIT;
EOF

# ------------------------------------------------------------------ setup
echo "=== setting up database \"$DBNAME\" ($ROWS rows) ==="
createdb "$DBNAME" 2>/dev/null || echo "    (database already exists, reusing)"
psql -d "$DBNAME" -q -f "$WORKDIR/schema.sql" || exit 2

psql -d "$DBNAME" -tA -c \
    "SELECT 'max_pred_locks_per_transaction=' || current_setting('max_pred_locks_per_transaction')
          || ' max_connections='               || current_setting('max_connections')"

# ------------------------------------------------------------------- run
echo "=== starting anchor session (serializable snapshot held ${ANCHOR_SLEEP}s at a time) ==="
pgbench -d "$DBNAME" -c 1 -T "$DURATION" --no-vacuum \
        -f "$WORKDIR/anchor.sql" > "$WORKDIR/anchor.log" 2>&1 &
ANCHOR_PID=$!

sleep 1

echo "=== running write-skew load: $CLIENTS clients, ${DURATION}s ==="
pgbench -d "$DBNAME" -c "$CLIENTS" -j "$THREADS" -T "$DURATION" \
        --max-tries=3 --no-vacuum \
        -f "$WORKDIR/write_skew_pair.sql" > "$WORKDIR/load.log" 2>&1
echo
cat "$WORKDIR/load.log"

# ---------------------------------------------------------------- verdict
echo
if grep -q "out of shared memory" "$WORKDIR/anchor.log" "$WORKDIR/load.log"; then
    echo "=============================================================="
    echo "RESULT: FAILED - predicate lock table exhausted"
    echo "=============================================================="
    grep -h -A1 "out of shared memory" "$WORKDIR/load.log" | head -4
    echo
    echo "Summarized SIREAD locks were not reclaimed while the workload"
    echo "was running.  This is the unpatched behaviour."
    exit 1
else
    echo "=============================================================="
    echo "RESULT: PASSED - no shared memory exhaustion"
    echo "=============================================================="
    echo
    echo "Summarized SIREAD locks were reclaimed as SxactGlobalXmin"
    echo "advanced.  This is the patched behaviour."
    exit 0
fi