#!/usr/bin/env bash
# v02-matrix.sh <bindir> <workdir>
#
# What the progress views show for a backend while it runs commands that nest
# one progress report inside another (file_fdw scans, COPY in triggers,
# REPACK/CLUSTER/VACUUM FULL rebuilding indexes), and what is left behind
# after errors.
#
# It creates a throw-away cluster in <workdir> with the binaries in <bindir>
# (file_fdw must be installed there), writes the data files there (about
# 130 MB), and runs each cell in its own new session while a second session
# samples pg_stat_progress_{analyze,copy,create_index,vacuum,repack} for that
# backend every ~0.5 ms, calling pg_stat_clear_snapshot() between samples.
# The cluster is stopped at the end.
#
# Two kinds of cells:
#   cell <name> <sql>        one psql -c; sampling stops when the backend
#                            goes idle (same as cf6958-matrix.sh).
#   series <name> <sql>...   one psql session runs the statements in order;
#                            before statement i it sets application_name to
#                            <name>/i, so each sample is tagged with the
#                            step it was taken in.  Idle samples are kept and
#                            tagged "idle:".  Errors do not stop the session.
#
# Output per cell: "count N <label>" per label; then the samples as runs of
# the same label in time order ("run i: <label> xN"), with the first..last
# bytes_processed and progress counter seen in that run and how many times
# bytes_processed went down inside the run; stderr; "SERVER CRASH" plus the
# first TRAP line if the backend died.  Labels: A = pg_stat_progress_analyze
# (phase, relid), C = pg_stat_progress_copy (command, type, relid),
# I = pg_stat_progress_create_index (command, phase, relid),
# V = pg_stat_progress_vacuum (phase, relid), R = pg_stat_progress_repack
# (command, phase, relid); "C (command not set yet)" = a COPY row whose
# parameters are still zero, as between pgstat_progress_start_command() and
# the first update; "(none)" = no row in any of them.  The counter is
# tuples_processed (C), blocks_done (I), heap_tuples_scanned (R),
# heap_blks_scanned (V) or sample_blks_scanned (A).
#
# The last section measures TopMemoryContext of one session before and after
# 10,000 nested commands, 10,000 errors caught inside nested commands, and
# 10,000 top-level errors inside nested commands.
set -euo pipefail
trap 'kill $(jobs -p) 2>/dev/null; exit 143' TERM INT
BIN=$1 W=$2; rm -rf "$W"; mkdir -p "$W"
PGDATA=$W/data; PORT=$((61000 + RANDOM % 4000)); export PGHOST=$W PGPORT=$PORT PGDATABASE=postgres
"$BIN/initdb" -A trust -N -D "$PGDATA" >/dev/null
"$BIN/pg_ctl" -D "$PGDATA" -o "-p $PORT -k $W -c autovacuum=off" -l "$W/log" -w start >/dev/null
trap '"$BIN/pg_ctl" -D "$PGDATA" -w -m immediate stop >/dev/null 2>&1 || true' EXIT
echo "# $("$BIN/postgres" --version); asserts: $("$BIN/pg_config" --configure | grep -c enable-cassert)"
# big.txt: 1,500,000 rows x 8 columns; small.txt: 200,000 rows; bad.txt: one bad
# line at the end; three.txt: 3 rows; ints.txt: 20,000 rows; ints10k.txt:
# 10,000 rows; badsmall.txt: "1" then a bad line.
awk 'BEGIN{for(i=1;i<=1500000;i++) printf "%d,%d,%d,%d,%d,x%d,y%d,z%d\n",i,i%997,i%101,i*7%100003,i%13,i%5003,i%211,i}' >"$W/big.txt"
head -200000 "$W/big.txt" >"$W/small.txt"
{ head -100000 "$W/big.txt"; echo 'oops,1,1,1,1,a,b,c'; } >"$W/bad.txt"
printf '1\n2\n3\n' >"$W/three.txt"
seq 20000 >"$W/ints.txt"; seq 10000 >"$W/ints10k.txt"; printf '1\noops\n' >"$W/badsmall.txt"
COLS='a int, b int, c int, d int, e int, f text, g text, h text'
SQ="$BIN/psql -X -q -v ON_ERROR_STOP=1"
$SQ <<SQL
CREATE EXTENSION file_fdw;
CREATE SERVER f FOREIGN DATA WRAPPER file_fdw;
CREATE FOREIGN TABLE ft ($COLS) SERVER f OPTIONS (filename '$W/big.txt', format 'csv');
CREATE FOREIGN TABLE ft_small ($COLS) SERVER f OPTIONS (filename '$W/small.txt', format 'csv');
CREATE FOREIGN TABLE ft_prog ($COLS) SERVER f OPTIONS (program 'cat $W/big.txt', format 'csv');
CREATE FOREIGN TABLE ft_bad ($COLS) SERVER f OPTIONS (filename '$W/bad.txt', format 'csv');
CREATE FOREIGN TABLE ft3 (a int) SERVER f OPTIONS (filename '$W/three.txt');
CREATE TABLE t ($COLS);
CREATE TABLE parted ($COLS) PARTITION BY RANGE (a);
CREATE TABLE p_heap PARTITION OF parted FOR VALUES FROM (2000000) TO (3000000);
INSERT INTO p_heap SELECT g, 1, 1, 1, 1, 'x', 'y', 'z' FROM generate_series(2000000, 2300000) g;
CREATE FOREIGN TABLE p_ft PARTITION OF parted FOR VALUES FROM (1) TO (2000000)
  SERVER f OPTIONS (filename '$W/big.txt', format 'csv');
