#!/usr/bin/env bash
# Counts, not times, for heap_truncate_find_FKs() with the confrelid index
# (CF 7365): master (seqscan), v6/v7 (one range scan from the lowest to the
# highest OID in the list) and v8 (one index probe per relation in the list).
# Every number is from a single-user backend under perf stat (user-space
# instructions) plus the pg_constraint / confrelid index scan counters.
#
#   spread:  Alvaro's case.  An old table and its FK, K foreign keys from
#            other tables to tables created later, then a new table and its
#            FK.  100 x TRUNCATE of the four, next to 100 x TRUNCATE of four
#            tables without triggers (that skips the lookup).
#            K grows 1000 -> 10000 -> 100000 in the same cluster; the new
#            table is dropped and created again at each step, so it always
#            has the highest OID.
#   list:    TRUNCATE of N tables in one statement, each with a trigger (so
#            the FK lookup runs, and finds nothing), in a fresh cluster;
#            control: the same N tables without triggers.
#   cascade: TRUNCATE c1 CASCADE on a chain of N tables, each referencing
#            the previous one (N levels of references), in a fresh cluster;
#            then a check that it empties all N tables.
#

# usage:
#   MASTER_BIN=.. RANGE_BIN=.. PROBE_BIN=.. WORK=<scratch dir> bash fk-lookup-counts.sh.txt
#     spread-setup master|idx <K>        (idx: catalog shared by range and probe)
#     spread       master|range|probe <K> <run>
#     list-setup   master|idx <N>
#     list         master|range|probe <N> <run>
#     cascade-setup master|idx <N>
#     cascade      master|range|probe <N> <run>
set -euo pipefail
: "${WORK:?}" "${MASTER_BIN:?}" "${RANGE_BIN:?}" "${PROBE_BIN:?}"
mkdir -p "$WORK"
bin() { case $1 in master) echo "$MASTER_BIN";; range) echo "$RANGE_BIN";;
        probe|idx) echo "$PROBE_BIN";; *) exit 2;; esac; }
cat_of() { case $1 in master) echo master;; *) echo idx;; esac; }
instr() { awk -F, '$3 ~ /^instructions/ {print $1}' "$1"; }
val() { sed -n "s/^[[:space:]]*[0-9]*: $2 = \"\([^\"]*\)\".*/\1/p" "$1" | tail -1; }
sb() { "$1/postgres" --single -j -D "$2" -c stats_fetch_consistency=none postgres; }
newdb() {
  rm -rf "$1"; "$2/initdb" -D "$1" -U postgres -A trust >/dev/null
  printf 'fsync = off\nautovacuum = off\nmax_locks_per_transaction = 1024\n' >>"$1/postgresql.conf"
}
COUNTERS="SELECT seq_scan AS SX, seq_tup_read AS TX,
  (SELECT idx_tup_read FROM pg_stat_sys_indexes
     WHERE indexrelname = 'pg_constraint_confrelid_index') AS IX
  FROM pg_stat_sys_tables WHERE relname = 'pg_constraint';"
# measure <bin> <datadir> <out prefix> <reps> <statement>: instructions seq_scan
#   seq_tup_read confrelid_idx_tup_read for <reps> runs of <statement>
#   optional: $6 = SQL run first in the same session (temp tables);
#   MEASURE_ERRORS=<n> = each statement is expected to fail, n errors in all
measure() {
  local f=$3.sql c0 c1
  c0=${COUNTERS//X/0}; c1=${COUNTERS//X/1}
  { [ -n "${6:-}" ] && printf '%s\n\n' "$6"
    printf 'SELECT pg_stat_force_next_flush();\n\nSELECT 1;\n\n%s\n\n' "$c0"
    for _ in $(seq "$4"); do printf '%s\n\n' "$5"; done
    printf 'SELECT pg_stat_force_next_flush();\n\nSELECT 1;\n\n%s\n\n' "$c1"
  } >"$f"
  perf stat -x, -e instructions:u -o "$3.perf" -- \
    "$1/postgres" --single -j -D "$2" -c stats_fetch_consistency=none postgres <"$f" >"$3.out" 2>&1
  local ne; ne=$(grep -c 'ERROR:' "$3.out" || true)
  [ "$ne" = "${MEASURE_ERRORS:-0}" ] || { echo "errors=$ne"; grep -m3 'ERROR:' "$3.out"; exit 1; }

  local i0 i1; i0=$(val "$3.out" i0); i1=$(val "$3.out" i1)
  echo "$(instr "$3.perf") $(( $(val "$3.out" s1) - $(val "$3.out" s0) ))" \
       "$(( $(val "$3.out" t1) - $(val "$3.out" t0) )) $(( ${i1:-0} - ${i0:-0} ))"
}
fkrows() { sb "$1" "$2" <<<"SELECT count(*) AS n, count(*) FILTER (WHERE contype = 'f') AS f FROM pg_constraint;" 2>/dev/null |
           sed -n 's/.* [nf] = "\([0-9]*\)".*/\1/p' | paste -sd/; }

case ${1:?stage} in
spread-setup)
  C=${2:?}; K=${3:?}; B=$(bin "$C"); D=$WORK/spread-$C
  if [ ! -d "$D" ]; then
    newdb "$D" "$B"
    sb "$B" "$D" >/dev/null <<'SQL'
