#!/usr/bin/env bash
# Counts, not times, for the confrelid index (CF 7365), master vs master+v6.
#
#   upgrade:  David Rowley's pg_upgrade_catalog_repro.sql, at 1/10 of its
#             default scale, upgraded with pg_upgrade; pg_constraint scan
#             counters of the restored database and user-space instructions
#             of pg_upgrade and every process it starts (perf stat).
#   truncate: 100 x TRUNCATE of a referenced table, and 100 commits that
#             touch an ON COMMIT DELETE ROWS temp table with a trigger,
#             each next to a control that skips the FK lookup; run in the
#             upgraded cluster, single-user backend under perf stat.
#   notnull:  instructions to create 1000 NOT NULL constraints, minus the
#             same tables with nullable columns, in a fresh cluster.
#
# usage (each stage is independent; N runs = call it N times):
#   OLD_BIN=<master>/bin MASTER_BIN=<master>/bin V6_BIN=<v6>/bin \
#   REPRO_SQL=<path>/pg_upgrade_catalog_repro.sql WORK=<scratch dir> \
#   bash repro-counts.sh.txt setup
#   bash repro-counts.sh.txt upgrade  master|v6 <run>
#   bash repro-counts.sh.txt truncate master|v6 <run>     (after "upgrade" of that side)
#   bash repro-counts.sh.txt notnull  master|v6 <run>
# Scale: PARENTS=32 PARTS=100 PLAIN=397 NNCOLS=200 (script defaults are
# 323/100/3970/200); the old cluster is built once, with OLD_BIN.
set -euo pipefail
: "${WORK:?}" "${OLD_BIN:?}" "${MASTER_BIN:?}" "${V6_BIN:?}"
PARENTS=${PARENTS:-32} PARTS=${PARTS:-100} PLAIN=${PLAIN:-397} NNCOLS=${NNCOLS:-200}
PORT=${PORT:-55432}
mkdir -p "$WORK"
side_bin() { case $1 in master) echo "$MASTER_BIN";; v6) echo "$V6_BIN";; *) exit 2;; esac; }
instr() { awk -F, '$3 ~ /^instructions/ {print $1}' "$1"; }
# one value from a single-user backend's output: val <file> <column>
val() { sed -n "s/^[[:space:]]*[0-9]*: $2 = \"\([^\"]*\)\".*/\1/p" "$1" | tail -1; }

case ${1:?stage} in
setup)
  : "${REPRO_SQL:?}"
  n=$(grep -c '^CALL public\.build_catalog_bloat();' "$REPRO_SQL" || true)
  [ "$n" = 1 ] || { echo "expected one CALL line in $REPRO_SQL"; exit 1; }
  rm -rf "$WORK/old"
  "$OLD_BIN/initdb" -D "$WORK/old" -U postgres -A trust >/dev/null
  cat >>"$WORK/old/postgresql.conf" <<EOF
max_locks_per_transaction = 4096
fsync = off
synchronous_commit = off
autovacuum = off
listen_addresses = ''
EOF
  "$OLD_BIN/pg_ctl" -D "$WORK/old" -l "$WORK/old.log" -w -o "-p $PORT -k $WORK" start >/dev/null
  "$OLD_BIN/createdb" -h "$WORK" -p $PORT -U postgres test
  sed "s/^CALL public\.build_catalog_bloat();/CALL public.build_catalog_bloat($PARENTS, $PARTS, 10, $PLAIN, $NNCOLS);/" \
    "$REPRO_SQL" | "$OLD_BIN/psql" -h "$WORK" -p $PORT -U postgres -d test -X -q \
    -v ON_ERROR_STOP=1 -f - >"$WORK/setup.out" 2>&1
  echo "old cluster: scale $PARENTS/$PARTS/10/$PLAIN/$NNCOLS"
  "$OLD_BIN/psql" -h "$WORK" -p $PORT -U postgres -d test -X -qtA -F' ' -c \
    "SELECT coalesce(contype::text, 'all'), count(*) FROM pg_constraint
     GROUP BY ROLLUP (contype) ORDER BY 2 DESC" | tr '\n' ' '; echo
  "$OLD_BIN/psql" -h "$WORK" -p $PORT -U postgres -d test -X -qtA -c \
    "SELECT 'partitions: ' || count(*) FROM pg_class WHERE relpartbound IS NOT NULL
       AND relkind IN ('r', 'p')"
  "$OLD_BIN/pg_ctl" -D "$WORK/old" -m fast -w stop >/dev/null
  ;;
