#!/usr/bin/env bash
# David Rowley's pg_upgrade case at full scale (his script, his defaults:
# 323 partitioned tables x 100 partitions, 3,970 tables x 200 NOT NULL
# columns), upgraded from master to each build.
#
#   setup                 build the old cluster once, with master's binaries
#   upgrade <build> <run> pg_upgrade -b master -B <build>; every output line
#                         gets a timestamp, like "| ts" in David's mail
#
#   OLD_BIN=<master bindir> BINS="master=<dir> v10=<dir> ..." WORK=<dir>
#   REPRO_SQL=<path to pg_upgrade_catalog_repro.sql> bash david-repro.sh.txt ...
#
# With VACLOOP=1, the run also does what David tried as a workaround: VACUUM
# (FREEZE, SKIP_DATABASE_STATS) pg_constraint in the new cluster every 3 s
# while pg_upgrade restores the schema.
#
# The restore time is the one David quoted: from the line "Restoring global
# objects in the new cluster ... ok" to the line "Restoring database schemas
# in the new cluster ... ok".  pg_upgrade runs with --no-sync, so the fsync of
# the new cluster at the end does not enter the total either.
set -euo pipefail
: "${WORK:?}" "${OLD_BIN:?}" "${BINS:?}"
PORT=${PORT:-61540}
mkdir -p "$WORK"
bin() { local kv; for kv in $BINS; do [ "${kv%%=*}" = "$1" ] && { echo "${kv#*=}"; return; }; done; exit 2; }
stamp() { perl -MTime::HiRes=time -MPOSIX=strftime -ne \
  'BEGIN{$|=1} $t=time; printf "[%s.%03d] %s", strftime("%H:%M:%S", localtime $t), ($t-int $t)*1000, $_'; }

case ${1:?stage} in
setup)
  : "${REPRO_SQL:?}"
  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
  tr -d '\r' <"$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
  q() { "$OLD_BIN/psql" -h "$WORK" -p $PORT -U postgres -d test -X -qtA -F' ' -c "$1"; }
  echo "old cluster: $("$OLD_BIN/postgres" --version)"
  echo "pg_constraint by contype: $(q "SELECT coalesce(contype::text, 'all'), count(*) FROM pg_constraint
     GROUP BY ROLLUP (contype) ORDER BY 2 DESC" | tr '\n' ' ')"
  echo "partitions: $(q "SELECT count(*) FROM pg_class WHERE relpartbound IS NOT NULL")" \
       "partitioned tables: $(q "SELECT count(*) FROM pg_class WHERE relkind = 'p'")"
  "$OLD_BIN/pg_ctl" -D "$WORK/old" -m fast -w stop >/dev/null
  ;;
upgrade)
  S=${2:?build}; R=${3:?run}; B=$(bin "$S"); D=$WORK/new; U=$WORK/up-$S-$R
  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"
  # RSS of the servers pg_upgrade starts (old and new cluster) and their children only,
  # not of other clusters on the machine
  ( while sleep 5; do
      pids=$(for f in "$WORK/old/postmaster.pid" "$D/postmaster.pid"; do
               if [ -f "$f" ]; then p=$(head -1 "$f"); echo "$p"; ps --ppid "$p" -o pid= || :; fi; done)
      if [ -n "$pids" ]; then ps -o rss= -p "$(echo $pids | tr ' ' ',')" | awk '{s+=$1} END{print s+0}' || :; fi
    done ) >"$U/rss.log" &
  rp=$!
  vp=""
  if [ -n "${VACLOOP:-}" ]; then
    ( while sleep 3; do
        "$B/psql" -h "$U" -p $((PORT + 1)) -U postgres -d test -X -q \
          -c "VACUUM (FREEZE, SKIP_DATABASE_STATS) pg_constraint" >>"$U/vacloop.log" 2>&1 &&
          echo "VACUUM done $(date +%T)" >>"$U/vacloop.log" || :
      done ) &
    vp=$!
  fi
  t0=$(date +%s.%N)
  "$B/pg_upgrade" -b "$OLD_BIN" -B "$B" -d "$WORK/old" -D "$D" -U postgres \
    --no-sync -s "$U" -p $PORT -P $((PORT + 1)) 2>&1 | stamp >"$U/pg_upgrade.out" || true
  t1=$(date +%s.%N)
  kill $rp $vp 2>/dev/null || true
  grep -q 'Upgrade Complete' "$U/pg_upgrade.out" || { tail -20 "$U/pg_upgrade.out"; exit 1; }
  secs() { sed -n "s/^\[\([0-9:.]*\)\] $1.*/\1/p" "$U/pg_upgrade.out" |
           awk -F: '{print $1*3600 + $2*60 + $3}'; }
  g=$(secs 'Restoring global objects in the new cluster'); s=$(secs 'Restoring database schemas in the new cluster')
  restore=$(awk -v a="$g" -v b="$s" 'BEGIN{d=b-a; if (d<0) d+=86400; printf "%.1f", d}')
  total=$(awk -v a="$t0" -v b="$t1" 'BEGIN{printf "%.1f", b-a}')
  "$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"; }
  n=$(q "SELECT count(*) FROM pg_constraint")
  sz=$(q "SELECT pg_relation_size('pg_constraint') || ' ' ||
          coalesce((SELECT pg_relation_size(indexrelid)::text FROM pg_index
                    WHERE indrelid = 'pg_constraint'::regclass
                      AND indexrelid::regclass::text = 'pg_constraint_confrelid_index'), '-')")
  "$B/pg_ctl" -D "$D" -m fast -w stop >/dev/null
  set -- $sz
  [ -n "${VACLOOP:-}" ] && S="$S+vacfreeze($(grep -c '^VACUUM done' "$U/vacloop.log")vacuums)"
  echo "upgrade $S run=$R restore_schemas_s=$restore pg_upgrade_total_s=$total" \
       "pg_constraint_rows=$n pg_constraint_bytes=$1 confrelid_index_bytes=$2" \
       "peak_postgres_rss_kb=$(sort -n "$U/rss.log" | tail -1)"
  ;;
*) echo "stage?"; exit 2;;
esac