CREATE TABLE old_pk (id int PRIMARY KEY);

CREATE TABLE old_fk (id int REFERENCES old_pk);

CREATE TABLE plain1 (id int);

CREATE TABLE plain2 (id int);

CREATE TABLE plain3 (id int);

CREATE TABLE plain4 (id int);
SQL
  fi
  have=$(sb "$B" "$D" <<<"SELECT count(*) AS n FROM pg_class WHERE relkind = 'r' AND relname LIKE 'mid\_%';" 2>/dev/null |
         sed -n 's/.*n = "\([0-9]*\)".*/\1/p')
  # one referenced table per 100 foreign keys; each gets 100 FKs from its own table
  for j in $(seq $((have + 1)) $((K / 100))); do
    printf 'CREATE TABLE mid_%s (id int PRIMARY KEY);\n\n' "$j"
    printf 'CREATE TABLE ref_%s (%s);\n\n' "$j" \
      "$(seq -s, -f "c%g int REFERENCES mid_$j" 1 100 | sed 's/,$//')"
  done | sb "$B" "$D" >/dev/null
  sb "$B" "$D" >/dev/null <<'SQL'
DROP TABLE IF EXISTS new_fk, new_pk;

CREATE TABLE new_pk (id int PRIMARY KEY);

CREATE TABLE new_fk (id int REFERENCES new_pk);
SQL
  echo "spread-setup $C K=$K pg_constraint rows/fk=$(fkrows "$B" "$D")"
  ;;
spread)
  S=${2:?}; K=${3:?}; R=${4:?}; B=$(bin "$S"); D=$WORK/spread-$(cat_of "$S"); O=$WORK/sp-$S-$K-$R
  set -- $(measure "$B" "$D" "$O-fk" 100 'TRUNCATE old_pk, old_fk, new_pk, new_fk;') \
         $(measure "$B" "$D" "$O-plain" 100 'TRUNCATE plain1, plain2, plain3, plain4;')
  echo "spread $S K=$K run=$R x100 TRUNCATE old+new: instructions=$1 seq_scan=$2" \
       "seq_tup_read=$3 confrelid_idx_tup_read=$4; control: instructions=$5;" \
       "lookup per TRUNCATE=$(( ($1 - $5) / 100 ))"
  ;;
list-setup)
  C=${2:?}; N=${3:?}; B=$(bin "$C"); D=$WORK/list-$C-$N
  newdb "$D" "$B"
  { printf 'CREATE FUNCTION f() RETURNS trigger LANGUAGE plpgsql AS $$BEGIN RETURN NEW; END$$;\n\n'
    for i in $(seq "$N"); do
      printf 'CREATE TABLE t%s (id int);\n\nCREATE TRIGGER g BEFORE INSERT ON t%s FOR EACH ROW EXECUTE FUNCTION f();\n\n' "$i" "$i"
      printf 'CREATE TABLE p%s (id int);\n\n' "$i"
    done; } | sb "$B" "$D" >/dev/null
  echo "list-setup $C N=$N pg_constraint rows/fk=$(fkrows "$B" "$D")"
  ;;
list)
  S=${2:?}; N=${3:?}; R=${4:?}; B=$(bin "$S"); D=$WORK/list-$(cat_of "$S")-$N; O=$WORK/li-$S-$N-$R
  reps=100; [ "$N" -ge 1000 ] && reps=10
  set -- $(measure "$B" "$D" "$O-trg" $reps "TRUNCATE $(seq -s, -f 't%g' 1 "$N" | sed 's/,$//');") \
         $(measure "$B" "$D" "$O-plain" $reps "TRUNCATE $(seq -s, -f 'p%g' 1 "$N" | sed 's/,$//');")
  echo "list $S N=$N run=$R x$reps TRUNCATE: instructions=$1 seq_tup_read=$3" \
       "confrelid_idx_tup_read=$4; control: instructions=$5;" \
       "per TRUNCATE: total=$(( $1 / reps )) lookup=$(( ($1 - $5) / reps ))"
  ;;