CREATE TABLE t3 (a int);
CREATE FUNCTION read_ft() RETURNS trigger LANGUAGE plpgsql AS
  \$\$ BEGIN PERFORM count(*) FROM ft_small; RETURN NEW; END \$\$;
CREATE TRIGGER t3_read_ft BEFORE INSERT ON t3 FOR EACH ROW EXECUTE FUNCTION read_ft();
CREATE TABLE t4 (a int);
CREATE FUNCTION copy_ft() RETURNS trigger LANGUAGE plpgsql AS
  \$\$ BEGIN COPY (SELECT a FROM ft_small) TO '/dev/null'; RETURN NEW; END \$\$;
CREATE TRIGGER t4_copy_ft BEFORE INSERT ON t4 FOR EACH ROW EXECUTE FUNCTION copy_ft();
-- new in v02
CREATE TABLE heapx (a int, b int);
INSERT INTO heapx SELECT g, g % 1000 FROM generate_series(1, 1000000) g;
CREATE TABLE small_heap (a int);
INSERT INTO small_heap VALUES (1), (2), (3);
CREATE TABLE tbad (a int);
CREATE TABLE t6 (a int);
CREATE FUNCTION copy_ft3() RETURNS trigger LANGUAGE plpgsql AS
  \$\$ BEGIN COPY (SELECT a FROM ft3) TO '/dev/null'; RETURN NEW; END \$\$;
CREATE TRIGGER t6_copy_ft3 BEFORE INSERT ON t6 FOR EACH ROW EXECUTE FUNCTION copy_ft3();
CREATE TABLE t6h (a int);
CREATE FUNCTION copy_heap() RETURNS trigger LANGUAGE plpgsql AS
  \$\$ BEGIN COPY (SELECT a FROM small_heap) TO '/dev/null'; RETURN NEW; END \$\$;
CREATE TRIGGER t6h_copy_heap BEFORE INSERT ON t6h FOR EACH ROW EXECUTE FUNCTION copy_heap();
CREATE TABLE t7 (a int);
CREATE FUNCTION copy_bad_caught() RETURNS trigger LANGUAGE plpgsql AS
  \$\$ BEGIN
       BEGIN COPY tbad FROM '$W/badsmall.txt';
       EXCEPTION WHEN others THEN NULL;
       END;
       RETURN NEW;
     END \$\$;