upgrade)
  S=${2:?side}; R=${3:?run}; B=$(side_bin "$S"); D=$WORK/new-$S; U=$WORK/up-$S
  rm -rf "$D" "$U"; mkdir -p "$U"
  "$B/initdb" -D "$D" -U postgres -A trust >/dev/null
  echo "max_locks_per_transaction = 4096" >>"$D/postgresql.conf"
  cd "$U"
  perf stat -x, -e instructions:u -o "$U/perf.csv" -- \
    "$B/pg_upgrade" -b "$OLD_BIN" -B "$B" -d "$WORK/old" -D "$D" -U postgres \
    --no-sync -s "$U" -p $PORT -P $((PORT + 1)) >"$U/pg_upgrade.out" 2>&1 \
    || { tail -20 "$U/pg_upgrade.out"; exit 1; }
  "$B/pg_ctl" -D "$D" -l "$U/server.log" -w -o "-p $PORT -k $U -c listen_addresses=''" start >/dev/null
  q() { "$B/psql" -h "$U" -p $PORT -U postgres -d test -X -qtA -F' ' -c "$1"; }
  c=$(q "SELECT seq_scan, seq_tup_read, coalesce(idx_scan, 0) FROM pg_stat_sys_tables
         WHERE relname = 'pg_constraint'")
  i=$(q "SELECT coalesce((SELECT idx_scan || ' ' || idx_tup_read FROM pg_stat_sys_indexes
         WHERE indexrelname = 'pg_constraint_confrelid_index'), '- -')")
  n=$(q "SELECT count(*) FROM pg_constraint")
  a=""
  if [ "$S" = v6 ]; then
    q "CREATE EXTENSION amcheck" >/dev/null
    a=$(q "SELECT bt_index_parent_check('pg_constraint_confrelid_index', true)" >/dev/null &&
        q "SELECT count(*) FROM pg_index i, LATERAL bt_index_check(i.indexrelid, true)
           WHERE i.indrelid = 'pg_constraint'::regclass" | sed 's/^/amcheck_ok_indexes=/')
    q "DROP EXTENSION amcheck" >/dev/null
  fi
  "$B/pg_ctl" -D "$D" -m fast -w stop >/dev/null
  set -- $c; set -- "$@" $i
  echo "upgrade $S run=$R instructions=$(instr "$U/perf.csv") pg_constraint_rows=$n" \
       "seq_scan=$1 seq_tup_read=$2 idx_scan=$3 confrelid_idx_scan=$4" \
       "confrelid_idx_tup_read=$5 $a"
  ;;
truncate)
  S=${2:?side}; R=${3:?run}; B=$(side_bin "$S"); D=$WORK/new-$S; T=$WORK/tr-$S
  [ -d "$D" ] || { echo "run upgrade $S first"; exit 1; }
  mkdir -p "$T"
  sb() { "$B/postgres" --single -j -D "$D" -c stats_fetch_consistency=none test; }
  sb >/dev/null <<'SQL'
DROP TABLE IF EXISTS tr_fk, tr_pk, tr_plain;

CREATE TABLE tr_pk (id int PRIMARY KEY);

CREATE TABLE tr_fk (id int REFERENCES tr_pk);