cascade-setup)
  C=${2:?}; N=${3:?}; B=$(bin "$C"); D=$WORK/cascade-$C-$N
  newdb "$D" "$B"
  { printf 'CREATE TABLE c1 (id int PRIMARY KEY);\n\n'
    for i in $(seq 2 "$N"); do
      printf 'CREATE TABLE c%s (id int PRIMARY KEY, p int REFERENCES c%s);\n\n' "$i" $((i - 1))
    done; } | sb "$B" "$D" >/dev/null
  echo "cascade-setup $C N=$N pg_constraint rows/fk=$(fkrows "$B" "$D")"
  ;;
cascade)
  S=${2:?}; N=${3:?}; R=${4:?}; B=$(bin "$S"); D=$WORK/cascade-$(cat_of "$S")-$N; O=$WORK/ca-$S-$N-$R
  set -- $(measure "$B" "$D" "$O" 3 'TRUNCATE c1 CASCADE;')
  # one row in every table, TRUNCATE c1 CASCADE, rows left must be 0
  left=$({ printf 'INSERT INTO c1 VALUES (1);\n\n'
           for i in $(seq 2 "$N"); do printf 'INSERT INTO c%s VALUES (1, 1);\n\n' "$i"; done
           printf 'TRUNCATE c1 CASCADE;\n\nSELECT %s AS left_rows;\n\n' \
             "$(seq -s+ -f '(SELECT count(*) FROM c%g)' 1 "$N")"; } |
         sb "$B" "$D" 2>/dev/null | sed -n 's/.*left_rows = "\([0-9]*\)".*/\1/p')
  echo "cascade $S N=$N run=$R x3 TRUNCATE c1 CASCADE (chain of $N tables):" \
       "instructions=$1 seq_scan=$2 seq_tup_read=$3 confrelid_idx_tup_read=$4;" \
       "per TRUNCATE=$(( $1 / 3 )); rows left after one more=$left"
  ;;
wide-setup)
  C=${2:?}; N=${3:?}; B=$(bin "$C"); D=$WORK/wide-$C-$N
  newdb "$D" "$B"
  { printf 'CREATE TABLE w0 (id int PRIMARY KEY);\n\n'
    for i in $(seq "$N"); do
      printf 'CREATE TABLE w%s (id int PRIMARY KEY, p int REFERENCES w0);\n\n' "$i"
    done; } | sb "$B" "$D" >/dev/null
  echo "wide-setup $C N=$N pg_constraint rows/fk=$(fkrows "$B" "$D")"
  ;;
wide)
  S=${2:?}; N=${3:?}; R=${4:?}; B=$(bin "$S"); D=$WORK/wide-$(cat_of "$S")-$N; O=$WORK/wi-$S-$N-$R
  set -- $(measure "$B" "$D" "$O" 3 'TRUNCATE w0 CASCADE;')
  left=$({ printf 'INSERT INTO w0 VALUES (1);\n\n'
           for i in $(seq "$N"); do printf 'INSERT INTO w%s VALUES (1, 1);\n\n' "$i"; done
           printf 'TRUNCATE w0 CASCADE;\n\nSELECT %s AS left_rows;\n\n' \
             "$(seq -s+ -f '(SELECT count(*) FROM w%g)' 0 "$N")"; } |
         sb "$B" "$D" 2>/dev/null | sed -n 's/.*left_rows = "\([0-9]*\)".*/\1/p')
  echo "wide $S N=$N run=$R x3 TRUNCATE w0 CASCADE ($N tables reference w0):" \
       "instructions=$1 seq_scan=$2 seq_tup_read=$3 confrelid_idx_tup_read=$4;" \
       "per TRUNCATE=$(( $1 / 3 )); rows left after one more=$left"
  ;;
part-setup)
  # pp partitioned in P partitions, referenced by Q tables: every FK has a
  # row per partition with conparentid set, which is what makes
  # heap_truncate_find_FKs() climb to the parent and restart
  C=${2:?}; P=${3:?}; Q=${4:?}; B=$(bin "$C"); D=$WORK/part-$C-$P-$Q
  newdb "$D" "$B"
  { printf 'CREATE TABLE pp (id int PRIMARY KEY) PARTITION BY RANGE (id);\n\n'
    for i in $(seq "$P"); do
      printf 'CREATE TABLE pp_%s PARTITION OF pp FOR VALUES FROM (%s) TO (%s);\n\n' \
        "$i" $(( i * 10 )) $(( i * 10 + 10 ))
    done
    for j in $(seq "$Q"); do
      printf 'CREATE TABLE pr_%s (id int, p int REFERENCES pp);\n\n' "$j"
    done; } | sb "$B" "$D" >/dev/null
  echo "part-setup $C P=$P Q=$Q pg_constraint rows/fk=$(fkrows "$B" "$D")"
  ;;