CREATE TRIGGER t7_copy_bad BEFORE INSERT ON t7 FOR EACH ROW EXECUTE FUNCTION copy_bad_caught();
CREATE TABLE t8 (a int);
CREATE FUNCTION copy_bad() RETURNS trigger LANGUAGE plpgsql AS
  \$\$ BEGIN COPY tbad FROM '$W/badsmall.txt'; RETURN NEW; END \$\$;
CREATE TRIGGER t8_copy_bad BEFORE INSERT ON t8 FOR EACH ROW EXECUTE FUNCTION copy_bad();
CREATE TABLE heapf (a int);
INSERT INTO heapf SELECT generate_series(1, 300000);
CREATE FUNCTION f_catch(x int) RETURNS int LANGUAGE plpgsql IMMUTABLE AS
  \$\$ BEGIN RETURN 1 / (x % 2); EXCEPTION WHEN division_by_zero THEN RETURN 0; END \$\$;
CREATE TABLE rp (a int, b int, c text);
INSERT INTO rp SELECT g, (g::bigint * 7919 % 2000000)::int, md5(g::text) FROM generate_series(1, 2000000) g;
CREATE INDEX rp_a ON rp (a);
CREATE INDEX rp_b ON rp (b);
CREATE INDEX rp_c ON rp (c);
CREATE FUNCTION observe(app text, series bool) RETURNS SETOF text
LANGUAGE plpgsql AS \$\$
DECLARE
  p int;
  t0 timestamptz := clock_timestamp();
  st text; an text; rowl text; lab text; b bigint; c bigint;
  seen bool := false;
  labs text[] := '{}'; ns int[] := '{}';
  b0 bigint[] := '{}'; b1 bigint[] := '{}'; c0 bigint[] := '{}'; c1 bigint[] := '{}';
  bd int[] := '{}';
  k int := 0; maxb bigint := 0;