CREATE TABLE tr_plain (id int);
SQL
  # workload <name> <setup sql> <one statement, run 100 times>
  workload() {
    local f=$T/$1.sql
    { printf 'SELECT pg_stat_force_next_flush();\n\n'
      printf 'SELECT 1;\n\n'
      printf '%s\n\n' "$2"
      printf 'SELECT pg_stat_force_next_flush();\n\n'
      printf 'SELECT 1;\n\n'
      printf 'SELECT seq_scan AS s0, seq_tup_read AS t0 FROM pg_stat_sys_tables WHERE relname = %s;\n\n' "'pg_constraint'"
      for _ in $(seq 100); do printf '%s\n\n' "$3"; done
      printf 'SELECT pg_stat_force_next_flush();\n\n'
      printf 'SELECT 1;\n\n'
      printf 'SELECT seq_scan AS s1, seq_tup_read AS t1 FROM pg_stat_sys_tables WHERE relname = %s;\n\n' "'pg_constraint'"
    } >"$f"
    perf stat -x, -e instructions:u -o "$T/$1.perf" -- \
      "$B/postgres" --single -j -D "$D" -c stats_fetch_consistency=none test <"$f" >"$T/$1.out" 2>&1
    echo "$(instr "$T/$1.perf") $(( $(val "$T/$1.out" s1) - $(val "$T/$1.out" s0) )) $(( $(val "$T/$1.out" t1) - $(val "$T/$1.out" t0) ))"
  }
  set -- $(workload trunc_fk 'SELECT 1;' 'TRUNCATE tr_pk, tr_fk;') \
         $(workload trunc_plain 'SELECT 1;' 'TRUNCATE tr_plain;') \
         $(workload commit_trg 'CREATE TEMP TABLE tt (a int) ON COMMIT DELETE ROWS;

CREATE FUNCTION pg_temp.f() RETURNS trigger LANGUAGE plpgsql AS $$BEGIN RETURN NEW; END$$;

CREATE TRIGGER t BEFORE INSERT ON tt FOR EACH ROW EXECUTE FUNCTION pg_temp.f();' 'INSERT INTO tt VALUES (1);') \
         $(workload commit_plain 'CREATE TEMP TABLE tt (a int) ON COMMIT DELETE ROWS;' 'INSERT INTO tt VALUES (1);')
  rows=$(sb <<<'SELECT count(*) AS n FROM pg_constraint;' 2>/dev/null | sed -n 's/.*n = "\([0-9]*\)".*/\1/p')
  sb >/dev/null <<<'DROP TABLE tr_fk, tr_pk, tr_plain;'
  echo "truncate $S run=$R pg_constraint_rows=$rows" \
       "x100 TRUNCATE pk,fk: instructions=$1 seq_scan=$2 seq_tup_read=$3;" \
       "x100 TRUNCATE plain (control): instructions=$4 seq_scan=$5 seq_tup_read=$6;" \
       "x100 commit, on-commit temp w/ trigger: instructions=$7 seq_scan=$8 seq_tup_read=$9;" \
       "x100 commit, on-commit temp plain (control): instructions=${10} seq_scan=${11} seq_tup_read=${12}"
  ;;
notnull)
  S=${2:?side}; R=${3:?run}; B=$(side_bin "$S"); D=$WORK/nn-$S
  rm -rf "$D"; "$B/initdb" -D "$D" -U postgres -A trust >/dev/null
  # 5 tables x 200 columns, NOT NULL or nullable, in one transaction each
  mk() {
    local f=$D.$1.sql
    for t in 1 2 3 4 5; do
      printf 'CREATE TABLE nn%s (%s);\n\n' $t \
        "$(seq -s, -f "c%g int $2" 1 200 | sed 's/,$//')"
    done >"$f"
    perf stat -x, -e instructions:u -o "$D.$1.perf" -- \
      "$B/postgres" --single -j -D "$D" postgres <"$f" >"$D.$1.out" 2>&1
    "$B/postgres" --single -j -D "$D" postgres >/dev/null 2>&1 \
      <<<'DROP TABLE nn1, nn2, nn3, nn4, nn5;'
    instr "$D.$1.perf"
  }
  nn=$(mk notnull 'NOT NULL'); nl=$(mk nullable '')
  echo "notnull $S run=$R fresh initdb: 1000 NOT NULL cols instructions=$nn," \
       "1000 nullable cols instructions=$nl, difference=$((nn - nl))"
  rm -rf "$D" "$D".*
  ;;
*) echo "stage?"; exit 2;;
esac