part)
  S=${2:?}; P=${3:?}; Q=${4:?}; R=${5:?}; B=$(bin "$S"); D=$WORK/part-$(cat_of "$S")-$P-$Q
  O=$WORK/pa-$S-$P-$Q-$R
  set -- $(measure "$B" "$D" "$O-all" 3 'TRUNCATE pp CASCADE;') \
         $(measure "$B" "$D" "$O-one" 3 'TRUNCATE pp_1 CASCADE;')
  left=$({ printf 'INSERT INTO pp VALUES (10);\n\n'
           for j in $(seq "$Q"); do printf 'INSERT INTO pr_%s VALUES (1, 10);\n\n' "$j"; done
           printf 'TRUNCATE pp_1 CASCADE;\n\nSELECT %s AS left_rows;\n\n' \
             "(SELECT count(*) FROM pp)+$(seq -s+ -f '(SELECT count(*) FROM pr_%g)' 1 "$Q")"; } |
         sb "$B" "$D" 2>/dev/null | sed -n 's/.*left_rows = "\([0-9]*\)".*/\1/p')
  echo "part $S P=$P Q=$Q run=$R x3 TRUNCATE pp CASCADE: instructions=$1 seq_scan=$2" \
       "seq_tup_read=$3 confrelid_idx_tup_read=$4; x3 TRUNCATE pp_1 CASCADE:" \
       "instructions=$5 seq_scan=$6 seq_tup_read=$7 confrelid_idx_tup_read=$8;" \
       "rows left after one more=$left"
  ;;
oncommit-setup)
  C=${2:?}; B=$(bin "$C"); D=$WORK/oc-$C
  newdb "$D" "$B"
  sb "$B" "$D" >/dev/null <<'SQL'
CREATE FUNCTION f() RETURNS trigger LANGUAGE plpgsql AS $$BEGIN RETURN NEW; END$$;
SQL
  echo "oncommit-setup $C pg_constraint rows/fk=$(fkrows "$B" "$D")"
  ;;
oncommit)
  # N ON COMMIT DELETE ROWS temp tables with a trigger, all truncated at
  # every commit that touches temp tables; per commit = (100 - 10) commits
  S=${2:?}; N=${3:?}; R=${4:?}; B=$(bin "$S"); D=$WORK/oc-$(cat_of "$S"); O=$WORK/oc-$S-$N-$R
  mk=$(for i in $(seq "$N"); do
         printf 'CREATE TEMP TABLE tt%s (a int) ON COMMIT DELETE ROWS;\n\n' "$i"
         printf 'CREATE TRIGGER g BEFORE INSERT ON tt%s FOR EACH ROW EXECUTE FUNCTION f();\n\n' "$i"
       done)
  set -- $(measure "$B" "$D" "$O-100" 100 'INSERT INTO tt1 VALUES (1);' "$mk") \
         $(measure "$B" "$D" "$O-10" 10 'INSERT INTO tt1 VALUES (1);' "$mk")
  echo "oncommit $S N=$N run=$R: x100 instructions=$1 seq_scan=$2 seq_tup_read=$3" \
       "confrelid_idx_tup_read=$4; x10 instructions=$5;" \
       "per commit=$(( ($1 - $5) / 90 ))"
  ;;
errpath-setup)
  # e1..eN with a trigger each; eN referenced by ex from outside the list
  C=${2:?}; N=${3:?}; B=$(bin "$C"); D=$WORK/err-$C-$N
  newdb "$D" "$B"
  { printf 'CREATE FUNCTION f() RETURNS trigger LANGUAGE plpgsql AS $$BEGIN RETURN NEW; END$$;\n\n'
    for i in $(seq "$N"); do
      printf 'CREATE TABLE e%s (id int PRIMARY KEY);\n\nCREATE TRIGGER g BEFORE INSERT ON e%s FOR EACH ROW EXECUTE FUNCTION f();\n\n' "$i" "$i"
    done
    printf 'CREATE TABLE ex (id int REFERENCES e%s);\n\n' "$N"; } | sb "$B" "$D" >/dev/null
  echo "errpath-setup $C N=$N pg_constraint rows/fk=$(fkrows "$B" "$D")"
  ;;
errpath)
  # TRUNCATE e1..eN fails on eN, the last one heap_truncate_check_FKs() tries
  S=${2:?}; N=${3:?}; R=${4:?}; B=$(bin "$S"); D=$WORK/err-$(cat_of "$S")-$N; O=$WORK/er-$S-$N-$R
  set -- $(MEASURE_ERRORS=10 measure "$B" "$D" "$O" 10 \
           "TRUNCATE $(seq -s, -f 'e%g' 1 "$N" | sed 's/,$//');")
  echo "errpath $S N=$N run=$R x10 failing TRUNCATE: instructions=$1 seq_scan=$2" \
       "seq_tup_read=$3 confrelid_idx_tup_read=$4; per TRUNCATE=$(( $1 / 10 ))"
  ;;
*) echo "stage?"; exit 2;;
esac