BEGIN
  LOOP
    PERFORM pg_stat_clear_snapshot();
    SELECT pid INTO p FROM pg_stat_activity
     WHERE application_name LIKE app || '/%' AND pid <> pg_backend_pid();
    EXIT WHEN p IS NOT NULL;
    IF clock_timestamp() - t0 > interval '30 s' THEN
      RAISE EXCEPTION 'session % never showed up', app;
    END IF;
    PERFORM pg_sleep(0.0005);
  END LOOP;
  LOOP
    PERFORM pg_stat_clear_snapshot();
    st := NULL;
    SELECT a.state, a.application_name INTO st, an FROM pg_stat_activity a WHERE a.pid = p;
    EXIT WHEN st IS NULL;
    IF series THEN
      EXIT WHEN an = app || '/end';
    ELSIF st = 'active' THEN
      seen := true;
    ELSIF seen THEN
      EXIT;
    ELSE
      PERFORM pg_sleep(0.0005);
      CONTINUE;
    END IF;
    SELECT concat_ws(' | ',
      (SELECT 'A ' || phase || ' relid=' || relid::regclass ||
              CASE WHEN child_tables_total > 0
                   THEN ' children ' || child_tables_done || '/' || child_tables_total
                   ELSE '' END
         FROM pg_stat_progress_analyze WHERE pid = p),
      (SELECT 'C ' || coalesce(command || ' ' || type, '(command not set yet)') ||
              ' relid=' || relid::regclass
         FROM pg_stat_progress_copy WHERE pid = p),
      (SELECT 'I ' || command || ' ' || phase || ' relid=' || relid::regclass
         FROM pg_stat_progress_create_index WHERE pid = p),
      (SELECT 'V ' || phase || ' relid=' || relid::regclass
         FROM pg_stat_progress_vacuum WHERE pid = p),
      (SELECT 'R ' || command || ' ' || phase || ' relid=' || relid::regclass
         FROM pg_stat_progress_repack WHERE pid = p)),
      (SELECT bytes_processed FROM pg_stat_progress_copy WHERE pid = p),
      coalesce((SELECT tuples_processed FROM pg_stat_progress_copy WHERE pid = p),
               (SELECT blocks_done FROM pg_stat_progress_create_index WHERE pid = p),
               (SELECT heap_tuples_scanned FROM pg_stat_progress_repack WHERE pid = p),
               (SELECT heap_blks_scanned FROM pg_stat_progress_vacuum WHERE pid = p),
               (SELECT sample_blks_scanned FROM pg_stat_progress_analyze WHERE pid = p))
      INTO rowl, b, c;
    lab := CASE WHEN series THEN 'step ' || split_part(an, '/', 2) || ': ' ELSE '' END ||
           CASE WHEN st = 'active' THEN '' ELSE st || ': ' END ||
           coalesce(nullif(rowl, ''), '(none)');
    maxb := greatest(maxb, b);
    IF k = 0 OR labs[k] <> lab THEN
      k := k + 1;
      labs[k] := lab; ns[k] := 1; b0[k] := b; b1[k] := b; c0[k] := c; c1[k] := c; bd[k] := 0;
    ELSE
      ns[k] := ns[k] + 1;
      IF b < b1[k] THEN bd[k] := bd[k] + 1; END IF;
      b1[k] := b; c1[k] := c;
    END IF;
    PERFORM pg_sleep(0.0005);
  END LOOP;
  RETURN QUERY
    SELECT 'count ' || sum(n) || '  ' || l
      FROM unnest(labs, ns) WITH ORDINALITY u(l, n, o) GROUP BY l ORDER BY min(o);
  RETURN NEXT 'runs: ' || k;
  FOR i IN 1 .. k LOOP
    CONTINUE WHEN k > 60 AND i > 40 AND i <= k - 15;
    IF k > 60 AND i = 41 THEN RETURN NEXT '...'; END IF;
    RETURN NEXT 'run ' || i || ': ' || labs[i] || ' x' || ns[i] ||
      CASE WHEN b0[i] IS NOT NULL
           THEN ' [bytes ' || b0[i] || '..' || b1[i] ||
                CASE WHEN bd[i] > 0 THEN ', went down ' || bd[i] || 'x' ELSE '' END || ']'
           ELSE '' END ||
      CASE WHEN c0[i] IS NOT NULL THEN ' [n ' || c0[i] || '..' || c1[i] || ']' ELSE '' END;
  END LOOP;
  RETURN NEXT 'max bytes_processed seen: ' || maxb;
END \$\$;
CREATE FUNCTION topmem() RETURNS text LANGUAGE sql AS
  \$\$ SELECT 'TopMemoryContext total ' || total_bytes || ' used ' || used_bytes
       FROM pg_backend_memory_contexts WHERE name = 'TopMemoryContext' \$\$;
SQL
crashcheck() {
  if grep -q 'terminated by signal' "$W/log"; then
    echo "   SERVER CRASH: $(grep -m1 -E '^TRAP|terminated by signal' "$W/log")"
    : >"$W/log.seen"; cat "$W/log" >>"$W/log.seen"; : >"$W/log"
    for _ in $(seq 600); do $SQ -c 'SELECT 1' >/dev/null 2>&1 && break; sleep 0.1; done
  fi
}
report() {			# report <name>
  sed 's/^/   /' "$W/obs.$1"
  echo "   bytes written to the client: $(cat "$W/out.$1")"
  grep -v '^$' "$W/err.$1" | sed 's/^/   stderr: /' || true
  crashcheck
}
cell() {			# cell <name> <sql>
  local name=$1 sql=$2
  $SQ -At -c "SELECT * FROM observe('$name', false)" >"$W/obs.$name" 2>&1 & local o=$!
  PGAPPNAME=$name/1 PGOPTIONS='-c default_statistics_target=1000' \
    $BIN/psql -X -q -c "$sql" 2>"$W/err.$name" | wc -c >"$W/out.$name" || true
  wait $o || true
  echo "== $name: $sql"
  report "$name"
}
series() {			# series <name> <sql>...
  local name=$1 i=0; shift
  : >"$W/in.$name"
  echo "== $name (one session):"
  for s in "$@"; do
    i=$((i + 1))
    echo "SET application_name TO '$name/$i';" >>"$W/in.$name"
    echo "$s" >>"$W/in.$name"
    echo "   step $i: $s"
  done
  echo "SET application_name TO '$name/end';" >>"$W/in.$name"
  $SQ -At -c "SELECT * FROM observe('$name', true)" >"$W/obs.$name" 2>&1 & local o=$!
  PGAPPNAME=$name/0 PGOPTIONS='-c default_statistics_target=1000' \
    $BIN/psql -X -q -f "$W/in.$name" 2>"$W/err.$name" | wc -c >"$W/out.$name" || true
  wait $o || true
  report "$name"
}
PROBE='SELECT pg_sleep(0.3);'
echo "######## P1: the cf6958 matrix, unchanged"
cell analyze_ft          'ANALYZE ft'
cell analyze_partitioned 'ANALYZE parted'
cell vacuum_analyze_parted 'VACUUM ANALYZE parted'
cell analyze_all         'ANALYZE'
cell copy_to_select_ft   'COPY (SELECT a FROM ft) TO STDOUT'
cell copy_from_trigger   "COPY t3 FROM '$W/three.txt'"
cell copy_from_trigger_copy_to "COPY t4 FROM '$W/three.txt'"
cell copy_to_two_scans   'COPY (SELECT a FROM ft UNION ALL SELECT a FROM ft_small) TO STDOUT'
cell insert_select_ft    'INSERT INTO t SELECT * FROM ft'
cell ctas_ft             'CREATE TABLE ctas AS SELECT * FROM ft'
cell select_ft           'SELECT count(*) FROM ft'
cell explain_then_select 'EXPLAIN SELECT * FROM ft_small; SELECT count(*) FROM ft'
cell select_program      'SELECT count(*) FROM ft_prog'
cell copy_to_program     'COPY (SELECT a FROM ft_prog) TO STDOUT'
cell copy_to_plain       'COPY (SELECT a FROM t) TO STDOUT'
cell select_two_scans    'SELECT (SELECT count(*) FROM ft) + (SELECT count(*) FROM ft_small)'
cell select_rescan       'SELECT count(*) FROM generate_series(1, 3) g(i), LATERAL (SELECT a FROM ft WHERE a > i OFFSET 0) s'
cell select_error        'SELECT count(*) FROM ft_bad'
echo "######## P2: COPY TO over a file_fdw scan; TO PROGRAM keeps EndCopy waiting 1 s"
cell copy_to_ft_program_sleep "COPY (SELECT a FROM ft) TO PROGRAM 'cat >/dev/null; sleep 1'"
echo "######## P3: one file_fdw scan per statement, same session"
series alternate 'SELECT count(*) FROM ft;' 'SELECT count(*) FROM ft;' \
  'SELECT count(*) FROM ft;' 'SELECT count(*) FROM ft;'
series explain_first 'EXPLAIN SELECT * FROM ft_small;' 'SELECT count(*) FROM ft;'
echo "######## P4: two file_fdw scans that fail, then other commands in the same session"
series two_scans_error \
  'SELECT (SELECT count(*) FROM ft_bad) + (SELECT count(*) FROM ft_small);' \
  'CREATE INDEX heapx_a ON heapx (a);' 'VACUUM (DISABLE_PAGE_SKIPPING) rp;' 'ANALYZE heapx;' \
  "COPY heapx TO '/dev/null';" 'SELECT count(*) FROM ft;' 'DROP INDEX heapx_a;'
series control_new_session \
  'CREATE INDEX heapx_a ON heapx (a);' 'VACUUM (DISABLE_PAGE_SKIPPING) rp;' 'ANALYZE heapx;' \
  "COPY heapx TO '/dev/null';" 'SELECT count(*) FROM ft;' 'DROP INDEX heapx_a;'
series two_scans_cancel 'SET statement_timeout = 200;' \
  'SELECT (SELECT count(*) FROM ft) + (SELECT count(*) FROM ft_small);' \
  'RESET statement_timeout;' 'CREATE INDEX heapx_a ON heapx (a);' 'DROP INDEX heapx_a;'
echo "######## P5: generic plan with two file_fdw scans (merge join)"
series generic_plan 'SET plan_cache_mode = force_generic_plan;' \
  'SET enable_hashjoin = off;' 'SET enable_nestloop = off;' \
  'PREPARE q AS SELECT count(*) FROM ft a JOIN ft_small b ON a.a = b.a;' \
  'EXECUTE q;' 'EXECUTE q;' 'EXECUTE q;' 'EXECUTE q;' "$PROBE"
echo "######## P6: three levels; COPY FROM whose trigger runs COPY (SELECT ...) TO"
cell copy_from_trigger_copy_ft3 "COPY t6 FROM '$W/ints.txt'"
cell copy_from_trigger_copy_heap "COPY t6h FROM '$W/ints.txt'"
echo "######## P7: errors inside nested commands"
cell caught_error_in_trigger "COPY t7 FROM '$W/ints.txt'"
cell index_function_catches 'CREATE INDEX heapf_f ON heapf (f_catch(a))'
series toplevel_error_in_trigger "COPY t8 FROM '$W/three.txt';" "$PROBE" \
  'CREATE INDEX heapx_a ON heapx (a);' "$PROBE" 'DROP INDEX heapx_a;'
series caught_error_no_outer \
  "DO \$\$BEGIN COPY tbad FROM '$W/badsmall.txt'; EXCEPTION WHEN others THEN NULL; END\$\$;" \
  "$PROBE" 'CREATE INDEX heapx_a ON heapx (a);' "$PROBE" 'DROP INDEX heapx_a;'
series toplevel_error_copy_to_ft_bad "COPY (SELECT a FROM ft_bad) TO '/dev/null';" "$PROBE" \
  'CREATE INDEX heapx_a ON heapx (a);' "$PROBE" 'DROP INDEX heapx_a;'
echo "######## P9: REPACK, CLUSTER and VACUUM FULL of a 2M-row table with 3 indexes"
cell repack        'REPACK rp'
cell cluster       'CLUSTER rp USING rp_a'
cell vacuum_full   'VACUUM FULL rp'
echo "######## P8: TopMemoryContext of one session"
{
  echo "SELECT 'start: ' || topmem();"
  echo "COPY t6h FROM '$W/ints10k.txt';"
  echo "SELECT 'after 10000 nested COPY TO inside COPY FROM: ' || topmem();"
  echo "COPY t7 FROM '$W/ints10k.txt';"
  echo "SELECT 'after 10000 errors caught inside COPY FROM: ' || topmem();"
  echo "DO \$\$BEGIN FOR i IN 1..10000 LOOP BEGIN COPY tbad FROM '$W/badsmall.txt';"
  echo "  EXCEPTION WHEN others THEN NULL; END; END LOOP; END\$\$;"
  echo "SELECT 'after 10000 errors caught in a DO block: ' || topmem();"
  for _ in $(seq 10000); do echo "COPY t8 FROM '$W/three.txt';"; done
  echo "SELECT 'after 10000 top-level errors in a nested COPY: ' || topmem();"
} >"$W/in.mem"
$BIN/psql -X -q -At -f "$W/in.mem" 2>"$W/err.mem" | sed 's/^/   /'
echo "   errors: $(grep -c ERROR "$W/err.mem")"
crashcheck
echo "######## server log: TRAP / PANIC / terminated"
grep -hE 'TRAP|PANIC|terminated by signal' "$W/log" "$W/log.seen" 2>/dev/null | sort | uniq -c || true